Update handler.py

handler.py

@@ -6,7 +6,8 @@ import base64
 
 class EndpointHandler():
     def __init__(self, path=""):
-        # Use the specific GLM-OCR architecture
+        # We MUST use trust_remote_code=True because the architecture 
+        # is defined in the files you downloaded, not in the library itself.
         self.processor = AutoProcessor.from_pretrained(path, trust_remote_code=True)
         self.model = AutoModelForVision2Seq.from_pretrained(
             path, 
@@ -17,17 +18,27 @@ class EndpointHandler():
         self.model.eval()
 
     def __call__(self, data):
-        # Decode the image sent from Google Sheets
-        inputs = data.pop("inputs", data)
-        image_data = base64.b64decode(inputs)
-        image = Image.open(io.BytesIO(image_data)).convert("RGB")
+        # Handle the input format from Google Sheets
+        inputs_data = data.pop("inputs", data)
         
-        # Format for GLM-OCR
-        prompt = "Identify Date, Vendor, and list every Item with description, qty, and price. Return as JSON."
-        inputs = self.processor(text=prompt, images=image, return_tensors="pt").to(self.model.device)
+        # If the data comes in as a string (base64)
+        if isinstance(inputs_data, str):
+            image_bytes = base64.b64decode(inputs_data)
+        else:
+            # Handle direct bytes if necessary
+            image_bytes = inputs_data
+            
+        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+        
+        # The prompt that tells the AI what to look for
+        prompt = "Extract all line items. Return a JSON array of objects with: date, vendor, description, qty, price, total."
+        
+        # Process the image and text
+        model_inputs = self.processor(text=prompt, images=image, return_tensors="pt").to(self.model.device)
         
         with torch.no_grad():
-            generated_ids = self.model.generate(**inputs, max_new_tokens=1024)
+            generated_ids = self.model.generate(**model_inputs, max_new_tokens=1024)
         
         result = self.processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+        
         return [{"generated_text": result}]