Update inference.py

inference.py

@@ -37,36 +37,36 @@ 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
-    )
+    # response = processor.decode(   # past code
+    #     generated_ids,
+    #     skip_special_tokens=True
+    # )
 
-    return response.strip()
+    # return response.strip()
 
 
     
-    # response = processor.decode(
-    # 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 parsed
+    response = processor.decode(
+    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 parsed
 # import json
 # from model_loader import get_model
 # from processor_utils import load_input