main.py

from fastapi import FastAPI, Request
from fastapi.responses import FileResponse, HTMLResponse
from transformers import pipeline
from fastapi.staticfiles import StaticFiles

app = FastAPI()

pipe_flan = pipeline("text2text-generation", model="google/flan-t5-small")

def t5(input):
    output = pipe_flan(input)
    return {"output": output[0]["generated_text"]}

@app.post("/infer_t5")
def infer_endpoint(data: dict):
    """Receive input and generate text."""
    input_text = data.get("input")
    if input_text is None:
        return {"error": "No input text detected."}
    else:
        result = t5(input_text)
        return result

# Route to render and present the index.html file
@app.get("/", response_class=HTMLResponse)
async def index(request: Request):
    return FileResponse(await request.url.file("/app/static/index.html"))

# Serve the static files
app.mount("/static", StaticFiles(directory="static"), name="static")