from fastapi import FastAPI
from fastapi.responses import HTMLResponse
from transformers import pipeline

# Create a new FastAPI app instance
app = FastAPI()

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

@app.get("/")
def home():
    html_content = """
    <!DOCTYPE html>
    <html lang="en">
    <head>
        <meta charset="UTF-8">
        <meta name="viewport" content="width=device-width, initial-scale=1.0">
        <title>Text to Text Generation</title>
    </head>
    <body>
        <h1>Welcome to Text to Text Generation</h1>
        <iframe
            src="https://rohitjakkam-text2text-using-docker.hf.space"
            frameborder="0"
            width="850"
            height="450"
        ></iframe>
    </body>
    </html>
    """
    return HTMLResponse(content=html_content)

@app.get("/generate")
def generate(text: str):
    # Use pipeline to generate text from given text as input parameter
    output = pipe(text)

    # Return the generated text in JSON response
    return {"output": output[0]['generated_text']}