download nltk if not detected (#42)

README.md

@@ -89,7 +89,9 @@ pip install git+https://github.com/Storia-AI/sage.git@main
     export PINECONE_INDEX_NAME=...
     ```
 
-3. For reranking, we support <a href="https://developer.nvidia.com/blog/enhancing-rag-pipelines-with-re-ranking/">NVIDIA</a>, <a href="https://docs.voyageai.com/docs/reranker">Voyage</a>, <a href="https://cohere.com/rerank">Cohere</a>, and <a href="https://jina.ai/reranker/">Jina</a>. According to [our experiments](benchmark/retrieval/README.md), NVIDIA performs best. Export the API key of the desired provider:
+3. For reranking, we support <a href="https://developer.nvidia.com/blog/enhancing-rag-pipelines-with-re-ranking/">NVIDIA</a>, <a href="https://docs.voyageai.com/docs/reranker">Voyage</a>, <a href="https://cohere.com/rerank">Cohere</a>, and <a href="https://jina.ai/reranker/">Jina</a>. According to [our experiments](benchmark/retrieval/README.md), NVIDIA performs best. Note: for NVIDIA you should use the `nvidia/nv-rerankqa-mistral-4b-v3` reranker. 
+
+Export the API key of the desired provider:
     ```
     export NVIDIA_API_KEY=...  # or
     export VOYAGE_API_KEY=...  # or
@@ -102,6 +104,10 @@ pip install git+https://github.com/Storia-AI/sage.git@main
     export ANTHROPIC_API_KEY=...
     ```
 
+For easier configuration, adapt the entries within the sample `.sage-env` (change the API keys names based on your desired setup) and run:
+```
+source .sage-env
+```
 </details>
 
 ### Optional

sage/.sage-env

@@ -0,0 +1,10 @@
+# Embeddings
+export OPENAI_API_KEY=
+# Vector store
+export PINECONE_API_KEY=
+# Reranking
+export NVIDIA_API_KEY=
+# Generation LLM
+export ANTHROPIC_API_KEY=
+# Github issues
+export GITHUB_TOKEN=

sage/.sample-env

@@ -1,3 +0,0 @@
-OPENAI_API_KEY=
-PINECONE_API_KEY=
-GITHUB_TOKEN=

sage/index.py

@@ -42,6 +42,7 @@ def main():
     if args.embedding_provider == "marqo" and args.vector_store_provider != "marqo":
         parser.error("When using the marqo embedder, the vector store type must also be marqo.")
 
+    
     ######################
     # Step 1: Embeddings #
     ######################

sage/vector_store.py

@@ -12,6 +12,7 @@ from langchain_community.vectorstores import Marqo
 from langchain_community.vectorstores import Pinecone as LangChainPinecone
 from langchain_core.documents import Document
 from langchain_core.embeddings import Embeddings
+from nltk.data import find
 from pinecone import Pinecone, ServerlessSpec
 from pinecone_text.sparse import BM25Encoder
 
@@ -20,6 +21,12 @@ from sage.data_manager import DataManager
 
 Vector = Tuple[Dict, List[float]]  # (metadata, embedding)
 
+def is_punkt_downloaded():
+    try:
+        find('tokenizers/punkt_tab')
+        return True
+    except LookupError:
+        return False
 
 class VectorStore(ABC):
     """Abstract class for a vector store."""
@@ -69,6 +76,13 @@ class PineconeVectorStore(VectorStore):
         if alpha < 1.0:
             if bm25_cache and os.path.exists(bm25_cache):
                 logging.info("Loading BM25 encoder from cache.")
+                # We need nltk tokenizers for bm25 tokenization
+                if is_punkt_downloaded():
+                    print("punkt is already downloaded")
+                else:
+                    print("punkt is not downloaded")
+                    # Optionally download it
+                    nltk.download('punkt_tab')
                 self.bm25_encoder = BM25Encoder()
                 self.bm25_encoder.load(path=bm25_cache)
             else: