Update main.py

main.py

@@ -1,8 +1,20 @@
 from fastapi import FastAPI
+from transformers import pipeline
 
 app = FastAPI()
 
-@app.get("/")
+pipe_flan = pipeline("text2text-generation", model="google/flan-t5-small")
 
-def read_root():
-    return {"Hello": "World!"}
+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