Use custom app

Dockerfile

@@ -17,6 +17,6 @@ RUN wget -q -O model.gguf "https://huggingface.co/dicta-il/dictalm2.0-instruct-G
 
 COPY --chown=user . .
 
-# Run llama-cpp-python server
+# Run FastAPI server
 # Port 7860 is required for HF Spaces
-CMD ["python", "-m", "llama_cpp.server", "--model", "model.gguf", "--host", "0.0.0.0", "--port", "7860", "--n_ctx", "4096", "--chat_format", "chatml"]
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

app.py

@@ -0,0 +1,73 @@
+from fastapi import FastAPI, Request, HTTPException
+import uvicorn
+from llama_cpp import Llama
+
+app = FastAPI()
+
+print("Loading Llama model...")
+llm = Llama(
+    model_path="model.gguf",
+    n_ctx=4096,
+    n_gpu_layers=0,
+)
+print("Model loaded successfully!")
+
+@app.post("/v1/chat/completions")
+async def chat_completions(request: Request):
+    try:
+        data = await request.json()
+    except:
+        raise HTTPException(status_code=400, detail="Invalid JSON body")
+        
+    messages = data.get("messages", [])
+    max_tokens = data.get("max_tokens", 1024)
+    temperature = data.get("temperature", 0.7)
+    
+    # Construct Dicta-LM Prompt
+    prompt = "<s>"
+    for m in messages:
+        if m["role"] == "user":
+            prompt += f"[INST] {m['content']} [/INST]\n"
+        elif m["role"] == "assistant":
+            prompt += f"{m['content']}</s> "
+            
+    try:
+        response = llm(
+            prompt=prompt,
+            max_tokens=max_tokens,
+            temperature=temperature,
+            stop=["</s>", "[INST]", "[/INST]"],
+            echo=False
+        )
+        
+        return {
+            "id": response["id"],
+            "object": "chat.completion",
+            "created": response["created"],
+            "model": "dictalm2.0-instruct",
+            "choices": [
+                {
+                    "index": 0,
+                    "message": {
+                        "role": "assistant",
+                        "content": response["choices"][0]["text"].strip()
+                    },
+                    "finish_reason": response["choices"][0]["finish_reason"]
+                }
+            ],
+            "usage": response["usage"]
+        }
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.get("/v1/models")
+def get_models():
+    return {
+        "data": [
+            {"id": "dictalm2.0-instruct", "object": "model", "owned_by": "dicta"}
+        ]
+    }
+
+@app.get("/")
+def health_check():
+    return {"status": "running"}

requirements.txt

@@ -1,3 +1,5 @@
-llama-cpp-python[server]==0.2.90
+llama-cpp-python
 huggingface_hub
 jinja2
+fastapi
+uvicorn