Update working_yolo_pipeline.py

working_yolo_pipeline.py

@@ -23,8 +23,8 @@ import re
 
 import torch.nn as nn
 from TorchCRF import CRF
-from transformers import LayoutLMv3TokenizerFast, LayoutLMv3Model, LayoutLMv3Config
-# from transformers import LayoutLMv3Tokenizer, LayoutLMv3Model, LayoutLMv3Config
+# from transformers import LayoutLMv3TokenizerFast, LayoutLMv3Model, LayoutLMv3Config
+from transformers import LayoutLMv3Tokenizer, LayoutLMv3Model, LayoutLMv3Config
 from typing import List, Dict, Any, Optional, Union, Tuple
 from ultralytics import YOLO
 import glob
@@ -75,51 +75,18 @@ except Exception as e:
 
 from typing import Optional
 
-# def sanitize_text(text: Optional[str]) -> str:
-#     """Removes surrogate characters and other invalid code points that cause UTF-8 encoding errors."""
-#     if not isinstance(text, str) or text is None:
-#         return ""
-    
-#     # Matches all surrogates (\ud800-\udfff) and common non-characters (\ufffe, \uffff).
-#     # This specifically removes '\udefd' which is causing your error.
-#     surrogates_and_nonchars = re.compile(r'[\ud800-\udfff\ufffe\uffff]')
-    
-#     # Replace the invalid characters with a standard space.
-#     # We strip afterward in the calling function.
-#     return surrogates_and_nonchars.sub(' ', text)
-
-
-
-
-# Robust sanitize_text: removes surrogates/non-characters, normalizes line breaks,
-# collapses multiple spaces, and removes remaining invalid bytes via utf-8 ignore.
-import unicodedata
-
 def sanitize_text(text: Optional[str]) -> str:
+    """Removes surrogate characters and other invalid code points that cause UTF-8 encoding errors."""
     if not isinstance(text, str) or text is None:
         return ""
-
-    # 1) Normalize common unicode forms (NFKC keeps compatibility forms reasonable)
-    try:
-        text = unicodedata.normalize("NFC", text)
-    except Exception:
-        pass
-
-    # 2) Remove surrogate codepoint range and common non-characters
-    #    \ud800-\udfff are surrogate halves; also remove \ufffe and \uffff
-    text = re.sub(r'[\uD800-\uDFFF\uFFFE\uFFFF]', ' ', text)
-
-    # 3) Remove other control chars except common whitespace (newline/tab)
-    text = re.sub(r'[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]', ' ', text)
-
-    # 4) Normalize newlines to single space, collapse repeated whitespace
-    text = re.sub(r'[\r\n]+', ' ', text)
-    text = re.sub(r'\s+', ' ', text).strip()
-
-    # 5) Final safety: encode/decode ignoring errors (this strips any remaining bad bytes)
-    cleaned = text.encode('utf-8', 'ignore').decode('utf-8', 'ignore')
-    return cleaned
-
+    
+    # Matches all surrogates (\ud800-\udfff) and common non-characters (\ufffe, \uffff).
+    # This specifically removes '\udefd' which is causing your error.
+    surrogates_and_nonchars = re.compile(r'[\ud800-\udfff\ufffe\uffff]')
+    
+    # Replace the invalid characters with a standard space.
+    # We strip afterward in the calling function.
+    return surrogates_and_nonchars.sub(' ', text)
 
 
 
@@ -750,105 +717,74 @@ def pixmap_to_numpy(pix: fitz.Pixmap) -> np.ndarray:
 
 
 
-# def extract_native_words_and_convert(fitz_page, scale_factor: float = 2.0) -> list:
-#     # 1. Get raw data
-#     try:
-#         raw_word_data = fitz_page.get_text("words")
-#     except Exception as e:
-#         print(f"  ❌ PyMuPDF extraction failed completely: {e}")
-#         return []
+def extract_native_words_and_convert(fitz_page, scale_factor: float = 2.0) -> list:
+    # 1. Get raw data
+    try:
+        raw_word_data = fitz_page.get_text("words")
+    except Exception as e:
+        print(f"  ❌ PyMuPDF extraction failed completely: {e}")
+        return []
 
-#     # ==============================================================================
-#     # --- DEBUGGING BLOCK: CHECK FIRST 50 NATIVE WORDS (SAFE PRINT) ---
-#     # ==============================================================================
-#     print(f"\n[DEBUG] Native Extraction (Page {fitz_page.number + 1}): Checking first 50 words...")
+    # ==============================================================================
+    # --- DEBUGGING BLOCK: CHECK FIRST 50 NATIVE WORDS (SAFE PRINT) ---
+    # ==============================================================================
+    print(f"\n[DEBUG] Native Extraction (Page {fitz_page.number + 1}): Checking first 50 words...")
     
-#     debug_count = 0
-#     for item in raw_word_data:
-#         if debug_count >= 150: break
+    debug_count = 0
+    for item in raw_word_data:
+        if debug_count >= 50: break
         
-#         word_text = item[4]
+        word_text = item[4]
         
-#         # --- SAFE PRINTING LOGIC ---
-#         # We encode/decode to ignore surrogates just for the print statement
-#         # This prevents the "UnicodeEncodeError" that was crashing your script
-#         safe_text = word_text.encode('utf-8', 'ignore').decode('utf-8')
+        # --- SAFE PRINTING LOGIC ---
+        # We encode/decode to ignore surrogates just for the print statement
+        # This prevents the "UnicodeEncodeError" that was crashing your script
+        safe_text = word_text.encode('utf-8', 'ignore').decode('utf-8')
         
-#         # Get hex codes (handling potential errors in 'ord')
-#         try:
-#             unicode_points = [f"\\u{ord(c):04x}" for c in word_text]
-#         except:
-#             unicode_points = ["ERROR"]
+        # Get hex codes (handling potential errors in 'ord')
+        try:
+            unicode_points = [f"\\u{ord(c):04x}" for c in word_text]
+        except:
+            unicode_points = ["ERROR"]
             
-#         print(f"  Word {debug_count}: '{safe_text}' -> Codes: {unicode_points}")
-#         debug_count += 1
-#     print("----------------------------------------------------------------------\n")
-#     # ==============================================================================
-
-#     converted_ocr_output = []
-#     DEFAULT_CONFIDENCE = 99.0
-
-#     for x1, y1, x2, y2, word, *rest in raw_word_data:
-#         # --- FIX: ROBUST SANITIZATION ---
-#         # 1. Encode to UTF-8 ignoring errors (strips surrogates)
-#         # 2. Decode back to string
-#         cleaned_word_bytes = word.strip()
-#         # cleaned_word_bytes = word.encode('utf-8', 'ignore')
-#         # cleaned_word = cleaned_word_bytes.decode('utf-8')
-#         # cleaned_word = word.encode('utf-8', 'ignore').decode('utf-8').strip()
-        
-#         # cleaned_word = cleaned_word.strip()
-#         if not cleaned_word: continue
-        
-#         x1_pix = int(x1 * scale_factor)
-#         y1_pix = int(y1 * scale_factor)
-#         x2_pix = int(x2 * scale_factor)
-#         y2_pix = int(y2 * scale_factor)
-        
-#         converted_ocr_output.append({
-#             'type': 'text',
-#             'word': cleaned_word, 
-#             'confidence': DEFAULT_CONFIDENCE,
-#             'bbox': [x1_pix, y1_pix, x2_pix, y2_pix],
-#             'y0': y1_pix, 'x0': x1_pix
-#         })
-        
-#     return converted_ocr_output
+        print(f"  Word {debug_count}: '{safe_text}' -> Codes: {unicode_points}")
+        debug_count += 1
+    print("----------------------------------------------------------------------\n")
+    # ==============================================================================
 
-
-
-
-
-
-
-def extract_native_words_and_convert(fitz_page, scale_factor: float = 2.0) -> list:
-    raw_word_data = fitz_page.get_text("words")
     converted_ocr_output = []
     DEFAULT_CONFIDENCE = 99.0
 
     for x1, y1, x2, y2, word, *rest in raw_word_data:
-        if not word.strip(): continue
+        # --- FIX: ROBUST SANITIZATION ---
+        # 1. Encode to UTF-8 ignoring errors (strips surrogates)
+        # 2. Decode back to string
+        cleaned_word_bytes = word.encode('utf-8', 'ignore')
+        cleaned_word = cleaned_word_bytes.decode('utf-8')
+        cleaned_word = word.encode('utf-8', 'ignore').decode('utf-8').strip()
+        
+        # cleaned_word = cleaned_word.strip()
+        if not cleaned_word: continue
+        
         x1_pix = int(x1 * scale_factor)
         y1_pix = int(y1 * scale_factor)
         x2_pix = int(x2 * scale_factor)
         y2_pix = int(y2 * scale_factor)
+        
         converted_ocr_output.append({
             'type': 'text',
-            'word': word,
+            'word': cleaned_word, 
             'confidence': DEFAULT_CONFIDENCE,
             'bbox': [x1_pix, y1_pix, x2_pix, y2_pix],
             'y0': y1_pix, 'x0': x1_pix
         })
+        
     return converted_ocr_output
 
 
 
 
 
-
-
-
-
 #===================================================================================================
 #===================================================================================================
 #===================================================================================================
@@ -1038,375 +974,111 @@ def post_process_json_with_inference(json_data, classifier):
 
 
 
