diff --git "a/gemini" "b/gemini" --- "a/gemini" +++ "b/gemini" @@ -1,2864 +1,6 @@ - -import fitz # PyMuPDF -import numpy as np -import cv2 -import torch -import torch -import torch.serialization - -_original_torch_load = torch.load - -def patched_torch_load(*args, **kwargs): - # FORCE classic behavior - kwargs["weights_only"] = False - return _original_torch_load(*args, **kwargs) - -torch.load = patched_torch_load - - - - - -#================================================================================== -#RAPID OCR -#================================================================================== - -from rapidocr import RapidOCR, OCRVersion - -# Initialize RapidOCR (v5 is generally the most accurate current version) -# We use return_word_box=True to get word-level precision similar to Tesseract's image_to_data -ocr_engine = RapidOCR(params={ - "Det.ocr_version": OCRVersion.PPOCRV5, - "Rec.ocr_version": OCRVersion.PPOCRV5, - "Cls.ocr_version": OCRVersion.PPOCRV4, -}) - - - -#================================================================================== -#RAPID OCR -#================================================================================== - - - -import json -import argparse -import os -import re - -import torch.nn as nn -from TorchCRF import CRF -# 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 -import pytesseract -from PIL import Image -from scipy.signal import find_peaks -from scipy.ndimage import gaussian_filter1d -import sys -import io -import base64 -import tempfile -import time -import shutil -from sklearn.feature_extraction.text import CountVectorizer -from sklearn.metrics.pairwise import cosine_similarity -import logging -from transformers import TrOCRProcessor -from optimum.onnxruntime import ORTModelForVision2Seq - - - -# ============================================================================ -# --- TR-OCR/ORT MODEL INITIALIZATION --- -# ============================================================================ - -logging.basicConfig(level=logging.WARNING) - -processor = None -ort_model = None - -try: - MODEL_NAME = 'breezedeus/pix2text-mfr-1.5' - processor = TrOCRProcessor.from_pretrained(MODEL_NAME) - - # Initialize the model for ONNX Runtime - # NOTE: Set use_cache=False to avoid caching warnings/issues if reloading - ort_model = ORTModelForVision2Seq.from_pretrained(MODEL_NAME, use_cache=False) - - print("✅ ORTModelForVision2Seq and TrOCRProcessor initialized successfully for equation conversion.") -except Exception as e: - print(f"❌ Error initializing TrOCR/ORT model. Equations will not be converted: {e}") - processor = None - ort_model = None - - - -#===================================================================================================================== -#===================================================================================================================== - - - -# ============================================================================ -# --- CUSTOM MODEL DEFINITIONS (ADD THIS BLOCK) --- -# ============================================================================ -import torch -import torch.nn as nn -import torch.nn.functional as F -from collections import Counter -import pickle - -# --- CONSTANTS FOR CUSTOM MODEL --- -MODEL_FILE = "model_enhanced.pt" # Ensure this file is in your directory -VOCAB_FILE = "vocabs_enhanced.pkl" # Ensure this file is in your directory -DEVICE = torch.device("cpu") # Use "cuda" if available -MAX_CHAR_LEN = 16 -EMBED_DIM = 128 -CHAR_EMBED_DIM = 50 -CHAR_CNN_OUT = 50 -BBOX_DIM = 128 -HIDDEN_SIZE = 768 -SPATIAL_FEATURE_DIM = 64 -POSITIONAL_DIM = 128 -INFERENCE_CHUNK_SIZE = 450 - -LABELS = [ - "O", "B-QUESTION", "I-QUESTION", "B-OPTION", "I-OPTION", - "B-ANSWER", "I-ANSWER", "B-IMAGE", "I-IMAGE", - "B-SECTION HEADING", "I-SECTION HEADING", "B-PASSAGE", "I-PASSAGE" -] -IDX2LABEL = {i: l for i, l in enumerate(LABELS)} - -# --- CRF DEPENDENCY --- -try: - from torch_crf import CRF -except ImportError: - try: - from TorchCRF import CRF - except ImportError: - # Minimal fallback if CRF library is missing (though you should install it) - class CRF(nn.Module): - def __init__(self, *args, **kwargs): super().__init__() - -# --- MODEL CLASSES --- -class Vocab: - def __init__(self, min_freq=1, unk_token="", pad_token=""): - self.min_freq = min_freq - self.unk_token = unk_token - self.pad_token = pad_token - self.freq = Counter() - self.itos = [] - self.stoi = {} - def __len__(self): return len(self.itos) - def __getitem__(self, token): return self.stoi.get(token, self.stoi.get(self.unk_token, 0)) - -class CharCNNEncoder(nn.Module): - def __init__(self, char_vocab_size, char_emb_dim, out_dim, kernel_sizes=(2, 3, 4, 5)): - super().__init__() - self.char_emb = nn.Embedding(char_vocab_size, char_emb_dim, padding_idx=0) - self.convs = nn.ModuleList([nn.Conv1d(char_emb_dim, out_dim, kernel_size=k) for k in kernel_sizes]) - self.out_dim = out_dim * len(kernel_sizes) - def forward(self, char_ids): - B, L, C = char_ids.size() - emb = self.char_emb(char_ids.view(B * L, C)).transpose(1, 2) - outs = [torch.max(torch.relu(conv(emb)), dim=2)[0] for conv in self.convs] - return torch.cat(outs, dim=1).view(B, L, -1) - -class SpatialAttention(nn.Module): - def __init__(self, hidden_dim): - super().__init__() - self.query = nn.Linear(hidden_dim, hidden_dim) - self.key = nn.Linear(hidden_dim, hidden_dim) - self.value = nn.Linear(hidden_dim, hidden_dim) - self.scale = hidden_dim ** 0.5 - def forward(self, x, mask): - Q, K, V = self.query(x), self.key(x), self.value(x) - scores = torch.matmul(Q, K.transpose(-2, -1)) / self.scale - mask_expanded = mask.unsqueeze(1).expand_as(scores) - scores = scores.masked_fill(~mask_expanded, float('-inf')) - attn_weights = F.softmax(scores, dim=-1).masked_fill(torch.isnan(scores), 0.0) - return torch.matmul(attn_weights, V) - -class MCQTagger(nn.Module): - def __init__(self, vocab_size, char_vocab_size, n_labels): - super().__init__() - self.word_emb = nn.Embedding(vocab_size, EMBED_DIM, padding_idx=0) - self.char_enc = CharCNNEncoder(char_vocab_size, CHAR_EMBED_DIM, CHAR_CNN_OUT) - self.bbox_proj = nn.Sequential(nn.Linear(4, BBOX_DIM), nn.ReLU(), nn.Dropout(0.1), nn.Linear(BBOX_DIM, BBOX_DIM)) - self.spatial_proj = nn.Sequential(nn.Linear(11, SPATIAL_FEATURE_DIM), nn.ReLU(), nn.Dropout(0.1)) - self.context_proj = nn.Sequential(nn.Linear(8, 32), nn.ReLU(), nn.Dropout(0.1)) - self.positional_encoding = nn.Embedding(512, POSITIONAL_DIM) - in_dim = (EMBED_DIM + self.char_enc.out_dim + BBOX_DIM + SPATIAL_FEATURE_DIM + 32 + POSITIONAL_DIM) - self.bilstm = nn.LSTM(in_dim, HIDDEN_SIZE // 2, num_layers=3, batch_first=True, bidirectional=True, dropout=0.3) - self.spatial_attention = SpatialAttention(HIDDEN_SIZE) - self.ff = nn.Sequential(nn.Linear(HIDDEN_SIZE * 2, HIDDEN_SIZE), nn.ReLU(), nn.Dropout(0.3), nn.Linear(HIDDEN_SIZE, n_labels)) - self.crf = CRF(n_labels) - self.dropout = nn.Dropout(p=0.5) - def forward(self, words, chars, bboxes, spatial_feats, context_feats, mask): - B, L = words.size() - wemb = self.word_emb(words) - cenc = self.char_enc(chars) - benc = self.bbox_proj(bboxes) - senc = self.spatial_proj(spatial_feats) - cxt_enc = self.context_proj(context_feats) - pos = torch.arange(L, device=words.device).unsqueeze(0).expand(B, -1) - pos_enc = self.positional_encoding(pos.clamp(max=511)) - enc_in = self.dropout(torch.cat([wemb, cenc, benc, senc, cxt_enc, pos_enc], dim=-1)) - lengths = mask.sum(dim=1).cpu() - packed_in = nn.utils.rnn.pack_padded_sequence(enc_in, lengths, batch_first=True, enforce_sorted=False) - packed_out, _ = self.bilstm(packed_in) - lstm_out, _ = nn.utils.rnn.pad_packed_sequence(packed_out, batch_first=True) - attn_out = self.spatial_attention(lstm_out, mask) - emissions = self.ff(torch.cat([lstm_out, attn_out], dim=-1)) - return self.crf.viterbi_decode(emissions, mask=mask) - -# --- INJECT DEPENDENCIES FOR PICKLE LOADING --- -import sys -from types import ModuleType -train_mod = ModuleType("train_model") -sys.modules["train_model"] = train_mod -train_mod.Vocab = Vocab -train_mod.MCQTagger = MCQTagger -train_mod.CharCNNEncoder = CharCNNEncoder -train_mod.SpatialAttention = SpatialAttention - - - - - - -# ============================================================================ -# --- CUSTOM FEATURE EXTRACTORS --- -# ============================================================================ -def extract_spatial_features(tokens, idx): - curr = tokens[idx] - f = [] - # Vertical distance to next - if idx < len(tokens)-1: f.append(min((tokens[idx+1]['y0'] - curr['y1'])/100.0, 1.0)) - else: f.append(0.0) - # Vertical distance from prev - if idx > 0: f.append(min((curr['y0'] - tokens[idx-1]['y1'])/100.0, 1.0)) - else: f.append(0.0) - # Geometry - f.extend([curr['x0']/1000.0, (curr['x1']-curr['x0'])/1000.0, (curr['y1']-curr['y0'])/1000.0]) - f.extend([(curr['x0']+curr['x1'])/2000.0, (curr['y0']+curr['y1'])/2000.0, curr['x0']/1000.0]) - # Aspect ratio - f.append(min(((curr['x1']-curr['x0'])/max((curr['y1']-curr['y0']),1.0))/10.0, 1.0)) - # Alignment check - if idx > 0: f.append(float(abs(curr['x0'] - tokens[idx-1]['x0']) < 5)) - else: f.append(0.0) - # Area - f.append(min(((curr['x1']-curr['x0'])*(curr['y1']-curr['y0']))/(1000.0**2), 1.0)) - return f - -def extract_context_features(tokens, idx, window=3): - f = [] - def check_p(i): - t = str(tokens[i]['word']).lower().strip() # Changed 'text' to 'word' to match pipeline - return [float(bool(re.match(r'^q?\.?\d+[.:]', t))), float(bool(re.match(r'^[a-dA-D][.)]', t))), float(t.isupper() and len(t)>2)] - - prev_res = [0.0, 0.0, 0.0] - for i in range(max(0, idx-window), idx): - res = check_p(i) - prev_res = [max(prev_res[j], res[j]) for j in range(3)] - f.extend(prev_res) - next_res = [0.0, 0.0, 0.0] - for i in range(idx+1, min(len(tokens), idx+window+1)): - res = check_p(i) - next_res = [max(next_res[j], res[j]) for j in range(3)] - f.extend(next_res) - dq, dopt = 1.0, 1.0 - for i in range(idx+1, min(len(tokens), idx+window+1)): - t = str(tokens[i]['word']).lower().strip() - if re.match(r'^q?\.?\d+[.:]', t): dq = min(dq, (i-idx)/window) - if re.match(r'^[a-dA-D][.)]', t): dopt = min(dopt, (i-idx)/window) - f.extend([dq, dopt]) - return f - -#====================================================================================================================================================== -#====================================================================================================================================================== - -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) - - - - - -def get_latex_from_base64(base64_string: str) -> str: - """ - Decodes a Base64 image string and uses the pre-initialized TrOCR/ORT model - to recognize the formula. It cleans the output by removing spaces and - crucially, replacing double backslashes with single backslashes for correct LaTeX. - """ - if ort_model is None or processor is None: - return "[MODEL_ERROR: Model not initialized]" - - try: - # 1. Decode Base64 to Image - image_data = base64.b64decode(base64_string) - # We must ensure the image is RGB format for the model input - image = Image.open(io.BytesIO(image_data)).convert('RGB') - - # 2. Preprocess the image - pixel_values = processor(images=image, return_tensors="pt").pixel_values - - # 3. Text Generation (OCR) - generated_ids = ort_model.generate(pixel_values) - raw_generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True) - - if not raw_generated_text: - return "[OCR_WARNING: No formula found]" - - latex_string = raw_generated_text[0] - - # --- 4. Post-processing and Cleanup --- - - # # A. Remove all spaces/line breaks - # cleaned_latex = re.sub(r'\s+', '', latex_string) - cleaned_latex = re.sub(r'[\r\n]+', '', latex_string) - - # B. CRITICAL FIX: Replace double backslashes (\\) with single backslashes (\). - # This corrects model output that already over-escaped the LaTeX commands. - # Python literal: '\\\\' is replaced with '\\'. - #cleaned_latex = cleaned_latex.replace('\\\\', '\\') - - return cleaned_latex - - - except Exception as e: - # Catch any unexpected errors - print(f" ❌ TR-OCR Recognition failed: {e}") - return f"[TR_OCR_ERROR: Recognition failed: {e}]" - - - - - - - -# ============================================================================ -# --- CONFIGURATION AND CONSTANTS --- -# ============================================================================ - - -# NOTE: Update these paths to match your environment before running! -WEIGHTS_PATH = 'best.pt' -DEFAULT_LAYOUTLMV3_MODEL_PATH = "98.pth" - -# DIRECTORY CONFIGURATION -OCR_JSON_OUTPUT_DIR = './ocr_json_output_final' -FIGURE_EXTRACTION_DIR = './figure_extraction' -TEMP_IMAGE_DIR = './temp_pdf_images' - -# Detection parameters -# CONF_THRESHOLD = 0.2 -TARGET_CLASSES = ['figure', 'equation'] -IOU_MERGE_THRESHOLD = 0.4 -IOA_SUPPRESSION_THRESHOLD = 0.7 -LINE_TOLERANCE = 15 - -# Similarity -SIMILARITY_THRESHOLD = 0.10 -RESOLUTION_MARGIN = 0.05 - -# Global counters for sequential numbering across the entire PDF -GLOBAL_FIGURE_COUNT = 0 -GLOBAL_EQUATION_COUNT = 0 - -# LayoutLMv3 Labels -ID_TO_LABEL = { - 0: "O", - 1: "B-QUESTION", 2: "I-QUESTION", - 3: "B-OPTION", 4: "I-OPTION", - 5: "B-ANSWER", 6: "I-ANSWER", - 7: "B-SECTION_HEADING", 8: "I-SECTION_HEADING", - 9: "B-PASSAGE", 10: "I-PASSAGE" -} -NUM_LABELS = len(ID_TO_LABEL) - - -# ============================================================================ -# --- PERFORMANCE OPTIMIZATION: OCR CACHE --- -# ============================================================================ - -class OCRCache: - """Caches OCR results per page to avoid redundant Tesseract runs.""" - - def __init__(self): - self.cache = {} - - def get_key(self, pdf_path: str, page_num: int) -> str: - return f"{pdf_path}:{page_num}" - - def has_ocr(self, pdf_path: str, page_num: int) -> bool: - return self.get_key(pdf_path, page_num) in self.cache - - def get_ocr(self, pdf_path: str, page_num: int) -> Optional[list]: - return self.cache.get(self.get_key(pdf_path, page_num)) - - def set_ocr(self, pdf_path: str, page_num: int, ocr_data: list): - self.cache[self.get_key(pdf_path, page_num)] = ocr_data - - def clear(self): - self.cache.clear() - - -# Global OCR cache instance -_ocr_cache = OCRCache() - - -# ============================================================================ -# --- PHASE 1: YOLO/OCR PREPROCESSING FUNCTIONS --- -# ============================================================================ - -def calculate_iou(box1, box2): - x1_a, y1_a, x2_a, y2_a = box1 - x1_b, y1_b, x2_b, y2_b = box2 - x_left = max(x1_a, x1_b) - y_top = max(y1_a, y1_b) - x_right = min(x2_a, x2_b) - y_bottom = min(y2_a, y2_b) - intersection_area = max(0, x_right - x_left) * max(0, y_bottom - y_top) - box_a_area = (x2_a - x1_a) * (y2_a - y1_a) - box_b_area = (x2_b - x1_b) * (y2_b - y1_b) - union_area = float(box_a_area + box_b_area - intersection_area) - return intersection_area / union_area if union_area > 0 else 0 - - -def calculate_ioa(box1, box2): - x1_a, y1_a, x2_a, y2_a = box1 - x1_b, y1_b, x2_b, y2_b = box2 - x_left = max(x1_a, x1_b) - y_top = max(y1_a, y1_b) - x_right = min(x2_a, x2_b) - y_bottom = min(y2_a, y2_b) - intersection_area = max(0, x_right - x_left) * max(0, y_bottom - y_top) - box_a_area = (x2_a - x1_a) * (y2_a - y1_a) - return intersection_area / box_a_area if box_a_area > 0 else 0 - - -def filter_nested_boxes(detections, ioa_threshold=0.80): - """ - Removes boxes that are inside larger boxes (Containment Check). - Prioritizes keeping the LARGEST box (the 'parent' container). - """ - if not detections: - return [] - - # 1. Calculate Area for all detections - for d in detections: - x1, y1, x2, y2 = d['coords'] - d['area'] = (x2 - x1) * (y2 - y1) - - # 2. Sort by Area Descending (Largest to Smallest) - # This ensures we process the 'container' first - detections.sort(key=lambda x: x['area'], reverse=True) - - keep_indices = [] - is_suppressed = [False] * len(detections) - - for i in range(len(detections)): - if is_suppressed[i]: continue - - keep_indices.append(i) - box_a = detections[i]['coords'] - - # Compare with all smaller boxes - for j in range(i + 1, len(detections)): - if is_suppressed[j]: continue - - box_b = detections[j]['coords'] - - # Calculate Intersection - x_left = max(box_a[0], box_b[0]) - y_top = max(box_a[1], box_b[1]) - x_right = min(box_a[2], box_b[2]) - y_bottom = min(box_a[3], box_b[3]) - - if x_right < x_left or y_bottom < y_top: - intersection = 0 - else: - intersection = (x_right - x_left) * (y_bottom - y_top) - - # Calculate IoA (Intersection over Area of the SMALLER box) - area_b = detections[j]['area'] - - if area_b > 0: - ioa_small = intersection / area_b - - # If the small box is > 90% inside the big box, suppress the small one. - if ioa_small > ioa_threshold: - is_suppressed[j] = True - # print(f" [Suppress] Removed nested object inside larger '{detections[i]['class']}'") - - return [detections[i] for i in keep_indices] - - -def merge_overlapping_boxes(detections, iou_threshold): - if not detections: return [] - detections.sort(key=lambda d: d['conf'], reverse=True) - merged_detections = [] - is_merged = [False] * len(detections) - for i in range(len(detections)): - if is_merged[i]: continue - current_box = detections[i]['coords'] - current_class = detections[i]['class'] - merged_x1, merged_y1, merged_x2, merged_y2 = current_box - for j in range(i + 1, len(detections)): - if is_merged[j] or detections[j]['class'] != current_class: continue - other_box = detections[j]['coords'] - iou = calculate_iou(current_box, other_box) - if iou > iou_threshold: - merged_x1 = min(merged_x1, other_box[0]) - merged_y1 = min(merged_y1, other_box[1]) - merged_x2 = max(merged_x2, other_box[2]) - merged_y2 = max(merged_y2, other_box[3]) - is_merged[j] = True - merged_detections.append({ - 'coords': (merged_x1, merged_y1, merged_x2, merged_y2), - 'y1': merged_y1, 'class': current_class, 'conf': detections[i]['conf'] - }) - return merged_detections - - -def merge_yolo_into_word_data(raw_word_data: list, yolo_detections: list, scale_factor: float) -> list: - """ - Filters out raw words that are inside YOLO boxes and replaces them with - a single solid 'placeholder' block for the column detector. - """ - if not yolo_detections: - return raw_word_data - - # 1. Convert YOLO boxes (Pixels) to PDF Coordinates (Points) - pdf_space_boxes = [] - for det in yolo_detections: - x1, y1, x2, y2 = det['coords'] - pdf_box = ( - x1 / scale_factor, - y1 / scale_factor, - x2 / scale_factor, - y2 / scale_factor - ) - pdf_space_boxes.append(pdf_box) - - # 2. Filter out raw words that are inside YOLO boxes - cleaned_word_data = [] - for word_tuple in raw_word_data: - wx1, wy1, wx2, wy2 = word_tuple[1], word_tuple[2], word_tuple[3], word_tuple[4] - w_center_x = (wx1 + wx2) / 2 - w_center_y = (wy1 + wy2) / 2 - - is_inside_yolo = False - for px1, py1, px2, py2 in pdf_space_boxes: - if px1 <= w_center_x <= px2 and py1 <= w_center_y <= py2: - is_inside_yolo = True - break - - if not is_inside_yolo: - cleaned_word_data.append(word_tuple) - - # 3. Add the YOLO boxes themselves as "Solid Words" - for i, (px1, py1, px2, py2) in enumerate(pdf_space_boxes): - dummy_entry = (f"BLOCK_{i}", px1, py1, px2, py2) - cleaned_word_data.append(dummy_entry) - - return cleaned_word_data - - -# ============================================================================ -# --- MISSING HELPER FUNCTION --- -# ============================================================================ - -def preprocess_image_for_ocr(img_np): - """ - Converts image to grayscale and applies Otsu's Binarization - to separate text from background clearly. - """ - # 1. Convert to Grayscale if needed - if len(img_np.shape) == 3: - gray = cv2.cvtColor(img_np, cv2.COLOR_BGR2GRAY) - else: - gray = img_np - - # 2. Apply Otsu's Thresholding (Automatic binary threshold) - # This makes text solid black and background solid white - _, thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU) - - return thresh - - -def calculate_vertical_gap_coverage(word_data: list, sep_x: int, page_height: float, gutter_width: int = 10) -> float: - """ - Calculates what percentage of the page's vertical text span is 'cleanly split' by the separator. - A valid column split should split > 65% of the page verticality. - """ - if not word_data: - return 0.0 - - # Determine the vertical span of the actual text content - y_coords = [w[2] for w in word_data] + [w[4] for w in word_data] # y1 and y2 - min_y, max_y = min(y_coords), max(y_coords) - total_text_height = max_y - min_y - - if total_text_height <= 0: - return 0.0 - - # Create a boolean array representing the Y-axis (1 pixel per unit) - gap_open_mask = np.ones(int(total_text_height) + 1, dtype=bool) - - zone_left = sep_x - (gutter_width / 2) - zone_right = sep_x + (gutter_width / 2) - offset_y = int(min_y) - - for _, x1, y1, x2, y2 in word_data: - # Check if this word horizontally interferes with the separator - if x2 > zone_left and x1 < zone_right: - y_start_idx = max(0, int(y1) - offset_y) - y_end_idx = min(len(gap_open_mask), int(y2) - offset_y) - if y_end_idx > y_start_idx: - gap_open_mask[y_start_idx:y_end_idx] = False - - open_pixels = np.sum(gap_open_mask) - coverage_ratio = open_pixels / len(gap_open_mask) - - return coverage_ratio - - -def calculate_x_gutters(word_data: list, params: Dict, page_height: float) -> List[int]: - """ - Calculates X-axis histogram and validates using BRIDGING DENSITY and Vertical Coverage. - """ - if not word_data: return [] - - x_points = [] - # Use only word_data elements 1 (x1) and 3 (x2) - for item in word_data: - x_points.extend([item[1], item[3]]) - - if not x_points: return [] - max_x = max(x_points) - - # 1. Determine total text height for ratio calculation - y_coords = [item[2] for item in word_data] + [item[4] for item in word_data] - min_y, max_y = min(y_coords), max(y_coords) - total_text_height = max_y - min_y - if total_text_height <= 0: return [] - - # Histogram Setup - bin_size = params.get('cluster_bin_size', 5) - smoothing = params.get('cluster_smoothing', 1) - min_width = params.get('cluster_min_width', 20) - threshold_percentile = params.get('cluster_threshold_percentile', 85) - - num_bins = int(np.ceil(max_x / bin_size)) - hist, bin_edges = np.histogram(x_points, bins=num_bins, range=(0, max_x)) - smoothed_hist = gaussian_filter1d(hist.astype(float), sigma=smoothing) - inverted_signal = np.max(smoothed_hist) - smoothed_hist - - peaks, properties = find_peaks( - inverted_signal, - height=np.max(inverted_signal) - np.percentile(smoothed_hist, threshold_percentile), - distance=min_width / bin_size - ) - - if not peaks.size: return [] - separator_x_coords = [int(bin_edges[p]) for p in peaks] - final_separators = [] - - for x_coord in separator_x_coords: - # --- CHECK 1: BRIDGING DENSITY (The "Cut Through" Check) --- - # Calculate the total vertical height of words that physically cross this line. - bridging_height = 0 - bridging_count = 0 - - for item in word_data: - wx1, wy1, wx2, wy2 = item[1], item[2], item[3], item[4] - - # Check if this word physically sits on top of the separator line - if wx1 < x_coord and wx2 > x_coord: - word_h = wy2 - wy1 - bridging_height += word_h - bridging_count += 1 - - # Calculate Ratio: How much of the page's text height is blocked by these crossing words? - bridging_ratio = bridging_height / total_text_height - - # THRESHOLD: If bridging blocks > 8% of page height, REJECT. - # This allows for page numbers or headers (usually < 5%) to cross, but NOT paragraphs. - if bridging_ratio > 0.08: - print( - f" ❌ Separator X={x_coord} REJECTED: Bridging Ratio {bridging_ratio:.1%} (>15%) cuts through text.") - continue - - # --- CHECK 2: VERTICAL GAP COVERAGE (The "Clean Split" Check) --- - # The gap must exist cleanly for > 65% of the text height. - coverage = calculate_vertical_gap_coverage(word_data, x_coord, page_height, gutter_width=min_width) - - if coverage >= 0.80: - final_separators.append(x_coord) - print(f" -> Separator X={x_coord} ACCEPTED (Coverage: {coverage:.1%}, Bridging: {bridging_ratio:.1%})") - else: - print(f" ❌ Separator X={x_coord} REJECTED (Coverage: {coverage:.1%}, Bridging: {bridging_ratio:.1%})") - - return sorted(final_separators) - -#====================================================================================================================================== - - -def get_word_data_for_detection(page: fitz.Page, pdf_path: str, page_num: int, - top_margin_percent=0.10, bottom_margin_percent=0.10) -> list: - """ - Retrieves word data using PyMuPDF native extraction or RapidOCR fallback. - """ - # 1. Attempt Native Extraction - word_data = page.get_text("words") - - if len(word_data) > 5: - word_data = [(w[4], w[0], w[1], w[2], w[3]) for w in word_data] - else: - # 2. Check Cache - if _ocr_cache.has_ocr(pdf_path, page_num): - cached_data = _ocr_cache.get_ocr(pdf_path, page_num) - if cached_data and len(cached_data) > 0: - return cached_data - - # 3. OCR Fallback (RapidOCR) - try: - zoom_level = 2.0 - pix = page.get_pixmap(matrix=fitz.Matrix(zoom_level, zoom_level)) - img_np = np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.height, pix.width, pix.n) - - if pix.n == 3: - img_np = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR) - elif pix.n == 4: - img_np = cv2.cvtColor(img_np, cv2.COLOR_RGBA2BGR) - - # CRITICAL FIX: Use return_word_box=True and access word_results - ocr_result = ocr_engine(img_np, return_word_box=True) - - full_word_data = [] - - # Check if we got valid results - if ocr_result and ocr_result.word_results: - scale_adjustment = 1.0 / zoom_level - - - # Flatten the per-line word results - flat_results = sum(ocr_result.word_results, ()) - - for text, score, bbox in flat_results: - text = str(text).strip() - - if text: - # Convert Polygon to BBox - xs = [p[0] for p in bbox] - ys = [p[1] for p in bbox] - - x1 = min(xs) * scale_adjustment - y1 = min(ys) * scale_adjustment - x2 = max(xs) * scale_adjustment - y2 = max(ys) * scale_adjustment - - full_word_data.append((text, x1, y1, x2, y2)) - - word_data = full_word_data - - if len(word_data) > 0: - _ocr_cache.set_ocr(pdf_path, page_num, word_data) - - except Exception as e: - print(f" ❌ RapidOCR Error in detection phase: {e}") - import traceback - traceback.print_exc() - return [] - - # 4. Apply Margin Filtering - page_height = page.rect.height - y_min = page_height * top_margin_percent - y_max = page_height * (1 - bottom_margin_percent) - - return [d for d in word_data if d[2] >= y_min and d[4] <= y_max] - -#========================================================================================================================================= -#============================================================================================================================================= - - - - - - - - -def pixmap_to_numpy(pix: fitz.Pixmap) -> np.ndarray: - img_data = pix.samples - img = np.frombuffer(img_data, dtype=np.uint8).reshape(pix.height, pix.width, pix.n) - if pix.n == 4: - img = cv2.cvtColor(img, cv2.COLOR_RGBA2BGR) - elif pix.n == 3: - img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) - return img - - - - - - - - -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...") - - debug_count = 0 - for item in raw_word_data: - if debug_count >= 50: break - - 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') - - # 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.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 - - - - - -#=================================================================================================== -#=================================================================================================== -#=================================================================================================== - - - -import pandas as pd -import pickle -import os -import time -import json -from sklearn.feature_extraction.text import TfidfVectorizer -import numpy as np -from collections import defaultdict - -# --- Model File Paths (Required for the Classifier to load) --- -VECTORIZER_FILE = 'tfidf_vectorizer_conditional.pkl' -SUBJECT_MODEL_FILE = 'subject_classifier_model_conditional.pkl' -CONDITIONAL_CONCEPT_MODELS_FILE = 'conditional_concept_models.pkl' - - -# --- Hierarchical Classifier Class (Dependency for the helper function) --- - -class HierarchicalClassifier: - """ - A two-stage classification system based on conditional training. - Loads the vectorizer, subject classifier, and conditional concept models. - """ - - def __init__(self): - self.vectorizer = None - self.subject_model = None - self.conditional_concept_models = {} - self.is_ready = False - - def load_models(self): - """Loads the vectorizer, subject model, and conditional concept models.""" - try: - start_time = time.time() - # 1. Load the TF-IDF Vectorizer - with open(VECTORIZER_FILE, 'rb') as f: - self.vectorizer = pickle.load(f) - - # 2. Load the Level 1 (Subject) Classifier - with open(SUBJECT_MODEL_FILE, 'rb') as f: - self.subject_model = pickle.load(f) - - # 3. Load the dictionary of conditional Level 2 (Concept) Models - with open(CONDITIONAL_CONCEPT_MODELS_FILE, 'rb') as f: - conditional_data = pickle.load(f) - - # Extract just the models for easy access - for subject, data in conditional_data.items(): - self.conditional_concept_models[subject] = data['model'] - - print(f"Loaded models successfully in {time.time() - start_time:.2f} seconds.") - self.is_ready = True - - except FileNotFoundError as e: - print(f"Error: Required model file not found: {e.filename}.") - self.is_ready = False - except Exception as e: - print(f"An error occurred while loading models: {e}") - self.is_ready = False - - return self.is_ready - - def predict_subject(self, text_chunk): - """Predicts the top Subject (Level 1).""" - if not self.is_ready: - return "Unknown", 0.0 - - # Vectorize the input - text_vector = self.vectorizer.transform([text_chunk]).astype(np.float64) - - if hasattr(self.subject_model, 'predict_proba'): - probabilities = self.subject_model.predict_proba(text_vector)[0] - classes = self.subject_model.classes_ - - top_index = np.argmax(probabilities) - return classes[top_index], probabilities[top_index] - else: - return self.subject_model.predict(text_vector)[0], 1.0 - - def predict_concept_hierarchical(self, text_chunk, predicted_subject): - """ - Predicts the top Concept (Level 2) using the specialized conditional model. - """ - if not self.is_ready: - return "Unknown", 0.0 - - concept_model = self.conditional_concept_models.get(predicted_subject) - - if concept_model is None or len(getattr(concept_model, 'classes_', [])) <= 1: - return "No_Conditional_Model_Found", 0.0 - - # Vectorize the input - text_vector = self.vectorizer.transform([text_chunk]).astype(np.float64) - - if hasattr(concept_model, 'predict_proba'): - probabilities = concept_model.predict_proba(text_vector)[0] - classes = concept_model.classes_ - - top_index = np.argmax(probabilities) - return classes[top_index], probabilities[top_index] - else: - return concept_model.predict(text_vector)[0], 1.0 - - -# -------------------------------------------------------------------------------------- -# --- The Requested Helper Function --- - -def post_process_json_with_inference(json_data, classifier): - """ - Takes JSON data, runs hierarchical inference on all question/option text, - and adds 'predicted_subject' and 'predicted_concept' tags to each entry. - - Args: - json_data (list): The list of dictionaries containing question entries. - classifier (HierarchicalClassifier): An initialized and loaded classifier object. - - Returns: - list: The modified list of dictionaries with classification tags added. - """ - if not classifier.is_ready: - print("Classifier not ready. Skipping inference.") - return json_data - - # This print statement can be removed for silent pipeline integration - print("\n--- Starting Subject/Concept Detection ---") - - for entry in json_data: - # Only process entries that have a 'question' field - if 'question' not in entry: - continue - - # 1. Combine Question Text and Option Text for robust feature extraction - full_text = entry.get('question', '') - - options = entry.get('options', {}) - for option_key, option_value in options.items(): - # Use the text component of the option if available - option_text = option_value if isinstance(option_value, str) else option_key - full_text += " " + option_text.replace('\n', ' ') - - # Clean up text (remove multiple spaces and surrounding whitespace) - full_text = ' '.join(full_text.split()).strip() - - # Handle empty text - if not full_text: - entry['predicted_subject'] = {'label': 'Empty_Text', 'confidence': 0.0} - entry['predicted_concept'] = {'label': 'Empty_Text', 'confidence': 0.0} - continue - - # 2. STAGE 1: Predict Subject - subj_label, subj_conf = classifier.predict_subject(full_text) - - # 3. STAGE 2: Predict Concept (Conditional on predicted subject) - conc_label, conc_conf = classifier.predict_concept_hierarchical(full_text, subj_label) - - # 4. Add results to the JSON entry - entry['predicted_subject'] = { - 'label': subj_label, - 'confidence': round(subj_conf, 4) - } - entry['predicted_concept'] = { - 'label': conc_label, - 'confidence': round(conc_conf, 4) - } - - # This print statement can be removed for silent pipeline integration - # print("--- JSON Post-Processing Complete ---") - - return json_data - - - - - -#=================================================================================================== -#=================================================================================================== -#=================================================================================================== - - - - - - - -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=0.2, imgsz=640, verbose=False) - - - - - relevant_detections = [] - - THRESHOLDS = { - 'figure': 0.75, - 'equation': 0.20 - } - - - - - if results and results[0].boxes: - for box in results[0].boxes: - class_id = int(box.cls[0]) - class_name = model.names[class_id] - conf = float(box.conf[0]) - - # Logic: Check if class is in our list AND meets its specific threshold - if class_name in THRESHOLDS: - if conf >= THRESHOLDS[class_name]: - 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': conf - }) - - 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) - - 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 (Standardizing to Zoom 4.0) - ocr_zoom = 4.0 - pix_ocr = fitz_page.get_pixmap(matrix=fitz.Matrix(ocr_zoom, ocr_zoom)) - - # Convert PyMuPDF Pixmap to OpenCV format (BGR) - 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. Run RapidOCR - # FIX 1: Capture the object (Removes the "cannot unpack" error) - ocr_out = ocr_engine(img_ocr_np) - - # FIX 2: Use 'is not None' (Removes the "ambiguous truth value" error) - if ocr_out is not None and ocr_out.boxes is not None: - # Calculate scaling from OCR image (4.0) to your pipeline standard (scale_factor=2.0) - scale_adjustment = scale_factor / ocr_zoom - - # FIX 3: Zip the attributes to restore your expected (box, text, score) format - for box, text, score in zip(ocr_out.boxes, ocr_out.txts, ocr_out.scores): - # Sanitize and clean text - cleaned_text = sanitize_text(str(text)).strip() - - if cleaned_text: - # 3. Coordinate Mapping (Convert 4-point polygon to x1, y1, x2, y2) - xs = [p[0] for p in box] - ys = [p[1] for p in box] - - x1 = int(min(xs) * scale_adjustment) - y1 = int(min(ys) * scale_adjustment) - x2 = int(max(xs) * scale_adjustment) - y2 = int(max(ys) * scale_adjustment) - - raw_ocr_output.append({ - 'type': 'text', - 'word': cleaned_text, - 'confidence': float(score) * 100, # Converting 0-1.0 to 0-100 scale - 'bbox': [x1, y1, x2, y2], - 'y0': y1, - 'x0': x1 - }) - except Exception as e: - print(f" ❌ RapidOCR Fallback Error: {e}") - # === END OF RAPIDOCR BLOCK ========================== - # === END OF RAPIDOCR BLOCK ==================================================================================================================================== -#=========================================================================================================================================================================== - # === 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 run_single_pdf_preprocessing(pdf_path: str, preprocessed_json_path: str) -> Optional[str]: - global GLOBAL_FIGURE_COUNT, GLOBAL_EQUATION_COUNT - - GLOBAL_FIGURE_COUNT = 0 - GLOBAL_EQUATION_COUNT = 0 - _ocr_cache.clear() - - print("\n" + "=" * 80) - print("--- 1. STARTING OPTIMIZED YOLO/OCR PREPROCESSING PIPELINE ---") - print("=" * 80) - - if not os.path.exists(pdf_path): - print(f"❌ FATAL ERROR: Input PDF not found at {pdf_path}.") - return None - - os.makedirs(os.path.dirname(preprocessed_json_path), exist_ok=True) - os.makedirs(FIGURE_EXTRACTION_DIR, exist_ok=True) - - model = YOLO(WEIGHTS_PATH) - pdf_name = os.path.splitext(os.path.basename(pdf_path))[0] - - try: - doc = fitz.open(pdf_path) - print(f"✅ Opened PDF: {pdf_name} ({doc.page_count} pages)") - except Exception as e: - print(f"❌ ERROR loading PDF file: {e}") - return None - - all_pages_data = [] - total_pages_processed = 0 - mat = fitz.Matrix(2.0, 2.0) - - 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] - - final_word_ids = [None] # CLS token (index 0) - final_word_ids.extend(manual_word_ids) - - if sequence_length > 1: - final_word_ids.append(None) # SEP token - - final_word_ids.extend([None] * (512 - len(final_word_ids))) - word_ids = final_word_ids[:512] # Final array for mapping - - # 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) - - 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 - - # 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 - 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. ***") - - doc.close() - print("\n" + "=" * 80) - print("--- LAYOUTLMV3 INFERENCE COMPLETE ---") - print("=" * 80) - return final_page_predictions - - - - - - - - -# ============================================================================ -# --- PHASE 2 REPLACEMENT: CUSTOM INFERENCE PIPELINE --- -# ============================================================================ -def run_custom_inference_and_get_raw_words(preprocessed_json_path: str) -> List[Dict[str, Any]]: - print("\n" + "=" * 80) - print("--- 2. STARTING CUSTOM MODEL INFERENCE PIPELINE ---") - print("=" * 80) - - # 1. Load Resources - if not os.path.exists(MODEL_FILE) or not os.path.exists(VOCAB_FILE): - print("❌ Error: Missing custom model or vocab files.") - return [] - - try: - print(" -> Loading Vocab and Model...") - with open(VOCAB_FILE, "rb") as f: - word_vocab, char_vocab = pickle.load(f) - - model = MCQTagger(len(word_vocab), len(char_vocab), len(LABELS)).to(DEVICE) - - # Load state dict safe - state_dict = torch.load(MODEL_FILE, map_location=DEVICE) - model.load_state_dict(state_dict if isinstance(state_dict, dict) else state_dict.state_dict()) - model.eval() - print("✅ Custom Model loaded successfully.") - except Exception as e: - print(f"❌ Error loading custom model: {e}") - return [] - - # 2. Load Preprocessed Data - try: - with open(preprocessed_json_path, 'r', encoding='utf-8') as f: - preprocessed_data = json.load(f) - print(f"✅ Loaded preprocessed data for {len(preprocessed_data)} pages.") - except Exception: - print("❌ Error loading preprocessed JSON.") - return [] - - final_page_predictions = [] - scale_factor = 2.0 # The pipeline scales PDF points to 2.0 for YOLO. We need to reverse this. - - for page_data in preprocessed_data: - page_num = page_data['page_number'] - raw_items = page_data['data'] - - if not raw_items: continue - - # --- A. ADAPTER: Convert Pipeline Data format to Custom Model format --- - # Pipeline Data: {'word': 'Text', 'bbox': [x1, y1, x2, y2]} (scaled by 2.0) - # Custom Model Needed: {'word': 'Text', 'x0': x, 'y0': y, 'x1': x, 'y1': y} (PDF points) - - tokens_for_inference = [] - for item in raw_items: - bbox = item['bbox'] - # Revert scale to get native PDF coordinates - x0 = bbox[0] / scale_factor - y0 = bbox[1] / scale_factor - x1 = bbox[2] / scale_factor - y1 = bbox[3] / scale_factor - - tokens_for_inference.append({ - 'word': str(item['word']), # Ensure string - 'x0': x0, 'y0': y0, 'x1': x1, 'y1': y1, - 'original_bbox': bbox # Keep for output - }) - - # --- B. FEATURE EXTRACTION --- - for i in range(len(tokens_for_inference)): - tokens_for_inference[i]['spatial_features'] = extract_spatial_features(tokens_for_inference, i) - tokens_for_inference[i]['context_features'] = extract_context_features(tokens_for_inference, i) - - # --- C. BATCH INFERENCE --- - page_raw_predictions = [] - - # Process in chunks - for i in range(0, len(tokens_for_inference), INFERENCE_CHUNK_SIZE): - chunk = tokens_for_inference[i : i + INFERENCE_CHUNK_SIZE] - - # Prepare Tensors - w_ids = torch.LongTensor([[word_vocab[t['word']] for t in chunk]]).to(DEVICE) - - c_ids_list = [] - for t in chunk: - chars = [char_vocab[c] for c in t['word'][:MAX_CHAR_LEN]] - chars += [0] * (MAX_CHAR_LEN - len(chars)) - c_ids_list.append(chars) - c_ids = torch.LongTensor([c_ids_list]).to(DEVICE) - - bboxes = torch.FloatTensor([[[t['x0']/1000.0, t['y0']/1000.0, t['x1']/1000.0, t['y1']/1000.0] for t in chunk]]).to(DEVICE) - s_feats = torch.FloatTensor([[t['spatial_features'] for t in chunk]]).to(DEVICE) - c_feats = torch.FloatTensor([[t['context_features'] for t in chunk]]).to(DEVICE) - mask = torch.ones(w_ids.size(), dtype=torch.bool).to(DEVICE) - - # Predict - with torch.no_grad(): - preds = model(w_ids, c_ids, bboxes, s_feats, c_feats, mask)[0] - - # --- D. FORMAT OUTPUT --- - for t, p in zip(chunk, preds): - label = IDX2LABEL[p] - # Create the exact dictionary structure expected by the rest of the pipeline - page_raw_predictions.append({ - "word": t['word'], - "bbox": t['original_bbox'], # Pass back the scaled bbox the pipeline uses - "predicted_label": label, - "page_number": page_num - }) - - if page_raw_predictions: - final_page_predictions.append({ - "page_number": page_num, - "data": page_raw_predictions - }) - print(f" -> Page {page_num} Inference Complete: {len(page_raw_predictions)} labeled words.") - - return final_page_predictions - - - -# ============================================================================ -# --- PHASE 3: BIO TO STRUCTURED JSON DECODER --- -# ============================================================================ - - - - - - - - -def convert_bio_to_structured_json_relaxed(input_path: str, output_path: str) -> Optional[List[Dict[str, Any]]]: - print("\n" + "=" * 80) - print("--- 3. STARTING BIO TO STRUCTURED JSON DECODING ---") - print(f"Source: {input_path}") - print("=" * 80) - - start_time = time.time() - - try: - with open(input_path, 'r', encoding='utf-8') as f: - predictions_by_page = json.load(f) - print(f"✅ Successfully loaded raw predictions ({len(predictions_by_page)} pages found)") - except Exception as e: - print(f"❌ Error loading raw prediction file: {e}") - return None - - predictions = [] - for page_item in predictions_by_page: - if isinstance(page_item, dict) and 'data' in page_item: - predictions.extend(page_item['data']) - - total_words = len(predictions) - print(f"📋 Total words to process: {total_words}") - - structured_data = [] - current_item = None - current_option_key = None - current_passage_buffer = [] - current_text_buffer = [] - first_question_started = False - last_entity_type = None - just_finished_i_option = False - is_in_new_passage = False - - def finalize_passage_to_item(item, passage_buffer): - if passage_buffer: - passage_text = re.sub(r'\s{2,}', ' ', ' '.join(passage_buffer)).strip() - print(f" ↳ [Buffer] Finalizing passage ({len(passage_buffer)} words) into current item") - if item.get('passage'): - item['passage'] += ' ' + passage_text - else: - item['passage'] = passage_text - passage_buffer.clear() - - # Iterate through every predicted word - for idx, item in enumerate(predictions): - word = item['word'] - label = item['predicted_label'] - entity_type = label[2:].strip() if label.startswith(('B-', 'I-')) else None - current_text_buffer.append(word) - - previous_entity_type = last_entity_type - is_passage_label = (entity_type == 'PASSAGE') - - # --- LOGGING: Track progress every 500 words or on B- labels --- - if label.startswith('B-'): - print(f"[Word {idx}/{total_words}] Found Label: {label} | Word: '{word}'") - - if not first_question_started: - if label != 'B-QUESTION' and not is_passage_label: - just_finished_i_option = False - is_in_new_passage = False - continue - if is_passage_label: - current_passage_buffer.append(word) - last_entity_type = 'PASSAGE' - just_finished_i_option = False - is_in_new_passage = False - continue - - if label == 'B-QUESTION': - print(f"🔍 Detection: New Question Started at word {idx}") - if not first_question_started: - header_text = ' '.join(current_text_buffer[:-1]).strip() - if header_text or current_passage_buffer: - print(f" -> Creating METADATA item for text found before first question") - metadata_item = {'type': 'METADATA', 'passage': ''} - finalize_passage_to_item(metadata_item, current_passage_buffer) - if header_text: metadata_item['text'] = header_text - structured_data.append(metadata_item) - first_question_started = True - current_text_buffer = [word] - - if current_item is not None: - finalize_passage_to_item(current_item, current_passage_buffer) - current_item['text'] = ' '.join(current_text_buffer[:-1]).strip() - structured_data.append(current_item) - print(f" -> Saved Question. Total structured items so far: {len(structured_data)}") - current_text_buffer = [word] - - current_item = { - 'question': word, 'options': {}, 'answer': '', 'passage': '', 'text': '' - } - current_option_key = None - last_entity_type = 'QUESTION' - just_finished_i_option = False - is_in_new_passage = False - continue - - if current_item is not None: - if is_in_new_passage: - if 'new_passage' not in current_item: - current_item['new_passage'] = word - else: - current_item['new_passage'] += f' {word}' - - if label.startswith('B-') or (label.startswith('I-') and entity_type != 'PASSAGE'): - print(f" ↳ [State] Exiting new_passage mode at label {label}") - is_in_new_passage = False - - if label.startswith(('B-', 'I-')): - last_entity_type = entity_type - continue - - is_in_new_passage = False - - if label.startswith('B-'): - if entity_type in ['QUESTION', 'OPTION', 'ANSWER', 'SECTION_HEADING']: - finalize_passage_to_item(current_item, current_passage_buffer) - current_passage_buffer = [] - - last_entity_type = entity_type - - if entity_type == 'PASSAGE': - if previous_entity_type == 'OPTION' and just_finished_i_option: - print(f" ↳ [State] Transitioning to new_passage (Option -> Passage boundary)") - current_item['new_passage'] = word - is_in_new_passage = True - else: - current_passage_buffer.append(word) - - elif entity_type == 'OPTION': - current_option_key = word - current_item['options'][current_option_key] = word - just_finished_i_option = False - - elif entity_type == 'ANSWER': - current_item['answer'] = word - current_option_key = None - just_finished_i_option = False - - elif entity_type == 'QUESTION': - current_item['question'] += f' {word}' - just_finished_i_option = False - - elif label.startswith('I-'): - if entity_type == 'QUESTION': - current_item['question'] += f' {word}' - elif entity_type == 'PASSAGE': - if previous_entity_type == 'OPTION' and just_finished_i_option: - current_item['new_passage'] = word - is_in_new_passage = True - else: - if not current_passage_buffer: last_entity_type = 'PASSAGE' - current_passage_buffer.append(word) - elif entity_type == 'OPTION' and current_option_key is not None: - current_item['options'][current_option_key] += f' {word}' - just_finished_i_option = True - elif entity_type == 'ANSWER': - current_item['answer'] += f' {word}' - - just_finished_i_option = (entity_type == 'OPTION') - - elif label == 'O': - # if last_entity_type == 'QUESTION': - # current_item['question'] += f' {word}' - # just_finished_i_option = False - pass - - # Final wrap up - if current_item is not None: - print(f"🏁 Finalizing the very last item...") - finalize_passage_to_item(current_item, current_passage_buffer) - current_item['text'] = ' '.join(current_text_buffer).strip() - structured_data.append(current_item) - - # Clean up and regex replacement - for item in structured_data: - item['text'] = re.sub(r'\s{2,}', ' ', item['text']).strip() - if 'new_passage' in item: - item['new_passage'] = re.sub(r'\s{2,}', ' ', item['new_passage']).strip() - - print(f"💾 Saving {len(structured_data)} items to {output_path}") - try: - with open(output_path, 'w', encoding='utf-8') as f: - json.dump(structured_data, f, indent=2, ensure_ascii=False) - print(f"✅ Decoding Complete. Total time: {time.time() - start_time:.2f}s") - except Exception as e: - print(f"⚠️ Error saving final JSON: {e}") - - return structured_data - -def create_query_text(entry: Dict[str, Any]) -> str: - """Combines question and options into a single string for similarity matching.""" - query_parts = [] - if entry.get("question"): - query_parts.append(entry["question"]) - - for key in ["options", "options_text"]: - options = entry.get(key) - if options and isinstance(options, dict): - for value in options.values(): - if value and isinstance(value, str): - query_parts.append(value) - return " ".join(query_parts) - - -def calculate_similarity(doc1: str, doc2: str) -> float: - """Calculates Cosine Similarity between two text strings.""" - if not doc1 or not doc2: - return 0.0 - - def clean_text(text): - return re.sub(r'^\s*[\(\d\w]+\.?\s*', '', text, flags=re.MULTILINE) - - clean_doc1 = clean_text(doc1) - clean_doc2 = clean_text(doc2) - corpus = [clean_doc1, clean_doc2] - - try: - vectorizer = CountVectorizer(stop_words='english', lowercase=True, token_pattern=r'(?u)\b\w\w+\b') - tfidf_matrix = vectorizer.fit_transform(corpus) - if tfidf_matrix.shape[1] == 0: - return 0.0 - vectors = tfidf_matrix.toarray() - # Handle cases where vectors might be empty or too short - if len(vectors) < 2: - return 0.0 - score = cosine_similarity(vectors[0:1], vectors[1:2])[0][0] - return score - except Exception: - return 0.0 - - - - - - -def process_context_linking(data: List[Dict[str, Any]]) -> List[Dict[str, Any]]: - """ - Links questions to passages based on 'passage' flow vs 'new_passage' priority. - Includes 'Decay Logic': If 2 consecutive questions fail to match the active passage, - the passage context is dropped to prevent false positives downstream. - """ - print("\n" + "=" * 80) - print("--- STARTING CONTEXT LINKING (WITH DECAY LOGIC) ---") - print("=" * 80) - - if not data: return [] - - # --- PHASE 1: IDENTIFY PASSAGE DEFINERS --- - passage_definer_indices = [] - for i, entry in enumerate(data): - if entry.get("passage") and entry["passage"].strip(): - passage_definer_indices.append(i) - if entry.get("new_passage") and entry["new_passage"].strip(): - if i not in passage_definer_indices: - passage_definer_indices.append(i) - - # --- PHASE 2: CONTEXT TRANSFER & LINKING --- - current_passage_text = None - current_new_passage_text = None - - # NEW: Counter to track consecutive linking failures - consecutive_failures = 0 - MAX_CONSECUTIVE_FAILURES = 2 - - for i, entry in enumerate(data): - item_type = entry.get("type", "Question") - - # A. UNCONDITIONALLY UPDATE CONTEXTS (And Reset Decay Counter) - if entry.get("passage") and entry["passage"].strip(): - current_passage_text = entry["passage"] - consecutive_failures = 0 # Reset because we have fresh explicit context - # print(f" [Flow] Updated Standard Context from Item {i}") - - if entry.get("new_passage") and entry["new_passage"].strip(): - current_new_passage_text = entry["new_passage"] - # We don't necessarily reset standard failures here as this is a local override - - # B. QUESTION LINKING - if entry.get("question") and item_type != "METADATA": - combined_query = create_query_text(entry) - - # Skip if query is too short (noise) - if len(combined_query.strip()) < 5: - continue - - # Calculate scores - score_old = calculate_similarity(current_passage_text, combined_query) if current_passage_text else 0.0 - score_new = calculate_similarity(current_new_passage_text, - combined_query) if current_new_passage_text else 0.0 - - # ------------------------------------------------------------------ - # 🛑 CRITICAL FIX APPLIED HERE 🛑 - # The original line: q_preview = entry['question'][:30] + '...' - - # 1. Capture the raw preview string (which might contain the bad surrogate) - q_preview_raw = entry['question'][:30] + '...' - - # 2. Safely clean the string by encoding to UTF-8 and ignoring errors, - # then decoding back. This removes the invalid surrogate character. - q_preview = q_preview_raw.encode('utf-8', errors='ignore').decode('utf-8') - # ------------------------------------------------------------------ - - # RESOLUTION LOGIC - linked = False - - # 1. Prefer New Passage if significantly better - if current_new_passage_text and (score_new > score_old + RESOLUTION_MARGIN) and ( - score_new >= SIMILARITY_THRESHOLD): - entry["passage"] = current_new_passage_text - print(f" [Linker] 🚀 Q{i} ('{q_preview}') -> NEW PASSAGE (Score: {score_new:.3f})") - linked = True - # Note: We do not reset 'consecutive_failures' for the standard passage here, - # because we matched the *new* passage, not the standard one. - - # 2. Otherwise use Standard Passage if it meets threshold - elif current_passage_text and (score_old >= SIMILARITY_THRESHOLD): - entry["passage"] = current_passage_text - print(f" [Linker] ✅ Q{i} ('{q_preview}') -> STANDARD PASSAGE (Score: {score_old:.3f})") - linked = True - consecutive_failures = 0 # Success! Reset the kill switch. - - if not linked: - # 3. DECAY LOGIC - if current_passage_text: - consecutive_failures += 1 - # This is the line that was failing (or similar logging lines) - print( - f" [Linker] ⚠️ Q{i} NOT LINKED. (Failures: {consecutive_failures}/{MAX_CONSECUTIVE_FAILURES})") - - if consecutive_failures >= MAX_CONSECUTIVE_FAILURES: - print(f" [Linker] 🗑️ Context dropped due to {consecutive_failures} consecutive misses.") - current_passage_text = None - consecutive_failures = 0 - else: - print(f" [Linker] ⚠️ Q{i} NOT LINKED (No active context).") - - # --- PHASE 3: CLEANUP AND INTERPOLATION --- - print(" [Linker] Running Cleanup & Interpolation...") - - # 3A. Self-Correction (Remove weak links) - for i in passage_definer_indices: - entry = data[i] - if entry.get("question") and entry.get("type") != "METADATA": - passage_to_check = entry.get("passage") or entry.get("new_passage") - if passage_to_check: - self_sim = calculate_similarity(passage_to_check, create_query_text(entry)) - if self_sim < SIMILARITY_THRESHOLD: - entry["passage"] = "" - if "new_passage" in entry: entry["new_passage"] = "" - print(f" [Cleanup] Removed weak link for Q{i}") - - # 3B. Interpolation (Fill gaps) - # We only interpolate if the gap is strictly 1 question wide to avoid undoing the decay logic - for i in range(1, len(data) - 1): - current_entry = data[i] - is_gap = current_entry.get("question") and not current_entry.get("passage") - if is_gap: - prev_p = data[i - 1].get("passage") - next_p = data[i + 1].get("passage") - if prev_p and next_p and (prev_p == next_p) and prev_p.strip(): - current_entry["passage"] = prev_p - print(f" [Linker] 🥪 Q{i} Interpolated from neighbors.") - - return data - - - - - - - - - - -def correct_misaligned_options(structured_data: List[Dict[str, Any]]) -> List[Dict[str, Any]]: - print("\n" + "=" * 80) - print("--- 5. STARTING POST-PROCESSING: OPTION ALIGNMENT CORRECTION ---") - print("=" * 80) - tag_pattern = re.compile(r'(EQUATION\d+|FIGURE\d+)') - corrected_count = 0 - for item in structured_data: - if item.get('type') in ['METADATA']: continue - options = item.get('options') - if not options or len(options) < 2: continue - option_keys = list(options.keys()) - for i in range(len(option_keys) - 1): - current_key = option_keys[i] - next_key = option_keys[i + 1] - current_value = options[current_key].strip() - next_value = options[next_key].strip() - is_current_empty = current_value == current_key - content_in_next = next_value.replace(next_key, '', 1).strip() - tags_in_next = tag_pattern.findall(content_in_next) - has_two_tags = len(tags_in_next) == 2 - if is_current_empty and has_two_tags: - tag_to_move = tags_in_next[0] - options[current_key] = f"{current_key} {tag_to_move}".strip() - options[next_key] = f"{next_key} {tags_in_next[1]}".strip() - corrected_count += 1 - print(f"✅ Option alignment correction finished. Total corrections: {corrected_count}.") - return structured_data - - - -def get_base64_for_file(filepath: str) -> Optional[str]: - """Reads a file and returns its Base64 encoded string without the data URI prefix.""" - try: - with open(filepath, "rb") as image_file: - # Return raw base64 string - return base64.b64encode(image_file.read()).decode('utf-8') - except Exception as e: - print(f"Error reading and encoding file {filepath}: {e}") - return None - - - - - - -def embed_images_as_base64_in_memory(structured_data: List[Dict[str, Any]], figure_extraction_dir: str) -> List[ Dict[str, Any]]: - print("\n" + "=" * 80) - print("--- 4. STARTING IMAGE EMBEDDING (Base64) / EQUATION TO LATEX CONVERSION ---") - print("=" * 80) - if not structured_data: - return [] - - image_files = glob.glob(os.path.join(figure_extraction_dir, "*.png")) - image_lookup = {} - tag_regex = re.compile(r'(figure|equation)(\d+)', re.IGNORECASE) - - for filepath in image_files: - filename = os.path.basename(filepath) - match = re.search(r'_(figure|equation)(\d+)\.png$', filename, re.IGNORECASE) - if match: - key = f"{match.group(1).upper()}{match.group(2)}" - image_lookup[key] = filepath - - print(f" -> Found {len(image_lookup)} image components.") - - final_structured_data = [] - - for item in structured_data: - text_fields = [item.get('question', ''), item.get('passage', '')] - if 'options' in item: - for opt_val in item['options'].values(): - text_fields.append(opt_val) - if 'new_passage' in item: - text_fields.append(item['new_passage']) - - unique_tags_to_embed = set() - for text in text_fields: - if not text: continue - for match in tag_regex.finditer(text): - tag = match.group(0).upper() - if tag in image_lookup: - unique_tags_to_embed.add(tag) - - # List of tags that were successfully converted to LaTeX - tags_converted_to_latex = set() - - for tag in sorted(list(unique_tags_to_embed)): - filepath = image_lookup[tag] - base_key = tag.replace(' ', '').lower() # e.g., figure1 or equation1 - - if 'EQUATION' in tag: - # Equation to LaTeX conversion - base64_code = get_base64_for_file(filepath) # This reads the file for conversion - if base64_code: - latex_output = get_latex_from_base64(base64_code) - if not latex_output.startswith('[P2T_ERROR') and not latex_output.startswith('[P2T_WARNING'): - # *** CORE CHANGE: Store the clean LaTeX output directly *** - item[base_key] = latex_output - tags_converted_to_latex.add(tag) - print(f" ✅ Embedded Clean LaTeX for {tag}") - else: - # On failure, embed the error message - item[base_key] = latex_output - print(f" ⚠️ Failed to convert {tag} to LaTeX. Embedding error message.") - else: - item[base_key] = "[FILE_ERROR: Could not read image file]" - print(f" ❌ File read error for {tag}.") - - elif 'FIGURE' in tag: - # Figure to Base64 conversion - base64_code = get_base64_for_file(filepath) - item[base_key] = base64_code - print(f" ✅ Embedded Base64 for {tag}") - - final_structured_data.append(item) - - print(f"✅ Image embedding complete.") - return final_structured_data - - - - - - - - - -# ============================================================================ -# --- MAIN FUNCTION --- -# ============================================================================ - - - - - - - - - - - -def classify_question_type(item: Dict[str, Any]) -> str: - """ - Classifies a question as 'MCQ', 'DESCRIPTIVE', or 'INTEGER' based on its options. - - Args: - item: Dictionary containing question data with 'options' field - - Returns: - str: 'MCQ' if options exist and are non-empty, 'DESCRIPTIVE' otherwise - """ - # Check if options exist and have meaningful content - options = item.get('options', {}) - - if not options: - return 'DESCRIPTIVE' - - # Check if options dict has keys and at least one non-empty value - has_valid_options = False - for key, value in options.items(): - # Check if the value is more than just the key itself (e.g., "A" vs "A Some text") - if value and isinstance(value, str): - # Remove the key from value and check if there's remaining content - remaining_text = value.replace(key, '').strip() - if remaining_text and len(remaining_text) > 0: - has_valid_options = True - break - - return 'MCQ' if has_valid_options else 'DESCRIPTIVE' - - -def add_question_type_validation(structured_data: List[Dict[str, Any]]) -> List[Dict[str, Any]]: - """ - Adds 'question_type' field to all question entries in the structured data. - - Args: - structured_data: List of dictionaries containing question data - - Returns: - List[Dict[str, Any]]: Modified list with 'question_type' field added - """ - print("\n" + "=" * 80) - print("--- ADDING QUESTION TYPE VALIDATION ---") - print("=" * 80) - - mcq_count = 0 - descriptive_count = 0 - metadata_count = 0 - - for item in structured_data: - item_type = item.get('type', 'Question') - - # Skip metadata entries - if item_type == 'METADATA': - metadata_count += 1 - item['question_type'] = 'METADATA' - continue - - # Classify the question - question_type = classify_question_type(item) - item['question_type'] = question_type - - if question_type == 'MCQ': - mcq_count += 1 - else: - descriptive_count += 1 - - print(f" ✅ Classification Complete:") - print(f" - MCQ Questions: {mcq_count}") - print(f" - Descriptive/Integer Questions: {descriptive_count}") - print(f" - Metadata Entries: {metadata_count}") - print(f" - Total Entries: {len(structured_data)}") - - return structured_data - - - - -import time -import traceback -import glob - - - - - - - -# def run_document_pipeline(input_pdf_path: str, layoutlmv3_model_path: str, structured_intermediate_output_path: Optional[str] = None) -> Optional[List[Dict[str, Any]]]: -# if not os.path.exists(input_pdf_path): -# print(f"❌ ERROR: File not found: {input_pdf_path}") -# return None - -# print("\n" + "#" * 80) -# print("### STARTING OPTIMIZED FULL DOCUMENT ANALYSIS PIPELINE ###") -# print(f"Input: {input_pdf_path}") -# print("#" * 80) - -# overall_start = time.time() -# pdf_name = os.path.splitext(os.path.basename(input_pdf_path))[0] -# temp_pipeline_dir = os.path.join(tempfile.gettempdir(), f"pipeline_run_{pdf_name}_{os.getpid()}") -# os.makedirs(temp_pipeline_dir, exist_ok=True) - -# preprocessed_json_path = os.path.join(temp_pipeline_dir, f"{pdf_name}_preprocessed.json") -# raw_output_path = os.path.join(temp_pipeline_dir, f"{pdf_name}_raw_predictions.json") - -# if structured_intermediate_output_path is None: -# structured_intermediate_output_path = os.path.join( -# temp_pipeline_dir, f"{pdf_name}_structured_intermediate.json" -# ) - -# final_result = None -# try: -# # --- Phase 1: Preprocessing --- -# print(f"\n[Step 1/5] Preprocessing (YOLO + Masking)...") -# p1_start = time.time() -# preprocessed_json_path_out = run_single_pdf_preprocessing(input_pdf_path, preprocessed_json_path) -# if not preprocessed_json_path_out: -# print("❌ FAILED at Step 1: Preprocessing returned None.") -# return None -# print(f"✅ Step 1 Complete ({time.time() - p1_start:.2f}s)") - -# # --- Phase 2: Inference --- -# print(f"\n[Step 2/5] Inference (LayoutLMv3)...") -# p2_start = time.time() -# page_raw_predictions_list = run_inference_and_get_raw_words( -# input_pdf_path, layoutlmv3_model_path, preprocessed_json_path_out -# ) -# if not page_raw_predictions_list: -# print("❌ FAILED at Step 2: Inference returned no data.") -# return None - -# with open(raw_output_path, 'w', encoding='utf-8') as f: -# json.dump(page_raw_predictions_list, f, indent=4) -# print(f"✅ Step 2 Complete ({time.time() - p2_start:.2f}s)") - -# # --- Phase 3: Decoding --- -# print(f"\n[Step 3/5] Decoding (BIO to Structured JSON)...") -# p3_start = time.time() -# structured_data_list = convert_bio_to_structured_json_relaxed( -# raw_output_path, structured_intermediate_output_path -# ) -# if not structured_data_list: -# print("❌ FAILED at Step 3: BIO conversion failed.") -# return None - -# print("... Correcting misalignments and linking context ...") -# structured_data_list = correct_misaligned_options(structured_data_list) -# structured_data_list = process_context_linking(structured_data_list) -# print(f"✅ Step 3 Complete ({time.time() - p3_start:.2f}s)") - -# # --- Phase 4: Base64 & LaTeX --- -# print(f"\n[Step 4/5] Finalizing Layout (Base64 Images & LaTeX)...") -# p4_start = time.time() -# final_result = embed_images_as_base64_in_memory(structured_data_list, FIGURE_EXTRACTION_DIR) -# if not final_result: -# print("❌ FAILED at Step 4: Final formatting failed.") -# return None -# print(f"✅ Step 4 Complete ({time.time() - p4_start:.2f}s)") - -# # --- Phase 4.5: Question Type Classification --- -# print(f"\n[Step 4.5/5] Adding Question Type Classification...") -# p4_5_start = time.time() -# final_result = add_question_type_validation(final_result) -# print(f"✅ Step 4.5 Complete ({time.time() - p4_5_start:.2f}s)") - -# # --- Phase 5: Hierarchical Tagging --- -# print(f"\n[Step 5/5] AI Classification (Subject/Concept Tagging)...") -# p5_start = time.time() -# classifier = HierarchicalClassifier() -# if classifier.load_models(): -# final_result = post_process_json_with_inference(final_result, classifier) -# print(f"✅ Step 5 Complete: Tags added ({time.time() - p5_start:.2f}s)") -# else: -# print("⚠️ WARNING: Classifier models failed to load. Skipping tags.") - -# # ============================================================ -# # 🔧 NEW STEP: FILTER OUT METADATA ENTRIES -# # ============================================================ -# print(f"\n[Post-Processing] Removing METADATA entries...") -# initial_count = len(final_result) -# final_result = [item for item in final_result if item.get('type') != 'METADATA'] -# removed_count = initial_count - len(final_result) -# print(f"✅ Removed {removed_count} METADATA entries. {len(final_result)} questions remain.") -# # ============================================================ - -# except Exception as e: -# print(f"\n‼️ FATAL PIPELINE EXCEPTION:") -# print(f"Error Message: {str(e)}") -# traceback.print_exc() -# return None - -# # finally: -# # print(f"\nCleaning up temporary files in {temp_pipeline_dir}...") -# # try: -# # for f in glob.glob(os.path.join(temp_pipeline_dir, '*')): -# # os.remove(f) -# # os.rmdir(temp_pipeline_dir) -# # print("🧹 Cleanup successful.") -# # except Exception as e: -# # print(f"⚠️ Cleanup failed: {e}") - -# total_time = time.time() - overall_start -# print("\n" + "#" * 80) -# print(f"### PIPELINE COMPLETE | Total Time: {total_time:.2f}s ###") -# print("#" * 80) - -# return final_result - - - -def run_document_pipeline(input_pdf_path: str, layoutlmv3_model_path: str, structured_intermediate_output_path: Optional[str] = None) -> Optional[List[Dict[str, Any]]]: - if not os.path.exists(input_pdf_path): - print(f"❌ ERROR: File not found: {input_pdf_path}") - return None - - print("\n" + "#" * 80) - print("### STARTING OPTIMIZED FULL DOCUMENT ANALYSIS PIPELINE ###") - print(f"Input: {input_pdf_path}") - print("#" * 80) - - overall_start = time.time() - pdf_name = os.path.splitext(os.path.basename(input_pdf_path))[0] - temp_pipeline_dir = os.path.join(tempfile.gettempdir(), f"pipeline_run_{pdf_name}_{os.getpid()}") - os.makedirs(temp_pipeline_dir, exist_ok=True) - - preprocessed_json_path = os.path.join(temp_pipeline_dir, f"{pdf_name}_preprocessed.json") - raw_output_path = os.path.join(temp_pipeline_dir, f"{pdf_name}_raw_predictions.json") - - if structured_intermediate_output_path is None: - structured_intermediate_output_path = os.path.join( - temp_pipeline_dir, f"{pdf_name}_structured_intermediate.json" - ) - - final_result = None - try: - # --- Phase 1: Preprocessing --- - print(f"\n[Step 1/5] Preprocessing (YOLO + Masking)...") - p1_start = time.time() - preprocessed_json_path_out = run_single_pdf_preprocessing(input_pdf_path, preprocessed_json_path) - if not preprocessed_json_path_out: - print("❌ FAILED at Step 1: Preprocessing returned None.") - return None - print(f"✅ Step 1 Complete ({time.time() - p1_start:.2f}s)") - - # --- Phase 2: Inference (MODIFIED) --- - print(f"\n[Step 2/5] Inference (Custom Model)...") - p2_start = time.time() - - # ------------------------------------------------------------------------- - # --- COMMENTED OUT OLD LAYOUTLMV3 CALL FOR REVERSION --- - # page_raw_predictions_list = run_inference_and_get_raw_words( - # input_pdf_path, layoutlmv3_model_path, preprocessed_json_path_out - # ) - # ------------------------------------------------------------------------- - - # --- NEW CUSTOM MODEL CALL --- - # Note: We only pass the JSON path because the custom function - # doesn't need to re-read the PDF or use the layoutlmv3 model path. - page_raw_predictions_list = run_custom_inference_and_get_raw_words( - preprocessed_json_path_out - ) - # ----------------------------- - - if not page_raw_predictions_list: - print("❌ FAILED at Step 2: Inference returned no data.") - return None - - with open(raw_output_path, 'w', encoding='utf-8') as f: - json.dump(page_raw_predictions_list, f, indent=4) - print(f"✅ Step 2 Complete ({time.time() - p2_start:.2f}s)") - - # --- Phase 3: Decoding --- - print(f"\n[Step 3/5] Decoding (BIO to Structured JSON)...") - p3_start = time.time() - structured_data_list = convert_bio_to_structured_json_relaxed( - raw_output_path, structured_intermediate_output_path - ) - if not structured_data_list: - print("❌ FAILED at Step 3: BIO conversion failed.") - return None - - print("... Correcting misalignments and linking context ...") - structured_data_list = correct_misaligned_options(structured_data_list) - structured_data_list = process_context_linking(structured_data_list) - print(f"✅ Step 3 Complete ({time.time() - p3_start:.2f}s)") - - # --- Phase 4: Base64 & LaTeX --- - print(f"\n[Step 4/5] Finalizing Layout (Base64 Images & LaTeX)...") - p4_start = time.time() - final_result = embed_images_as_base64_in_memory(structured_data_list, FIGURE_EXTRACTION_DIR) - if not final_result: - print("❌ FAILED at Step 4: Final formatting failed.") - return None - print(f"✅ Step 4 Complete ({time.time() - p4_start:.2f}s)") - - # --- Phase 4.5: Question Type Classification --- - print(f"\n[Step 4.5/5] Adding Question Type Classification...") - p4_5_start = time.time() - final_result = add_question_type_validation(final_result) - print(f"✅ Step 4.5 Complete ({time.time() - p4_5_start:.2f}s)") - - # --- Phase 5: Hierarchical Tagging --- - print(f"\n[Step 5/5] AI Classification (Subject/Concept Tagging)...") - p5_start = time.time() - classifier = HierarchicalClassifier() - if classifier.load_models(): - final_result = post_process_json_with_inference(final_result, classifier) - print(f"✅ Step 5 Complete: Tags added ({time.time() - p5_start:.2f}s)") - else: - print("⚠️ WARNING: Classifier models failed to load. Skipping tags.") - - # ============================================================ - # 🔧 NEW STEP: FILTER OUT METADATA ENTRIES - # ============================================================ - print(f"\n[Post-Processing] Removing METADATA entries...") - initial_count = len(final_result) - final_result = [item for item in final_result if item.get('type') != 'METADATA'] - removed_count = initial_count - len(final_result) - print(f"✅ Removed {removed_count} METADATA entries. {len(final_result)} questions remain.") - # ============================================================ - - except Exception as e: - print(f"\n‼️ FATAL PIPELINE EXCEPTION:") - print(f"Error Message: {str(e)}") - traceback.print_exc() - return None - - # finally: - # print(f"\nCleaning up temporary files in {temp_pipeline_dir}...") - # try: - # for f in glob.glob(os.path.join(temp_pipeline_dir, '*')): - # os.remove(f) - # os.rmdir(temp_pipeline_dir) - # print("🧹 Cleanup successful.") - # except Exception as e: - # print(f"⚠️ Cleanup failed: {e}") - - total_time = time.time() - overall_start - print("\n" + "#" * 80) - print(f"### PIPELINE COMPLETE | Total Time: {total_time:.2f}s ###") - print("#" * 80) - - return final_result - - - - - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Complete Pipeline") - parser.add_argument("--input_pdf", type=str, required=True, help="Input PDF") - parser.add_argument("--layoutlmv3_model_path", type=str, default=DEFAULT_LAYOUTLMV3_MODEL_PATH, help="Model Path") - - # --- ADDED ARGUMENT FOR DEBUGGING --- - parser.add_argument("--raw_preds_path", type=str, default='BIO_debug.json', - help="Debug path for raw BIO tag predictions (JSON).") - # ------------------------------------ - args = parser.parse_args() - - pdf_name = os.path.splitext(os.path.basename(args.input_pdf))[0] - final_output_path = os.path.abspath(f"{pdf_name}_final_output_embedded.json") - - # --- CALCULATE RAW PREDICTIONS OUTPUT PATH (Kept commented as per original script) --- - # raw_predictions_output_path = os.path.abspath( - # args.raw_preds_path if args.raw_preds_path else f"{pdf_name}_raw_predictions_debug.json") - # --------------------------------------------- - - # --- UPDATED FUNCTION CALL --- - final_json_data = run_document_pipeline( - args.input_pdf, - args.layoutlmv3_model_path ) - # ----------------------------- - - # 🛑 CRITICAL FINAL FIX: AGGRESSIVE CUSTOM JSON SAVING 🛑 - if final_json_data: - # 1. Dump the Python object to a standard JSON string. - # This converts the in-memory double backslash ('\\') into a quadruple backslash ('\\\\') - # in the raw json_str string content. - json_str = json.dumps(final_json_data, indent=2, ensure_ascii=False) - - # 2. **AGGRESSIVE UNDO ESCAPING:** We assume we have quadruple backslashes and - # replace them with the double backslashes needed for the LaTeX command to work. - # This operation essentially replaces four literal backslashes with two literal backslashes. - # final_output_content = json_str.replace('\\\\\\\\', '\\\\') - - # 3. Write the corrected string content to the file. - with open(final_output_path, 'w', encoding='utf-8') as f: - f.write(json_str) - - print(f"\n✅ Final Data Saved: {final_output_path}") - else: - print("\n❌ Pipeline Failed.") - sys.exit(1) +gradio +pymupdf +langchain-community +langchain-huggingface +faiss-cpu +sentence-transformers \ No newline at end of file