Add bm25 service (primeqa#402)

* Add BM25 retriever and indexer to components and services

* add get_engine_type api

* save engine_type  in information.json

* Bump version: 0.10.0 → 0.11.0

Author: bhavani105

.bumpversion.cfg

@@ -1,5 +1,5 @@
 [bumpversion]
-current_version = 0.10.0
+current_version = 0.11.0
 commit = True
 
 [bumpversion:file:VERSION]

VERSION

@@ -1 +1 @@
-0.10.0
+0.11.0

primeqa/ir/README.md

@@ -169,7 +169,7 @@ The following command builds an index for BM25 retrieval.
 python primeqa/ir/run_ir.py \
     --do_index \
     --engine_type BM25 \
-    --corpus_path <document_collection> \
+    --collection <document_collection> \
     --index_path <index_dir>
     --threads <num_threads>
 ```

primeqa/ir/sparse/bm25_engine.py

@@ -12,11 +12,12 @@ class BM25Engine:
     def __init__(self, config: BM25Config):
         self.config = config
         logger.info(f"Running BM25")
+        logger.info(config)
 
     def do_index(self):
         logger.info("Running BM25 indexing")
         indexer = PyseriniIndexer()
-        rc = indexer.index_collection(self.config.corpus_path, self.config.index_location, 
+        rc = indexer.index_collection(self.config.collection, self.config.index_location, 
                     self.config.fieldnames, self.config.overwrite, 
                     self.config.threads, self.config.additional_indexing_args )
         logger.info(f"BM25 Indexing finished with rc: {rc}")

primeqa/ir/sparse/config.py

@@ -4,11 +4,11 @@
 @dataclass
 class IndexingArguments():
 
-    index_path: str = field(default=None, metadata={"help":"Path to the index directory location"})
+    index_location: str = field(default=None, metadata={"help":"Path to the index directory location"})
 
     overwrite: bool = field(default=False, metadata={"help": "Overwrite existing directory"})
 
-    corpus_path: str = field(default=None, metadata={"help":"Path to a corpus tsv or json file or directory"})
+    collection: str = field(default=None, metadata={"help":"Path to a corpus tsv or json file or directory"})
 
     fieldnames: list = field(default=None, metadata={"help":"fields names to use to identify document_id, title, text if corpus tsv has no headings"})
 

primeqa/ir/sparse/indexer.py

@@ -29,8 +29,8 @@ def _run_command(self, cmd):
         rc = process.wait()
         return rc
 
-    def _preprocess_corpus(self, corpus_path, tmpdirname, fieldnames=None):
-        reader = corpus_reader(corpus_path, fieldnames=fieldnames)
+    def _preprocess_corpus(self, collection, tmpdirname, fieldnames=None):
+        reader = corpus_reader(collection, fieldnames=fieldnames)
         outf = open( os.path.join(tmpdirname,"corpus_pyserini_fmt.jsonl"), 'w' )
         num_docs = 0
         for passage in tqdm(reader):
@@ -53,7 +53,7 @@ def _preprocess_corpus(self, corpus_path, tmpdirname, fieldnames=None):
 
 
         Args:
-            corpus_path (str) : path to file or directory of documents in tsv or jsonl format.
+            collection (str) : path to file or directory of documents in tsv or jsonl format.
             index_path (str) : output directory path where the index is written
             fieldnames ( List, Optional): column headers to be assigned to tsv without headers
             overwrite (bool, Optional): overwrite an existing directory, defaults to false
@@ -64,7 +64,7 @@ def _preprocess_corpus(self, corpus_path, tmpdirname, fieldnames=None):
 
 
         """
-    def index_collection(self, corpus_path: str, index_path: str, fieldnames=None, overwrite=False, 
+    def index_collection(self, collection: str, index_path: str, fieldnames=None, overwrite=False, 
             threads=1, additional_index_cmd_args='--storePositions --storeDocvectors --storeRaw' ):
         if not overwrite and os.path.exists(index_path) and os.listdir(index_path) :
             raise ValueError(f"Index path not empty '{index_path}' and overwrite not specified")
@@ -73,7 +73,7 @@ def index_collection(self, corpus_path: str, index_path: str, fieldnames=None, o
         # create temporary subdirectory for the corpus
         with tempfile.TemporaryDirectory(prefix='tmp',dir=index_path) as tmpdirname:
             # convert corpus documents to pyserini jsonl
-            num_docs = self._preprocess_corpus(corpus_path, tmpdirname, fieldnames=fieldnames)
+            num_docs = self._preprocess_corpus(collection, tmpdirname, fieldnames=fieldnames)
             # build index command
             cmd1 = f'python -m pyserini.index.lucene -collection JsonCollection ' + \
                 f'-generator DefaultLuceneDocumentGenerator ' + \
@@ -90,5 +90,6 @@ def index_collection(self, corpus_path: str, index_path: str, fieldnames=None, o
         logger.info(f"Index {index_path} contains {searcher.num_docs} documents")
         assert(searcher.num_docs == num_docs)
         logging.info(f"Index available at {index_path}")
+        searcher.close()
         return rc
 

primeqa/ir/sparse/retriever.py

@@ -6,55 +6,7 @@
 
 logger = logging.getLogger(__name__)
 
-class BaseRetriever(metaclass=ABCMeta):
-    """ 
-        Base class for Retriever
-    """
-
-    @abstractmethod
-    def retrieve(self, query: str, topK: Optional[int] = 10):
-        """
-
-        Run queries against the index to retrieve ranked list of documents
-        Return documents that are most relevant to the query.
-
-        Args:
-             query: search
-             top_k: number of hits to return, defaults to 10
-
-
-        Returns:
-             List of hits, each hit is a dict containing :
-             {
-                "rank": i,
-                "score": hit.score,
-                "doc_id": docid,
-                "title": title,
-                "text": text 
-            }
-                
-
-        """
-        pass
-
-    @abstractmethod
-    def batch_retrieve(self,  queries: List[str], qids: List[str], topK: int = 10, threads: int = 1):
-        """
-           Run a batch of queries 
-
-           Args:
-                queries:  list of query strings
-                qids:     list of qid strings corresponding to queries
-                top_k:    number of hits to return, defaults to 10
-                threads:  maximum number of threads to use
-                
-            Returns:
-                Dict of qid to hits
-                
-        """
-        pass
-
-class PyseriniRetriever(BaseRetriever):
+class PyseriniRetriever:
     def __init__(self, index_location: str, use_bm25: bool = True, k1: float = float(0.9), b: float = float(0.4)):
         """
         Initialize Pyserini retriever

primeqa/pipelines/components/base.py

@@ -54,6 +54,19 @@ class IndexerComponent(Component):
     @abstractmethod
     def index(self, collection: Union[List[dict], str], *args, **kwargs):
         pass
+    
+    @abstractmethod    
+    def get_engine_type() -> str:
+        """
+        Return this retriever engine type. Must match with the retriever tha will be used to query the index.
+
+        Raises:
+            NotImplementedError:
+
+        Returns:
+            str: engine type
+        """
+        raise NotImplementedError
 
 
 @dataclass(init=False, repr=False, eq=False)
@@ -91,3 +104,16 @@ def __hash__(self) -> int:
     @abstractmethod
     def retrieve(self, input_texts: List[str], *args, **kwargs):
         pass
+    
+    @abstractmethod
+    def get_engine_type() -> str:
+        """
+        Return this retriever engine type. Must match with the indexer used to generate the index.
+
+        Raises:
+            NotImplementedError:
+
+        Returns:
+            str: engine type
+        """
+        raise NotImplementedError

primeqa/pipelines/components/indexer/dense.py

@@ -139,3 +139,6 @@ def index(self, collection: Union[List[dict], str], *args, **kwargs):
             collection,
             overwrite="overwrite" in kwargs and kwargs["overwrite"],
         )
+    
+    def get_engine_type(self):
+        return "ColBERT"

primeqa/pipelines/components/indexer/sparse.py

@@ -1,14 +1,19 @@
 from typing import Union, List
-from dataclasses import dataclass
+from dataclasses import dataclass, field
+import json
 
 from primeqa.pipelines.components.base import IndexerComponent
+from primeqa.ir.sparse.indexer import PyseriniIndexer
+
 
 
 @dataclass
 class BM25Indexer(IndexerComponent):
     """_summary_
 
     Args:
+        index_root (str): Path to root directory where index to be stored.
+        index_name (str): Index name.
 
     Important:
     1. Each field has metadata property which can carry additional information for other downstream usages.
@@ -17,9 +22,49 @@ class BM25Indexer(IndexerComponent):
         b. exclude_from_hash (bool,optional): If set to True, that parameter is not considered while building the hash representation for the object. Defaults to False.
 
     """
-
+    
+    num_workers: int = field(
+        default=1,
+        metadata={
+            "name": "Number of worker threads",
+        },
+    )
+    
+    additional_index_args: str = field(
+        default='--storePositions --storeDocvectors --storeRaw',
+        metadata={
+            "name": "Additional index arguments",
+        },
+    )
+    
+    def __post_init__(self):
+        self._indexer = None
+    
+    def __hash__(self) -> int:
+            return hash(
+            f"{self.__class__.__name__}::{json.dumps({k: v.default for k, v in self.__class__.__dataclass_fields__.items() if not 'exclude_from_hash' in v.metadata or not v.metadata['exclude_from_hash']}, sort_keys=True)}"
+        )
+            
     def load(self, *args, **kwargs):
-        pass
+        self._index_path=f"{self.index_root}/{self.index_name}"
+        self._indexer = PyseriniIndexer()
 
     def index(self, collection: Union[List[dict], str], *args, **kwargs):
-        pass
+        if not isinstance(collection, str):
+                raise TypeError(
+                "Pyserini indexer expects path to `documents.tsv` as value for `collection` argument."
+            )
+                
+        self._indexer.index_collection(collection = collection, index_path=self._index_path, 
+            fieldnames=None, 
+            overwrite="overwrite" in kwargs and kwargs["overwrite"],
+            threads=kwargs["num_workers"] if "num_workers" in kwargs else 1, 
+            additional_index_cmd_args=kwargs["additional_index_args"] if "additional_index_args" in kwargs 
+                else '--storePositions --storeDocvectors --storeRaw' )
+    
+    def get_engine_type(self) -> str:
+        return "BM25"
+        
+
+        
+        

primeqa/pipelines/components/retriever/dense.py

@@ -127,3 +127,6 @@ def retrieve(self, input_texts: List[str], *args, **kwargs):
             [(result[0], result[-1]) for result in results_per_query]
             for results_per_query in ranking_results.data.values()
         ]
+        
+    def get_engine_type(self):
+        return "ColBERT"

primeqa/pipelines/components/retriever/sparse.py

@@ -1,7 +1,9 @@
 from typing import List
 from dataclasses import dataclass, field
+import json
 
 from primeqa.pipelines.components.base import RetrieverComponent
+from primeqa.ir.sparse.retriever import PyseriniRetriever
 
 
 @dataclass
@@ -29,22 +31,55 @@ class BM25Retriever(RetrieverComponent):
             "description": "Path to root directory where index is stored",
         },
     )
-    index_name: str = field(
+
+    max_num_documents: int = field(
+        default=5,
         metadata={
-            "name": "Index name",
+            "name": "Maximum number of retrieved documents",
+            "range": [1, 100, 1],
+            "api_support": True,
+            "exclude_from_hash": True,
         },
     )
-    max_num_documents: int = field(
-        default=5,
-        metadata={"name": "Maximum number of documents", "range": [1, 100, 1]},
+    
+    num_workers: int = field(
+        default=1,
+        metadata={
+            "name": "Num worker threads",
+            "range": [1, 100, 1],
+            "exclude_from_hash": True,
+        },
     )
 
     def __post_init__(self):
         # Placeholder variables
+        self._index_path=f"{self.index_root}/{self.index_name}"
         self._searcher = None
+        
+    def __hash__(self) -> int:
+        # Step 1: Identify all fields to be included in the hash
+        hashable_fields = [
+            k
+            for k, v in self.__class__.__dataclass_fields__.items()
+            if not "exclude_from_hash" in v.metadata
+            or not v.metadata["exclude_from_hash"]
+        ]
+
+        # Step 2: Run
+        return hash(
+            f"{self.__class__.__name__}::{json.dumps({k: v for k, v in vars(self).items() if k in hashable_fields}, sort_keys=True)}"
+        )
 
     def load(self, *args, **kwargs):
-        pass
+        self._searcher = PyseriniRetriever(self._index_path)
 
     def retrieve(self, input_texts: List[str], *args, **kwargs):
-        pass
+        qids = [str(idx) for  idx, query in enumerate(input_texts) ]
+        hits = self._searcher.batch_retrieve(input_texts, qids, topK=self.max_num_documents, threads=self.num_workers)
+        return [
+            [(result['doc_id'], result['score']) for result in results_per_query]
+            for results_per_query in hits.values()
+        ]
+    
+    def get_engine_type(self):
+        return "BM25"

primeqa/services/constants.py

@@ -2,6 +2,7 @@
 
 ATTR_INDEX_ID = "index_id"
 ATTR_STATUS = "status"
+ATTR_ENGINE_TYPE  ="engine_type"
 
 
 class IndexStatus(str, Enum):

primeqa/services/factories.py

@@ -12,19 +12,23 @@
 from primeqa.pipelines.components.reader.extractive import ExtractiveReader
 
 from primeqa.pipelines.components.retriever.dense import ColBERTRetriever
+from primeqa.pipelines.components.retriever.sparse import BM25Retriever
 
 from primeqa.pipelines.components.indexer.dense import ColBERTIndexer
+from primeqa.pipelines.components.indexer.sparse import BM25Indexer
 
 READERS_REGISTRY = {
     ExtractiveReader.__name__: ExtractiveReader,
 }
 
 RETRIEVERS_REGISTRY = {
     ColBERTRetriever.__name__: ColBERTRetriever,
+    BM25Retriever.__name__: BM25Retriever,
 }
 
 INDEXERS_REGISTRY = {
     ColBERTIndexer.__name__: ColBERTIndexer,
+    BM25Indexer.__name__: BM25Indexer,
 }