from fastapi import FastAPI
from pydantic import BaseModel
from llama_cpp import Llama

app = FastAPI()

# Your specific requested model
MODEL_ID = "Qwen/Qwen2.5-3B-Instruct-GGUF"
MODEL_FILE = "qwen2.5-3b-instruct-q4_k_m.gguf"

# Loading the model. 
# n_ctx=2048 helps keep memory usage stable on free tier.
llm = Llama.from_pretrained(
    repo_id=MODEL_ID,
    filename=MODEL_FILE,
    n_ctx=2048,
    verbose=False
)

class Request(BaseModel):
    prompt: str

@app.get("/")
def home():
    return {"status": "Model is loaded and ready!"}

@app.post("/generate")
def generate(req: Request):
    output = llm.create_chat_completion(
        messages=[
            {"role": "user", "content": req.prompt}
        ],
        temperature=0.7,
        max_tokens=128
    )
    return {
        "response": output["choices"][0]["message"]["content"]
    }