Update app.py

app.py

@@ -3,39 +3,42 @@ import torch
 from transformers import AutoModel, AutoTokenizer
 from PIL import Image
 
-model_id = "zai-org/GLM-OCR" 
+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(
     model_id, 
     trust_remote_code=True, 
-    device_map="auto"
+    device_map="auto",
+    torch_dtype=torch.float16
 ).eval()
 
 def extract_all_tables(image):
     if image is None: return "Please upload an image."
     
-    # We explicitly tell the model to find multiple tables
-    prompt = (
-        "Extract all text and all tables from this image. "
-        "Format every table in Markdown. Do not skip any data. "
-        "Provide the full content of the page."
-    )
+    prompt = "Extract all text and tables from this image. Format tables in Markdown. Do not skip any data."
     
+    # 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)}
+    
+    # 2. Use the standard generate method
     with torch.no_grad():
-        # repetition_penalty helps prevent the model from getting stuck
-        # max_new_tokens is high to ensure both tables fit
-        response, _ = model.chat(
-            tokenizer, 
-            query=prompt, 
-            images=[image], 
-            history=[], 
+        outputs = model.generate(
+            **inputs,
             max_new_tokens=4096,
+            do_sample=False,
             repetition_penalty=1.1
         )
+        
+    # 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
 
-# UI with Compatibility Fix
+# UI with Copy Functionality
 with gr.Blocks() as demo:
     gr.Markdown("# 📊 Accountancy Paper Extractor")
     
@@ -51,18 +54,17 @@ with gr.Blocks() as demo:
                 interactive=True,
                 elem_id="result-text"
             )
-            # Standard Gradio version of a copy functionality
             copy_btn = gr.Button("📋 Copy to Clipboard")
-            # This uses a simple JavaScript snippet to copy the text
+            
+            # JavaScript for copying text
             copy_btn.click(None, None, None, js="""
                 () => {
                     const text = document.querySelector('#result-text textarea').value;
                     navigator.clipboard.writeText(text);
-                    alert('Copied to clipboard!');
+                    alert('Copied!');
                 }
             """)
 
     btn.click(extract_all_tables, inputs=img_input, outputs=text_output)
 
-# Fixed theme warning by moving it to launch
-demo.launch(theme=gr.themes.Default())
+demo.launch()