-# def preprocess_and_ocr_page(original_img: np.ndarray, model, pdf_path: str,
-#                             page_num: int, fitz_page: fitz.Page,
-#                             pdf_name: str) -> Tuple[List[Dict[str, Any]], Optional[int]]:
-#     """
-#     OPTIMIZED FLOW:
-#     1. Run YOLO to find Equations/Tables.
-#     2. Mask raw text with YOLO boxes.
-#     3. Run Column Detection on the MASKED data.
-#     4. Proceed with OCR (Native or High-Res Tesseract Fallback) and Output.
-#     """
-#     global GLOBAL_FIGURE_COUNT, GLOBAL_EQUATION_COUNT
-
-#     start_time_total = time.time()
+def preprocess_and_ocr_page(original_img: np.ndarray, model, pdf_path: str,
+                            page_num: int, fitz_page: fitz.Page,
+                            pdf_name: str) -> Tuple[List[Dict[str, Any]], Optional[int]]:
+    """
+    OPTIMIZED FLOW:
+    1. Run YOLO to find Equations/Tables.
+    2. Mask raw text with YOLO boxes.
+    3. Run Column Detection on the MASKED data.
+    4. Proceed with OCR (Native or High-Res Tesseract Fallback) and Output.
+    """
+    global GLOBAL_FIGURE_COUNT, GLOBAL_EQUATION_COUNT
 
-#     if original_img is None:
-#         print(f"  ❌ Invalid image for page {page_num}.")
-#         return None, None
+    start_time_total = time.time()
 
-#     # ====================================================================
-#     # --- STEP 1: YOLO DETECTION ---
-#     # ====================================================================
-#     start_time_yolo = time.time()
-#     results = model.predict(source=original_img, conf=CONF_THRESHOLD, imgsz=640, verbose=False)
+    if original_img is None:
+        print(f"  ❌ Invalid image for page {page_num}.")
+        return None, None
 
-#     relevant_detections = []
-#     if results and results[0].boxes:
-#         for box in results[0].boxes:
-#             class_id = int(box.cls[0])
-#             class_name = model.names[class_id]
-#             if class_name in TARGET_CLASSES:
-#                 x1, y1, x2, y2 = box.xyxy[0].cpu().numpy().astype(int)
-#                 relevant_detections.append(
-#                     {'coords': (x1, y1, x2, y2), 'y1': y1, 'class': class_name, 'conf': float(box.conf[0])}
-#                 )
+    # ====================================================================
+    # --- STEP 1: YOLO DETECTION ---
+    # ====================================================================
+    start_time_yolo = time.time()
+    results = model.predict(source=original_img, conf=CONF_THRESHOLD, imgsz=640, verbose=False)
 
-#     merged_detections = merge_overlapping_boxes(relevant_detections, IOU_MERGE_THRESHOLD)
-#     print(f"    [LOG] YOLO found {len(merged_detections)} objects in {time.time() - start_time_yolo:.3f}s.")
+    relevant_detections = []
+    if results and results[0].boxes:
+        for box in results[0].boxes:
+            class_id = int(box.cls[0])
+            class_name = model.names[class_id]
+            if class_name in TARGET_CLASSES:
+                x1, y1, x2, y2 = box.xyxy[0].cpu().numpy().astype(int)
+                relevant_detections.append(
+                    {'coords': (x1, y1, x2, y2), 'y1': y1, 'class': class_name, 'conf': float(box.conf[0])}
+                )
 
-#     # ====================================================================
-#     # --- STEP 2: PREPARE DATA FOR COLUMN DETECTION (MASKING) ---
-#     # ====================================================================
-#     # Note: This uses the updated 'get_word_data_for_detection' which has its own optimizations
-#     raw_words_for_layout = get_word_data_for_detection(
-#         fitz_page, pdf_path, page_num,
-#         top_margin_percent=0.10, bottom_margin_percent=0.10
-#     )
+    merged_detections = merge_overlapping_boxes(relevant_detections, IOU_MERGE_THRESHOLD)
+    print(f"    [LOG] YOLO found {len(merged_detections)} objects in {time.time() - start_time_yolo:.3f}s.")
 
-#     masked_word_data = merge_yolo_into_word_data(raw_words_for_layout, merged_detections, scale_factor=2.0)
+    # ====================================================================
+    # --- STEP 2: PREPARE DATA FOR COLUMN DETECTION (MASKING) ---
+    # ====================================================================
+    # Note: This uses the updated 'get_word_data_for_detection' which has its own optimizations
+    raw_words_for_layout = get_word_data_for_detection(
+        fitz_page, pdf_path, page_num,
+        top_margin_percent=0.10, bottom_margin_percent=0.10
+    )
 
-#     # ====================================================================
-#     # --- STEP 3: COLUMN DETECTION ---
-#     # ====================================================================
-#     page_width_pdf = fitz_page.rect.width
-#     page_height_pdf = fitz_page.rect.height
+    masked_word_data = merge_yolo_into_word_data(raw_words_for_layout, merged_detections, scale_factor=2.0)
 
-#     column_detection_params = {
-#         'cluster_bin_size': 2, 'cluster_smoothing': 2,
-#         'cluster_min_width': 10, 'cluster_threshold_percentile': 85,
-#     }
+    # ====================================================================
+    # --- STEP 3: COLUMN DETECTION ---
+    # ====================================================================
+    page_width_pdf = fitz_page.rect.width
+    page_height_pdf = fitz_page.rect.height
 
-#     separators = calculate_x_gutters(masked_word_data, column_detection_params, page_height_pdf)
+    column_detection_params = {
+        'cluster_bin_size': 2, 'cluster_smoothing': 2,
+        'cluster_min_width': 10, 'cluster_threshold_percentile': 85,
+    }
 
-#     page_separator_x = None
-#     if separators:
-#         central_min = page_width_pdf * 0.35
-#         central_max = page_width_pdf * 0.65
-#         central_separators = [s for s in separators if central_min <= s <= central_max]
+    separators = calculate_x_gutters(masked_word_data, column_detection_params, page_height_pdf)
 
-#         if central_separators:
-#             center_x = page_width_pdf / 2
-#             page_separator_x = min(central_separators, key=lambda x: abs(x - center_x))
-#             print(f"      ✅ Column Split Confirmed at X={page_separator_x:.1f}")
-#         else:
-#             print("      ⚠️ Gutter found off-center. Ignoring.")
-#     else:
-#         print("      -> Single Column Layout Confirmed.")
+    page_separator_x = None
+    if separators:
+        central_min = page_width_pdf * 0.35
+        central_max = page_width_pdf * 0.65
+        central_separators = [s for s in separators if central_min <= s <= central_max]
 
-#     # ====================================================================
-#     # --- STEP 4: COMPONENT EXTRACTION (Save Images) ---
-#     # ====================================================================
-#     start_time_components = time.time()
-#     component_metadata = []
-#     fig_count_page = 0
-#     eq_count_page = 0
+        if central_separators:
+            center_x = page_width_pdf / 2
+            page_separator_x = min(central_separators, key=lambda x: abs(x - center_x))
+            print(f"      ✅ Column Split Confirmed at X={page_separator_x:.1f}")
+        else:
+            print("      ⚠️ Gutter found off-center. Ignoring.")
+    else:
+        print("      -> Single Column Layout Confirmed.")
 
-#     for detection in merged_detections:
-#         x1, y1, x2, y2 = detection['coords']
-#         class_name = detection['class']
+    # ====================================================================
+    # --- STEP 4: COMPONENT EXTRACTION (Save Images) ---
+    # ====================================================================
+    start_time_components = time.time()
+    component_metadata = []
+    fig_count_page = 0
+    eq_count_page = 0
 
-#         if class_name == 'figure':
-#             GLOBAL_FIGURE_COUNT += 1
-#             counter = GLOBAL_FIGURE_COUNT
-#             component_word = f"FIGURE{counter}"
-#             fig_count_page += 1
-#         elif class_name == 'equation':
-#             GLOBAL_EQUATION_COUNT += 1
-#             counter = GLOBAL_EQUATION_COUNT
-#             component_word = f"EQUATION{counter}"
-#             eq_count_page += 1
-#         else:
-#             continue
+    for detection in merged_detections:
+        x1, y1, x2, y2 = detection['coords']
+        class_name = detection['class']
 
-#         component_crop = original_img[y1:y2, x1:x2]
-#         component_filename = f"{pdf_name}_page{page_num}_{class_name}{counter}.png"
-#         cv2.imwrite(os.path.join(FIGURE_EXTRACTION_DIR, component_filename), component_crop)
+        if class_name == 'figure':
+            GLOBAL_FIGURE_COUNT += 1
+            counter = GLOBAL_FIGURE_COUNT
+            component_word = f"FIGURE{counter}"
+            fig_count_page += 1
+        elif class_name == 'equation':
+            GLOBAL_EQUATION_COUNT += 1
+            counter = GLOBAL_EQUATION_COUNT
+            component_word = f"EQUATION{counter}"
+            eq_count_page += 1
+        else:
+            continue
 
