rsdt

app.py

@@ -1,51 +1,50 @@
-from fastapi import FastAPI
-from pydantic import BaseModel
-from llama_cpp import Llama
 import os
 import requests
+from fastapi import FastAPI, Request
+from fastapi.responses import JSONResponse
+from llama_cpp import Llama
 
-app = FastAPI()
-
-# === Constants ===
-MODEL_REPO="nilbot/gemma-2b-it-Q4_K.gguf"
-MODEL_FILE="gemma-2b-it-Q4_K.gguf"
-MODEL_URL = f"https://huggingface.co/{MODEL_REPO}/resolve/main/{MODEL_FILE}"
-MODEL_DIR = "./models"
-MODEL_PATH = os.path.join(MODEL_DIR, MODEL_FILE)
-
+REPO_ID = "TheBloke/TinyLlama-1.1B-Chat-v1.0-GGUF"
+MODEL_FILENAME = "tinyllama-1.1b-chat-v1.0.Q4_K_M.gguf"
+MODEL_PATH = f"./models/{MODEL_FILENAME}"
 HF_TOKEN = os.getenv("HF_TOKEN")
 
-# === Create model directory ===
-os.makedirs(MODEL_DIR, exist_ok=True)
+# Ensure models folder exists
+os.makedirs("./models", exist_ok=True)
+os.chmod("./models", 0o777)  # ensure write access
 
-# === Manual download of GGUF ===
+# Download model if missing
 if not os.path.exists(MODEL_PATH):
-    print("📦 Downloading GGUF model manually from Hugging Face...")
+    print("📦 Downloading TinyLlama Q4_K_M model...")
+    url = f"https://huggingface.co/{REPO_ID}/resolve/main/{MODEL_FILENAME}"
     headers = {"Authorization": f"Bearer {HF_TOKEN}"}
-    response = requests.get(MODEL_URL, headers=headers, stream=True)
-    if response.status_code != 200:
-        raise RuntimeError(f"❌ Failed to download model. Status: {response.status_code}")
+    r = requests.get(url, headers=headers, stream=True)
+    if r.status_code != 200:
+        raise RuntimeError(f"❌ Download failed: {r.status_code} {r.text[:200]}")
     with open(MODEL_PATH, "wb") as f:
-        for chunk in response.iter_content(chunk_size=8192):
+        for chunk in r.iter_content(8192):
             f.write(chunk)
-    print(f"✅ Model downloaded to {MODEL_PATH}")
+    print("✅ Model downloaded")
+
+# Load into llama-cpp
+print("🔧 Loading TinyLlama model...")
+llm = Llama(
+    model_path=MODEL_PATH,
+    n_ctx=512,
+    n_threads=os.cpu_count() or 1
+)
 
-# === Load model ===
-print("🔧 Loading GGUF model...")
-llm = Llama(model_path=MODEL_PATH, n_ctx=512, n_threads=os.cpu_count())
+app = FastAPI()
 
-# === Inference ===
-class PromptRequest(BaseModel):
-    prompt: str
-    max_tokens: int = 256
-    temperature: float = 0.7
+@app.get("/")
+async def root():
+    return {"status": "🟢 TinyLlama-1.1B Q4_K_M is ready"}
 
 @app.post("/prompt")
-def generate_prompt(req: PromptRequest):
-    output = llm(
-        prompt=req.prompt,
-        max_tokens=req.max_tokens,
-        temperature=req.temperature,
-        stop=["</s>"],
-    )
-    return {"response": output["choices"][0]["text"].strip()}
+async def prompt(req: Request):
+    body = await req.json()
+    prompt = body.get("prompt") or ""
+    if not prompt:
+        return JSONResponse(status_code=400, content={"error": "Missing 'prompt' field"})
+    resp = llm(prompt, max_tokens=512, stop=["</s>"])
+    return {"response": resp["choices"][0]["text"].strip()}