Retrieval Extremely Slow for Small Corpus

[aaaaaanya]

Good afternoon, I tried to implement the BM25s save and reloading of index and vocabulary as per the example given here: examples/save_and_reload_end_to_end.py. Despite using the "numba" backend and confirming that the index and vocab are loaded from disk in under a second, the actual retrieval step is taking 8+ seconds per query.
Here is my code for reference:

from pathlib import Path
import pandas as pd
import bm25s
import time
import Stemmer
from bm25s.tokenization import Tokenizer


class BM25sRetriever:
    def __init__(self,past_df: pd.DataFrame,k1: float = 1.5, b: float = 0.75,backend: str = "numba",stemmer_lang: str = "english"):

         # Validate required columns in the DataFrame
        required_columns = ["ticket_no", "ticket_description"]
        for col in required_columns:
            if col not in past_df.columns:
                raise ValueError(f"Missing required column '{col}' in the DataFrame.")
        
        self.df=past_df
        corpus = self.df.rename(columns={"ticket_no": "id", "ticket_description": "text"})[["id", "text"]].to_dict(orient="records")
        self.tokenizer = Tokenizer(splitter=lambda x: x.split())
        self.tokens = self.tokenizer.tokenize(self.df["ticket_description"].tolist(),return_as="tuple")
        self.retriever = bm25s.BM25(corpus=corpus, backend=backend,k1=k1,b=b)
        # Initialize stemmer and tokenizer
        stemmer = Stemmer.Stemmer(stemmer_lang)
        #self.reloaded_tokenizer = None
        #self.reloaded_retriever = None
        


    def build_bm25s_index(self):
        """
        Build a BM25S index from a list of texts using batching.
        Returns the retriever and tokenizer.
        """
        st= time.time()
        texts = self.df["ticket_description"].tolist()
        if not texts:
            raise ValueError("The input DataFrame is empty or does not contain 'ticket_description' column.")
        
        
        #tokens = bm25s.tokenize(texts,return_ids=True)
        self.retriever.index(self.tokens)
        et = time.time()
        print(f"BM25S index built in {et - st:.4f} seconds")
        return self.retriever, self.tokenizer
    

    def save_bm25s_index(self,save_dir: str):
        """
        Save the BM25S index and tokenizer/vocab to disk.
        """
        st= time.time()
        save_dir = Path(save_dir)
        save_dir.mkdir(parents=True, exist_ok=True)
        self.retriever.save(save_dir)
        self.tokenizer.save_vocab(save_dir)
        #self.tokenizer.save_stopwords(save_dir)
        et= time.time()
        print(f"BM25S index and vocab saved in {et - st:.4f} seconds")
        print(f"BM25S index and vocab saved to {save_dir}")


    def load_bm25s_index(self, save_dir: str, mmap: bool = True):
        """
        Loads a BM25S retriever and tokenizer from disk.
        """
        st=time.time()
        save_dir = Path(save_dir)
        reloaded_retriever = bm25s.BM25.load(save_dir, mmap=mmap, load_corpus=True)
        reloaded_tokenizer = Tokenizer(splitter=lambda x: x.split())
        reloaded_tokenizer.load_vocab(save_dir)
        print(f"Loaded BM25S index and vocab from {save_dir}")
        et=time.time()
        print(f"BM25S index and vocab loaded in {et - st:.4f} seconds")
        return reloaded_tokenizer,reloaded_retriever

            
        
    def bm25_retrieve(self, query, save_dir, top_k=10):
        """
        Retrieve top_k results for a query. Builds and saves index if not present.
        """
        
        save_dir = Path(save_dir)
        index_file = save_dir / "vocab.index.json"
        vocab_file = save_dir / "vocab.tokenizer.json"

        # Check if index and vocab exist
        if not (index_file.exists() and vocab_file.exists()):
            print("BM25S index or vocab not found. Building and saving index...")
            self.build_bm25s_index()
            self.save_bm25s_index(save_dir)      
        else:
            print("BM25S index and vocab found. Loading...")
        reloaded_tokenizer,reloaded_retriever= self.load_bm25s_index(save_dir)


        # Tokenize the query and retrieves
        query_tokens = reloaded_tokenizer.tokenize([query], update_vocab=False,return_as="tuple")
        #query_tokens = bm25s.tokenize(query,return_ids=True)
        
        
        reloaded_retriever.activate_numba_scorer()
        st= time.time()

        
        results, scores = reloaded_retriever.retrieve(query_tokens.ids, k=top_k)
        #results, scores = self.retriever.retrieve(query_tokens, k=top_k)
        et = time.time()

        docs = pd.DataFrame(list(results[0, :]))
        df_out = docs.rename(columns={"id": "ticket_no", "text": "ticket_description"})
        df_out["bm25s_score"] = scores[0, :]

        print(f"BM25S retrieval completed in {et - st:.4f} seconds")
        return df_out[["ticket_no", "ticket_description", "bm25s_score"]]
                    
                    
        

The retrieve function in particular i.e. this line:
results, scores = reloaded_retriever.retrieve(query_tokens.ids, k=top_k)
Is taking particularly long.
Is there a reason why this is taking so unusually long despite following the required steps?

[xhluca]

you are using reloaded_retriever.activate_numba_scorer() and loading the index/retriever every time you are retrieving something. this will take a lot of time to run (especiall the activate part). can you do that only at the beginning, so you don't need to load & activate the retriever every time you are retrieving something?

Also, I would generally appreciate if you can provide a minimally reproducible example. i.e.:

1. avoid using pandas throughout the code. i'd like to have a function named preprocess(...) that returns corpus, queries where you would handle the pandas processing. this isolates the time it takes for pandas to
2. share code I can just copy & run. here i don't have access to your pandas df so i can't even run your code. i cannot help without the sufficient code to run your script.
3. avoid classes in a bug report/Q&A. bm25s is designed to be used as is; there's generally no need to wrap it inside a class. it is already a class. write functions that preprocess your df, and use bm25 directly after that.