Update app.py

app.py

@@ -1,25 +1,25 @@
 import gradio as gr
 import torch
 import json
-import uuid
 from PIL import Image
-
-from transformers import AutoProcessor, AutoModelForVision2Seq
+from transformers import AutoProcessor, AutoModel
 
 MODEL_ID = "Qwen/Qwen2.5-VL-7B-Instruct"
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
+# Processor
 processor = AutoProcessor.from_pretrained(
     MODEL_ID,
     trust_remote_code=True
 )
 
-model = AutoModelForVision2Seq.from_pretrained(
+# Model (REMOTE CODE LOAD — critical)
+model = AutoModel.from_pretrained(
     MODEL_ID,
-    device_map="auto",
+    trust_remote_code=True,
     torch_dtype=torch.float16 if device == "cuda" else torch.float32,
-    trust_remote_code=True
+    device_map="auto"
 )
 
 model.eval()
@@ -30,9 +30,9 @@ You are a universal document understanding AI.
 Return ONLY valid JSON.
 
 Extract:
-- document type
+- document_type
 - key-value fields
-- tables with rows & columns
+- tables with rows and columns
 
 Be document-agnostic.
 Do not hallucinate.
@@ -45,23 +45,34 @@ Do not hallucinate.
     ).to(model.device)
 
     with torch.no_grad():
-        output = model.generate(**inputs, max_new_tokens=2048)
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=2048,
+            temperature=0.0
+        )
 
-    text = processor.decode(output[0], skip_special_tokens=True)
+    text = processor.decode(outputs[0], skip_special_tokens=True)
 
     try:
         start = text.find("{")
         end = text.rfind("}") + 1
         return json.loads(text[start:end])
-    except:
-        return {"error": "Failed to parse output"}
+    except Exception:
+        return {
+            "error": "Model output could not be parsed",
+            "raw_output": text
+        }
 
 with gr.Blocks() as demo:
-    gr.Markdown("# 📄 DocAI – Universal Document Intelligence")
+    gr.Markdown("# 📄 DocAI — Universal Document Intelligence")
 
-    img = gr.Image(type="pil")
-    out = gr.JSON()
+    image = gr.Image(type="pil", label="Upload document")
+    output = gr.JSON(label="Extracted JSON")
 
-    gr.Button("Extract").click(extract_document, img, out)
+    gr.Button("Extract").click(
+        extract_document,
+        inputs=image,
+        outputs=output
+    )
 
 demo.launch()