Update app.py

app.py

@@ -1,36 +1,29 @@
+import subprocess
 from fastapi import FastAPI
-from pydantic import BaseModel, Field
-from llama_cpp import Llama
+from pydantic import BaseModel
+import json
 import os
 
-MODEL_PATH = "model/model.gguf"
+app = FastAPI()
 
-# GPU auto-detection: use cuBLAS if available, fallback to CPU
-GPU_LAYERS = 0
-if os.environ.get("CUDA_VISIBLE_DEVICES") not in [None, "", "None"]:
-    GPU_LAYERS = 20  # enables GPU acceleration for Qwen 0.5B
+MODEL_PATH = "/app/model/qwen2.5-0.5b-instruct-q4_k_m.gguf"
+LLAMA = "/app/llama.cpp/llama-cli"  # CLI mode → fastest for small CPUs
 
-llm = Llama(
-    model_path=MODEL_PATH,
-    n_ctx=2048,
-    n_threads=2,
-    n_gpu_layers=GPU_LAYERS,
-    n_batch=64,
-    verbose=False
-)
-
-api = FastAPI()
 
 class Query(BaseModel):
-    prompt: str = Field(..., min_length=1)
+    prompt: str
+    max_tokens: int = 128
+
+
+@app.post("/generate")
+def generate_text(data: Query):
+    cmd = [
+        LLAMA,
+        "-m", MODEL_PATH,
+        "-p", data.prompt,
+        "--n-predict", str(data.max_tokens),
+        "--temp", "0.2"
+    ]
 
-@api.post("/generate")
-def generate(q: Query):
-    output = llm(
-        q.prompt,
-        max_tokens=256,
-        temperature=0.2,
-        top_p=0.9,
-        stop=["</s>"]
-    )
-    return {"reply": output["choices"][0]["text"].strip()}
+    out = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
+    return {"output": out.stdout.strip()}