from fastapi import FastAPI
from transformers import pipeline

# Create a FastAPI instance
app = FastAPI(docs_url='/')

# Initialize the text generation pipeline
pipe = pipeline('text2text-generation', model='google/flan-t5-small')

# The GET endpoint of the application,
# which corresponds to the text generation
# functionality. The generate() method takes 
# a prompt as input and returns the generated
# text as output in the form of a JSON object.
@app.get("/generate")
def generate(prompt: str):
    output = pipe(prompt)
    return {'output': output[0]['generated_text']}