Update main.py

main.py

@@ -3,9 +3,11 @@ from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 from llama_cpp import Llama
 from huggingface_hub import hf_hub_download
+import os
 
 app = FastAPI()
 
+# CORS (allow all for now)
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -13,36 +15,68 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# Download Aevis-Medical-SLM
-try:
-    print("📥 Downloading model...")
-    model_path = hf_hub_download(
-        repo_id="devray11/Aevis-Medical-SLM",
-        filename="DeepSeek-R1-Distill-Llama-8B.Q4_K_M.gguf"
-    )
-    # Initialize Model (Optimized for 2-core CPU)
-    llm = Llama(model_path=model_path, n_ctx=1024, n_threads=2)
-    print("✅ Model Loaded Successfully")
-except Exception as e:
-    print(f"❌ Error: {e}")
-    llm = None
+# Global model variable
+llm = None
+
+def load_model():
+    global llm
+    try:
+        print("📥 Downloading model from Hugging Face...")
+
+        model_path = hf_hub_download(
+            repo_id="devray11/Aevis-Medical-SLM",
+            filename="DeepSeek-R1-Distill-Llama-8B.Q4_K_M.gguf"
+        )
+
+        print("⚙️ Initializing model...")
+
+        llm = Llama(
+            model_path=model_path,
+            n_ctx=512,          # Reduced for low RAM
+            n_threads=2,        # HF free CPU = 2 cores
+            n_batch=128,
+            use_mmap=True,
+            use_mlock=False
+        )
+
+        print("✅ Model Loaded Successfully")
+
+    except Exception as e:
+        print(f"❌ Model Load Error: {e}")
+        llm = None
+
+
+# Load model at startup
+load_model()
+
 
 class Query(BaseModel):
     prompt: str
 
+
 @app.post("/generate")
 async def generate(query: Query):
-    if not llm:
+    global llm
+
+    if llm is None:
         return {"error": "Model not initialized"}
-    
-    output = llm(
-        f"### Instruction:\n{query.prompt}\n\n### Response:\n",
-        max_tokens=256,
-        stop=["###"],
-        echo=False
-    )
-    return {"response": output["choices"][0]["text"]}
+
+    try:
+        output = llm(
+            f"### Instruction:\n{query.prompt}\n\n### Response:\n",
+            max_tokens=128,     # Reduced for speed
+            stop=["###"],
+            echo=False
+        )
+
+        return {
+            "response": output["choices"][0]["text"].strip()
+        }
+
+    except Exception as e:
+        return {"error": str(e)}
+
 
 @app.get("/")
 def health():
-    return {"status": "Aevis API is running"}
+    return {"status": "Aevis API is running 🚀"}