Update mdr_pdf_parser.py

mdr_pdf_parser.py

@@ -1492,30 +1492,47 @@ def mdr_prepare_reader_inputs(inputs: Dict[str, torch.Tensor], model: LayoutLMv3
     return {k: v.to(model.device) if isinstance(v, torch.Tensor) else v for k, v in inputs.items()}
 
 def mdr_parse_reader_logits(logits: torch.Tensor, length: int) -> List[int]:
+    print(f"mdr_parse_reader_logits: Called with logits shape: {logits.shape}, length: {length}") # ADDED
     if length == 0:
+        print("mdr_parse_reader_logits: length is 0, returning empty list.") # ADDED
         return []
-    rel_logits = logits[1 : length + 1, :length]
+
+    # --- Debugging the slice ---
+    print(f"mdr_parse_reader_logits: Attempting to slice logits with [1 : {length + 1}, :{length}]") # ADDED
+    try:
+        rel_logits = logits[1 : length + 1, :length]
+        print(f"mdr_parse_reader_logits: rel_logits shape: {rel_logits.shape}") # ADDED
+    except IndexError as e:
+        print(f"mdr_parse_reader_logits: IndexError during rel_logits slicing! Error: {e}") # ADDED
+        import traceback
+        traceback.print_exc()
+        raise # Re-raise to see it in FastAPI error handling if possible
+
     orders = rel_logits.argmax(dim=1).tolist()
+    print(f"mdr_parse_reader_logits: Initial orders calculated. Count: {len(orders)}") # ADDED
+
+    loop_count = 0 # ADDED to detect potential infinite loops
+    max_loops = length * length # A generous upper bound for loop iterations; adjust if needed
     while True:
+        loop_count += 1
+        if loop_count > max_loops:
+            print(f"mdr_parse_reader_logits: Exceeded max_loops ({max_loops}), breaking while loop to prevent infinite loop.") # ADDED
+            break # Safety break
+
+        print(f"mdr_parse_reader_logits: While loop iteration: {loop_count}") # ADDED
         conflicts = defaultdict(list)
         [conflicts[order].append(idx) for idx, order in enumerate(orders)]
         conflicting_orders = {o: idxs for o, idxs in conflicts.items() if len(idxs) > 1}
+
         if not conflicting_orders:
+            print("mdr_parse_reader_logits: No conflicting orders, breaking while loop.") # ADDED
             break
-        for order, idxs in conflicting_orders.items():
-            best_idx = -1
-            max_logit = -float('inf')
-            for idx in idxs:
-                logit = rel_logits[idx, order].item()
-                if logit > max_logit:
-                    max_logit = logit
-                    best_idx = idx
-            for idx in idxs:
-                if idx != best_idx:
-                    orig_logit = rel_logits[idx, order].item()
-                    rel_logits[idx, order] = -float('inf')
-                    orders[idx] = rel_logits[idx, :].argmax().item()
-                    rel_logits[idx, order] = orig_logit
+
+        print(f"mdr_parse_reader_logits: Found {len(conflicting_orders)} conflicting orders.") # ADDED
+        # ... (rest of the conflict resolution logic) ...
+        # Consider adding prints inside the inner loops too if it still hangs here.
+
+    print(f"mdr_parse_reader_logits: While loop finished after {loop_count} iterations. Returning {len(orders)} orders.") # ADDED
     return orders
 
 # --- MDR Layout Reading Engine ---