Update inference.py

inference.py

@@ -1,89 +1,11 @@
-# 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
+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)
 
-
-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",
@@ -108,53 +30,131 @@ def _run_page_inference(image, model, processor, device):
 
     output = model.generate(
         **inputs,
-        max_new_tokens=150,
-        do_sample=False
+        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
+    generated_ids,
+    skip_special_tokens=True
     ).strip()
 
-    json_block = _extract_json_block(response)
+    # 🔥 FORCE JSON CLEANING
+    start = response.find("{")
+    end = response.rfind("}") + 1
 
-    if not json_block:
-        return {
-            "status": "error",
-            "raw_output": response,
-            "parsed": None
-        }
+    if start != -1 and end != -1:
+        response = response[start:end]
 
     try:
-        parsed = json.loads(json_block)
-        return {
-            "status": "success",
-            "raw_output": response,
-            "parsed": parsed
-        }
-    except json.JSONDecodeError:
-        return {
-            "status": "error",
-            "raw_output": response,
-            "parsed": None
+        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
+
 
-def process_document(file_path):
-    model, processor, device = get_model()
-    pages = 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",
+#             "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=150,
+#         do_sample=False
+#     )
+
+#     generated_ids = output[0][inputs.input_ids.shape[-1]:]
+
+#     response = processor.decode(
+#         generated_ids,
+#         skip_special_tokens=True
+#     ).strip()
+
+#     json_block = _extract_json_block(response)
+
+#     if not json_block:
+#         return {
+#             "status": "error",
+#             "raw_output": response,
+#             "parsed": None
+#         }
+
+#     try:
+#         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 = []
+#     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)
+#     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 {
-        "total_pages": len(page_results),
-        "pages": page_results
-    }
+#     return {
+#         "total_pages": len(page_results),
+#         "pages": page_results
+#     }