Update app.py

app.py

@@ -86,6 +86,9 @@ from bs4 import BeautifulSoup
 from transformers import AutoTokenizer, AutoModelForCausalLM, T5Tokenizer, T5ForConditionalGeneration
 import torch
 import re
+from fastapi.responses import JSONResponse
+from fastapi.requests import Request
+from fastapi import status
 
 app = FastAPI()
 
@@ -102,7 +105,6 @@ class GenerateResponse(BaseModel):
     reasoning_content: str
     generated_text: str
 
-
 # --- Utility Functions ---
 
 def clean_text(text: str) -> str:
@@ -110,7 +112,6 @@ def clean_text(text: str) -> str:
     text = re.sub(r"\b\d+\s*likes?,?\s*\d*\s*replies?$", "", text, flags=re.IGNORECASE).strip()
     return text
 
-
 # --- Scraping Endpoint ---
 
 @app.get("/scrape", response_model=ThreadResponse)
@@ -128,7 +129,6 @@ def scrape(url: str):
             return ThreadResponse(question=question, replies=replies)
     return ThreadResponse(question="", replies=[])
 
-
 # --- Load DeepSeek-R1-Distill-Qwen-1.5B Model & Tokenizer ---
 
 deepseek_model_name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
@@ -137,7 +137,6 @@ deepseek_model = AutoModelForCausalLM.from_pretrained(deepseek_model_name)
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 deepseek_model = deepseek_model.to(device)
 
-
 # --- Load T5-Large Model & Tokenizer ---
 
 t5_model_name = "google-t5/t5-large"
@@ -145,7 +144,6 @@ t5_tokenizer = T5Tokenizer.from_pretrained(t5_model_name)
 t5_model = T5ForConditionalGeneration.from_pretrained(t5_model_name)
 t5_model = t5_model.to(device)
 
-
 # --- Generation Functions ---
 
 def generate_deepseek(prompt: str) -> (str, str):
@@ -167,9 +165,7 @@ def generate_deepseek(prompt: str) -> (str, str):
     else:
         return "", generated_text.strip()
 
-
 def generate_t5(prompt: str) -> (str, str):
-    # T5 expects prompt with task prefix, e.g. "summarize: ..."
     inputs = t5_tokenizer.encode(prompt, return_tensors="pt", max_length=512, truncation=True).to(device)
     outputs = t5_model.generate(
         inputs,
@@ -181,14 +177,12 @@ def generate_t5(prompt: str) -> (str, str):
     )
     generated_text = t5_tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    # Optional reasoning parsing if </think> is used
     if "</think>" in generated_text:
         reasoning_content, content = generated_text.split("</think>", 1)
         return reasoning_content.strip(), content.strip()
     else:
         return "", generated_text.strip()
 
-
 # --- API Endpoints ---
 
 @app.post("/generate/{model_name}", response_model=GenerateResponse)
@@ -201,6 +195,19 @@ async def generate(
     elif model_name == "t5-large":
         reasoning, text = generate_t5(request.prompt)
     else:
-        return {"reasoning_content": "", "generated_text": f"Error: Unknown model '{model_name}'."}
+        return GenerateResponse(reasoning_content="", generated_text=f"Error: Unknown model '{model_name}'.")
 
     return GenerateResponse(reasoning_content=reasoning, generated_text=text)
+
+# --- Global Exception Handler ---
+
+@app.exception_handler(Exception)
+async def global_exception_handler(request: Request, exc: Exception):
+    print(f"Exception: {exc}")
+    return JSONResponse(
+        status_code=status.HTTP_200_OK,
+        content={
+            "reasoning_content": "",
+            "generated_text": f"Error: {str(exc)}"
+        }
+    )