from transformers import pipeline
from fastapi import FastAPI
from payload import SomeText

app = FastAPI()

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

@app.get("/api/check-heartbeat")
def get_heartbeat():
    return {"Status": "Running. Try out the endpoints in swagger"}

@app.post("/api/generate", summary="Generate text from prompt", tags=["Generate"], response_model=SomeText)
def inference(input_prompt: SomeText):
    output = pipe_flan(input_prompt)
    return {"output": output[0]["generated_text"]}