Update inference.py

inference.py

@@ -2,7 +2,7 @@ import torch
 from model_loader import model, processor, device
 from processor_utils import load_input
 from prompt import get_prompt
-
+import json
 def process_document(file_path):
     image = load_input(file_path)
 
@@ -37,9 +37,33 @@ def process_document(file_path):
 
     generated_ids = output[0][inputs.input_ids.shape[-1]:]
 
+    # response = processor.decode(   # past code
+    #     generated_ids,
+    #     skip_special_tokens=True
+    # )
+
+    # return response.strip()
+
+
+    
     response = processor.decode(
-        generated_ids,
-        skip_special_tokens=True
-    )
+    generated_ids,
+    skip_special_tokens=True
+    ).strip()
+
+    # 🔥 FORCE JSON CLEANING
+    start = response.find("{")
+    end = response.rfind("}") + 1
+
+    if start != -1 and end != -1:
+        response = response[start:end]
+
+    try:
+        parsed = json.loads(response)
+    except:
+        parsed = {
+        "error": "Invalid JSON",
+        "raw": response
+        }
 
-    return response.strip()
+    return parsed