Update app.py

app.py

@@ -1,6 +1,6 @@
 """
-PDF Document Processing Gradio App for HuggingFace Spaces
-Built on DOLPHIN model for document parsing and analysis
+DOLPHIN PDF Document AI - Final Version
+Optimized for HuggingFace Spaces NVIDIA T4 Small deployment
 """
 
 import gradio as gr
@@ -25,28 +25,22 @@ try:
     from mdx_math import MathExtension
     MATH_EXTENSION_AVAILABLE = True
 except ImportError:
-    # mdx_math is not available in standard PyPI, gracefully continue without it
     pass
 
 
 class DOLPHIN:
     def __init__(self, model_id_or_path):
-        """Initialize the Hugging Face model optimized for HF Spaces
-        
-        Args:
-            model_id_or_path: Path to local model or Hugging Face model ID
-        """
+        """Initialize the Hugging Face model optimized for T4 Small"""
         self.processor = AutoProcessor.from_pretrained(model_id_or_path)
         self.model = VisionEncoderDecoderModel.from_pretrained(
             model_id_or_path,
-            torch_dtype=torch.float16,  # Use half precision for memory efficiency
+            torch_dtype=torch.float16,
             device_map="auto" if torch.cuda.is_available() else None
         )
         self.model.eval()
         
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
         if not torch.cuda.is_available():
-            # Keep full precision on CPU
             self.model = self.model.float()
         
         self.tokenizer = self.processor.tokenizer
@@ -62,7 +56,6 @@ class DOLPHIN:
             images = image
             prompts = prompt if isinstance(prompt, list) else [prompt] * len(images)
         
-        # Prepare image
         batch_inputs = self.processor(images, return_tensors="pt", padding=True)
         batch_pixel_values = batch_inputs.pixel_values
         
@@ -71,7 +64,6 @@ class DOLPHIN:
         else:
             batch_pixel_values = batch_pixel_values.to(self.device)
         
