upload models and embeddings to hf hub

.gitignore

@@ -17,6 +17,7 @@ venv/
 
 .chroma/
 embeddings/
+chromadb.tar.gz
 data/
 models/
 other_models/

app/scripts/download_assets.py

@@ -8,6 +8,7 @@ import sys
 import tarfile
 import zipfile
 from pathlib import Path
+from src.settings import settings
 
 try:
     from huggingface_hub import snapshot_download, hf_hub_download
@@ -84,7 +85,8 @@ def download_chromadb(chromadb_repo: str, chromadb_dir: Path, hf_token: str | No
     print(f"Checking ChromaDB in {chromadb_dir}...")
     
     # Check if ChromaDB directory already has content
-    if chromadb_dir.exists() and any(chromadb_dir.iterdir()):
+    expected_chroma_path = chromadb_dir / settings.chroma_db
+    if expected_chroma_path.exists() and any(expected_chroma_path.iterdir()):
         print(f"ChromaDB directory already contains files. Skipping download.")
         print(f"To force re-download, delete {chromadb_dir} and restart.")
         return

app/src/settings.py

@@ -6,14 +6,6 @@ from pydantic_settings import BaseSettings
 load_dotenv(find_dotenv())
 
 
-SRC_DIR: str = os.path.dirname(__file__)
-DATA_DIR: str = os.path.join(SRC_DIR, "../../data")
-MODELS_DIR: str = os.path.join(SRC_DIR, "../../models")
-CHROMA_DIR: str = os.path.join(SRC_DIR, "../../.chroma")
-CHROMA_DB: str = os.path.join(CHROMA_DIR, "bge-small-finetuned-chroma")
-CHROMA_COLLECTION: str =  "bge_small_finetuned_astra_collection" # use bge_small_finetuned_astra_collection for the bege-large model and embeddings
-
-
 class Settings(BaseSettings):
     cohere_api_key: str = ""
     groq_api_key: str = ""
@@ -23,6 +15,16 @@ class Settings(BaseSettings):
     hf_models_repo: str = os.getenv("HF_MODELS_REPO", "")
     hf_chromadb_repo: str = os.getenv("HF_CHROMADB_REPO", "")
     hf_token: str = os.getenv("HF_TOKEN", "")
+    
+    chroma_db: str = os.getenv("CHROMA_DB", "")
+    chroma_collection: str = os.getenv("CHROMA_COLLECTION", "")
 
 
 settings = Settings()
+
+SRC_DIR: str = os.path.dirname(__file__)
+DATA_DIR: str = os.path.join(SRC_DIR, "../../data")
+MODELS_DIR: str = os.path.join(SRC_DIR, "../../models")
+CHROMA_DIR: str = os.path.join(SRC_DIR, "../../.chroma")
+CHROMA_DB: str = os.path.join(CHROMA_DIR, settings.chroma_db)
+CHROMA_COLLECTION: str =  settings.chroma_collection