update

api.py

@@ -1,8 +1,9 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from llama_cpp import Llama
 import os
 import subprocess
+import threading
 
 MODEL_PATH = "/data/phi-3.gguf"
 MODEL_URL = (
@@ -11,44 +12,94 @@ MODEL_URL = (
     "phi-3-mini-4k-instruct.Q4_K_M.gguf"
 )
 
-app = FastAPI(title="llama.cpp Code API")
+app = FastAPI(title="llama.cpp Phi-3 API")
 
-def download_model():
-    if not os.path.exists(MODEL_PATH):
-        os.makedirs("/data", exist_ok=True)
-        subprocess.run([
-            "curl", "-L", "-o", MODEL_PATH, MODEL_URL
-        ], check=True)
+llm = None
+lock = threading.Lock()
 
-download_model()
 
-llm = Llama(
-    model_path=MODEL_PATH,
-    n_ctx=4096,
-    n_threads=2,        # HF CPU safe
-    n_batch=256,
-    n_gpu_layers=0,
-    verbose=False,
-)
+def ensure_model():
+    if os.path.exists(MODEL_PATH) and os.path.getsize(MODEL_PATH) > 100_000_000:
+        return
+
+    os.makedirs("/data", exist_ok=True)
+
+    result = subprocess.run(
+        ["curl", "-L", "--fail", "--retry", "3", "-o", MODEL_PATH, MODEL_URL],
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        text=True,
+    )
+
+    if result.returncode != 0 or not os.path.exists(MODEL_PATH):
+        raise RuntimeError(f"Model download failed: {result.stderr}")
+
+    if os.path.getsize(MODEL_PATH) < 100_000_000:
+        raise RuntimeError("Downloaded model file is corrupted or incomplete")
+
+
+def get_llm():
+    global llm
+    with lock:
+        if llm is None:
+            ensure_model()
+            llm = Llama(
+                model_path=MODEL_PATH,
+                n_ctx=4096,
+                n_threads=2,        # HF Space CPU safe
+                n_batch=256,
+                n_gpu_layers=0,
+                use_mmap=True,
+                use_mlock=False,
+                verbose=False,
+            )
+    return llm
+
 
 class GenerateRequest(BaseModel):
     prompt: str
     max_tokens: int = 512
     temperature: float = 0.1
+    top_p: float = 0.9
+
 
 @app.get("/")
 def root():
-    return {"message": "llama.cpp Phi-3 API ready"}
+    return {
+        "status": "ok",
+        "model_loaded": llm is not None,
+        "model_file_exists": os.path.exists(MODEL_PATH),
+    }
+
 
 @app.post("/generate")
 def generate(req: GenerateRequest):
-    output = llm(
-        f"<|user|>\n{req.prompt}\n<|assistant|>",
+    try:
+        model = get_llm()
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+    prompt = (
+        "<|system|>\n"
+        "You are an expert software engineer.\n"
+        "<|user|>\n"
+        f"{req.prompt}\n"
+        "<|assistant|>\n"
+    )
+
+    output = model(
+        prompt,
         max_tokens=req.max_tokens,
         temperature=req.temperature,
-        stop=["<|user|>"]
+        top_p=req.top_p,
+        stop=["<|user|>", "<|system|>"],
+        echo=False,
     )
-    return {"response": output["choices"][0]["text"].strip()}
+
+    return {
+        "response": output["choices"][0]["text"].strip()
+    }
+
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)