Update main.py

main.py

@@ -1,89 +1,120 @@
-from fastapi import FastAPI, UploadFile, File, Form
-from fastapi.responses import JSONResponse
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-import uvicorn
-import tempfile
 import os
+import io
+from fastapi import FastAPI, UploadFile, File, Form
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import JSONResponse, HTMLResponse
+from huggingface_hub import InferenceClient
+from PyPDF2 import PdfReader
+from docx import Document
 from PIL import Image
-import torch
+from io import BytesIO
+
+# Load Hugging Face Token securely
+HUGGINGFACE_TOKEN = os.getenv("HF_TOKEN")
 
 app = FastAPI()
 
-# Load tokenizers fast but not full models immediately
-tokenizers = {
-    "qwen": AutoTokenizer.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct", trust_remote_code=True),
-    "deepseek": AutoTokenizer.from_pretrained("deepseek-ai/DeepSeek-V2-Chat", trust_remote_code=True),
-    "llama": AutoTokenizer.from_pretrained("meta-llama/Llama-2-70b-chat-hf", trust_remote_code=True),
-}
-
-models = {}
-
-def load_model(name):
-    if name not in models:
-        if name == "qwen":
-            models[name] = AutoModelForCausalLM.from_pretrained(
-                "Qwen/Qwen2.5-VL-7B-Instruct",
-                device_map="auto",
-                trust_remote_code=True,
-                torch_dtype=torch.float16
-            )
-        elif name == "deepseek":
-            models[name] = AutoModelForCausalLM.from_pretrained(
-                "deepseek-ai/DeepSeek-V2-Chat",
-                device_map="auto",
-                trust_remote_code=True,
-                torch_dtype=torch.float16
-            )
-        elif name == "llama":
-            models[name] = AutoModelForCausalLM.from_pretrained(
-                "meta-llama/Llama-2-70b-chat-hf",
-                device_map="auto",
-                trust_remote_code=True,
-                torch_dtype=torch.float16
-            )
-    return models[name]
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
 
-@app.post("/api/summarize")
-async def summarize(file: UploadFile = File(...)):
-    ext = os.path.splitext(file.filename)[1].lower()
-    temp_path = os.path.join(tempfile.gettempdir(), file.filename)
-    with open(temp_path, "wb") as f:
-        f.write(await file.read())
+# Initialize Hugging Face clients
+summary_client = InferenceClient(model="facebook/bart-large-cnn", token=HUGGINGFACE_TOKEN)
+qa_client = InferenceClient(model="deepset/roberta-base-squad2", token=HUGGINGFACE_TOKEN)
+image_caption_client = InferenceClient(model="nlpconnect/vit-gpt2-image-captioning", token=HUGGINGFACE_TOKEN)
+
+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()
 
-    # For now: Just simulate basic summarization
-    summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
-    with open(temp_path, 'r', errors='ignore') as f:
-        text = f.read()
+def extract_text_from_docx(content: bytes) -> str:
+    doc = Document(io.BytesIO(content))
+    return "\n".join(para.text for para in doc.paragraphs).strip()
 
-    if len(text) > 1024:
-        text = text[:1024]
+def process_uploaded_file(file: UploadFile) -> str:
+    content = file.file.read()
+    extension = file.filename.split('.')[-1].lower()
+    if extension == "pdf":
+        return extract_text_from_pdf(content)
+    elif extension == "docx":
+        return extract_text_from_docx(content)
+    elif extension == "txt":
+        return content.decode("utf-8").strip()
+    else:
+        raise ValueError("Unsupported file type.")
 
-    summary = summarizer(text, max_length=150, min_length=40, do_sample=False)[0]['summary_text']
-    return JSONResponse({"result": summary})
+@app.get("/", response_class=HTMLResponse)
+async def serve_homepage():
+    with open("index.html", "r", encoding="utf-8") as f:
+        return HTMLResponse(content=f.read(), status_code=200)
+
+@app.post("/api/summarize")
+async def summarize_document(file: UploadFile = File(...)):
+    try:
+        text = process_uploaded_file(file)
+        if len(text) < 20:
+            return {"result": "Document too short to summarize."}
+        summary = summary_client.summarization(text[:3000])
+        return {"result": summary}
+    except Exception as e:
+        return JSONResponse(status_code=500, content={"error": str(e)})
 
 @app.post("/api/caption")
-async def caption(file: UploadFile = File(...)):
-    image = Image.open(await file.read())
-    # For now: Use a simple vision model, because Qwen2.5 VL loading takes a lot of time
-    captioner = pipeline("image-to-text", model="nlpconnect/vit-gpt2-image-captioning")
-    caption = captioner(image)[0]['generated_text']
-    return JSONResponse({"result": caption})
+async def caption_image(file: UploadFile = File(...)):
+    try:
+        image_bytes = await file.read()
+        image_pil = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+        image_pil.thumbnail((1024, 1024))
+        img_byte_arr = BytesIO()
+        image_pil.save(img_byte_arr, format='JPEG')
+        img_byte_arr = img_byte_arr.getvalue()
+        result = image_caption_client.image_to_text(img_byte_arr)
+
+        if isinstance(result, dict):
+            caption = result.get("generated_text") or result.get("caption") or "No caption found."
+        elif isinstance(result, list) and result:
+            caption = result[0].get("generated_text", "No caption found.")
+        elif isinstance(result, str):
+            caption = result
+        else:
+            caption = "No caption found."
+
+        return {"result": caption}
+    except Exception as e:
+        return JSONResponse(status_code=500, content={"error": str(e)})
 
 @app.post("/api/qa")
-async def question_answer(file: UploadFile = File(...), question: str = Form(...)):
-    temp_path = os.path.join(tempfile.gettempdir(), file.filename)
-    with open(temp_path, "wb") as f:
-        f.write(await file.read())
+async def question_answering(file: UploadFile = File(...), question: str = Form(...)):
+    try:
+        content_type = file.content_type
+        if content_type.startswith("image/"):
+            image_bytes = await file.read()
+            image_pil = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+            image_pil.thumbnail((1024, 1024))
+            img_byte_arr = BytesIO()
+            image_pil.save(img_byte_arr, format='JPEG')
+            img_byte_arr = img_byte_arr.getvalue()
+            result = image_caption_client.image_to_text(img_byte_arr)
+            context = result.get("generated_text") if isinstance(result, dict) else result
+        else:
+            text = process_uploaded_file(file)
+            if len(text) < 20:
+                return {"result": "Document too short to answer questions."}
+            context = text[:3000]
 
-    # For now: pick deepseek model for QA
-    tokenizer = tokenizers["deepseek"]
-    model = load_model("deepseek")
+        if not context:
+            return {"result": "No context available to answer."}
 
-    inputs = tokenizer(question, return_tensors="pt").to(model.device)
-    outputs = model.generate(**inputs, max_new_tokens=100)
-    answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        answer = qa_client.question_answering(question=question, context=context)
+        return {"result": answer.get("answer", "No answer found.")}
 
-    return JSONResponse({"result": answer})
+    except Exception as e:
+        return JSONResponse(status_code=500, content={"error": str(e)})
 
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    import uvicorn
+    uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True)