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

app = FastAPI()

pipe = pipeline(
    "text2text-generation",
    model="google/flan-t5-small"
)

class Prompt(BaseModel):
    input: str

@app.post("/generate")
def generate(prompt: Prompt):
    output = pipe(prompt.input)
    return {
        "output": output[0]["generated_text"]
    }