Update llm.py

llm.py

@@ -34,26 +34,13 @@ if torch.cuda.is_available():
 def execute_llm(model, processor, image, prompt: str):
 
     if not prompt:
-        prompt = """
-        Extract all text from image.
-        Return ONLY valid JSON.
-        """
-
-    messages = [{
-        "role": "user",
-        "content": [
-            {"type": "image", "image": image},
-            {"type": "text", "text": prompt}
-        ]
-    }]
+        prompt = "Extract all text from the given image and return ONLY valid JSON."
 
     with torch.inference_mode():
 
-        inputs = processor.apply_chat_template(
-            messages,
-            add_generation_prompt=True,
-            tokenize=True,
-            return_dict=True,
+        inputs = processor(
+            text=prompt,
+            images=image,
             return_tensors="pt"
         )
 
@@ -65,11 +52,13 @@ def execute_llm(model, processor, image, prompt: str):
             do_sample=False
         )
 
-        return processor.decode(
-            outputs[0][inputs["input_ids"].shape[-1]:],
+        result = processor.decode(
+            outputs[0],
             skip_special_tokens=True
         )
 
+    return result
+
 async def call_llm(image, prompt: str = ""):
     print("call llm")