Update main.py

main.py

@@ -1,9 +1,9 @@
-from fastapi import FastAPI
-from fastapi.responses import FileResponse
+from fastapi import FastAPI, Request
+from fastapi.responses import FileResponse, JSONResponse
 from fastapi.staticfiles import StaticFiles
 from transformers import pipeline
 from pydantic import BaseModel
-from typing import Union
+from typing import Optional, Any
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
 
@@ -14,24 +14,28 @@ app.mount("/static", StaticFiles(directory="static"), name="static")
 
 pipe_flan = pipeline("text2text-generation", model="google/flan-t5-small")
 
-def t5(input):
+def t5(input: str) -> dict[str, str]:
     output = pipe_flan(input)
-    return {"output": output[0]["generated_text"]}
+    return {"output": output[0].get("generated_text", "")}
 
 @app.post("/infer_t5")
-def infer_endpoint(data: dict):
+async def infer_endpoint(data: dict = Depends(parse_raw)):
     """Receive input and generate text."""
-    input_text = data.get("input")
+    try:
+        input_text = data.get("input")
 
-    # Validate that the input is a string
-    assert isinstance(input_text, str), "Input must be a string."
+        # Validate that the input is a string
+        assert isinstance(input_text, str), "Input must be a string."
+
+        if input_text is None:
+            return {"error": "No input text detected."}
 
-    if input_text is None:
-        return {"error": "No input text detected."}
-    else:
         result = t5(input_text)
         return result
 
+    except AssertionError as e:
+        return JSONResponse({"error": f"Invalid Input Format: {e}"}, status_code=400)
+
 @app.get("/infer_t5")
 def get_default_inference_endpoint():
     return {"message": "Use POST method to submit input data"}
@@ -55,27 +59,30 @@ except Exception as e:
     model = AutoModelForCausalLM.from_pretrained("miqudev/miqu-1-70b", config=model_config).to(device)
     tokenizer = AutoTokenizer.from_pretrained("miqudev/miqu-1-70b")
 
-def miuk_answer(query: str):
+def miuk_answer(query: str) -> str:
     query_tokens = tokenizer.encode(query, return_tensors="pt")
     query_tokens = query_tokens.to(device)
     answer = model.generate(query_tokens, max_length=128, temperature=1, pad_token_id=tokenizer.pad_token_id)
     return tokenizer.decode(answer[:, 0]).replace(" ", "")
 
 @app.post("/infer_miku")
-def infer_endpoint(data: dict):
+async def infer_endpoint(data: dict = Depends(parse_raw)):
     """Receive input and generate text."""
-    input_text = data.get("input")
+    try:
+        input_text = data.get("input")
+
+        # Validate that the input is a string
+        assert isinstance(input_text, str), "Input must be a string."
 
-    # Validate that the input is a string
-    assert isinstance(input_text, str), "Input must be a string."
+        if input_text is None:
+            return {"error": "No input text detected."}
 
-    if input_text is None:
-        return {"error": "No input text detected."}
-    else:
         result = {"output": miuk_answer(input_text)}
         return result
 
+    except AssertionError as e:
+        return JSONResponse({"error": f"Invalid Input Format: {e}"}, status_code=400)
+
 @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"}