improved word tokenization for BM25Retriever

Author: GFJHogue

Dockerfile

@@ -27,6 +27,9 @@ RUN poetry config virtualenvs.in-project true
 # Install dependencies without dev dependencies
 RUN poetry install --no-root --without dev
 
+# Download NLTK data
+RUN poetry run python -m nltk.downloader punkt_tab
+
 # Adjust PATH to include the virtual environment's bin directory
 ENV PATH="${VENV_PATH}/bin:${PATH}"
 

poetry.lock

@@ -46,6 +46,7 @@ pydantic = "^2.10.5"
 pyyaml = "^6.0.2"
 tavily-python = "^0.5.0"
 openpyxl = "^3.1.5"
+nltk = "^3.9.1"
 
 [tool.poetry.group.dev.dependencies]
 ruff = "^0.7.1"

pyproject.toml

@@ -11,6 +11,7 @@
 from langchain_core.embeddings import Embeddings
 from langchain_core.language_models.chat_models import BaseChatModel
 from langchain_core.retrievers import BaseRetriever
+from nltk.tokenize import word_tokenize
 
 chroma_settings = chromadb.config.Settings(anonymized_telemetry=False)
 
@@ -38,7 +39,10 @@ def create_bm25_chroma_ensemble_retriever(
         loader = CSVLoader(file_path=reactome_csvs_dir / csv_file_name)
         data = loader.load()
         bm25_retriever = BM25Retriever.from_documents(
-            data, preprocess_func=lambda text: text.casefold().split()
+            data,
+            preprocess_func=lambda text: word_tokenize(
+                text.casefold(), language="english"
+            ),
         )
         bm25_retriever.k = 10