from fastapi import FastAPI
from pydantic import BaseModel
from transformers import pipeline

app = FastAPI()

generator = pipeline(
    "text-generation",
    model="meta-llama/Llama-3.1-8B-Instruct",  # or your HF model
    device=0  # remove if CPU
)

class Request(BaseModel):
    prompt: str
    max_new_tokens: int = 150

@app.post("/generate")
def generate(req: Request):
    output = generator(
        req.prompt,
        max_new_tokens=req.max_new_tokens,
        do_sample=True,
    )
    return {"result": output[0]["generated_text"]}