from fastapi import FastAPI
from transformers import pipeline

app = FastAPI()


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


@app.get("/generate")
def generate(query: str):
    output = pipe(query, max_length=50, do_sample=True, top_k=50)   
    return {"output": output[0]["generated_text"], "query": query, "model": "google/flan-t5-small"}