from fastapi import FastAPI
from transformers import pipeline

# Create a FastAPI instance
app = FastAPI()

# The root endpoint of the application
@app.get("/")
def root():
    return {'message': 'Welcome to the Text Generation API!'}

# Initialize the text generation pipeline
pipe = pipeline('text2test-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']}