Optimize Hugging Face Space: add eager model loading, reduce max tokens, fix stop tokens, limit CPU threads

Dockerfile

@@ -6,7 +6,8 @@ ENV PYTHONUNBUFFERED=1 \
     PYTHONDONTWRITEBYTECODE=1 \
     PORT=7860 \
     HOME=/home/user \
-    USE_OLLAMA=false
+    USE_OLLAMA=false \
+    LLAMA_THREADS=2
 
 # Install system dependencies
 RUN apt-get update && apt-get install -y --no-install-recommends \

config.py

@@ -5,6 +5,6 @@ VECTOR_DB_PATH = "vector_store/faiss_index"
 EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
 CHUNK_SIZE = 500
 CHUNK_OVERLAP = 50
-MAX_TOKENS = 512
+MAX_TOKENS = 150
 TEMPERATURE = 0.7
 USE_OLLAMA = os.getenv("USE_OLLAMA", "True").lower() == "true"

llm/inference.py

@@ -117,7 +117,7 @@ def _generate_response_ollama(prompt: str) -> str:
         "options": {
             "num_predict": MAX_TOKENS,
             "temperature": TEMPERATURE,
-            "stop": ["Question:", "<|im_end|>", "<|im_start|>"]
+            "stop": ["Question:", "<|im_end|>", "<|im_start|>", "<|endoftext|>", "<|end_of_text|>"]
         }
     }
     
@@ -144,7 +144,7 @@ def generate_response(prompt: str) -> str:
             prompt,
             max_tokens=MAX_TOKENS,
             temperature=TEMPERATURE,
-            stop=["Question:", "<|im_end|>", "<|im_start|>"]
+            stop=["Question:", "<|im_end|>", "<|im_start|>", "<|endoftext|>", "<|end_of_text|>"]
         )
         text = output["choices"][0]["text"]
         return text.strip()

llm/model_loader.py

@@ -36,7 +36,8 @@ def get_llm() -> Llama:
             )
         try:
             cpu_count = os.cpu_count()
-            threads = max(1, min(4, cpu_count if cpu_count else 2))
+            default_threads = max(1, min(4, cpu_count if cpu_count else 2))
+            threads = int(os.getenv("LLAMA_THREADS", str(default_threads)))
             _llm_instance = Llama(
                 model_path=MODEL_PATH,
                 n_ctx=4096,

main.py

@@ -1,4 +1,5 @@
 import traceback
+from contextlib import asynccontextmanager
 
 from fastapi import FastAPI, Request
 from fastapi.responses import JSONResponse
@@ -8,8 +9,39 @@ from routes.upload import router as upload_router
 from routes.health import router as health_router
 
 from utils.logger import logger
+from llm.model_loader import get_llm
+from embeddings.embedding_model import get_embedding_model
+from embeddings.vector_store import load_vector_store
 
-app = FastAPI(title="AI Assistant")
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    # Eagerly load models on startup
+    logger.info("Eagerly loading LLM model on startup...")
+    try:
+        get_llm()
+        logger.info("LLM model loaded successfully!")
+    except Exception as e:
+        logger.error(f"Error loading LLM model on startup: {e}")
+
+    logger.info("Eagerly loading embedding model on startup...")
+    try:
+        get_embedding_model()
+        logger.info("Embedding model loaded successfully!")
+    except Exception as e:
+        logger.error(f"Error loading embedding model on startup: {e}")
+
+    logger.info("Eagerly loading vector store on startup...")
+    try:
+        load_vector_store()
+        logger.info("Vector store loaded successfully!")
+    except Exception as e:
+        logger.error(f"Error loading vector store on startup: {e}")
+
+    yield
+
+
+app = FastAPI(title="AI Assistant", lifespan=lifespan)
 
 app.include_router(chat_router)
 app.include_router(upload_router)