from fastapi import FastAPI
from transformers import pipeline

app = FastAPI()


@app.get('/')
def home():
    return {"hello": "world"}


@app.get('/ask')
def ask(prompt: str):
    pipe = pipeline("text2text-generation", model="google/flan-t5-base")
    output = pipe(prompt)
    return output[0]['generated_text']