from fastapi import FastAPI, Request
from transformers import pipeline

app = FastAPI()

# Corrigido com from_tf=True
generator = pipeline(
    "text-generation",
    model="pierreguillou/gpt2-small-portuguese",
    max_length=300,
    from_tf=True
)

@app.post("/gerar")
async def gerar_historia(req: Request):
    body = await req.json()
    relato = body.get("relato", "")
    estilo = body.get("estilo", "")
    prompt = f"Crie uma história no estilo {estilo}. Baseie-se neste relato: {relato}"

    try:
        resultado = generator(prompt, max_new_tokens=200)[0]["generated_text"]
        return {"historia": resultado}
    except Exception as e:
        return {"erro": str(e)}