Update inference.py

inference.py

@@ -1,11 +1,89 @@
-import torch
-from model_loader import model, processor, device
+# 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)
+
+#     messages = [
+#         {
+#             "role": "user",
+#             "content": [
+#                 {"type": "image", "image": image},
+#                 {"type": "text", "text": get_prompt()}
+#             ]
+#         }
+#     ]
+
+#     text = processor.apply_chat_template(
+#         messages,
+#         tokenize=False,
+#         add_generation_prompt=True
+#     )
+
+#     inputs = processor(
+#         text=[text],
+#         images=[image],
+#         return_tensors="pt"
+#     ).to(device)
+
+#     output = model.generate(
+#         **inputs,
+#         max_new_tokens=1500,
+#         do_sample=False,   #  if it is true there will be extra text with output
+#         # temperature=0.1   # temp is not required
+#     )
+
+#     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
+#     ).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
 from prompt import get_prompt
-import json
-def process_document(file_path):
-    image = load_input(file_path)
 
+
+def _extract_json_block(text):
+    start = text.find("{")
+    end = text.rfind("}") + 1
+
+    if start == -1 or end == 0:
+        return None
+
+    return text[start:end]
+
+
+def _run_page_inference(image, model, processor, device):
     messages = [
         {
             "role": "user",
@@ -30,40 +108,53 @@ def process_document(file_path):
 
     output = model.generate(
         **inputs,
-        max_new_tokens=1500,
-        do_sample=False,   #  if it is true there will be extra text with output
-        # temperature=0.1   # temp is not required
+        max_new_tokens=150,
+        do_sample=False
     )
 
     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
+    json_block = _extract_json_block(response)
 
-    if start != -1 and end != -1:
-        response = response[start:end]
+    if not json_block:
+        return {
+            "status": "error",
+            "raw_output": response,
+            "parsed": None
+        }
 
     try:
-        parsed = json.loads(response)
-    except:
-        parsed = {
-        "error": "Invalid JSON",
-        "raw": response
+        parsed = json.loads(json_block)
+        return {
+            "status": "success",
+            "raw_output": response,
+            "parsed": parsed
         }
+    except json.JSONDecodeError:
+        return {
+            "status": "error",
+            "raw_output": response,
+            "parsed": None
+        }
+
+
+def process_document(file_path):
+    model, processor, device = get_model()
+    pages = load_input(file_path)
+
+    page_results = []
+
+    for page_number, image in enumerate(pages, start=1):
+        result = _run_page_inference(image, model, processor, device)
+        result["page_number"] = page_number
+        page_results.append(result)
 
-    return parsed
+    return {
+        "total_pages": len(page_results),
+        "pages": page_results
+    }