Update inference.py

inference.py

@@ -25,22 +25,26 @@ def process_document(image):
 
     text = processor.apply_chat_template(
         messages,
-        tokenize=False,
-        add_generation_prompt=True
+        tokenize=False,  # so that this can return string output
+        add_generation_prompt=True   # if true it will add extra on start and end
     )
+    print(f"The text of inference is {text}")
 
     inputs = processor(
         text=[text],
         images=[image],
         return_tensors="pt"
     ).to(device)
-
+    print(f"The inputs of inference is {inputs}")
+    
     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
     )
+    print(f"The output of inference is {output}")
+    
 
     generated_ids = output[0][inputs.input_ids.shape[-1]:]
 
@@ -73,95 +77,4 @@ def process_document(image):
         "raw": response
         }
 
-    return parsed
-# import json
-# from model_loader import get_model
-# from processor_utils import load_input
-# from prompt import get_prompt
-
-
-# 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 = []
-
-#     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 parsed