Rename train_layoutlmv3.py to inference.py

inference.py

@@ -0,0 +1,394 @@
+ import os
+import json
+import torch
+import torch.nn as nn
+import pdfplumber
+import argparse
+import time
+import re
+from typing import List, Dict, Any, Optional
+from transformers import LayoutLMv3TokenizerFast, LayoutLMv3Model
+from TorchCRF import CRF
+
+# --- Configuration (Must match training) ---
+BASE_MODEL_ID = "microsoft/layoutlmv3-base"
+MAX_BBOX_DIMENSION = 1000
+LABELS = ["O", "B-QUESTION", "I-QUESTION", "B-OPTION", "I-OPTION", "B-ANSWER", "I-ANSWER", "B-SECTION_HEADING", "I-SECTION_HEADING", "B-PASSAGE", "I-PASSAGE"]
+LABEL2ID = {l: i for i, l in enumerate(LABELS)}
+ID2LABEL = {i: l for l, i in LABEL2ID.items()}
+
+# -------------------------
+# Part 1: Model Architecture 
+# (Must be identical to training script)
+# -------------------------
+class LayoutLMv3CRF(nn.Module):
+    def __init__(self, num_labels):
+        super().__init__()
+        self.layoutlm = LayoutLMv3Model.from_pretrained(BASE_MODEL_ID)
+        hidden_size = self.layoutlm.config.hidden_size
+        
+        self.classifier = nn.Sequential(
+            nn.Linear(hidden_size, hidden_size),
+            nn.GELU(),
+            nn.LayerNorm(hidden_size),
+            nn.Dropout(0.1),
+            nn.Linear(hidden_size, num_labels)
+        )
+        self.crf = CRF(num_labels)
+
+    def forward(self, input_ids, bbox, attention_mask, labels=None):
+        # Note: Your training script did not use pixel_values, so we omit them here too
+        outputs = self.layoutlm(input_ids=input_ids, bbox=bbox, attention_mask=attention_mask)
+        sequence_output = outputs.last_hidden_state
+        emissions = self.classifier(sequence_output)
+
+        if labels is not None:
+            log_likelihood = self.crf(emissions, labels, mask=attention_mask.bool())
+            return -log_likelihood.mean()
+        else:
+            return self.crf.viterbi_decode(emissions, mask=attention_mask.bool())
+
+# -------------------------
+# Part 2: PDF Extraction
+# -------------------------
+def extract_pdf_data(pdf_path):
+    """
+    Extracts words and normalized bounding boxes (0-1000) from a PDF.
+    """
+    extracted_pages = []
+    print(f"📄 Extracting content from: {pdf_path}")
+    
+    with pdfplumber.open(pdf_path) as pdf:
+        for page_idx, page in enumerate(pdf.pages):
+            width, height = page.width, page.height
+            words_data = page.extract_words()
+            
+            page_tokens = []
+            page_bboxes = []
+            
+            for w in words_data:
+                text = w['text']
+                # Normalize bbox to 0-1000 scale
+                x0 = int((w['x0'] / width) * 1000)
+                top = int((w['top'] / height) * 1000)
+                x1 = int((w['x1'] / width) * 1000)
+                bottom = int((w['bottom'] / height) * 1000)
+                
+                # Clamp
+                box = [
+                    max(0, min(x0, 1000)),
+                    max(0, min(top, 1000)),
+                    max(0, min(x1, 1000)),
+                    max(0, min(bottom, 1000))
+                ]
+                
+                page_tokens.append(text)
+                page_bboxes.append(box)
+                
+            extracted_pages.append({
+                "page_id": page_idx,
+                "tokens": page_tokens,
+                "bboxes": page_bboxes
+            })
+            
+    print(f"✅ Extracted {len(extracted_pages)} pages.")
+    return extracted_pages
+
+# -------------------------
+# Part 3: Inference Logic
+# -------------------------
+def run_inference(model, tokenizer, pages_data, device):
+    results = []
+    model.eval()
+    
+    print("🧠 Running Inference...")
+    
+    for page in pages_data:
+        tokens = page['tokens']
+        bboxes = page['bboxes']
+        
+        if not tokens:
+            continue
+
+        # Tokenize
+        encoding = tokenizer(
+            tokens, 
+            boxes=bboxes, 
+            return_tensors="pt", 
+            padding="max_length", 
+            truncation=True, 
+            max_length=512,
+            return_offsets_mapping=True 
+        )
+        
+        input_ids = encoding.input_ids.to(device)
+        bbox = encoding.bbox.to(device)
+        attention_mask = encoding.attention_mask.to(device)
+        
+        # Predict
+        with torch.no_grad():
+            # returns list of lists (batch_size=1)
+            preds = model(input_ids=input_ids, bbox=bbox, attention_mask=attention_mask)
+            pred_tags = preds[0] # Take first item in batch
+
+        # Align sub-word predictions back to original words
+        word_ids = encoding.word_ids()
+        aligned_data = []
+        
+        previous_word_idx = None
+        
+        for i, word_idx in enumerate(word_ids):
+            # Special tokens (None) or padding (masked) are skipped
+            if word_idx is None:
+                continue
+                
+            # If we are at the start of a new word (or the only token for that word)
+            if word_idx != previous_word_idx:
+                # Get the label ID
+                label_id = pred_tags[i]
+                label_str = ID2LABEL[label_id]
+                
+                # Retrieve original word text
+                original_word = tokens[word_idx]
+                
+                aligned_data.append({
+                    "word": original_word,
+                    "predicted_label": label_str
+                })
+            
+            previous_word_idx = word_idx
+            
+        results.append({
+            "page": page['page_id'],
+            "data": aligned_data
+        })
+        
+    return results
+
+# -------------------------
+# Part 4: User's Conversion Function
+# -------------------------
+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')
+
+        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':
+                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:
+        if 'text' in item:
+            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
+
+# -------------------------
+# Part 5: Main Execution
+# -------------------------
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--pdf_path", type=str, required=True, help="Path to the PDF file")
+    parser.add_argument("--model_path", type=str, required=True, help="Path to the .pth checkpoint")
+    parser.add_argument("--output_json", type=str, default="final_output.json", help="Path for final structured JSON")
+    args = parser.parse_args()
+
+    # 1. Setup Device
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"⚙️ Using device: {device}")
+
+    # 2. Load Model
+    print(f"🔄 Loading model from {args.model_path}...")
+    model = LayoutLMv3CRF(num_labels=len(LABELS))
+    # Load state dict
+    state_dict = torch.load(args.model_path, map_location=device)
+    model.load_state_dict(state_dict)
+    model.to(device)
+    
+    tokenizer = LayoutLMv3TokenizerFast.from_pretrained(BASE_MODEL_ID)
+
+    # 3. Extract PDF Data
+    pages_data = extract_pdf_data(args.pdf_path)
+
+    # 4. Run Inference
+    raw_predictions = run_inference(model, tokenizer, pages_data, device)
+
+    # 5. Save Intermediate (BIO tagged format)
+    intermediate_path = "temp_inference_bio.json"
+    with open(intermediate_path, "w", encoding="utf-8") as f:
+        json.dump(raw_predictions, f, indent=2, ensure_ascii=False)
+    print(f"💾 Intermediate BIO predictions saved to {intermediate_path}")
+
+    # 6. Convert to Structured JSON
+    convert_bio_to_structured_json_relaxed(intermediate_path, args.output_json)

train_layoutlmv3.py

@@ -1 +0,0 @@
-