Update app.py

app.py

@@ -1,13 +1,19 @@
 import gradio as gr
 import torch
-from transformers import AutoModel, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer
 from PIL import Image
 
 model_id = "zai-org/GLM-OCR"
 
-# Load with trust_remote_code=True to get the model's specific processing logic
-tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
-model = AutoModel.from_pretrained(
+# 1. Force use_fast=False to load the custom Python methods
+tokenizer = AutoTokenizer.from_pretrained(
+    model_id, 
+    trust_remote_code=True, 
+    use_fast=False
+)
+
+# 2. Load the model with appropriate precision
+model = AutoModelForCausalLM.from_pretrained(
     model_id, 
     trust_remote_code=True, 
     device_map="auto",
@@ -15,53 +21,62 @@ model = AutoModel.from_pretrained(
 ).eval()
 
 def extract_all_tables(image):
-    if image is None: return "Please upload an image."
-    
-    prompt = "Extract all text and tables from this image. Format tables in Markdown. Do not skip any data."
+    if image is None: 
+        return "Please upload an image."
     
-    # 1. Manually build the input using the model's internal helper
-    # This replaces the 'model.chat' method that was failing
-    inputs = tokenizer.build_chat_input(prompt, images=[image], history=[])
-    inputs = {k: v.to(model.device) for k, v in inputs.items() if torch.is_tensor(v)}
+    # Precise prompt for accountancy data
+    prompt = "Extract all text and all tables from this image. Format tables strictly in Markdown. Preserve all currency symbols."
     
-    # 2. Use the standard generate method
-    with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=4096,
-            do_sample=False,
-            repetition_penalty=1.1
-        )
+    # 3. Use the model's internal processing logic manually 
+    # This avoids the build_chat_input attribute error entirely
+    try:
+        # Most GLM-OCR models use this specific conversation template
+        query = f"<|user|>\n<|image_1|>\n{prompt}<|assistant|>\n"
+        
+        inputs = tokenizer(query, return_tensors="pt").to(model.device)
+        # Handle the image processing separately as expected by the vision encoder
+        pixel_values = tokenizer.process_images([image]).to(model.device).to(torch.float16)
+        
+        with torch.no_grad():
+            outputs = model.generate(
+                input_ids=inputs.input_ids,
+                pixel_values=pixel_values,
+                max_new_tokens=4096,
+                do_sample=False,
+                repetition_penalty=1.1
+            )
+            
+        # 4. Decode the result, skipping the input tokens
+        response = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
+        return response.strip()
         
-    # 3. Decode only the new parts (the answer)
-    input_len = inputs['input_ids'].shape[1]
-    response = tokenizer.decode(outputs[0][input_len:], skip_special_tokens=True)
-    return response
+    except Exception as e:
+        return f"Processing Error: {str(e)}\n\nTip: If this persists, try cropping the image to just the tables."
 
-# UI with Copy Functionality
+# UI with JS-based Copy Button for compatibility
 with gr.Blocks() as demo:
-    gr.Markdown("# 📊 Accountancy Paper Extractor")
+    gr.Markdown("# 📊 Precision Accountancy Table Extractor")
     
     with gr.Row():
         with gr.Column():
-            img_input = gr.Image(type="pil", label="Upload Paper")
-            btn = gr.Button("Extract All Tables", variant="primary")
+            img_input = gr.Image(type="pil", label="Upload Document")
+            btn = gr.Button("Extract Markdown", variant="primary")
         
         with gr.Column():
             text_output = gr.Textbox(
                 label="Extracted Result", 
-                lines=20,
+                lines=25, 
                 interactive=True,
-                elem_id="result-text"
+                elem_id="result_box"
             )
             copy_btn = gr.Button("📋 Copy to Clipboard")
             
-            # JavaScript for copying text
+            # JavaScript to handle clipboard copy
             copy_btn.click(None, None, None, js="""
                 () => {
-                    const text = document.querySelector('#result-text textarea').value;
+                    const text = document.querySelector('#result_box textarea').value;
                     navigator.clipboard.writeText(text);
-                    alert('Copied!');
+                    alert('Text copied to clipboard!');
                 }
             """)