Update app.py

app.py

@@ -38,7 +38,7 @@ def load_model():
             "openbmb/MiniCPM-V-2_6", 
             trust_remote_code=True,
             torch_dtype=torch.float16,
-            device_map="auto"  # Will move to GPU when @spaces.GPU is active
+            device_map="auto"
         )
         return _model, _tokenizer
     except Exception as e:
@@ -195,8 +195,45 @@ def extract_data_from_image(image, extraction_prompt, model, tokenizer):
             "extracted_data": None
         }
 
+def safe_merge_field(combined_data, field, value, page_num, extracted_pages):
+    """Safely merge field data with type checking"""
+    try:
+        if field in combined_data and value:
+            # Handle nested dictionaries (like insurance)
+            if isinstance(value, dict) and isinstance(combined_data[field], dict):
+                for sub_field, sub_value in value.items():
+                    if sub_field in combined_data[field] and sub_value and not combined_data[field][sub_field]:
+                        combined_data[field][sub_field] = sub_value
+                        if page_num not in extracted_pages:
+                            extracted_pages.append(page_num)
+            # Handle simple fields
+            elif not isinstance(value, (dict, list)) and not combined_data[field]:
+                combined_data[field] = value
+                if page_num not in extracted_pages:
+                    extracted_pages.append(page_num)
+    except Exception as e:
+        print(f"Warning: Error merging field {field}: {e}")
+
+def safe_merge_confidence(combined_confidence, field, score):
+    """Safely merge confidence scores with type checking"""
+    try:
+        # Handle nested confidence scores (like primary_insurance)
+        if isinstance(score, dict):
+            if field not in combined_confidence:
+                combined_confidence[field] = {}
+            for sub_field, sub_score in score.items():
+                if (sub_field not in combined_confidence[field] and 
+                    isinstance(sub_score, (int, float)) and sub_score > 0):
+                    combined_confidence[field][sub_field] = sub_score
+        # Handle simple confidence scores
+        elif isinstance(score, (int, float)) and score > 0:
+            if field not in combined_confidence:
+                combined_confidence[field] = score
+    except Exception as e:
+        print(f"Warning: Error merging confidence for {field}: {e}")
+
 def combine_page_data(pages_data):
