TREC CaST

.github/workflows/push.yml

@@ -14,7 +14,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: [3.7, 3.8, 3.9, '3.10']
+        python-version: ['3.8', '3.9', '3.10']
         os: ['ubuntu-latest', 'windows-latest', 'macOs-latest']
         architecture: ['x64']
 

ir_datasets/datasets/base.py

@@ -303,6 +303,9 @@ def queries_iter(self):
             if operator(query):
                 yield query
 
+    def queries_cls(self):
+        return self._queries_handler.queries_cls()
+
     def queries_handler(self):
         return self
 

ir_datasets/datasets/kilt.py

@@ -66,28 +66,26 @@ def __init__(self, streamer, count_hint=None):
 
     @ir_datasets.util.use_docstore
     def docs_iter(self):
-        with self._streamer.stream() as stream:
-            for doc in stream:
-                doc = json.loads(doc)
-                yield KiltDoc(
-                    doc['wikipedia_id'],
-                    doc['wikipedia_title'],
-                    ''.join(strip_markup(t) for t in doc['text']),
-                    tuple(doc['text']),
-                    tuple(KiltDocAnchor(
-                        a['text'],
-                        a['href'],
-                        a['paragraph_id'],
-                        a['start'],
-                        a['end']) for a in doc['anchors']),
-                    tuple(doc['categories'].split(',')),
-                    doc.get('wikidata_info', {}).get('wikidata_id', ''),
-                    str(doc['history']['revid']),
-                    doc['history']['timestamp'],
-                    str(doc['history']['parentid']),
-                    str(doc['history']['pageid']),
-                    doc['history']['url'],
-                )
+        for doc in self.docs_kilt_raw_iter():
+            yield KiltDoc(
+                doc['wikipedia_id'],
+                doc['wikipedia_title'],
+                ''.join(strip_markup(t) for t in doc['text']),
+                tuple(doc['text']),
+                tuple(KiltDocAnchor(
+                    a['text'],
+                    a['href'],
+                    a['paragraph_id'],
+                    a['start'],
+                    a['end']) for a in doc['anchors']),
+                tuple(doc['categories'].split(',')),
+                doc.get('wikidata_info', {}).get('wikidata_id', ''),
+                str(doc['history']['revid']),
+                doc['history']['timestamp'],
+                str(doc['history']['parentid']),
+                str(doc['history']['pageid']),
+                doc['history']['url'],
+            )
 
     def docs_cls(self):
         return KiltDoc
@@ -112,6 +110,11 @@ def docs_namespace(self):
     def docs_lang(self):
         return 'en'
 
+    def docs_kilt_raw_iter(self):
+        with self._streamer.stream() as stream:
+            for doc in stream:
+                yield json.loads(doc)
+
 
 def _init():
     base_path = ir_datasets.util.home_path()/NAME