Add download model simplifier

app.py

@@ -58,31 +58,16 @@ os.environ["TOKENIZERS_PARALLELISM"] = "false"
 # os.makedirs(project_dir, exist_ok=True)
 # huggingface_cache_dir = os.path.join(project_dir, "huggingface_models")
 # os.environ["HF_HOME"] = huggingface_cache_dir  # Use this folder for HF cache
-# 2b) Setup Hugging Face Cloud project model cache
-hf_cache_dir = "/home/user/.cache/huggingface"
-# Model storage location
-hf_cache_dir = "/home/user/.cache/huggingface"
-model_cache_dir = "/app/model_cache"
-os.environ["HF_HOME"] = hf_cache_dir
-os.environ["SENTENCE_TRANSFORMERS_HOME"] = hf_cache_dir
-# 3. Download (or load from cache) the SentenceTransformer model 
-from huggingface_hub import snapshot_download
-print("⏳ Checking or downloading the all-MiniLM-L6-v2 model from huggingface_hub...")
-# st.write("⏳ Checking or downloading the all-MiniLM-L6-v2 model from huggingface_hub...")
-# a) First, try loading from our copied cache
-if os.path.exists(model_cache_dir) and os.listdir(model_cache_dir):  # Check if model folder exists and is not empty
-    print(f"✅ Found cached model at {model_cache_dir}")
-    model_loc = model_cache_dir
-# b) Else, try loading backup from snapshot_download
+# 2. Setup Hugging Face Cloud project model cache
+MODEL_CACHE_DIR = "/app/model_cache"
+# Check if the required model files exist
+if os.path.exists(os.path.join(MODEL_CACHE_DIR, "config.json")):
+    print(f"✅ Found cached model at {MODEL_CACHE_DIR}")
+    model_loc = MODEL_CACHE_DIR
 else:
-    print(f"❌ Model not found in {model_cache_dir}. This should not happen!")
-    print("⚠️ Retrying with snapshot_download...")
-    model_loc = snapshot_download(
-        repo_id="sentence-transformers/all-MiniLM-L6-v2",
-        cache_dir=hf_cache_dir,
-        local_files_only=True # Change to `False` for fallback to online download
-    )
-# 4. Load the model to application
+    print(f"❌ Model not found in {MODEL_CACHE_DIR}. Critical error!")
+    exit(1)  # Exit since the model is missing
+# 3. Load the model to application
 from sentence_transformers import SentenceTransformer
 print("📥 **Loading Embedding Model...**")
 # st.write("📥 **Loading Embedding Model...**")

download_model.py

@@ -2,20 +2,19 @@ import os
 import shutil
 from huggingface_hub import snapshot_download
 
-# Define the target cache directory
+# Dir setup
+MODEL_REPO = "sentence-transformers/all-MiniLM-L6-v2"
 MODEL_CACHE_DIR = "/app/model_cache"
 
-# Download model
+# Download snapshots
 print("⏳ Downloading the SentenceTransformer model...")
-model_path = snapshot_download(repo_id="sentence-transformers/all-MiniLM-L6-v2", cache_dir=MODEL_CACHE_DIR)
+model_path = snapshot_download(repo_id=MODEL_REPO, cache_dir=MODEL_CACHE_DIR)
 
-# Find the snapshot folder
-snapshots_dir = os.path.join(model_path, "snapshots")
-if os.path.exists(snapshots_dir):
-    snapshot_subdirs = os.listdir(snapshots_dir)
-    if snapshot_subdirs:
-        snapshot_dir = os.path.join(snapshots_dir, snapshot_subdirs[0])
-        # Move all files to the main model cache directory
-        for filename in os.listdir(snapshot_dir):
-            shutil.move(os.path.join(snapshot_dir, filename), MODEL_CACHE_DIR)
-print(f"✅ Model downloaded and stored in {MODEL_CACHE_DIR}")
+# Ensure the model structure is correct
+snapshot_folders = os.path.join(model_path, "snapshots")
+if os.path.exists(snapshot_folders):
+    snapshot_dir = os.path.join(snapshot_folders, os.listdir(snapshot_folders)[0])  # Get first snapshot folder
+    for filename in os.listdir(snapshot_dir):
+        shutil.move(os.path.join(snapshot_dir, filename), MODEL_CACHE_DIR)  # Move files to /app/model_cache
+# Complete
+print(f"✅ Model downloaded and extracted to {MODEL_CACHE_DIR}")