-    """Combine extracted data from multiple pages into final medical record"""
+    """Combine extracted data from multiple pages into final medical record - FIXED VERSION"""
     combined_data = {
         "date_of_receipt": "",
         "patient_first_name": "",
@@ -238,24 +275,25 @@ def combine_page_data(pages_data):
     
     # Combine data from all pages
     for page_num, page_data in enumerate(pages_data, 1):
-        if page_data["page_data"]["status"] == "success":
-            extracted = page_data["page_data"]["extracted_data"]
-            
-            # If we got JSON data, merge it
-            if isinstance(extracted, dict) and "data" in extracted:
-                page_info = extracted["data"]
-                
-                # Merge non-empty fields (first non-empty value wins)
-                for field, value in page_info.items():
-                    if field in combined_data and value and not combined_data[field]:
-                        combined_data[field] = value
-                        combined_data["extracted_page_numbers"].append(page_num)
+        try:
+            if page_data.get("page_data", {}).get("status") == "success":
+                extracted = page_data["page_data"].get("extracted_data", {})
                 
-                # Merge confidence scores
-                if "confidence_scores" in extracted:
-                    for field, score in extracted["confidence_scores"].items():
-                        if field not in combined_confidence and score > 0:
-                            combined_confidence[field] = score
+                # If we got JSON data, merge it
+                if isinstance(extracted, dict) and "data" in extracted:
+                    page_info = extracted["data"]
+                    
+                    # Safely merge each field
+                    for field, value in page_info.items():
+                        safe_merge_field(combined_data, field, value, page_num, combined_data["extracted_page_numbers"])
+                    
+                    # Safely merge confidence scores
+                    if "confidence_scores" in extracted:
+                        for field, score in extracted["confidence_scores"].items():
+                            safe_merge_confidence(combined_confidence, field, score)
+        except Exception as e:
+            print(f"Warning: Error processing page {page_num}: {e}")
+            continue
     
     return {
         "data": combined_data,
@@ -282,7 +320,7 @@ def extract_efax_from_pdf(pdf_file, custom_prompt=None):
                 "pages_data": []
             }
         
-        # Step 1: Convert PDF to images (CPU operation - do this before GPU)
+        # Step 1: Convert PDF to images (CPU operation)
         print("Converting PDF to images...")
         images = pdf_to_images(pdf_file)
         
@@ -296,7 +334,7 @@ def extract_efax_from_pdf(pdf_file, custom_prompt=None):
         
         print(f"Converted {len(images)} pages. Starting GPU processing...")
         
-        # Step 2: Load model on GPU (happens once GPU session starts)
+        # Step 2: Load model on GPU
         model, tokenizer = load_model()
         
         # Step 3: Use custom prompt or default
@@ -314,7 +352,7 @@ def extract_efax_from_pdf(pdf_file, custom_prompt=None):
         
         print("GPU processing complete. Combining results...")
         
-        # Step 5: Combine data from all pages
+        # Step 5: Combine data from all pages (with error handling)
         combined_result = combine_page_data(pages_data)
         
         # Final result
@@ -331,6 +369,7 @@ def extract_efax_from_pdf(pdf_file, custom_prompt=None):
         return result
         
     except Exception as e:
+        print(f"Error in extract_efax_from_pdf: {e}")
         return {
             "status": "error",
             "error": str(e),
@@ -340,9 +379,9 @@ def extract_efax_from_pdf(pdf_file, custom_prompt=None):
 
 # Create Gradio Interface
 def create_gradio_interface():
-    with gr.Blocks(title="eFax PDF Data Extractor - Optimized ZeroGPU", theme=gr.themes.Soft()) as demo:
+    with gr.Blocks(title="eFax PDF Data Extractor - Fixed", theme=gr.themes.Soft()) as demo:
         gr.Markdown("# 🏥 eFax Medical Data Extraction API")
-        gr.Markdown("🚀 **Optimized GPU Usage** - Single 10-minute GPU session for entire document")
+        gr.Markdown("🚀 **Fixed Version** - Single 10-minute GPU session with proper error handling")
         
         with gr.Tab("📄 PDF Upload & Extraction"):
             with gr.Row():
@@ -361,13 +400,14 @@ def create_gradio_interface():
                             placeholder="Leave empty to use optimized medical data extraction prompt..."
                         )
                     
-                    extract_btn = gr.Button("🚀 Extract Medical Data (10min GPU)", variant="primary", size="lg")
+                    extract_btn = gr.Button("🚀 Extract Medical Data (Fixed)", variant="primary", size="lg")
                     
                     gr.Markdown("""
-                    ### ⚡ Optimized Processing
-                    - **Single GPU Session**: All pages processed in one 10-minute session
-                    - **No Timeouts**: Handles up to 20+ page documents
-                    - **Efficient**: PDF→Images (CPU) → All Processing (GPU) → Results
+                    ### ✅ Bug Fixes Applied
+                    - **Fixed**: Dict/int comparison error
+                    - **Added**: Safe type checking for all operations  
+                    - **Improved**: Error handling and logging
+                    - **Single GPU Session**: No more timeouts
                     """)
                 
                 with gr.Column():
@@ -376,7 +416,7 @@ def create_gradio_interface():
         
         with gr.Tab("🔌 API Usage"):
             gr.Markdown("""
-            ## Optimized API (No Timeout Issues)
+            ## Fixed API (No More Errors)
             
             ### Python Usage
             ```
@@ -396,36 +436,14 @@ def create_gradio_interface():
                 }
             )
             
-            # Now handles 13+ pages without timeout!
+            # Should work without dict/int comparison errors
             result = response.json()
-            medical_data = result["data"]["combined_extraction"]
+            if result["data"]["status"] == "success":
+                medical_data = result["data"]["combined_extraction"]
+                print("Patient:", medical_data["data"]["patient_first_name"])
             ```
             """)
         
-        with gr.Tab("⚡ Performance Info"):
-            gr.Markdown("""
-            ## Optimized ZeroGPU Performance
-            
-            ### Before Optimization (❌ Had Timeout Issues)
-            - GPU session per page = 13 × 30 seconds = 6.5 minutes
-            - Model loading repeated = wasted time
-            - Timeout around page 11/13
-            
-            ### After Optimization (✅ No Timeouts)
-            - **Single 10-minute GPU session** for entire document
-            - Model loads once, processes all pages
-            - Handles 15-20+ page documents easily
-            - PDF conversion on CPU (doesn't count toward GPU time)
-            
-            ### Processing Flow
-            1. **PDF → Images** (CPU, before GPU starts)
-            2. **🚀 GPU Session Starts** (10 minutes allocated)
-            3. **Load Model** (once, on GPU)
-            4. **Process All Pages** (GPU, sequential)
-            5. **GPU Session Ends**
-            6. **Combine Results** (CPU, after GPU)
-            """)
-        
         def process_with_status(pdf_file, custom_prompt):
             if pdf_file is None:
                 return "❌ No PDF file uploaded", {"error": "Please upload a PDF file"}
@@ -436,7 +454,7 @@ def create_gradio_interface():
                 result = extract_efax_from_pdf(pdf_file, custom_prompt if custom_prompt.strip() else None)
                 
                 if result["status"] == "success":
-                    yield f"✅ Successfully processed {result['total_pages']} pages in single GPU session", result
+                    yield f"✅ Successfully processed {result['total_pages']} pages", result
                 else:
                     yield f"❌ Error: {result.get('error', 'Unknown error')}", result
                     
@@ -463,4 +481,4 @@ if __name__ == "__main__":
         server_name="0.0.0.0", 
         server_port=7860,
         show_error=True
-    )
+    )