move punkt download to vector store where it's actually used

README.md

@@ -104,18 +104,9 @@ Export the API key of the desired provider:
     export ANTHROPIC_API_KEY=...
     ```
 
-For easier configuration, create a `.sage-env` file with the following contents (change the API keys names based on your desired setup):
-```
-# 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=
+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>
 

sage/.sample-env

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

sage/index.py

@@ -14,20 +14,10 @@ from sage.embedder import build_batch_embedder_from_flags
 from sage.github import GitHubIssuesChunker, GitHubIssuesManager
 from sage.vector_store import build_vector_store_from_args
 
-
-from nltk.data import find
-
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger()
 logger.setLevel(logging.INFO)
 
-def is_punkt_downloaded():
-    try:
-        find('tokenizers/punkt_tab')
-        return True
-    except LookupError:
-        return False
-
 
 def main():
     parser = configargparse.ArgParser(
@@ -53,14 +43,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.")
 
-    # We need nltk tokenizers for 
-    if is_punkt_downloaded():
-        print("punkt is already downloaded")
-    else:
-        print("punkt is not downloaded")
-        # Optionally download it
-        nltk.download('punkt_tab')
-
+    
     ######################
     # 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."""
@@ -71,6 +78,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: