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

app = FastAPI()

# Download model (replace filename with the exact GGUF file)
model_path = hf_hub_download(
    repo_id="mradermacher/Mythos-nano-i1-GGUF",
    filename="Mythos-nano-i1.Q4_K_M.gguf"
)

llm = Llama(
    model_path=model_path,
    n_ctx=2048
)

class Request(BaseModel):
    prompt: str
    max_tokens: int = 256

@app.post("/generate")
async def generate(req: Request):
    output = llm(
        req.prompt,
        max_tokens=req.max_tokens
    )

    return {
        "response": output["choices"][0]["text"]
    }