Update inference.py

inference.py

@@ -1,26 +1,24 @@
- import os
-import json
+import gradio as gr
 import torch
 import torch.nn as nn
 import pdfplumber
-import argparse
-import time
+import json
+import os
 import re
-from typing import List, Dict, Any, Optional
 from transformers import LayoutLMv3TokenizerFast, LayoutLMv3Model
 from TorchCRF import CRF
 
-# --- Configuration (Must match training) ---
+# --- Configuration ---
+# Ensure this filename matches exactly what you uploaded to the Space
+MODEL_FILENAME = "layoutlmv3_nonlinear_scratch.pth" 
 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)
-# -------------------------
+# ---------------------------------------------------------
+# 1. MODEL ARCHITECTURE
+# ---------------------------------------------------------
 class LayoutLMv3CRF(nn.Module):
     def __init__(self, num_labels):
         super().__init__()
@@ -37,7 +35,6 @@ class LayoutLMv3CRF(nn.Module):
         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)
@@ -48,150 +45,32 @@ class LayoutLMv3CRF(nn.Module):
         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
-            })
+# ---------------------------------------------------------
+# 2. MODEL LOADING
+# ---------------------------------------------------------
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+tokenizer = LayoutLMv3TokenizerFast.from_pretrained(BASE_MODEL_ID)
+model = None
+
+def load_model():
+    global model
+    if model is None:
+        print(f"🔄 Loading model from {MODEL_FILENAME}...")
+        if not os.path.exists(MODEL_FILENAME):
+            raise FileNotFoundError(f"Model file {MODEL_FILENAME} not found. Please upload it to the Space.")
             
-    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}")
-
+        model = LayoutLMv3CRF(num_labels=len(LABELS))
+        state_dict = torch.load(MODEL_FILENAME, map_location=device)
+        model.load_state_dict(state_dict)
+        model.to(device)
+        model.eval()
+        print("✅ Model loaded successfully.")
+    return model
+
+# ---------------------------------------------------------
+# 3. CONVERSION LOGIC (Your Custom Function)
+# ---------------------------------------------------------
+def convert_bio_to_structured_json(predictions):
     structured_data = []
     current_item = None
     current_option_key = None
@@ -205,15 +84,16 @@ def convert_bio_to_structured_json_relaxed(input_path: str, output_path: str) ->
     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
+            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):
+    # Flatten predictions list
+    flat_predictions = []
+    for page in predictions:
+        flat_predictions.extend(page['data'])
+
+    for idx, item in enumerate(flat_predictions):
         word = item['word']
         label = item['predicted_label']
         entity_type = label[2:].strip() if label.startswith(('B-', 'I-')) else None
@@ -235,11 +115,9 @@ def convert_bio_to_structured_json_relaxed(input_path: str, output_path: str) ->
                 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
@@ -251,7 +129,6 @@ def convert_bio_to_structured_json_relaxed(input_path: str, output_path: str) ->
                 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 = {
@@ -265,53 +142,38 @@ def convert_bio_to_structured_json_relaxed(input_path: str, output_path: str) ->
 
         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 '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
+                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)
-                
+                    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}'
+                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
@@ -322,73 +184,111 @@ def convert_bio_to_structured_json_relaxed(input_path: str, output_path: str) ->
                 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}'
-                
+                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
+    # Clean text
     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}")
+        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()
 
     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)
+# ---------------------------------------------------------
+# 4. PROCESSING PIPELINE
+# ---------------------------------------------------------
+def process_pdf(pdf_file):
+    if pdf_file is None:
+        return None, "Please upload a PDF file."
 
-    # 3. Extract PDF Data
-    pages_data = extract_pdf_data(args.pdf_path)
+    try:
+        model = load_model()
+        
+        # 1. Extract
+        extracted_pages = []
+        with pdfplumber.open(pdf_file.name) 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']
+                    x0 = int((w['x0'] / width) * 1000)
+                    top = int((w['top'] / height) * 1000)
+                    x1 = int((w['x1'] / width) * 1000)
+                    bottom = int((w['bottom'] / height) * 1000)
+                    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})
+
+        # 2. Inference
+        raw_predictions = []
+        for page in extracted_pages:
+            tokens = page['tokens']
+            bboxes = page['bboxes']
+            if not tokens: continue
+
+            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)
+            
+            with torch.no_grad():
+                preds = model(input_ids=input_ids, bbox=bbox, attention_mask=attention_mask)
+                pred_tags = preds[0]
 
-    # 4. Run Inference
-    raw_predictions = run_inference(model, tokenizer, pages_data, device)
+            word_ids = encoding.word_ids()
+            aligned_data = []
+            prev_word_idx = None
+            
+            for i, word_idx in enumerate(word_ids):
+                if word_idx is None: continue
+                if word_idx != prev_word_idx:
+                    label_str = ID2LABEL[pred_tags[i]]
+                    aligned_data.append({"word": tokens[word_idx], "predicted_label": label_str})
+                prev_word_idx = word_idx
+            raw_predictions.append({"data": aligned_data})
+
+        # 3. Structure
+        final_json = convert_bio_to_structured_json(raw_predictions)
+        
+        # Save to file for download
+        output_filename = "structured_output.json"
+        with open(output_filename, "w", encoding="utf-8") as f:
+            json.dump(final_json, f, indent=2, ensure_ascii=False)
+            
+        return output_filename, f"✅ Successfully processed {len(extracted_pages)} pages. Found {len(final_json)} structured items."
 
-    # 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}")
+    except Exception as e:
+        return None, f"❌ Error: {str(e)}"
+
+# ---------------------------------------------------------
+# 5. GRADIO INTERFACE
+# ---------------------------------------------------------
+iface = gr.Interface(
+    fn=process_pdf,
+    inputs=gr.File(label="Upload PDF", file_types=[".pdf"]),
+    outputs=[
+        gr.File(label="Download JSON Output"),
+        gr.Textbox(label="Status Log")
+    ],
+    title="LayoutLMv3 PDF Parser",
+    description="Upload a document to extract Questions, Options, and Passages into structured JSON.",
+    allow_flagging="never"
+)
 
-    # 6. Convert to Structured JSON
-    convert_bio_to_structured_json_relaxed(intermediate_path, args.output_json)
+if __name__ == "__main__":
+    iface.launch()