-        # Prepare prompt
         prompts = [f"<s>{p} <Answer/>" for p in prompts]
         batch_prompt_inputs = self.tokenizer(
             prompts,
@@ -82,14 +74,13 @@ class DOLPHIN:
         batch_prompt_ids = batch_prompt_inputs.input_ids.to(self.device)
         batch_attention_mask = batch_prompt_inputs.attention_mask.to(self.device)
         
-        # Generate text with memory-efficient settings
         with torch.no_grad():
             outputs = self.model.generate(
                 pixel_values=batch_pixel_values,
                 decoder_input_ids=batch_prompt_ids,
                 decoder_attention_mask=batch_attention_mask,
                 min_length=1,
-                max_length=2048,  # Reduced for memory efficiency
+                max_length=1024,  # Reduced for T4 Small
                 pad_token_id=self.tokenizer.pad_token_id,
                 eos_token_id=self.tokenizer.eos_token_id,
                 use_cache=True,
@@ -101,10 +92,8 @@ class DOLPHIN:
                 temperature=1.0
             )
         
-        # Process output
         sequences = self.tokenizer.batch_decode(outputs.sequences, skip_special_tokens=False)
         
-        # Clean prompt text from output
         results = []
         for i, sequence in enumerate(sequences):
             cleaned = sequence.replace(prompts[i], "").replace("<pad>", "").replace("</s>", "").strip()
@@ -120,24 +109,17 @@ def convert_pdf_to_images_gradio(pdf_file):
     try:
         import pymupdf
         
-        # Handle different file input types
         if isinstance(pdf_file, str):
-            # If it's a file path (Gradio 5.x behavior)
             pdf_document = pymupdf.open(pdf_file)
         else:
-            # If it's a file object with .read() method
             pdf_bytes = pdf_file.read()
             pdf_document = pymupdf.open(stream=pdf_bytes, filetype="pdf")
         
         images = []
         for page_num in range(len(pdf_document)):
             page = pdf_document[page_num]
-            
-            # Render page to image with high DPI for better quality
-            mat = pymupdf.Matrix(2.0, 2.0)  # 2x zoom for better quality
+            mat = pymupdf.Matrix(2.0, 2.0)
             pix = page.get_pixmap(matrix=mat)
-            
-            # Convert to PIL Image
             img_data = pix.tobytes("png")
             pil_image = Image.open(io.BytesIO(img_data)).convert("RGB")
             images.append(pil_image)
@@ -152,91 +134,60 @@ def convert_pdf_to_images_gradio(pdf_file):
 def process_pdf_document(pdf_file, model, progress=gr.Progress()):
     """Process uploaded PDF file page by page"""
     if pdf_file is None:
-        return "No PDF file uploaded", [], {}
+        return "No PDF file uploaded", ""
     
     try:
-        # Convert PDF to images
         progress(0.1, desc="Converting PDF to images...")
         images = convert_pdf_to_images_gradio(pdf_file)
         
         if not images:
-            return "Failed to convert PDF to images", [], {}
+            return "Failed to convert PDF to images", ""
         
-        # Process each page
         all_results = []
-        page_previews = []
         
         for page_idx, pil_image in enumerate(images):
             progress((page_idx + 1) / len(images) * 0.8 + 0.1, 
                     desc=f"Processing page {page_idx + 1}/{len(images)}...")
             
-            # Stage 1: Layout parsing
             layout_output = model.chat("Parse the reading order of this document.", pil_image)
             
-            # Stage 2: Element processing with memory optimization
             padded_image, dims = prepare_image(pil_image)
             recognition_results = process_elements_optimized(
                 layout_output, 
                 padded_image, 
                 dims, 
                 model, 
-                max_batch_size=4  # Smaller batch size for memory efficiency
+                max_batch_size=2  # Smaller batch for T4 Small
             )
             
-            # Convert to markdown
             try:
                 markdown_converter = MarkdownConverter()
                 markdown_content = markdown_converter.convert(recognition_results)
             except:
-                # Fallback markdown generation
                 markdown_content = generate_fallback_markdown(recognition_results)
             
-            # Store page results
             page_result = {
                 "page_number": page_idx + 1,
-                "layout_output": layout_output,
-                "elements": recognition_results,
                 "markdown": markdown_content
             }
             all_results.append(page_result)
-            
-            # Create page preview with results
-            page_preview = {
-                "image": pil_image,
-                "page_num": page_idx + 1,
-                "element_count": len(recognition_results),
-                "markdown_preview": markdown_content[:500] + "..." if len(markdown_content) > 500 else markdown_content
-            }
-            page_previews.append(page_preview)
         
         progress(1.0, desc="Processing complete!")
         
-        # Combine all markdown
         combined_markdown = "\n\n---\n\n".join([
             f"# Page {result['page_number']}\n\n{result['markdown']}" 
             for result in all_results
         ])
         
-        # Create summary JSON
-        summary_json = {
-            "total_pages": len(images),
-            "processing_status": "completed",
-            "pages": all_results,
-            "model_info": {
-                "device": model.device,
-                "total_elements": sum(len(page["elements"]) for page in all_results)
-            }
-        }
-        
-        return combined_markdown, page_previews, summary_json
+        return combined_markdown, "processing_complete"
         
     except Exception as e:
         error_msg = f"Error processing PDF: {str(e)}"
-        return error_msg, [], {"error": error_msg}
+        return error_msg, "error"
 
 
-def process_elements_optimized(layout_results, padded_image, dims, model, max_batch_size=4):
-    """Optimized element processing for memory efficiency"""
+def process_elements_optimized(layout_results, padded_image, dims, model, max_batch_size=2):
+    """Optimized element processing for T4 Small"""
     layout_results = parse_layout_string(layout_results)
     
     text_elements = []
@@ -245,7 +196,6 @@ def process_elements_optimized(layout_results, padded_image, dims, model, max_ba
     previous_box = None
     reading_order = 0
     
-    # Collect elements to process
     for bbox, label in layout_results:
         try:
             x1, y1, x2, y2, orig_x1, orig_y1, orig_x2, orig_y2, previous_box = process_coordinates(
@@ -255,7 +205,6 @@ def process_elements_optimized(layout_results, padded_image, dims, model, max_ba
             cropped = padded_image[y1:y2, x1:x2]
             if cropped.size > 0 and cropped.shape[0] > 3 and cropped.shape[1] > 3:
                 if label == "fig":
-                    # Convert to base64 for figure display
                     pil_crop = Image.fromarray(cv2.cvtColor(cropped, cv2.COLOR_BGR2RGB))
                     pil_crop = crop_margin(pil_crop)
                     
@@ -290,7 +239,6 @@ def process_elements_optimized(layout_results, padded_image, dims, model, max_ba
             print(f"Error processing element {label}: {str(e)}")
             continue
     
-    # Process elements in small batches
     recognition_results = figure_results.copy()
     
     if text_elements:
@@ -309,8 +257,8 @@ def process_elements_optimized(layout_results, padded_image, dims, model, max_ba
     return recognition_results
 
 
-def process_element_batch_optimized(elements, model, prompt, max_batch_size=4):
-    """Process elements in small batches for memory efficiency"""
+def process_element_batch_optimized(elements, model, prompt, max_batch_size=2):
+    """Process elements in small batches for T4 Small"""
     results = []
     batch_size = min(len(elements), max_batch_size)
     
@@ -319,7 +267,6 @@ def process_element_batch_optimized(elements, model, prompt, max_batch_size=4):
         crops_list = [elem["crop"] for elem in batch_elements]
         prompts_list = [prompt] * len(crops_list)
         
-        # Process batch
         batch_results = model.chat(prompts_list, crops_list)
         
         for j, result in enumerate(batch_results):
@@ -331,7 +278,6 @@ def process_element_batch_optimized(elements, model, prompt, max_batch_size=4):
                 "reading_order": elem["reading_order"],
             })
             
-        # Clear memory
         del crops_list, batch_elements
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
@@ -352,28 +298,6 @@ def generate_fallback_markdown(recognition_results):
     return markdown_content
 
 
-def create_page_gallery(page_previews):
-    """Create a gallery view of processed pages"""
-    if not page_previews:
-        return "No pages processed yet."
-    
-    gallery_html = "<div style='display: grid; grid-template-columns: repeat(auto-fit, minmax(300px, 1fr)); gap: 20px;'>"
-    
-    for preview in page_previews:
-        gallery_html += f"""
-        <div style='border: 1px solid #ddd; padding: 15px; border-radius: 8px;'>
-            <h3>Page {preview['page_num']}</h3>
-            <p><strong>Elements found:</strong> {preview['element_count']}</p>
-            <div style='max-height: 200px; overflow-y: auto; background: #f5f5f5; padding: 10px; border-radius: 4px; font-size: 12px;'>
-                {preview['markdown_preview']}
-            </div>
-        </div>
-        """
-    
-    gallery_html += "</div>"
-    return gallery_html
-
-
 # Initialize model
 model_path = "./hf_model"
 if not os.path.exists(model_path):
@@ -382,130 +306,210 @@ if not os.path.exists(model_path):
 try:
     dolphin_model = DOLPHIN(model_path)
     print(f"Model loaded successfully from {model_path}")
-    model_status = f"✅ Model loaded: {model_path} (Device: {dolphin_model.device})"
+    model_status = f"✅ Model ready (Device: {dolphin_model.device})"
 except Exception as e:
     print(f"Error loading model: {e}")
     dolphin_model = None
     model_status = f"❌ Model failed to load: {str(e)}"
 
 
+# Global state for managing tabs
+processed_markdown = ""
+show_results_tab = False
+
+
 def process_uploaded_pdf(pdf_file, progress=gr.Progress()):
     """Main processing function for uploaded PDF"""
+    global processed_markdown, show_results_tab
+    
     if dolphin_model is None:
-        return "Model not loaded", "Model not loaded", {}, "Model not loaded"
+        return "Model not loaded", gr.Tabs(visible=False)
     
     if pdf_file is None:
-        return "No PDF uploaded", "No PDF uploaded", {}, "No PDF uploaded"
+        return "No PDF uploaded", gr.Tabs(visible=False)
     
     try:
-        # Process the PDF
-        combined_markdown, page_previews, summary_json = process_pdf_document(
-            pdf_file, dolphin_model, progress
-        )
-        
-        # Create page gallery
-        gallery_html = create_page_gallery(page_previews)
-        
-        return combined_markdown, combined_markdown, summary_json, gallery_html
+        combined_markdown, status = process_pdf_document(pdf_file, dolphin_model, progress)
         
+        if status == "processing_complete":
+            processed_markdown = combined_markdown
+            show_results_tab = True
+            return "PDF processed successfully! Check the 'Document' tab above.", gr.Tabs(visible=True)
+        else:
+            show_results_tab = False
+            return combined_markdown, gr.Tabs(visible=False)
+            
     except Exception as e:
+        show_results_tab = False
         error_msg = f"Error processing PDF: {str(e)}"
-        return error_msg, error_msg, {"error": error_msg}, error_msg
+        return error_msg, gr.Tabs(visible=False)
+
+
+def get_processed_markdown():
+    """Return the processed markdown content"""
+    global processed_markdown
+    return processed_markdown if processed_markdown else "No document processed yet."
 
 
 def clear_all():
-    """Clear all inputs and outputs"""
-    return None, "", "", {}, ""
+    """Clear all data and hide results tab"""
+    global processed_markdown, show_results_tab
+    processed_markdown = ""
+    show_results_tab = False
+    return None, "Upload a PDF to get started", gr.Tabs(visible=False)
 
 
-# Create Gradio interface optimized for HuggingFace Spaces
+# Create Gradio interface
 with gr.Blocks(
-    title="DOLPHIN PDF Document AI", 
+    title="DOLPHIN PDF AI", 
     theme=gr.themes.Soft(),
     css="""
-    .main-container { max-width: 1200px; margin: 0 auto; }
-    .status-box { padding: 10px; border-radius: 5px; margin: 10px 0; }
-    .success { background-color: #d4edda; border: 1px solid #c3e6cb; }
-    .error { background-color: #f8d7da; border: 1px solid #f5c6cb; }
+    @import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700&display=swap');
+    
+    * {
+        font-family: 'Inter', -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif !important;
+    }
+    
+    .main-container { 
+        max-width: 1000px; 
+        margin: 0 auto; 
+    }
+    .upload-container { 
+        text-align: center; 
+        padding: 40px 20px;
+        border: 2px dashed #e0e0e0;
+        border-radius: 15px;
+        margin: 20px 0;
+    }
+    .upload-button {
+        font-size: 18px !important;
+        padding: 15px 30px !important;
+        margin: 20px 0 !important;
+        font-weight: 600 !important;
+    }
+    .status-message {
+        text-align: center;
+        padding: 15px;
+        margin: 10px 0;
+        border-radius: 8px;
+        font-weight: 500;
+    }
+    .chatbot-container {
+        max-height: 600px;
+    }
+    h1, h2, h3 {
+        font-weight: 700 !important;
+    }
     """
 ) as demo:
-    gr.Markdown("# 🐬 DOLPHIN PDF Document AI")
-    gr.Markdown(
-        "Upload a PDF document and process it page by page with the DOLPHIN model. "
-        "Optimized for HuggingFace Spaces deployment."
-    )
     
-    # Model status
-    gr.Markdown(f"**Model Status:** {model_status}")
-    
-    with gr.Row():
-        # Left column: Upload and controls
-        with gr.Column(scale=1):
-            gr.Markdown("### 📄 Upload PDF Document")
-            pdf_input = gr.File(
-                file_types=[".pdf"],
-                label="Select PDF File",
-                height=200
+    with gr.Tabs() as main_tabs:
+        # Home Tab
+        with gr.TabItem("🏠 Home", id="home"):
+            gr.Markdown(
+                "# Scholar Express\n"
+                "### Upload a research paper to get a web-friendly version, an AI chatbot, and a podcast summary. Because of our reliance on Generative AI, some errors are inevitable.\n"
+                f"**Status:** {model_status}"
             )
             
-            with gr.Row():
-                process_btn = gr.Button("🚀 Process PDF", variant="primary", size="lg")
-                clear_btn = gr.Button("🗑️ Clear All", variant="secondary")
-        
-        # Right column: Results tabs
-        with gr.Column(scale=2):
-            gr.Markdown("### 📊 Processing Results")
-            
-            with gr.Tabs():
-                with gr.TabItem("📖 Markdown Output"):
-                    markdown_output = gr.Markdown(
-                        label="Processed Document",
-                        latex_delimiters=[
-                            {"left": "$$", "right": "$$", "display": True},
-                            {"left": "$", "right": "$", "display": False}
-                        ],
-                        height=600
-                    )
+            with gr.Column(elem_classes="upload-container"):
+                gr.Markdown("## 📄 Upload Your PDF Document")
                 
-                with gr.TabItem("📝 Raw Markdown"):
-                    raw_markdown = gr.Code(
-                        label="Raw Markdown Text",
-                        language="markdown",
-                        lines=25
-                    )
+                pdf_input = gr.File(
+                    file_types=[".pdf"],
+                    label="",
+                    height=150,
+                    elem_id="pdf_upload"
+                )
                 
-                with gr.TabItem("🔍 Page Gallery"):
-                    page_gallery = gr.HTML(
-                        label="Page Overview"
-                    )
+                process_btn = gr.Button(
+                    "🚀 Process PDF", 
+                    variant="primary", 
+                    size="lg",
+                    elem_classes="upload-button"
+                )
                 
-                with gr.TabItem("🔧 JSON Details"):
-                    json_output = gr.JSON(
-                        label="Processing Details",
-                        height=600
-                    )
-    
-    # Progress bar
-    progress_bar = gr.HTML(visible=False)
+                clear_btn = gr.Button(
+                    "🗑️ Clear", 
+                    variant="secondary"
+                )
+            
+            # Status and progress
+            status_output = gr.Markdown(
+                "Upload a PDF to get started",
+                elem_classes="status-message"
+            )
+        
+        # Results Tab (initially hidden)
+        with gr.TabItem("📖 Document", id="results", visible=False) as results_tab:
+            gr.Markdown("## Processed Document")
+            
+            markdown_display = gr.Markdown(
+                value="",
+                latex_delimiters=[
+                    {"left": "$$", "right": "$$", "display": True},
+                    {"left": "$", "right": "$", "display": False}
+                ],
+                height=700
+            )
+        
+        # Chatbot Tab (initially hidden)
+        with gr.TabItem("💬 Chat", id="chat", visible=False) as chat_tab:
+            gr.Markdown("## Ask Questions About Your Document")
+            
+            chatbot = gr.Chatbot(
+                value=[],
+                height=500,
+                elem_classes="chatbot-container",
+                placeholder="Your conversation will appear here once you process a document..."
+            )
+            
+            with gr.Row():
+                msg_input = gr.Textbox(
+                    placeholder="Ask a question about the processed document...",
+                    scale=4,
+                    container=False
+                )
+                send_btn = gr.Button("Send", variant="primary", scale=1)
+            
+            gr.Markdown(
+                "*Chat functionality will be implemented in the next version*",
+                elem_id="chat-notice"
+            )
     
     # Event handlers
     process_btn.click(
         fn=process_uploaded_pdf,
         inputs=[pdf_input],
-        outputs=[markdown_output, raw_markdown, json_output, page_gallery],
+        outputs=[status_output, results_tab],
         show_progress=True
+    ).then(
+        fn=get_processed_markdown,
+        outputs=[markdown_display]
+    ).then(
+        fn=lambda: gr.TabItem(visible=True),
+        outputs=[chat_tab]
     )
     
     clear_btn.click(
         fn=clear_all,
-        outputs=[pdf_input, markdown_output, raw_markdown, json_output, page_gallery]
+        outputs=[pdf_input, status_output, results_tab]
+    ).then(
+        fn=lambda: gr.TabItem(visible=False),
+        outputs=[chat_tab]
     )
     
-    # Footer
-    gr.Markdown(
-        "---\n"
-        "**Note:** This app is optimized for NVIDIA T4 deployment on HuggingFace Spaces. "
-        "Processing time depends on document complexity and page count."
+    # Placeholder chat functionality
+    def placeholder_chat(message, history):
+        return history + [["Coming soon: AI-powered document Q&A", "This feature will allow you to ask questions about your processed PDF document."]]
+    
+    send_btn.click(
+        fn=placeholder_chat,
+        inputs=[msg_input, chatbot],
+        outputs=[chatbot]
+    ).then(
+        lambda: "",
+        outputs=[msg_input]
     )
 
 
@@ -515,7 +519,7 @@ if __name__ == "__main__":
         server_port=7860,
         share=False,
         show_error=True,
-        max_threads=2,  # Limit threads for memory efficiency
+        max_threads=1,  # Single thread for T4 Small
         inbrowser=False,
         quiet=True
     )