from fastapi import FastAPI
from transformers import pipeline

# create FastAPI app

app = FastAPI()

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


@app.get("/")
def home():
    """
        Home route for the FastAPI app.

        Returns:
            dict: A dictionary with a message indicating that it is a simple FastAPI app for text generation using T5.
    """
    return {"message": "This is a simple FastAPI app for text generation using T5"}


# create a route for text generation
@app.get("/generate/")
def generate_text(prompt: str):
    """
        Route for generating text using the T5 model.

        Args:
            prompt (str): The prompt for the text generation.

        Returns:
            dict: A dictionary with the generated text.
    """
    output = pipe(text=prompt)
    return {"generated_text": pipe(prompt)[0]["generated_text"]}