pre-downloading-all-models-at-startup

.gitignore

@@ -52,4 +52,5 @@ build/
 Thumbs.db
 
 # Gemini Plans
-gemini_plans/
+gemini_plans/
+llm_app_rework.md

Dockerfile

@@ -9,14 +9,22 @@ ENV HF_TOKEN=""
 COPY requirements.txt .
 RUN pip install --no-cache-dir -r requirements.txt
 
-
+# Pre-download all local models during build
 RUN --mount=type=secret,id=HF_TOKEN \
     export HF_TOKEN=$(cat /run/secrets/HF_TOKEN) && \
-    echo "--- Docker RUN: Downloading model using huggingface-cli..." && \
+    echo "--- Downloading Bielik-1.5B..." && \
     huggingface-cli download speakleash/Bielik-1.5B-v3.0-Instruct \
-        --local-dir ${MODEL_DIR} \
+        --local-dir ${MODEL_DIR}/bielik-1.5b \
+        --local-dir-use-symlinks=False && \
+    echo "--- Downloading Qwen2.5-3B..." && \
+    huggingface-cli download Qwen/Qwen2.5-3B-Instruct \
+        --local-dir ${MODEL_DIR}/qwen2.5-3b \
+        --local-dir-use-symlinks=False && \
+    echo "--- Downloading Gemma-2-2B..." && \
+    huggingface-cli download google/gemma-2-2b-it \
+        --local-dir ${MODEL_DIR}/gemma-2-2b \
         --local-dir-use-symlinks=False && \
-    echo "--- Docker RUN: Model download complete."
+    echo "--- All models downloaded."
 
 
 COPY . .

app/models/registry.py

@@ -16,18 +16,21 @@ from app.models.huggingface_inference_api import HuggingFaceInferenceAPI
 MODEL_CONFIG = {
     "bielik-1.5b": {
         "id": "speakleash/Bielik-1.5B-v3.0-Instruct",
+        "local_path": "bielik-1.5b",
         "type": "local",
         "polish_support": "excellent",
         "size": "1.5B",
     },
     "qwen2.5-3b": {
         "id": "Qwen/Qwen2.5-3B-Instruct",
+        "local_path": "qwen2.5-3b",
         "type": "local",
         "polish_support": "good",
         "size": "3B",
     },
     "gemma-2-2b": {
         "id": "google/gemma-2-2b-it",
+        "local_path": "gemma-2-2b",
         "type": "local",
         "polish_support": "medium",
         "size": "2B",
@@ -40,8 +43,8 @@ MODEL_CONFIG = {
     },
 }
 
-# For local model override (when model is pre-downloaded in container)
-LOCAL_MODEL_PATH = os.getenv("LOCAL_MODEL_PATH", "/app/pretrain_model")
+# Base path for pre-downloaded models in container
+LOCAL_MODEL_BASE = os.getenv("MODEL_DIR", "/app/pretrain_model")
 
 
 class ModelRegistry:
@@ -66,9 +69,14 @@ class ModelRegistry:
         model_type = config["type"]
         model_id = config["id"]
         
-        # Special case: local Bielik uses pre-downloaded path
-        if name == "bielik-1.5b" and os.path.exists(LOCAL_MODEL_PATH):
-            model_id = LOCAL_MODEL_PATH
+        # For local models, check if pre-downloaded version exists
+        if model_type == "local" and "local_path" in config:
+            local_path = os.path.join(LOCAL_MODEL_BASE, config["local_path"])
+            if os.path.exists(local_path):
+                print(f"Using pre-downloaded model at: {local_path}")
+                model_id = local_path
+            else:
+                print(f"Pre-downloaded model not found at {local_path}, will download from HuggingFace")
         
         if model_type == "local":
             return HuggingFaceLocal(