hmm need to sort the flow

Parry-Parry · Nov 13, 2024 · 9fcc7a7 · 9fcc7a7
1 parent 1fb36dd
commit 9fcc7a7
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/rankers/datasets/dataset.py b/rankers/datasets/dataset.py
@@ -134,7 +134,8 @@ def __getitem__(self, idx):
                     texts, scores = texts[:self.group_size], scores[:self.group_size]
                 else:
                     texts, scores = zip(*random.sample(list(zip(doc_id_a_text + doc_id_b_text, doc_id_a_scores + doc_id_b_scores)), self.group_size))
-            return (query, texts, scores)
+                return (query, texts, scores)
+            return (query, doc_id_a_text + doc_id_b_text, doc_id_a_scores + doc_id_b_scores)
         else:
             if len(doc_id_b_text) > (self.n_neg): doc_id_b_text = random.sample(doc_id_b_text, self.group_size)
             return (query, doc_id_a_text + doc_id_b_text)