Update main.py

main.py

@@ -1,5 +1,6 @@
 import os, io
 from pathlib import Path
+
 from fastapi import FastAPI, UploadFile, File, Form
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse, HTMLResponse, FileResponse
@@ -14,20 +15,20 @@ from io import BytesIO
 # CONFIGURATION
 # -----------------------------------------------------------------------------
 HUGGINGFACE_TOKEN = os.getenv("HF_TOKEN")
-PORT = int(os.getenv("PORT", 7860))
+PORT              = int(os.getenv("PORT", 7860))
 
 app = FastAPI(
-    title="AI‑Powered Web‑App API",
-    description="Backend for summarisation, captioning & QA",
-    version="1.2.2",
+    title       = "AI-Powered Web-App API",
+    description = "Backend for summarisation, captioning & QA",
+    version     = "1.2.3",               # <-- bumped
 )
 
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
+    allow_origins     = ["*"],
+    allow_credentials = True,
+    allow_methods     = ["*"],
+    allow_headers     = ["*"],
 )
 
 # -----------------------------------------------------------------------------
@@ -40,14 +41,31 @@ if static_dir.exists():
 # -----------------------------------------------------------------------------
 # HUGGING FACE INFERENCE CLIENTS
 # -----------------------------------------------------------------------------
-summary_client = InferenceClient("facebook/bart-large-cnn", token=HUGGINGFACE_TOKEN)
-qa_client = InferenceClient("deepset/roberta-base-squad2", token=HUGGINGFACE_TOKEN)
-image_caption_client = InferenceClient("nlpconnect/vit-gpt2-image-captioning", token=HUGGINGFACE_TOKEN)
+summary_client        = InferenceClient(
+    "facebook/bart-large-cnn",
+    token   = HUGGINGFACE_TOKEN,
+    timeout = 120,
+)
+
+# ➜ Upgraded QA model (higher accuracy than roberta-base)
+qa_client             = InferenceClient(
+    "deepset/roberta-large-squad2",
+    token   = HUGGINGFACE_TOKEN,
+    timeout = 120,
+)
+# If you need multilingual support, swap for:
+# qa_client = InferenceClient("deepset/xlm-roberta-large-squad2",
+#                             token=HUGGINGFACE_TOKEN, timeout=120)
+
+image_caption_client  = InferenceClient(
+    "nlpconnect/vit-gpt2-image-captioning",
+    token   = HUGGINGFACE_TOKEN,
+    timeout = 60,
+)
 
 # -----------------------------------------------------------------------------
 # UTILITIES
 # -----------------------------------------------------------------------------
-
 def extract_text_from_pdf(content: bytes) -> str:
     reader = PdfReader(io.BytesIO(content))
     return "\n".join(page.extract_text() or "" for page in reader.pages).strip()
@@ -58,7 +76,7 @@ def extract_text_from_docx(content: bytes) -> str:
 
 def process_uploaded_file(file: UploadFile) -> str:
     content = file.file.read()
-    ext = file.filename.split(".")[-1].lower()
+    ext     = file.filename.split(".")[-1].lower()
     if ext == "pdf":
         return extract_text_from_pdf(content)
     if ext == "docx":
@@ -70,7 +88,6 @@ def process_uploaded_file(file: UploadFile) -> str:
 # -----------------------------------------------------------------------------
 # ROUTES
 # -----------------------------------------------------------------------------
-
 @app.get("/", response_class=HTMLResponse)
 async def serve_index():
     return FileResponse("index.html")
@@ -90,51 +107,66 @@ async def summarize_document(file: UploadFile = File(...)):
         )
         return {"result": summary_txt}
     except Exception as exc:
-        return JSONResponse(status_code=500, content={"error": f"Summarisation failure: {exc}"})
+        return JSONResponse(status_code=500,
+                            content={"error": f"Summarisation failure: {exc}"})
+
 
-# -------------------- Image Caption -----------------------------------------
+# -------------------- Image Caption ------------------------------------------
 @app.post("/api/caption")
 async def caption_image(image: UploadFile = File(...)):
     """`image` field name matches frontend (was `file` before)."""
     try:
         img_bytes = await image.read()
-        img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
+        img       = Image.open(io.BytesIO(img_bytes)).convert("RGB")
         img.thumbnail((1024, 1024))
         buf = BytesIO(); img.save(buf, format="JPEG")
         result = image_caption_client.image_to_text(buf.getvalue())
         if isinstance(result, dict):
-            caption = result.get("generated_text") or result.get("caption") or "No caption found."
+            caption = (result.get("generated_text")
+                       or result.get("caption")
+                       or "No caption found.")
         elif isinstance(result, list):
             caption = result[0].get("generated_text", "No caption found.")
         else:
             caption = str(result)
         return {"result": caption}
     except Exception as exc:
-        return JSONResponse(status_code=500, content={"error": f"Caption failure: {exc}"})
+        return JSONResponse(status_code=500,
+                            content={"error": f"Caption failure: {exc}"})
+
 
-# -------------------- Question Answering ------------------------------------
+# -------------------- Question Answering -------------------------------------
 @app.post("/api/qa")
-async def question_answering(file: UploadFile = File(...), question: str = Form(...)):
+async def question_answering(file: UploadFile = File(...),
+                             question: str = Form(...)):
     try:
         if file.content_type.startswith("image/"):
             img_bytes = await file.read()
-            img = Image.open(io.BytesIO(img_bytes)).convert("RGB"); img.thumbnail((1024, 1024))
+            img       = Image.open(io.BytesIO(img_bytes)).convert("RGB")
+            img.thumbnail((1024, 1024))
             buf = BytesIO(); img.save(buf, format="JPEG")
-            res = image_caption_client.image_to_text(buf.getvalue())
-            context = res.get("generated_text") if isinstance(res, dict) else str(res)
+            res      = image_caption_client.image_to_text(buf.getvalue())
+            context  = (res.get("generated_text") if isinstance(res, dict)
+                        else str(res))
         else:
             context = process_uploaded_file(file)[:3000]
+
         if not context:
             return {"result": "No context – cannot answer."}
+
         answer = qa_client.question_answering(question=question, context=context)
         return {"result": answer.get("answer", "No answer found.")}
     except Exception as exc:
-        return JSONResponse(status_code=500, content={"error": f"QA failure: {exc}"})
+        return JSONResponse(status_code=500,
+                            content={"error": f"QA failure: {exc}"})
+
 
-# -------------------- Health -------------------------------------------------
+# -------------------- Health --------------------------------------------------
 @app.get("/api/health")
 async def health():
-    return {"status": "healthy", "hf_token_set": bool(HUGGINGFACE_TOKEN), "version": app.version}
+    return {"status": "healthy",
+            "hf_token_set": bool(HUGGINGFACE_TOKEN),
+            "version": app.version}
 
 # -----------------------------------------------------------------------------
 # ENTRYPOINT