-#         y_midpoint = (y1 + y2) // 2
-#         component_metadata.append({
-#             'type': class_name, 'word': component_word,
-#             'bbox': [int(x1), int(y1), int(x2), int(y2)],
-#             'y0': int(y_midpoint), 'x0': int(x1)
-#         })
-
-#     # ====================================================================
-#     # --- STEP 5: HYBRID OCR (Native Text + Cached Tesseract Fallback) ---
-#     # ====================================================================
-#     raw_ocr_output = []
-#     scale_factor = 2.0  # Pipeline standard scale
-
-#     try:
-#         # Try getting native text first
-#         # NOTE: extract_native_words_and_convert MUST ALSO BE UPDATED TO USE sanitize_text
-#         raw_ocr_output = extract_native_words_and_convert(fitz_page, scale_factor=scale_factor)
-#     except Exception as e:
-#         print(f"  ❌ Native text extraction failed: {e}")
-
-#     # If native text is missing, fall back to OCR
-#     if not raw_ocr_output:
-#         if _ocr_cache.has_ocr(pdf_path, page_num):
-#             print(f"  ⚡ Using cached Tesseract OCR for page {page_num}")
-#             cached_word_data = _ocr_cache.get_ocr(pdf_path, page_num)
-#             for word_tuple in cached_word_data:
-#                 word_text, x1, y1, x2, y2 = word_tuple
-
-#                 # Scale from PDF points to Pipeline Pixels (2.0)
-#                 x1_pix = int(x1 * scale_factor)
-#                 y1_pix = int(y1 * scale_factor)
-#                 x2_pix = int(x2 * scale_factor)
-#                 y2_pix = int(y2 * scale_factor)
-
-#                 raw_ocr_output.append({
-#                     'type': 'text', 'word': word_text, 'confidence': 95.0,
-#                     'bbox': [x1_pix, y1_pix, x2_pix, y2_pix],
-#                     'y0': y1_pix, 'x0': x1_pix
-#                 })
-#         else:
-#             # === START OF OPTIMIZED OCR BLOCK ===
-#             try:
-#                 # 1. Re-render Page at High Resolution (Zoom 4.0 = ~300 DPI)
-#                 ocr_zoom = 4.0
-#                 pix_ocr = fitz_page.get_pixmap(matrix=fitz.Matrix(ocr_zoom, ocr_zoom))
-
-#                 # Convert PyMuPDF Pixmap to OpenCV format
-#                 img_ocr_np = np.frombuffer(pix_ocr.samples, dtype=np.uint8).reshape(pix_ocr.height, pix_ocr.width,
-#                                                                                     pix_ocr.n)
-#                 if pix_ocr.n == 3:
-#                     img_ocr_np = cv2.cvtColor(img_ocr_np, cv2.COLOR_RGB2BGR)
-#                 elif pix_ocr.n == 4:
-#                     img_ocr_np = cv2.cvtColor(img_ocr_np, cv2.COLOR_RGBA2BGR)
-
-#                 # 2. Preprocess (Binarization)
-#                 processed_img = preprocess_image_for_ocr(img_ocr_np)
-
-#                 # 3. Run Tesseract with Optimized Configuration
-#                 custom_config = r'--oem 3 --psm 6'
-
-#                 hocr_data = pytesseract.image_to_data(
-#                     processed_img,
-#                     output_type=pytesseract.Output.DICT,
-#                     config=custom_config
-#                 )
-
-#                 for i in range(len(hocr_data['level'])):
-#                     text = hocr_data['text'][i] # Retrieve raw Tesseract text
-                    
-#                     # --- FIX: SANITIZE TEXT AND THEN STRIP ---
-#                     cleaned_text = sanitize_text(text).strip()
-                    
-#                     if cleaned_text and hocr_data['conf'][i] > -1:
-#                         # 4. Coordinate Mapping
-#                         scale_adjustment = scale_factor / ocr_zoom
-
-#                         x1 = int(hocr_data['left'][i] * scale_adjustment)
-#                         y1 = int(hocr_data['top'][i] * scale_adjustment)
-#                         w = int(hocr_data['width'][i] * scale_adjustment)
-#                         h = int(hocr_data['height'][i] * scale_adjustment)
-#                         x2 = x1 + w
-#                         y2 = y1 + h
-
-#                         raw_ocr_output.append({
-#                             'type': 'text',
-#                             'word': cleaned_text, # Use the sanitized word
-#                             'confidence': float(hocr_data['conf'][i]),
-#                             'bbox': [x1, y1, x2, y2],
-#                             'y0': y1,
-#                             'x0': x1
-#                         })
-#             except Exception as e:
-#                 print(f"  ❌ Tesseract OCR Error: {e}")
-#             # === END OF OPTIMIZED OCR BLOCK ===
-    
-#     # ====================================================================
-#     # --- STEP 6: OCR CLEANING AND MERGING ---
-#     # ====================================================================
-#     items_to_sort = []
-
-#     for ocr_word in raw_ocr_output:
-#         is_suppressed = False
-#         for component in component_metadata:
-#             # Do not include words that are inside figure/equation boxes
-#             ioa = calculate_ioa(ocr_word['bbox'], component['bbox'])
-#             if ioa > IOA_SUPPRESSION_THRESHOLD:
-#                 is_suppressed = True
-#                 break
-#         if not is_suppressed:
-#             items_to_sort.append(ocr_word)
-
-#     # Add figures/equations back into the flow as "words"
-#     items_to_sort.extend(component_metadata)
-
-#     # ====================================================================
-#     # --- STEP 7: LINE-BASED SORTING ---
-#     # ====================================================================
-#     items_to_sort.sort(key=lambda x: (x['y0'], x['x0']))
-#     lines = []
-
-#     for item in items_to_sort:
-#         placed = False
-#         for line in lines:
-#             y_ref = min(it['y0'] for it in line)
-#             if abs(y_ref - item['y0']) < LINE_TOLERANCE:
-#                 line.append(item)
-#                 placed = True
-#                 break
-#         if not placed and item['type'] in ['equation', 'figure']:
-#             for line in lines:
-#                 y_ref = min(it['y0'] for it in line)
-#                 if abs(y_ref - item['y0']) < 20:
-#                     line.append(item)
-#                     placed = True
-#                     break
-#         if not placed:
-#             lines.append([item])
-
-#     for line in lines:
-#         line.sort(key=lambda x: x['x0'])
-
-#     final_output = []
-#     for line in lines:
-#         for item in line:
-#             data_item = {"word": item["word"], "bbox": item["bbox"], "type": item["type"]}
-#             if 'tag' in item: data_item['tag'] = item['tag']
-#             final_output.append(data_item)
-
-#     return final_output, page_separator_x
-
-
-
-
-
-
-
-
-
-def preprocess_and_ocr_page(original_img: np.ndarray, model, pdf_path: str,
-                            page_num: int, fitz_page: fitz.Page,
-                            pdf_name: str) -> Tuple[List[Dict[str, Any]], Optional[int]]:
-    """
-    OPTIMIZED FLOW: 
-    1. Run YOLO to find Equations/Tables.
-    2. Mask raw text with YOLO boxes.
-    3. Run Column Detection on the MASKED data.
-    4. Proceed with OCR (Native or High-Res Tesseract Fallback) and Output.
-    """
-    global GLOBAL_FIGURE_COUNT, GLOBAL_EQUATION_COUNT
-
-    start_time_total = time.time()
-
-    if original_img is None:
-        print(f"  ❌ Invalid image for page {page_num}.")
-        return None, None
-
-    # ====================================================================
-    # --- STEP 1: YOLO DETECTION ---
-    # ====================================================================
-    start_time_yolo = time.time()
-    results = model.predict(source=original_img, conf=CONF_THRESHOLD, imgsz=640, verbose=False)
-    
-    relevant_detections = []
-    if results and results[0].boxes:
-        for box in results[0].boxes:
-            class_id = int(box.cls[0])
-            class_name = model.names[class_id]
-            if class_name in TARGET_CLASSES:
-                x1, y1, x2, y2 = box.xyxy[0].cpu().numpy().astype(int)
-                relevant_detections.append(
-                    {'coords': (x1, y1, x2, y2), 'y1': y1, 'class': class_name, 'conf': float(box.conf[0])}
-                )
-
-    merged_detections = merge_overlapping_boxes(relevant_detections, IOU_MERGE_THRESHOLD)
-    print(f"    [LOG] YOLO found {len(merged_detections)} objects in {time.time() - start_time_yolo:.3f}s.")
-
-    # ====================================================================
-    # --- STEP 2: PREPARE DATA FOR COLUMN DETECTION (MASKING) ---
-    # ====================================================================
-    # Note: This uses the updated 'get_word_data_for_detection' which has its own optimizations
-    raw_words_for_layout = get_word_data_for_detection(
-        fitz_page, pdf_path, page_num,
-        top_margin_percent=0.10, bottom_margin_percent=0.10
-    )
-
-    masked_word_data = merge_yolo_into_word_data(raw_words_for_layout, merged_detections, scale_factor=2.0)
-
-    # ====================================================================
-    # --- STEP 3: COLUMN DETECTION ---
-    # ====================================================================
-    page_width_pdf = fitz_page.rect.width
-    page_height_pdf = fitz_page.rect.height 
-    
-    column_detection_params = {
-        'cluster_bin_size': 2, 'cluster_smoothing': 2,
-        'cluster_min_width': 10, 'cluster_threshold_percentile': 85,
-    }
-
-    separators = calculate_x_gutters(masked_word_data, column_detection_params, page_height_pdf)
-    
-    page_separator_x = None
-    if separators:
-        central_min = page_width_pdf * 0.35
-        central_max = page_width_pdf * 0.65
-        central_separators = [s for s in separators if central_min <= s <= central_max]
-        
-        if central_separators:
-            center_x = page_width_pdf / 2
-            page_separator_x = min(central_separators, key=lambda x: abs(x - center_x))
-            print(f"      ✅ Column Split Confirmed at X={page_separator_x:.1f}")
-        else:
-            print("      ⚠️ Gutter found off-center. Ignoring.")
-    else:
-        print("      -> Single Column Layout Confirmed.")
-
-    # ====================================================================
-    # --- STEP 4: COMPONENT EXTRACTION (Save Images) ---
-    # ====================================================================
-    start_time_components = time.time()
-    component_metadata = []
-    fig_count_page = 0
-    eq_count_page = 0
-
-    for detection in merged_detections:
-        x1, y1, x2, y2 = detection['coords']
-        class_name = detection['class']
-
-        if class_name == 'figure':
-            GLOBAL_FIGURE_COUNT += 1
-            counter = GLOBAL_FIGURE_COUNT
-            component_word = f"FIGURE{counter}"
-            fig_count_page += 1
-        elif class_name == 'equation':
-            GLOBAL_EQUATION_COUNT += 1
-            counter = GLOBAL_EQUATION_COUNT
-            component_word = f"EQUATION{counter}"
-            eq_count_page += 1
-        else:
-            continue
-
-        component_crop = original_img[y1:y2, x1:x2]
-        component_filename = f"{pdf_name}_page{page_num}_{class_name}{counter}.png"
-        cv2.imwrite(os.path.join(FIGURE_EXTRACTION_DIR, component_filename), component_crop)
+        component_crop = original_img[y1:y2, x1:x2]
+        component_filename = f"{pdf_name}_page{page_num}_{class_name}{counter}.png"
+        cv2.imwrite(os.path.join(FIGURE_EXTRACTION_DIR, component_filename), component_crop)
 
         y_midpoint = (y1 + y2) // 2
         component_metadata.append({
@@ -1419,10 +1091,11 @@ def preprocess_and_ocr_page(original_img: np.ndarray, model, pdf_path: str,
     # --- STEP 5: HYBRID OCR (Native Text + Cached Tesseract Fallback) ---
     # ====================================================================
     raw_ocr_output = []
-    scale_factor = 2.0 # Pipeline standard scale
+    scale_factor = 2.0  # Pipeline standard scale
 
     try:
         # Try getting native text first
+        # NOTE: extract_native_words_and_convert MUST ALSO BE UPDATED TO USE sanitize_text
         raw_ocr_output = extract_native_words_and_convert(fitz_page, scale_factor=scale_factor)
     except Exception as e:
         print(f"  ❌ Native text extraction failed: {e}")
@@ -1434,13 +1107,13 @@ def preprocess_and_ocr_page(original_img: np.ndarray, model, pdf_path: str,
             cached_word_data = _ocr_cache.get_ocr(pdf_path, page_num)
             for word_tuple in cached_word_data:
                 word_text, x1, y1, x2, y2 = word_tuple
-                
+
                 # Scale from PDF points to Pipeline Pixels (2.0)
                 x1_pix = int(x1 * scale_factor)
                 y1_pix = int(y1 * scale_factor)
                 x2_pix = int(x2 * scale_factor)
                 y2_pix = int(y2 * scale_factor)
-                
+
                 raw_ocr_output.append({
                     'type': 'text', 'word': word_text, 'confidence': 95.0,
                     'bbox': [x1_pix, y1_pix, x2_pix, y2_pix],
@@ -1450,63 +1123,63 @@ def preprocess_and_ocr_page(original_img: np.ndarray, model, pdf_path: str,
             # === START OF OPTIMIZED OCR BLOCK ===
             try:
                 # 1. Re-render Page at High Resolution (Zoom 4.0 = ~300 DPI)
-                # We do this specifically for OCR accuracy, separate from the pipeline image
                 ocr_zoom = 4.0
                 pix_ocr = fitz_page.get_pixmap(matrix=fitz.Matrix(ocr_zoom, ocr_zoom))
-                
+
                 # Convert PyMuPDF Pixmap to OpenCV format
-                img_ocr_np = np.frombuffer(pix_ocr.samples, dtype=np.uint8).reshape(pix_ocr.height, pix_ocr.width, pix_ocr.n)
-                if pix_ocr.n == 3: img_ocr_np = cv2.cvtColor(img_ocr_np, cv2.COLOR_RGB2BGR)
-                elif pix_ocr.n == 4: img_ocr_np = cv2.cvtColor(img_ocr_np, cv2.COLOR_RGBA2BGR)
+                img_ocr_np = np.frombuffer(pix_ocr.samples, dtype=np.uint8).reshape(pix_ocr.height, pix_ocr.width,
+                                                                                    pix_ocr.n)
+                if pix_ocr.n == 3:
+                    img_ocr_np = cv2.cvtColor(img_ocr_np, cv2.COLOR_RGB2BGR)
+                elif pix_ocr.n == 4:
+                    img_ocr_np = cv2.cvtColor(img_ocr_np, cv2.COLOR_RGBA2BGR)
 
                 # 2. Preprocess (Binarization)
-                # Ensure 'preprocess_image_for_ocr' is defined at top of file!
                 processed_img = preprocess_image_for_ocr(img_ocr_np)
-                
+
                 # 3. Run Tesseract with Optimized Configuration
-                # --oem 3: Default LSTM engine
-                # --psm 6: Assume a single uniform block of text (Critical for lists/questions)
                 custom_config = r'--oem 3 --psm 6'
-                
+
                 hocr_data = pytesseract.image_to_data(
-                    processed_img, 
-                    output_type=pytesseract.Output.DICT, 
+                    processed_img,
+                    output_type=pytesseract.Output.DICT,
                     config=custom_config
                 )
-                
+
                 for i in range(len(hocr_data['level'])):
-                    text = hocr_data['text'][i].strip()
-                    if text and hocr_data['conf'][i] > -1:
-                        
+                    text = hocr_data['text'][i] # Retrieve raw Tesseract text
+                    
+                    # --- FIX: SANITIZE TEXT AND THEN STRIP ---
+                    cleaned_text = sanitize_text(text).strip()
+                    
+                    if cleaned_text and hocr_data['conf'][i] > -1:
                         # 4. Coordinate Mapping
-                        # We scanned at Zoom 4.0, but our pipeline expects Zoom 2.0.
-                        # Scale Factor = (Target 2.0) / (Source 4.0) = 0.5
-                        scale_adjustment = scale_factor / ocr_zoom 
-                        
+                        scale_adjustment = scale_factor / ocr_zoom
+
                         x1 = int(hocr_data['left'][i] * scale_adjustment)
                         y1 = int(hocr_data['top'][i] * scale_adjustment)
                         w = int(hocr_data['width'][i] * scale_adjustment)
                         h = int(hocr_data['height'][i] * scale_adjustment)
                         x2 = x1 + w
                         y2 = y1 + h
-                        
+
                         raw_ocr_output.append({
-                            'type': 'text', 
-                            'word': text, 
+                            'type': 'text',
+                            'word': cleaned_text, # Use the sanitized word
                             'confidence': float(hocr_data['conf'][i]),
-                            'bbox': [x1, y1, x2, y2], 
-                            'y0': y1, 
+                            'bbox': [x1, y1, x2, y2],
+                            'y0': y1,
                             'x0': x1
                         })
             except Exception as e:
                 print(f"  ❌ Tesseract OCR Error: {e}")
             # === END OF OPTIMIZED OCR BLOCK ===
-
+    
     # ====================================================================
     # --- STEP 6: OCR CLEANING AND MERGING ---
     # ====================================================================
     items_to_sort = []
-    
+
     for ocr_word in raw_ocr_output:
         is_suppressed = False
         for component in component_metadata:
@@ -1570,7 +1243,6 @@ def preprocess_and_ocr_page(original_img: np.ndarray, model, pdf_path: str,
 
 
 
-
 # def preprocess_and_ocr_page(original_img: np.ndarray, model, pdf_path: str,
 #                             page_num: int, fitz_page: fitz.Page,
 #                             pdf_name: str) -> Tuple[List[Dict[str, Any]], Optional[int]]:
@@ -1881,112 +1553,425 @@ def run_single_pdf_preprocessing(pdf_path: str, preprocessed_json_path: str) ->
     total_pages_processed = 0
     mat = fitz.Matrix(2.0, 2.0)
 
-    print("\n[STEP 1.2: ITERATING PAGES - IN-MEMORY PROCESSING]")
+    print("\n[STEP 1.2: ITERATING PAGES - IN-MEMORY PROCESSING]")
+
+    for page_num_0_based in range(doc.page_count):
+        page_num = page_num_0_based + 1
+        print(f"  -> Processing Page {page_num}/{doc.page_count}...")
+
+        fitz_page = doc.load_page(page_num_0_based)
+
+        try:
+            pix = fitz_page.get_pixmap(matrix=mat)
+            original_img = pixmap_to_numpy(pix)
+        except Exception as e:
+            print(f"  ❌ Error converting page {page_num} to image: {e}")
+            continue
+
+        final_output, page_separator_x = preprocess_and_ocr_page(
+            original_img,
+            model,
+            pdf_path,
+            page_num,
+            fitz_page,
+            pdf_name
+        )
+
+        if final_output is not None:
+            page_data = {
+                "page_number": page_num,
+                "data": final_output,
+                "column_separator_x": page_separator_x
+            }
+            all_pages_data.append(page_data)
+            total_pages_processed += 1
+        else:
+            print(f"  ❌ Skipped page {page_num} due to processing error.")
+
+    doc.close()
+
+    if all_pages_data:
+        try:
+            with open(preprocessed_json_path, 'w') as f:
+                json.dump(all_pages_data, f, indent=4)
+            print(f"\n  ✅ Combined structured OCR JSON saved to: {os.path.basename(preprocessed_json_path)}")
+        except Exception as e:
+            print(f"❌ ERROR saving combined JSON output: {e}")
+            return None
+    else:
+        print("❌ WARNING: No page data generated. Halting pipeline.")
+        return None
+
+    print("\n" + "=" * 80)
+    print(f"--- YOLO/OCR PREPROCESSING COMPLETE ({total_pages_processed} pages processed) ---")
+    print("=" * 80)
+
+    return preprocessed_json_path
+
+
+# ============================================================================
+# --- PHASE 2: LAYOUTLMV3 INFERENCE FUNCTIONS ---
+# ============================================================================
+
+class LayoutLMv3ForTokenClassification(nn.Module):
+    def __init__(self, num_labels: int = NUM_LABELS):
+        super().__init__()
+        self.num_labels = num_labels
+        config = LayoutLMv3Config.from_pretrained("microsoft/layoutlmv3-base", num_labels=num_labels)
+        self.layoutlmv3 = LayoutLMv3Model.from_pretrained("microsoft/layoutlmv3-base", config=config)
+        self.classifier = nn.Linear(config.hidden_size, num_labels)
+        self.crf = CRF(num_labels)
+        self.init_weights()
+
+    def init_weights(self):
+        nn.init.xavier_uniform_(self.classifier.weight)
+        if self.classifier.bias is not None: nn.init.zeros_(self.classifier.bias)
+
+    def forward(self, input_ids: torch.Tensor, bbox: torch.Tensor, attention_mask: torch.Tensor,
+                labels: Optional[torch.Tensor] = None):
+        outputs = self.layoutlmv3(input_ids=input_ids, bbox=bbox, attention_mask=attention_mask, return_dict=True)
+        sequence_output = outputs.last_hidden_state
+        emissions = self.classifier(sequence_output)
+        mask = attention_mask.bool()
+        if labels is not None:
+            loss = -self.crf(emissions, labels, mask=mask).mean()
+            return loss
+        else:
+            return self.crf.viterbi_decode(emissions, mask=mask)
+
+
+def _merge_integrity(all_token_data: List[Dict[str, Any]],
+                     column_separator_x: Optional[int]) -> List[List[Dict[str, Any]]]:
+    """Splits the token data objects into column chunks based on a separator."""
+    if column_separator_x is None:
+        print("    -> No column separator. Treating as one chunk.")
+        return [all_token_data]
+
+    left_column_tokens, right_column_tokens = [], []
+    for token_data in all_token_data:
+        bbox_raw = token_data['bbox_raw_pdf_space']
+        center_x = (bbox_raw[0] + bbox_raw[2]) / 2
+        if center_x < column_separator_x:
+            left_column_tokens.append(token_data)
+        else:
+            right_column_tokens.append(token_data)
+
+    chunks = [c for c in [left_column_tokens, right_column_tokens] if c]
+    print(f"    -> Data split into {len(chunks)} column chunk(s) using separator X={column_separator_x}.")
+    return chunks
+
+
+
+
+
+
+def run_inference_and_get_raw_words(pdf_path: str, model_path: str,
+                                    preprocessed_json_path: str,
+                                    column_detection_params: Optional[Dict] = None) -> List[Dict[str, Any]]:
+    print("\n" + "=" * 80)
+    print("--- 2. STARTING LAYOUTLMV3 INFERENCE PIPELINE (Raw Word Output) ---")
+    print("=" * 80)
+
+    tokenizer = LayoutLMv3Tokenizer.from_pretrained("microsoft/layoutlmv3-base")
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"  -> Using device: {device}")
+
+    try:
+        model = LayoutLMv3ForTokenClassification(num_labels=NUM_LABELS)
+        checkpoint = torch.load(model_path, map_location=device)
+        model_state = checkpoint.get('model_state_dict', checkpoint)
+        # Apply patch for layoutlmv3 compatibility with saved state_dict
+        fixed_state_dict = {key.replace('layoutlm.', 'layoutlmv3.'): value for key, value in model_state.items()}
+        model.load_state_dict(fixed_state_dict)
+        model.to(device)
+        model.eval()
+        print(f"✅ LayoutLMv3 Model loaded successfully from {os.path.basename(model_path)}.")
+    except Exception as e:
+        print(f"❌ FATAL ERROR during LayoutLMv3 model loading: {e}")
+        return []
+
+    try:
+        with open(preprocessed_json_path, 'r', encoding='utf-8') as f:
+            preprocessed_data = json.load(f)
+        print(f"✅ Loaded preprocessed data with {len(preprocessed_data)} pages.")
+    except Exception:
+        print("❌ Error loading preprocessed JSON.")
+        return []
+
+    try:
+        doc = fitz.open(pdf_path)
+    except Exception:
+        print("❌ Error loading PDF.")
+        return []
+
+    final_page_predictions = []
+    CHUNK_SIZE = 500
+
+    for page_data in preprocessed_data:
+        page_num_1_based = page_data['page_number']
+        page_num_0_based = page_num_1_based - 1
+        page_raw_predictions = []
+        print(f"\n  *** Processing Page {page_num_1_based} ({len(page_data['data'])} raw tokens) ***")
+
+        fitz_page = doc.load_page(page_num_0_based)
+        page_width, page_height = fitz_page.rect.width, fitz_page.rect.height
+        print(f"    -> Page dimensions: {page_width:.0f}x{page_height:.0f} (PDF points).")
+
+        all_token_data = []
+        scale_factor = 2.0
+
+        for item in page_data['data']:
+            raw_yolo_bbox = item['bbox']
+            bbox_pdf = [
+                int(raw_yolo_bbox[0] / scale_factor), int(raw_yolo_bbox[1] / scale_factor),
+                int(raw_yolo_bbox[2] / scale_factor), int(raw_yolo_bbox[3] / scale_factor)
+            ]
+            normalized_bbox = [
+                max(0, min(1000, int(1000 * bbox_pdf[0] / page_width))),
+                max(0, min(1000, int(1000 * bbox_pdf[1] / page_height))),
+                max(0, min(1000, int(1000 * bbox_pdf[2] / page_width))),
+                max(0, min(1000, int(1000 * bbox_pdf[3] / page_height)))
+            ]
+            all_token_data.append({
+                "word": item['word'],
+                "bbox_raw_pdf_space": bbox_pdf,
+                "bbox_normalized": normalized_bbox,
+                "item_original_data": item
+            })
+
+        if not all_token_data:
+            continue
+
+        column_separator_x = page_data.get('column_separator_x', None)
+        if column_separator_x is not None:
+            print(f"    -> Using SAVED column separator: X={column_separator_x}")
+        else:
+            print("    -> No column separator found. Assuming single chunk.")
+
+        token_chunks = _merge_integrity(all_token_data, column_separator_x)
+        total_chunks = len(token_chunks)
+
+        for chunk_idx, chunk_tokens in enumerate(token_chunks):
+            if not chunk_tokens: continue
+
+            # 1. Sanitize: Convert everything to strings and aggressively clean Unicode errors.
+            chunk_words = [
+                str(t['word']).encode('utf-8', errors='ignore').decode('utf-8')
+                for t in chunk_tokens
+            ]
+            chunk_normalized_bboxes = [t['bbox_normalized'] for t in chunk_tokens]
+
+            total_sub_chunks = (len(chunk_words) + CHUNK_SIZE - 1) // CHUNK_SIZE
+            for i in range(0, len(chunk_words), CHUNK_SIZE):
+                sub_chunk_idx = i // CHUNK_SIZE + 1
+                sub_words = chunk_words[i:i + CHUNK_SIZE]
+                sub_bboxes = chunk_normalized_bboxes[i:i + CHUNK_SIZE]
+                sub_tokens_data = chunk_tokens[i:i + CHUNK_SIZE]
+
+                print(f"      -> Chunk {chunk_idx + 1}/{total_chunks}, Sub-chunk {sub_chunk_idx}/{total_sub_chunks}: {len(sub_words)} words. Running Inference...")
+
+                # 2. Manual generation of word_ids
+                manual_word_ids = []
+                for current_word_idx, word in enumerate(sub_words):
+                    sub_tokens = tokenizer.tokenize(word)
+                    for _ in sub_tokens:
+                        manual_word_ids.append(current_word_idx)
+
+                encoded_input = tokenizer(
+                    sub_words,
+                    boxes=sub_bboxes,
+                    truncation=True,
+                    padding="max_length",
+                    max_length=512,
+                    is_split_into_words=True,
+                    return_tensors="pt"
+                )
+
+                # Check for empty sequence
+                if encoded_input['input_ids'].shape[0] == 0:
+                    print(f"        -> Warning: Sub-chunk {sub_chunk_idx} encoded to an empty sequence. Skipping.")
+                    continue
+
+                # 3. Finalize word_ids based on encoded output length
+                sequence_length = int(torch.sum(encoded_input['attention_mask']).item())
+                content_token_length = max(0, sequence_length - 2)
+
+                manual_word_ids = manual_word_ids[:content_token_length]
 
-    for page_num_0_based in range(doc.page_count):
-        page_num = page_num_0_based + 1
-        print(f"  -> Processing Page {page_num}/{doc.page_count}...")
+                final_word_ids = [None]  # CLS token (index 0)
+                final_word_ids.extend(manual_word_ids)
 
-        fitz_page = doc.load_page(page_num_0_based)
+                if sequence_length > 1:
+                    final_word_ids.append(None)  # SEP token
 
-        try:
-            pix = fitz_page.get_pixmap(matrix=mat)
-            original_img = pixmap_to_numpy(pix)
-        except Exception as e:
-            print(f"  ❌ Error converting page {page_num} to image: {e}")
-            continue
+                final_word_ids.extend([None] * (512 - len(final_word_ids)))
+                word_ids = final_word_ids[:512]  # Final array for mapping
 
-        final_output, page_separator_x = preprocess_and_ocr_page(
-            original_img,
-            model,
-            pdf_path,
-            page_num,
-            fitz_page,
-            pdf_name
-        )
+                # Inputs are already batched by the tokenizer as [1, 512] 
+                input_ids = encoded_input['input_ids'].to(device)
+                bbox = encoded_input['bbox'].to(device)
+                attention_mask = encoded_input['attention_mask'].to(device)
 
-        if final_output is not None:
-            page_data = {
-                "page_number": page_num,
-                "data": final_output,
-                "column_separator_x": page_separator_x
-            }
-            all_pages_data.append(page_data)
-            total_pages_processed += 1
-        else:
-            print(f"  ❌ Skipped page {page_num} due to processing error.")
+                with torch.no_grad():
+                    model_outputs = model(input_ids, bbox, attention_mask)
+
+                # --- Robust extraction: support several forward return types ---
+                # We'll try (in order):
+                # 1) model_outputs is (emissions_tensor, viterbi_list)  -> use emissions for logits, keep decoded
+                # 2) model_outputs has .logits attribute (HF ModelOutput)
+                # 3) model_outputs is tuple/list containing a logits tensor
+                # 4) model_outputs is a tensor (assume logits)
+                # 5) model_outputs is a list-of-lists of ints (viterbi decoded) -> use that directly (no logits)
+                logits_tensor = None
+                decoded_labels_list = None
+
+                # case 1: tuple/list with (emissions, viterbi)
+                if isinstance(model_outputs, (tuple, list)) and len(model_outputs) == 2:
+                    a, b = model_outputs
+                    # a might be tensor (emissions), b might be viterbi list
+                    if isinstance(a, torch.Tensor):
+                        logits_tensor = a
+                    if isinstance(b, list):
+                        decoded_labels_list = b
+
+                # case 2: HF ModelOutput with .logits
+                if logits_tensor is None and hasattr(model_outputs, 'logits') and isinstance(model_outputs.logits, torch.Tensor):
+                    logits_tensor = model_outputs.logits
+
+                # case 3: tuple/list - search for a 3D tensor (B, L, C)
+                if logits_tensor is None and isinstance(model_outputs, (tuple, list)):
+                    found_tensor = None
+                    for item in model_outputs:
+                        if isinstance(item, torch.Tensor):
+                            # prefer 3D (batch, seq, labels)
+                            if item.dim() == 3:
+                                logits_tensor = item
+                                break
+                            if found_tensor is None:
+                                found_tensor = item
+                    if logits_tensor is None and found_tensor is not None:
+                        # found_tensor may be (batch, seq, hidden) or (seq, hidden); we avoid guessing.
+                        # Keep found_tensor only if it matches num_labels dimension
+                        if found_tensor.dim() == 3 and found_tensor.shape[-1] == NUM_LABELS:
+                            logits_tensor = found_tensor
+                        elif found_tensor.dim() == 2 and found_tensor.shape[-1] == NUM_LABELS:
+                            logits_tensor = found_tensor.unsqueeze(0)
+
+                # case 4: model_outputs directly a tensor
+                if logits_tensor is None and isinstance(model_outputs, torch.Tensor):
+                    logits_tensor = model_outputs
+
+                # case 5: model_outputs is a decoded viterbi list (common for CRF-only forward)
+                if decoded_labels_list is None and isinstance(model_outputs, list) and model_outputs and isinstance(model_outputs[0], list):
+                    # assume model_outputs is already viterbi decoded: List[List[int]] with batch dim first
+                    decoded_labels_list = model_outputs
+
+                # If neither logits nor decoded exist, that's fatal
+                if logits_tensor is None and decoded_labels_list is None:
+                    # helpful debug info
+                    try:
+                        elem_shapes = [ (type(x), getattr(x, 'shape', None)) for x in model_outputs ] if isinstance(model_outputs, (list, tuple)) else [(type(model_outputs), getattr(model_outputs, 'shape', None))]
+                    except Exception:
+                        elem_shapes = str(type(model_outputs))
+                    raise RuntimeError(f"Model output of type {type(model_outputs)} did not contain a valid logits tensor or decoded viterbi. Contents: {elem_shapes}")
+
+                # If we have logits_tensor, normalize shape to [seq_len, num_labels]
+                if logits_tensor is not None:
+                    # If shape is [B, L, C] with B==1, squeeze batch
+                    if logits_tensor.dim() == 3 and logits_tensor.shape[0] == 1:
+                        preds_tensor = logits_tensor.squeeze(0)  # [L, C]
+                    else:
+                        preds_tensor = logits_tensor  # possibly [L, C] already
 
-    doc.close()
+                    # Safety: ensure we have at least seq_len x channels
+                    if preds_tensor.dim() != 2:
+                        # try to reshape or error
+                        raise RuntimeError(f"Unexpected logits tensor shape: {tuple(preds_tensor.shape)}")
+                    # We'll use preds_tensor[token_idx] to argmax
+                else:
+                    preds_tensor = None  # no logits available
 
-    if all_pages_data:
-        try:
-            with open(preprocessed_json_path, 'w') as f:
-                json.dump(all_pages_data, f, indent=4)
-            print(f"\n  ✅ Combined structured OCR JSON saved to: {os.path.basename(preprocessed_json_path)}")
-        except Exception as e:
-            print(f"❌ ERROR saving combined JSON output: {e}")
-            return None
-    else:
-        print("❌ WARNING: No page data generated. Halting pipeline.")
-        return None
+                # If decoded labels provided, make a token-level list-of-ints aligned to tokenizer tokens
+                decoded_token_labels = None
+                if decoded_labels_list is not None:
+                    # decoded_labels_list is batch-first; we used batch size 1
+                    # if multiple sequences returned, take first
+                    decoded_token_labels = decoded_labels_list[0] if isinstance(decoded_labels_list[0], list) else decoded_labels_list
 
-    print("\n" + "=" * 80)
-    print(f"--- YOLO/OCR PREPROCESSING COMPLETE ({total_pages_processed} pages processed) ---")
-    print("=" * 80)
+                # Now map token-level predictions -> word-level predictions using word_ids
+                word_idx_to_pred_id = {}
 
-    return preprocessed_json_path
+                if preds_tensor is not None:
+                    # We have logits. Use argmax of logits for each token id up to sequence_length
+                    for token_idx, word_idx in enumerate(word_ids):
+                        if token_idx >= sequence_length:
+                            break
+                        if word_idx is not None and word_idx < len(sub_words):
+                            if word_idx not in word_idx_to_pred_id:
+                                pred_id = torch.argmax(preds_tensor[token_idx]).item()
+                                word_idx_to_pred_id[word_idx] = pred_id
+                else:
+                    # No logits, but we have decoded_token_labels from CRF (one label per token)
+                    # We'll align decoded_token_labels to token positions.
+                    if decoded_token_labels is None:
+                        # should not happen due to earlier checks
+                        raise RuntimeError("No logits and no decoded labels available for mapping.")
+                    # decoded_token_labels length may be equal to content_token_length (no special tokens)
+                    # or equal to sequence_length; try to align intelligently:
+                    # Prefer using decoded_token_labels aligned to the tokenizer tokens (starting at token 1 for CLS)
+                    # If decoded length == content_token_length, then manual_word_ids maps sub-token -> word idx for content tokens only.
+                    # We'll iterate tokens and pick label accordingly.
+                    # Build token_idx -> decoded_label mapping:
+                    # We'll assume decoded_token_labels correspond to content tokens (no CLS/SEP). If decoded length == sequence_length, then shift by 0.
+                    decoded_len = len(decoded_token_labels)
+                    # Heuristic: if decoded_len == content_token_length -> alignment starts at token_idx 1 (skip CLS)
+                    if decoded_len == content_token_length:
+                        decoded_start = 1
+                    elif decoded_len == sequence_length:
+                        decoded_start = 0
+                    else:
+                        # fallback: prefer decoded_start=1 (most common)
+                        decoded_start = 1
+
+                    for tok_idx_in_decoded, label_id in enumerate(decoded_token_labels):
+                        tok_idx = decoded_start + tok_idx_in_decoded
+                        if tok_idx >= 512:
+                            break
+                        if tok_idx >= sequence_length:
+                            break
+                        # map this token to a word index if present
+                        word_idx = word_ids[tok_idx] if tok_idx < len(word_ids) else None
+                        if word_idx is not None and word_idx < len(sub_words):
+                            if word_idx not in word_idx_to_pred_id:
+                                # label_id may already be an int
+                                word_idx_to_pred_id[word_idx] = int(label_id)
+
+                # Finally convert mapped word preds -> page_raw_predictions entries
+                for current_word_idx in range(len(sub_words)):
+                    pred_id = word_idx_to_pred_id.get(current_word_idx, 0)  # default to 0
+                    predicted_label = ID_TO_LABEL[pred_id]
+                    original_token = sub_tokens_data[current_word_idx]
+                    page_raw_predictions.append({
+                        "word": original_token['word'],
+                        "bbox": original_token['bbox_raw_pdf_space'],
+                        "predicted_label": predicted_label,
+                        "page_number": page_num_1_based
+                    })
 
+        if page_raw_predictions:
+            final_page_predictions.append({
+                "page_number": page_num_1_based,
+                "data": page_raw_predictions
+            })
+            print(f"  *** Page {page_num_1_based} Finalized: {len(page_raw_predictions)} labeled words. ***")
 
-# ============================================================================
-# --- PHASE 2: LAYOUTLMV3 INFERENCE FUNCTIONS ---
-# ============================================================================
+    doc.close()
+    print("\n" + "=" * 80)
+    print("--- LAYOUTLMV3 INFERENCE COMPLETE ---")
+    print("=" * 80)
+    return final_page_predictions
 
-# class LayoutLMv3ForTokenClassification(nn.Module):
-#     def __init__(self, num_labels: int = NUM_LABELS):
-#         super().__init__()
-#         self.num_labels = num_labels
-#         config = LayoutLMv3Config.from_pretrained("microsoft/layoutlmv3-base", num_labels=num_labels)
-#         self.layoutlmv3 = LayoutLMv3Model.from_pretrained("microsoft/layoutlmv3-base", config=config)
-#         self.classifier = nn.Linear(config.hidden_size, num_labels)
-#         self.crf = CRF(num_labels)
-#         self.init_weights()
-
-#     def init_weights(self):
-#         nn.init.xavier_uniform_(self.classifier.weight)
-#         if self.classifier.bias is not None: nn.init.zeros_(self.classifier.bias)
-
-#     def forward(self, input_ids: torch.Tensor, bbox: torch.Tensor, attention_mask: torch.Tensor,
-#                 labels: Optional[torch.Tensor] = None):
-#         outputs = self.layoutlmv3(input_ids=input_ids, bbox=bbox, attention_mask=attention_mask, return_dict=True)
-#         sequence_output = outputs.last_hidden_state
-#         emissions = self.classifier(sequence_output)
-#         mask = attention_mask.bool()
-#         if labels is not None:
-#             loss = -self.crf(emissions, labels, mask=mask).mean()
-#             return loss
-#         else:
-#             return self.crf.viterbi_decode(emissions, mask=mask)
-
-
-# def _merge_integrity(all_token_data: List[Dict[str, Any]],
-#                      column_separator_x: Optional[int]) -> List[List[Dict[str, Any]]]:
-#     """Splits the token data objects into column chunks based on a separator."""
-#     if column_separator_x is None:
-#         print("    -> No column separator. Treating as one chunk.")
-#         return [all_token_data]
-
-#     left_column_tokens, right_column_tokens = [], []
-#     for token_data in all_token_data:
-#         bbox_raw = token_data['bbox_raw_pdf_space']
-#         center_x = (bbox_raw[0] + bbox_raw[2]) / 2
-#         if center_x < column_separator_x:
-#             left_column_tokens.append(token_data)
-#         else:
-#             right_column_tokens.append(token_data)
 
-#     chunks = [c for c in [left_column_tokens, right_column_tokens] if c]
-#     print(f"    -> Data split into {len(chunks)} column chunk(s) using separator X={column_separator_x}.")
-#     return chunks
 
 
 
@@ -2067,6 +2052,20 @@ def run_single_pdf_preprocessing(pdf_path: str, preprocessed_json_path: str) ->
 #                 "item_original_data": item
 #             })
 
+#         # ==============================================================================
+#         # --- DEBUGGING BLOCK: CHECK FIRST 50 TOKENS BEFORE INFERENCE ---
+#         # ==============================================================================
+#         print(f"\n[DEBUG] LayoutLMv3 Input (Page {page_num_1_based}): Checking first 50 tokens...")
+#         debug_count = 0
+#         for t in all_token_data:
+#              if debug_count >= 50: break
+#              w = t['word']
+#              unicode_points = [f"\\u{ord(c):04x}" for c in w]
+#              print(f"  Token {debug_count}: '{w}' -> Codes: {unicode_points}")
+#              debug_count += 1
+#         print("----------------------------------------------------------------------\n")
+#         # ==============================================================================
+
 #         if not all_token_data:
 #             continue
 
@@ -2144,19 +2143,12 @@ def run_single_pdf_preprocessing(pdf_path: str, preprocessed_json_path: str) ->
 #                     model_outputs = model(input_ids, bbox, attention_mask)
 
 #                 # --- Robust extraction: support several forward return types ---
-#                 # We'll try (in order):
-#                 # 1) model_outputs is (emissions_tensor, viterbi_list)  -> use emissions for logits, keep decoded
-#                 # 2) model_outputs has .logits attribute (HF ModelOutput)
-#                 # 3) model_outputs is tuple/list containing a logits tensor
-#                 # 4) model_outputs is a tensor (assume logits)
-#                 # 5) model_outputs is a list-of-lists of ints (viterbi decoded) -> use that directly (no logits)
 #                 logits_tensor = None
 #                 decoded_labels_list = None
 
 #                 # case 1: tuple/list with (emissions, viterbi)
 #                 if isinstance(model_outputs, (tuple, list)) and len(model_outputs) == 2:
 #                     a, b = model_outputs
-#                     # a might be tensor (emissions), b might be viterbi list
 #                     if isinstance(a, torch.Tensor):
 #                         logits_tensor = a
 #                     if isinstance(b, list):
@@ -2171,15 +2163,12 @@ def run_single_pdf_preprocessing(pdf_path: str, preprocessed_json_path: str) ->
 #                     found_tensor = None
 #                     for item in model_outputs:
 #                         if isinstance(item, torch.Tensor):
-#                             # prefer 3D (batch, seq, labels)
 #                             if item.dim() == 3:
 #                                 logits_tensor = item
 #                                 break
 #                             if found_tensor is None:
 #                                 found_tensor = item
 #                     if logits_tensor is None and found_tensor is not None:
-#                         # found_tensor may be (batch, seq, hidden) or (seq, hidden); we avoid guessing.
-#                         # Keep found_tensor only if it matches num_labels dimension
 #                         if found_tensor.dim() == 3 and found_tensor.shape[-1] == NUM_LABELS:
 #                             logits_tensor = found_tensor
 #                         elif found_tensor.dim() == 2 and found_tensor.shape[-1] == NUM_LABELS:
@@ -2191,12 +2180,10 @@ def run_single_pdf_preprocessing(pdf_path: str, preprocessed_json_path: str) ->
 
 #                 # case 5: model_outputs is a decoded viterbi list (common for CRF-only forward)
 #                 if decoded_labels_list is None and isinstance(model_outputs, list) and model_outputs and isinstance(model_outputs[0], list):
-#                     # assume model_outputs is already viterbi decoded: List[List[int]] with batch dim first
 #                     decoded_labels_list = model_outputs
 
 #                 # If neither logits nor decoded exist, that's fatal
 #                 if logits_tensor is None and decoded_labels_list is None:
-#                     # helpful debug info
 #                     try:
 #                         elem_shapes = [ (type(x), getattr(x, 'shape', None)) for x in model_outputs ] if isinstance(model_outputs, (list, tuple)) else [(type(model_outputs), getattr(model_outputs, 'shape', None))]
 #                     except Exception:
@@ -2205,32 +2192,25 @@ def run_single_pdf_preprocessing(pdf_path: str, preprocessed_json_path: str) ->
 
 #                 # If we have logits_tensor, normalize shape to [seq_len, num_labels]
 #                 if logits_tensor is not None:
-#                     # If shape is [B, L, C] with B==1, squeeze batch
 #                     if logits_tensor.dim() == 3 and logits_tensor.shape[0] == 1:
 #                         preds_tensor = logits_tensor.squeeze(0)  # [L, C]
 #                     else:
 #                         preds_tensor = logits_tensor  # possibly [L, C] already
 
-#                     # Safety: ensure we have at least seq_len x channels
 #                     if preds_tensor.dim() != 2:
-#                         # try to reshape or error
 #                         raise RuntimeError(f"Unexpected logits tensor shape: {tuple(preds_tensor.shape)}")
-#                     # We'll use preds_tensor[token_idx] to argmax
 #                 else:
 #                     preds_tensor = None  # no logits available
 
 #                 # If decoded labels provided, make a token-level list-of-ints aligned to tokenizer tokens
 #                 decoded_token_labels = None
 #                 if decoded_labels_list is not None:
-#                     # decoded_labels_list is batch-first; we used batch size 1
-#                     # if multiple sequences returned, take first
 #                     decoded_token_labels = decoded_labels_list[0] if isinstance(decoded_labels_list[0], list) else decoded_labels_list
 
 #                 # Now map token-level predictions -> word-level predictions using word_ids
 #                 word_idx_to_pred_id = {}
 
 #                 if preds_tensor is not None:
-#                     # We have logits. Use argmax of logits for each token id up to sequence_length
 #                     for token_idx, word_idx in enumerate(word_ids):
 #                         if token_idx >= sequence_length:
 #                             break
@@ -2239,26 +2219,14 @@ def run_single_pdf_preprocessing(pdf_path: str, preprocessed_json_path: str) ->
 #                                 pred_id = torch.argmax(preds_tensor[token_idx]).item()
 #                                 word_idx_to_pred_id[word_idx] = pred_id
 #                 else:
-#                     # No logits, but we have decoded_token_labels from CRF (one label per token)
-#                     # We'll align decoded_token_labels to token positions.
 #                     if decoded_token_labels is None:
-#                         # should not happen due to earlier checks
 #                         raise RuntimeError("No logits and no decoded labels available for mapping.")
-#                     # decoded_token_labels length may be equal to content_token_length (no special tokens)
-#                     # or equal to sequence_length; try to align intelligently:
-#                     # Prefer using decoded_token_labels aligned to the tokenizer tokens (starting at token 1 for CLS)
-#                     # If decoded length == content_token_length, then manual_word_ids maps sub-token -> word idx for content tokens only.
-#                     # We'll iterate tokens and pick label accordingly.
-#                     # Build token_idx -> decoded_label mapping:
-#                     # We'll assume decoded_token_labels correspond to content tokens (no CLS/SEP). If decoded length == sequence_length, then shift by 0.
 #                     decoded_len = len(decoded_token_labels)
-#                     # Heuristic: if decoded_len == content_token_length -> alignment starts at token_idx 1 (skip CLS)
 #                     if decoded_len == content_token_length:
 #                         decoded_start = 1
 #                     elif decoded_len == sequence_length:
 #                         decoded_start = 0
 #                     else:
-#                         # fallback: prefer decoded_start=1 (most common)
 #                         decoded_start = 1
 
 #                     for tok_idx_in_decoded, label_id in enumerate(decoded_token_labels):
@@ -2267,11 +2235,9 @@ def run_single_pdf_preprocessing(pdf_path: str, preprocessed_json_path: str) ->
 #                             break
 #                         if tok_idx >= sequence_length:
 #                             break
-#                         # map this token to a word index if present
 #                         word_idx = word_ids[tok_idx] if tok_idx < len(word_ids) else None
 #                         if word_idx is not None and word_idx < len(sub_words):
 #                             if word_idx not in word_idx_to_pred_id:
-#                                 # label_id may already be an int
 #                                 word_idx_to_pred_id[word_idx] = int(label_id)
 
 #                 # Finally convert mapped word preds -> page_raw_predictions entries
@@ -2300,196 +2266,6 @@ def run_single_pdf_preprocessing(pdf_path: str, preprocessed_json_path: str) ->
 #     return final_page_predictions
 
 
-class LayoutLMv3ForTokenClassification(nn.Module):
-    def __init__(self, num_labels: int = NUM_LABELS):
-        super().__init__()
-        self.num_labels = num_labels
-        config = LayoutLMv3Config.from_pretrained("microsoft/layoutlmv3-base", num_labels=num_labels)
-        self.layoutlmv3 = LayoutLMv3Model.from_pretrained("microsoft/layoutlmv3-base", config=config)
-        self.classifier = nn.Linear(config.hidden_size, num_labels)
-        self.crf = CRF(num_labels)
-        self.init_weights()
-
-    def init_weights(self):
-        nn.init.xavier_uniform_(self.classifier.weight)
-        if self.classifier.bias is not None: nn.init.zeros_(self.classifier.bias)
-
-    def forward(self, input_ids: torch.Tensor, bbox: torch.Tensor, attention_mask: torch.Tensor, labels: Optional[torch.Tensor] = None):
-        outputs = self.layoutlmv3(input_ids=input_ids, bbox=bbox, attention_mask=attention_mask, return_dict=True)
-        sequence_output = outputs.last_hidden_state
-        emissions = self.classifier(sequence_output)
-        mask = attention_mask.bool()
-        if labels is not None:
-            loss = -self.crf(emissions, labels, mask=mask).mean()
-            return loss
-        else:
-            return self.crf.viterbi_decode(emissions, mask=mask)
-
-def _merge_integrity(all_token_data: List[Dict[str, Any]],
-                      column_separator_x: Optional[int]) -> List[List[Dict[str, Any]]]:
-    """Splits the token data objects into column chunks based on a separator."""
-    if column_separator_x is None:
-        print("    -> No column separator. Treating as one chunk.")
-        return [all_token_data]
-
-    left_column_tokens, right_column_tokens = [], []
-    for token_data in all_token_data:
-        bbox_raw = token_data['bbox_raw_pdf_space']
-        center_x = (bbox_raw[0] + bbox_raw[2]) / 2
-        if center_x < column_separator_x:
-            left_column_tokens.append(token_data)
-        else:
-            right_column_tokens.append(token_data)
-
-    chunks = [c for c in [left_column_tokens, right_column_tokens] if c]
-    print(f"    -> Data split into {len(chunks)} column chunk(s) using separator X={column_separator_x}.")
-    return chunks
-
-def run_inference_and_get_raw_words(pdf_path: str, model_path: str,
-                                    preprocessed_json_path: str,
-                                    column_detection_params: Optional[Dict] = None) -> List[Dict[str, Any]]:
-    print("\n" + "=" * 80)
-    print("--- 2. STARTING LAYOUTLMV3 INFERENCE PIPELINE (Raw Word Output) ---")
-    print("=" * 80)
-
-    tokenizer = LayoutLMv3TokenizerFast.from_pretrained("microsoft/layoutlmv3-base")
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    print(f"  -> Using device: {device}")
-
-    try:
-        model = LayoutLMv3ForTokenClassification(num_labels=NUM_LABELS)
-        checkpoint = torch.load(model_path, map_location=device)
-        model_state = checkpoint.get('model_state_dict', checkpoint)
-        fixed_state_dict = {key.replace('layoutlm.', 'layoutlmv3.'): value for key, value in model_state.items()}
-        model.load_state_dict(fixed_state_dict)
-        model.to(device)
-        model.eval()
-        print(f"✅ LayoutLMv3 Model loaded successfully from {os.path.basename(model_path)}.")
-    except Exception as e:
-        print(f"❌ FATAL ERROR during LayoutLMv3 model loading: {e}")
-        return []
-
-    try:
-        with open(preprocessed_json_path, 'r', encoding='utf-8') as f:
-            preprocessed_data = json.load(f)
-        print(f"✅ Loaded preprocessed data with {len(preprocessed_data)} pages.")
-    except Exception:
-        print("❌ Error loading preprocessed JSON.")
-        return []
-
-    try:
-        doc = fitz.open(pdf_path)
-    except Exception:
-        print("❌ Error loading PDF.")
-        return []
-
-    final_page_predictions = []
-    CHUNK_SIZE = 500
-
-    for page_data in preprocessed_data:
-        page_num_1_based = page_data['page_number']
-        page_num_0_based = page_num_1_based - 1
-        page_raw_predictions = []
-        print(f"\n  *** Processing Page {page_num_1_based} ({len(page_data['data'])} raw tokens) ***")
-
-        fitz_page = doc.load_page(page_num_0_based)
-        page_width, page_height = fitz_page.rect.width, fitz_page.rect.height
-        print(f"    -> Page dimensions: {page_width:.0f}x{page_height:.0f} (PDF points).")
-
-        all_token_data = []
-        scale_factor = 2.0
-
-        for item in page_data['data']:
-            raw_yolo_bbox = item['bbox']
-            bbox_pdf = [
-                int(raw_yolo_bbox[0] / scale_factor), int(raw_yolo_bbox[1] / scale_factor),
-                int(raw_yolo_bbox[2] / scale_factor), int(raw_yolo_bbox[3] / scale_factor)
-            ]
-            normalized_bbox = [
-                max(0, min(1000, int(1000 * bbox_pdf[0] / page_width))),
-                max(0, min(1000, int(1000 * bbox_pdf[1] / page_height))),
-                max(0, min(1000, int(1000 * bbox_pdf[2] / page_width))),
-                max(0, min(1000, int(1000 * bbox_pdf[3] / page_height)))
-            ]
-            all_token_data.append({
-                "word": item['word'],
-                "bbox_raw_pdf_space": bbox_pdf,
-                "bbox_normalized": normalized_bbox,
-                "item_original_data": item
-            })
-
-        if not all_token_data: continue
-
-        column_separator_x = page_data.get('column_separator_x', None)
-        if column_separator_x is not None:
-            print(f"    -> Using SAVED column separator: X={column_separator_x}")
-        else:
-            print("    -> No column separator found. Assuming single chunk.")
-
-        token_chunks = _merge_integrity(all_token_data, column_separator_x)
-        total_chunks = len(token_chunks)
-
-        for chunk_idx, chunk_tokens in enumerate(token_chunks):
-            if not chunk_tokens: continue
-
-            chunk_words = [t['word'] for t in chunk_tokens]
-            chunk_normalized_bboxes = [t['bbox_normalized'] for t in chunk_tokens]
-
-            total_sub_chunks = (len(chunk_words) + CHUNK_SIZE - 1) // CHUNK_SIZE
-            for i in range(0, len(chunk_words), CHUNK_SIZE):
-                sub_chunk_idx = i // CHUNK_SIZE + 1
-                sub_words = chunk_words[i:i + CHUNK_SIZE]
-                sub_bboxes = chunk_normalized_bboxes[i:i + CHUNK_SIZE]
-                sub_tokens_data = chunk_tokens[i:i + CHUNK_SIZE]
-
-                print(f"      -> Chunk {chunk_idx + 1}/{total_chunks}, Sub-chunk {sub_chunk_idx}/{total_sub_chunks}: {len(sub_words)} words. Running Inference...")
-
-                encoded_input = tokenizer(
-                    [sub_words], boxes=[sub_bboxes],is_split_into_words=True, truncation=True, padding="max_length",
-                    max_length=512, return_tensors="pt"
-                )
-                input_ids = encoded_input['input_ids'].to(device)
-                bbox = encoded_input['bbox'].to(device)
-                attention_mask = encoded_input['attention_mask'].to(device)
-
-                with torch.no_grad():
-                    predictions_int_list = model(input_ids, bbox, attention_mask)
-
-                if not predictions_int_list: continue
-                predictions_int = predictions_int_list[0]
-                word_ids = encoded_input.word_ids(batch_index=0)
-                word_idx_to_pred_id = {}
-
-                for token_idx, word_idx in enumerate(word_ids):
-                    if word_idx is not None and word_idx < len(sub_words):
-                        if word_idx not in word_idx_to_pred_id:
-                            word_idx_to_pred_id[word_idx] = predictions_int[token_idx]
-
-                for current_word_idx in range(len(sub_words)):
-                    pred_id_or_tensor = word_idx_to_pred_id.get(current_word_idx, 0)
-                    pred_id = pred_id_or_tensor.item() if torch.is_tensor(pred_id_or_tensor) else pred_id_or_tensor
-                    predicted_label = ID_TO_LABEL[pred_id]
-                    original_token = sub_tokens_data[current_word_idx]
-                    page_raw_predictions.append({
-                        "word": original_token['word'],
-                        "bbox": original_token['bbox_raw_pdf_space'],
-                        "predicted_label": predicted_label,
-                        "page_number": page_num_1_based
-                    })
-
-        if page_raw_predictions:
-            final_page_predictions.append({
-                "page_number": page_num_1_based,
-                "data": page_raw_predictions
-            })
-            print(f"  *** Page {page_num_1_based} Finalized: {len(page_raw_predictions)} labeled words. ***")
-
-    doc.close()
-    print("\n" + "=" * 80)
-    print("--- LAYOUTLMV3 INFERENCE COMPLETE ---")
-    print("=" * 80)
-    return final_page_predictions
-