from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from transformers import pipeline

app = FastAPI()


# Load the text-to-text generation pipeline
pipe = pipeline("text2text-generation", model="google/flan-t5-small")

# Define the input data model
class GenerateRequest(BaseModel):
    text: str

@app.get("/")
def welcome():
    """
    Welcoming page.
    """
    return {"message": "Welcome to the Text-to-Text Generation API! Use the /generate endpoint to transform text."}

@app.post("/generate")
def generate_text(data: GenerateRequest):
    """
    Generate text from the input text using the text-to-text generation model.
    """
    try:
        # Perform text-to-text generation
        result = pipe(data.text, max_length=50, num_return_sequences=1)
        return {"input": data.text, "generated_text": result[0]["generated_text"]}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))