Update app.py

app.py

@@ -1,36 +1,38 @@
 from fastapi import FastAPI, UploadFile, File
 from fastapi.responses import RedirectResponse
 import gradio as gr
-from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import pipeline, AutoModelForSeq2SeqLM, AutoTokenizer
 import tempfile
 import os
 from PIL import Image
 import fitz  # PyMuPDF
 import docx
-import openpyxl
-from pptx import Presentation
 import easyocr
 
 app = FastAPI()
 
-# Initialize models with error handling
+# Lightweight model choices
+SUMMARIZATION_MODEL = "facebook/bart-large-cnn"  # 500MB
+IMAGE_CAPTIONING_MODEL = "Salesforce/blip-image-captioning-base"  # 300MB
+
+# Initialize models
 try:
-    # Load summarization model directly with tokenizer
-    tokenizer = AutoTokenizer.from_pretrained("FeruzaBoynazarovaas/my_awesome_billsum_model", use_fast=False)
-    model = AutoModelForSeq2SeqLM.from_pretrained("FeruzaBoynazarovaas/my_awesome_billsum_model")
     summarizer = pipeline(
-        "text2text-generation",
-        model=model,
-        tokenizer=tokenizer
+        "summarization",
+        model=SUMMARIZATION_MODEL,
+        device="cpu"
     )
 except Exception as e:
     print(f"Error loading summarizer: {e}")
-    # Fallback to a default model if custom fails
-    summarizer = pipeline("text2text-generation", model="t5-small")
+    summarizer = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6")  # Fallback 250MB model
+
+captioner = pipeline(
+    "image-to-text",
+    model=IMAGE_CAPTIONING_MODEL,
+    device="cpu"
+)
 
-# Other models (these should work fine)
-captioner = pipeline("image-to-text", model="nlpconnect/vit-gpt2-image-captioning")
-reader = easyocr.Reader(['en'])
+reader = easyocr.Reader(['en'])  # Lightweight OCR
 
 def extract_text_from_file(file_path: str, file_type: str):
     """Extract text from different document formats"""
@@ -41,26 +43,24 @@ def extract_text_from_file(file_path: str, file_type: str):
         elif file_type == "docx":
             doc = docx.Document(file_path)
             return "\n".join(p.text for p in doc.paragraphs)
-        elif file_type == "pptx":
-            prs = Presentation(file_path)
-            return "\n".join(shape.text for slide in prs.slides for shape in slide.shapes if hasattr(shape, "text"))
-        elif file_type == "xlsx":
-            wb = openpyxl.load_workbook(file_path)
-            return "\n".join(str(cell.value) for sheet in wb for row in sheet for cell in row)
         else:
-            return "Unsupported file format"
+            return "Unsupported file format (only PDF/DOCX supported in lightweight version)"
     except Exception as e:
         return f"Error reading file: {str(e)}"
 
 def process_document(file):
+    """Handle document summarization"""
     try:
         file_ext = os.path.splitext(file.name)[1][1:].lower()
+        if file_ext not in ["pdf", "docx"]:
+            return "Lightweight version only supports PDF and DOCX"
+            
         with tempfile.NamedTemporaryFile(delete=False, suffix=f".{file_ext}") as tmp:
             tmp.write(file.read())
             tmp_path = tmp.name
         
         text = extract_text_from_file(tmp_path, file_ext)
-        summary = summarizer(text, max_length=150, min_length=30, do_sample=False)[0]['generated_text']
+        summary = summarizer(text, max_length=130, min_length=30, do_sample=False)[0]['summary_text']
         
         os.unlink(tmp_path)
         return summary
@@ -68,11 +68,17 @@ def process_document(file):
         return f"Processing error: {str(e)}"
 
 def process_image(image):
+    """Handle image captioning and OCR"""
     try:
         img = Image.open(image)
+        
+        # Get caption
         caption = captioner(img)[0]['generated_text']
+        
+        # Get OCR text
         ocr_result = reader.readtext(img)
         ocr_text = " ".join([res[1] for res in ocr_result])
+        
         return {
             "caption": caption,
             "ocr_text": ocr_text if ocr_text else "No readable text found"
@@ -81,25 +87,29 @@ def process_image(image):
         return {"error": str(e)}
 
 # Gradio Interface
-with gr.Blocks() as demo:
-    gr.Markdown("# 📄 Document & Image Analysis")
+with gr.Blocks(title="Lightweight Document & Image Analysis") as demo:
+    gr.Markdown("## 📄 Lightweight Document & Image Analysis")
     
     with gr.Tab("Document Summarization"):
-        doc_input = gr.File(label="Upload Document")
+        gr.Markdown("Supports PDF and DOCX files (max 10MB)")
+        doc_input = gr.File(label="Upload Document", file_types=[".pdf", ".docx"])
         doc_output = gr.Textbox(label="Summary")
         doc_button = gr.Button("Summarize")
         
     with gr.Tab("Image Analysis"):
+        gr.Markdown("Get captions and extracted text from images")
         img_input = gr.Image(type="filepath", label="Upload Image")
-        caption_output = gr.Textbox(label="Image Caption")
-        ocr_output = gr.Textbox(label="Extracted Text")
+        with gr.Accordion("Results", open=False):
+            caption_output = gr.Textbox(label="Image Caption")
+            ocr_output = gr.Textbox(label="Extracted Text")
         img_button = gr.Button("Analyze")
     
     doc_button.click(process_document, inputs=doc_input, outputs=doc_output)
     img_button.click(process_image, inputs=img_input, outputs=[caption_output, ocr_output])
 
+# Mount Gradio app
 app = gr.mount_gradio_app(app, demo, path="/")
 
 @app.get("/")
-def redirect():
+def redirect_to_interface():
     return RedirectResponse(url="/")