memofy-api / app.py
memofy's picture
Update app.py
70315f4 verified
raw
history blame
619 Bytes
from fastapi import FastAPI, Request
from transformers import pipeline
app = FastAPI()
generator = pipeline("text-generation", model="pierreguillou/pt-gpt2-medium", max_length=300)
@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)}