Update app.py

app.py

@@ -1,61 +1,67 @@
 import gradio as gr
-from transformers import AutoModel, AutoTokenizer
 import torch
 from PIL import Image
+from transformers import AutoProcessor, GlmOcrForConditionalGeneration
 
-# Use the same model ID
+# 1. Load the model and processor correctly
 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", 
-    torch_dtype=torch.float16
+# Processor handles both the image and the text encoding
+processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
+model = GlmOcrForConditionalGeneration.from_pretrained(
+    model_id,
+    trust_remote_code=True,
+    device_map="auto",
+    torch_dtype=torch.bfloat16 # Use bfloat16 for better accuracy on T4/A10G GPUs
 ).eval()
 
 def process_accountancy_image(input_image):
     if input_image is None:
         return "Please upload an image."
 
-    # 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."
-    )
-
-    # 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("Extracts complex ledgers and financial tables into Markdown format.")
+    # 2. Prepare the input using the standard conversation format
+    prompt = "Extract all text and tables from this accountancy question. Format tables in Markdown."
+    
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "image", "image": input_image},
+                {"type": "text", "text": prompt},
+            ],
+        }
+    ]
+
+    # 3. Process inputs and generate
+    text = processor.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
+    inputs = processor(text=[text], images=[input_image], return_tensors="pt").to(model.device)
+
+    with torch.no_grad():
+        output_ids = model.generate(
+            **inputs,
+            max_new_tokens=2048,
+            do_sample=False # Keep it deterministic for accounting numbers
+        )
     
+    # 4. Decode the result
+    # We slice the output to remove the input tokens and show only the answer
+    generated_ids = [
+        output_ids[len(input_ids):] for input_ids, output_ids in zip(inputs.input_ids, output_ids)
+    ]
+    response = processor.batch_decode(generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)[0]
+    
+    return response
+
+# UI setup
+with gr.Blocks() as demo:
+    gr.Markdown("# 📊 Precision Accountancy OCR")
     with gr.Row():
         with gr.Column():
-            image_input = gr.Image(type="pil", label="Upload Accountancy Question")
-            submit_btn = gr.Button("Extract Text & Tables", variant="primary")
-        
+            img_in = gr.Image(type="pil", label="Upload Balance Sheet/Ledger")
+            btn = gr.Button("Extract Data", variant="primary")
         with gr.Column():
-            text_output = gr.Markdown(label="Extracted Result")
-
-    submit_btn.click(fn=process_accountancy_image, inputs=image_input, outputs=text_output)
+            out = gr.Markdown()
+            
+    btn.click(process_accountancy_image, inputs=img_in, outputs=out)
 
-# Launching with the fix for Gradio 6 theme warning
 demo.launch()