Update main.py

main.py

@@ -1,4 +1,4 @@
-from fastapi import FastAPI, Request
+from fastapi import FastAPI, Request, Depends
 from fastapi.responses import FileResponse, JSONResponse
 from fastapi.staticfiles import StaticFiles
 from transformers import pipeline
@@ -18,11 +18,14 @@ def t5(input: str) -> dict[str, str]:
     output = pipe_flan(input)
     return {"output": output[0].get("generated_text", "")}
 
+class ParseRaw(BaseModel):
+    raw: bytes
+
 @app.post("/infer_t5")
-async def infer_endpoint(data: dict = Depends(parse_raw)):
+async def infer_endpoint(data: ParseRaw = Depends(parse_raw)):
     """Receive input and generate text."""
     try:
-        input_text = data.get("input")
+        input_text = data.raw.decode("utf-8")
 
         # Validate that the input is a string
         assert isinstance(input_text, str), "Input must be a string."
@@ -66,10 +69,10 @@ def miuk_answer(query: str) -> str:
     return tokenizer.decode(answer[:, 0]).replace(" ", "")
 
 @app.post("/infer_miku")
-async def infer_endpoint(data: dict = Depends(parse_raw)):
+async def infer_endpoint(data: ParseRaw = Depends(parse_raw)):
     """Receive input and generate text."""
     try:
-        input_text = data.get("input")
+        input_text = data.raw.decode("utf-8")
 
         # Validate that the input is a string
         assert isinstance(input_text, str), "Input must be a string."
@@ -85,4 +88,7 @@ async def infer_endpoint(data: dict = Depends(parse_raw)):
 
 @app.get("/infer_miku")
 def get_default_inference_endpoint():
-    return {"message": "Use POST method to submit input data"}
+    return {"message": "Use POST method to submit input data"}
+
+def parse_raw(request: Request):
+    return ParseRaw(raw=request.body())