fixing doc_store

Parry-Parry · Nov 12, 2024 · e72b792 · e72b792
1 parent 464daaf
commit e72b792
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/rankers/datasets/dataset.py b/rankers/datasets/dataset.py
@@ -9,8 +9,8 @@
 from .corpus import Corpus
 
 class LazyTextLoader:
-    def __init__(self, docstore : callable) -> None:
-        self.docstore = docstore
+    def __init__(self, corpus : Union[Corpus, irds.Dataset]) -> None:
+        self.docstore = corpus.docs_store()
 
     def __getitem__(self, doc_id):
         if type == list: