Merge pull request #34 from allegro/data_loading_fix

Fix for data loading with non-ascending query ids
allegro · May 17, 2021 · 2923985 · 2923985
2 parents 5fc2eb7 + bbf3f87
commit 2923985
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/allrank/data/dataset_loading.py b/allrank/data/dataset_loading.py
@@ -106,7 +106,8 @@ def __init__(self, X, y, query_ids, transform=None):
         """
         X = X.toarray()
 
-        groups = np.cumsum(np.unique(query_ids, return_counts=True)[1])
+        _, indices, counts = np.unique(query_ids, return_index=True, return_counts=True)
+        groups = np.cumsum(counts[np.argsort(indices)])
 
         self.X_by_qid = np.split(X, groups)[:-1]
         self.y_by_qid = np.split(y, groups)[:-1]

diff --git a/setup.py b/setup.py
@@ -22,7 +22,7 @@
 
 setup(
     name="allRank",
-    version="1.4.2",
+    version="1.4.3",
     description="allRank is a framework for training learning-to-rank neural models",
     long_description=README,
     long_description_content_type="text/markdown",