Update app.py

app.py

@@ -3,50 +3,53 @@ from transformers import AutoModel, AutoTokenizer
 import torch
 from PIL import Image
 
-# 1. Load the High-Accuracy Model
-# Note: Using AutoModel instead of AutoModelForCausalLM for custom GLM architectures
+# Use the same model ID
 model_id = "zai-org/GLM-OCR"
 
+# 1. Load the tokenizer and model with remote code enabled
+# We use AutoModel because GLM-OCR has a custom vision-language architecture
 tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
 model = AutoModel.from_pretrained(
     model_id, 
     trust_remote_code=True, 
     device_map="auto", 
-    dtype=torch.float16 # Fixed the deprecated parameter
+    torch_dtype=torch.float16
 ).eval()
 
 def process_accountancy_image(input_image):
     if input_image is None:
         return "Please upload an image."
 
-    # 2. Define the Prompt for Accountancy
+    # Define the specific prompt for accountancy extraction
     prompt = (
         "Extract all text and tables from this image. "
         "If there is an accountancy table (Ledger, Balance Sheet, P&L), "
         "format it strictly as a Markdown table. Preserve all symbols and currency signs."
     )
 
-    # 3. Run Inference
-    # For GLM models, we use the specific chat input builder provided by their remote code
-    inputs = tokenizer.build_chat_input(prompt, images=[input_image], history=[])
-    inputs = {k: v.to(model.device) for k, v in inputs.items() if torch.is_tensor(v)}
-    
-    with torch.no_grad():
-        # GLM-OCR typically uses a custom generate method
-        outputs = model.generate(**inputs, max_new_tokens=2048)
-        # Decoding logic
-        response = tokenizer.decode(outputs[0][len(inputs['input_ids'][0]):], skip_special_tokens=True)
-
-    return response
-
-# 4. User Interface
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    # FIX: GLM-OCR uses a custom 'chat' method usually defined in its modeling script
+    # We will use the model's own chat functionality which handles the image processing internally
+    try:
+        # Most GLM vision models use this specific internal method for inference
+        with torch.no_grad():
+            response, _ = model.chat(
+                tokenizer,
+                query=prompt,
+                images=[input_image],
+                history=[]
+            )
+        return response
+    except Exception as e:
+        return f"Extraction Error: {str(e)}\n\nTry ensuring the image is clear and not too large."
+
+# Gradio Interface
+with gr.Blocks() as demo:
     gr.Markdown("# 📊 Professional Accountancy Table Extractor")
-    gr.Markdown("High-accuracy extraction for ledgers and financial questions.")
+    gr.Markdown("Extracts complex ledgers and financial tables into Markdown format.")
     
     with gr.Row():
         with gr.Column():
-            image_input = gr.Image(type="pil", label="Upload Image")
+            image_input = gr.Image(type="pil", label="Upload Accountancy Question")
             submit_btn = gr.Button("Extract Text & Tables", variant="primary")
         
         with gr.Column():
@@ -54,4 +57,5 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
 
     submit_btn.click(fn=process_accountancy_image, inputs=image_input, outputs=text_output)
 
+# Launching with the fix for Gradio 6 theme warning
 demo.launch()