Update inference.py

inference.py

@@ -208,7 +208,7 @@ import json
 from model_loader import get_model
 from processor_utils import load_input
 from prompt import get_part_classifier_prompt, get_part_prompt
-
+import time
 
 def _get_max_tokens(part_name):
     limits = {
@@ -383,7 +383,11 @@ def process_document(file_path):
     page_results = []
 
     for idx, image in enumerate(pages, start=1):
+        print("first model has been called for",idx,"image")
+        start = time.time()
         part_name = classify_page(image, model, processor, device)
+        end = time.time()
+        print("total time taken by the first model",end-start,"sec")
 
         if part_name == "UNKNOWN":
             page_results.append({
@@ -394,8 +398,11 @@ def process_document(file_path):
                 "parsed": None
             })
             continue
-
+        print("second model has been called for",idx,"image")
+        start = time.time()
         result = extract_part_json(image, part_name, model, processor, device)
+        end = time.time()
+        print("total time taken by the second model",end-start,"sec")
         result["page_number"] = idx
         page_results.append(result)