cranfield fix (#276)

* fixes cranfield query_ids * updated cranfield tests --------- Co-authored-by: Djoerd Hiemstra <[email protected]>
allenai · Nov 11, 2024 · b6da232 · b6da232
1 parent b9ae1f2
commit b6da232
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/ir_datasets/datasets/cranfield.py b/ir_datasets/datasets/cranfield.py
@@ -109,6 +109,7 @@ def queries_path(self):
     def queries_iter(self):
         with self.queries_dlc.stream() as stream:
             stream = io.TextIOWrapper(stream)
+            query_id = 1
             for lines in prefix_sentinel_splitter(stream, sentinel='.I '):
                 record = {'query_id': '', 'text': ''}
                 field = 'query_id'
@@ -118,7 +119,8 @@ def queries_iter(self):
                     else:
                         record[field] += line
                 record = {k: v.strip() for k, v in record.items()}
-                record['query_id'] = record['query_id'].lstrip('0') # remove leading 0s to match qrels
+                record['query_id'] = str(query_id) # overwrite query_id to match qrels
+                query_id += 1
                 yield GenericQuery(**record)
 
     def queries_cls(self):

diff --git a/test/integration/cranfield.py b/test/integration/cranfield.py
@@ -16,8 +16,8 @@ def test_docs(self):
     def test_queries(self):
         self._test_queries('cranfield', count=225, items={
             0: GenericQuery('1', 'what similarity laws must be obeyed when constructing aeroelastic models\nof heated high speed aircraft .'),
-            9: GenericQuery('18', 'are real-gas transport properties for air available over a wide range of\nenthalpies and densities .'),
-            224: GenericQuery('365', 'what design factors can be used to control lift-drag ratios at mach\nnumbers above 5 .'),
+            9: GenericQuery('10', 'are real-gas transport properties for air available over a wide range of\nenthalpies and densities .'),
+            224: GenericQuery('225', 'what design factors can be used to control lift-drag ratios at mach\nnumbers above 5 .'),
         })
 
     def test_qrels(self):