Update app.py

app.py

@@ -1,557 +1,4 @@
 
-# import os
-# import json
-# import pickle
-# from typing import List, Dict, Any, Tuple
-# from collections import Counter
-# import torch
-# import torch.nn as nn
-# import torch.nn.functional as F
-# import re
-# from tqdm import tqdm
-
-# # === GRADIO AND DEPENDENCIES ===
-# import gradio as gr
-# import fitz  # PyMuPDF
-# from PIL import Image, ImageEnhance
-# import pytesseract
-
-# try:
-#     # Attempt to import the actual CRF layer for correct Viterbi decoding
-#     from TorchCRF import CRF
-# except ImportError:
-#     # Placeholder for environments where it's not yet installed, enabling model definition
-#     class CRF:
-#         def __init__(self, *args, **kwargs):
-#             pass
-#         # Fallback to simple argmax decoding if the CRF module is missing
-#         def viterbi_decode(self, emissions, mask):
-#             return [list(torch.argmax(emissions[0], dim=-1).cpu().numpy())]
-
-
-# # ========== CONFIG (Must match Training Script) ==========
-# MODEL_FILE = "model_enhanced.pt"
-# VOCAB_FILE = "vocabs_enhanced.pkl"
-
-# DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-# MAX_CHAR_LEN = 16
-# EMBED_DIM = 100
-# CHAR_EMBED_DIM = 30
-# CHAR_CNN_OUT = 30
-# BBOX_DIM = 100
-# HIDDEN_SIZE = 512
-# BBOX_NORM_CONSTANT = 1000.0
-# INFERENCE_CHUNK_SIZE = 256
-
-# # ========== LABELS (Must match Training Script) ==========
-# # Including PASSAGE for the new structuring logic
-# # LABELS = ["O", "B-QUESTION", "I-QUESTION", "B-OPTION", "I-OPTION", "B-ANSWER", "I-ANSWER", "B-IMAGE", "I-IMAGE", "B-PASSAGE", "I-PASSAGE"]
-# # LABEL2IDX = {l: i for i, l in enumerate(LABELS)}
-# # IDX2LABEL = {i: l for i, l in enumerate(LABELS)}
-# LABELS = ["O", "B-QUESTION", "I-QUESTION", "B-OPTION", "I-OPTION", "B-ANSWER", "I-ANSWER", "B-PASSAGE", "I-PASSAGE", "B-SECTION HEADING", "I-SECTION HEADING"]
-# LABEL2IDX = {l: i for i, l in enumerate(LABELS)}
-# IDX2LABEL = {i: l for i, l in enumerate(LABELS)}
-
-# # =========================================================
-# # 1. Core Classes (Vocab, CharCNNEncoder, MCQTagger)
-# # =========================================================
-
-# class Vocab:
-#     def __init__(self, min_freq=1, unk_token="<UNK>", pad_token="<PAD>"):
-#         self.min_freq = min_freq
-#         self.unk_token = unk_token
-#         self.pad_token = pad_token
-#         self.freq = Counter()
-#         self.itos = []
-#         self.stoi = {}
-
-#     def add_sentence(self, toks):
-#         self.freq.update(toks)
-
-#     def build(self):
-#         items = [tok for tok, c in self.freq.items() if c >= self.min_freq]
-#         items = [self.pad_token, self.unk_token] + sorted(items)
-#         self.itos = items
-#         self.stoi = {s: i for i, s in enumerate(self.itos)}
-
-#     def __len__(self):
-#         return len(self.itos)
-
-#     def __getitem__(self, token: str) -> int:
-#         return self.stoi.get(token, self.stoi[self.unk_token])
-
-#     def __getstate__(self):
-#         return {
-#             'min_freq': self.min_freq,
-#             'unk_token': self.unk_token,
-#             'pad_token': self.pad_token,
-#             'itos': self.itos,
-#             'stoi': self.stoi,
-#         }
-
-#     def __setstate__(self, state):
-#         self.min_freq = state['min_freq']
-#         self.unk_token = state['unk_token']
-#         self.pad_token = state['pad_token']
-#         self.itos = state['itos']
-#         self.stoi = state['stoi']
-#         self.freq = Counter()
-
-
-# def load_vocabs(path: str) -> Tuple[Vocab, Vocab]:
-#     """Loads word and character vocabularies."""
-#     try:
-#         absolute_path = os.path.abspath(path)
-#         with open(absolute_path, "rb") as f:
-#             word_vocab, char_vocab = pickle.load(f)
-#         if len(word_vocab) <= 2:
-#             raise IndexError("CRITICAL: Word vocabulary size is too small.")
-#         return word_vocab, char_vocab
-#     except Exception as e:
-#         raise RuntimeError(f"Error loading vocabs from {path}: {e}")
-
-
-# class CharCNNEncoder(nn.Module):
-#     def __init__(self, char_vocab_size, char_emb_dim, out_dim, kernel_sizes=(3, 4, 5)):
-#         super().__init__()
-#         self.char_emb = nn.Embedding(char_vocab_size, char_emb_dim, padding_idx=0)
-#         convs = [nn.Conv1d(char_emb_dim, out_dim, kernel_size=k) for k in kernel_sizes]
-#         self.convs = nn.ModuleList(convs)
-#         self.out_dim = out_dim * len(convs)
-
-#     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]
-#         res = torch.cat(outs, dim=1)
-#         return res.view(B, L, -1)
-
-
-# class MCQTagger(nn.Module):
-#     def __init__(self, vocab_size, char_vocab_size, n_labels, bbox_dim=BBOX_DIM):
-#         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.Linear(4, bbox_dim)
-#         in_dim = EMBED_DIM + self.char_enc.out_dim + bbox_dim
-
-#         self.bilstm = nn.LSTM(in_dim, HIDDEN_SIZE // 2, num_layers=2, batch_first=True, bidirectional=True, dropout=0.3)
-#         self.ff = nn.Linear(HIDDEN_SIZE, n_labels)
-#         self.crf = CRF(n_labels)
-#         self.dropout = nn.Dropout(p=0.5)
-
-#     def forward_emissions(self, words, chars, bboxes, mask):
-#         wemb = self.word_emb(words)
-#         cenc = self.char_enc(chars)
-#         benc = self.bbox_proj(bboxes)
-#         enc_in = torch.cat([wemb, cenc, benc], dim=-1)
-#         enc_in = self.dropout(enc_in)
-#         lengths = mask.sum(dim=1).cpu()
-
-#         if lengths.max().item() == 0:
-#             B, L = enc_in.size(0), enc_in.size(1)
-#             # Return zero tensor if batch is empty
-#             return torch.zeros((B, L, len(LABELS)), device=enc_in.device)
-
-#         packed_in = nn.utils.rnn.pack_padded_sequence(enc_in, lengths, batch_first=True, enforce_sorted=False)
-#         packed_out, _ = self.bilstm(packed_in)
-#         padded_out, _ = nn.utils.rnn.pad_packed_sequence(packed_out, batch_first=True)
-
-#         return self.ff(padded_out)
-
-#     def forward(self, words, chars, bboxes, mask, labels=None, class_weights=None, alpha=0.7):
-#         emissions = self.forward_emissions(words, chars, bboxes, mask)
-#         return self.crf.viterbi_decode(emissions, mask=mask)
-
-
-# # =========================================================
-# # 2. PDF Processing Functions
-# # =========================================================
-
-# def ocr_fallback_page(page: fitz.Page, page_width: float, page_height: float) -> List[Dict[str, Any]]:
-#     """Renders a PyMuPDF page, runs Tesseract OCR, and tokenizes the result."""
-#     try:
-#         pix = page.get_pixmap(matrix=fitz.Matrix(3, 3))
-#         if pix.n - pix.alpha > 3:
-#             pix = fitz.Pixmap(fitz.csRGB, pix)
-
-#         img_pil = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
-
-#         # Preprocessing for Tesseract
-#         img_pil = img_pil.convert('L')
-#         img_pil = ImageEnhance.Contrast(img_pil).enhance(2.0)
-#         img_pil = ImageEnhance.Sharpness(img_pil).enhance(2.0)
-
-#         ocr_data = pytesseract.image_to_data(img_pil, output_type=pytesseract.Output.DICT)
-
-#         ocr_tokens = []
-#         for i in range(len(ocr_data['text'])):
-#             word = ocr_data['text'][i]
-#             conf = ocr_data['conf'][i]
-
-#             if word.strip() and int(conf) > 50:
-#                 left, top, width, height = (ocr_data[k][i] for k in ['left', 'top', 'width', 'height'])
-#                 scale = page_width / pix.width
-
-#                 raw_bbox = [
-#                     left * scale, top * scale, (left + width) * scale, (top + height) * scale
-#                 ]
-
-#                 normalized_bbox = [
-#                     (raw_bbox[0] / page_width) * BBOX_NORM_CONSTANT,
-#                     (raw_bbox[1] / page_height) * BBOX_NORM_CONSTANT,
-#                     (raw_bbox[2] / page_width) * BBOX_NORM_CONSTANT,
-#                     (raw_bbox[3] / page_height) * BBOX_NORM_CONSTANT
-#                 ]
-
-#                 ocr_tokens.append({
-#                     "word": word,
-#                     "raw_bbox": [int(b) for b in raw_bbox],
-#                     "normalized_bbox": [int(b) for b in normalized_bbox]
-#                 })
-
-#         return ocr_tokens
-
-#     except Exception as e:
-#         print(f"OCR fallback failed: {e}")
-#         return []
-
-
-# def extract_tokens_from_pdf_fitz_with_ocr(pdf_path: str) -> List[Dict[str, Any]]:
-#     """Extracts words and bboxes using PyMuPDF text layer and falls back to OCR."""
-#     all_tokens = []
-#     try:
-#         doc = fitz.open(pdf_path)
-#         for page_num in tqdm(range(len(doc)), desc="PDF Page Processing"):
-#             page = doc.load_page(page_num)
-#             page_width, page_height = page.rect.width, page.rect.height
-#             page_tokens = []
-
-#             # 1. Primary Extraction: PyMuPDF's word structure
-#             word_list = page.get_text("words", sort=True)
-
-#             if word_list:
-#                 for word_data in word_list:
-#                     word = word_data[4]
-#                     raw_bbox = word_data[:4]
-
-#                     normalized_bbox = [
-#                         (raw_bbox[0] / page_width) * BBOX_NORM_CONSTANT,
-#                         (raw_bbox[1] / page_height) * BBOX_NORM_CONSTANT,
-#                         (raw_bbox[2] / page_width) * BBOX_NORM_CONSTANT,
-#                         (raw_bbox[3] / page_height) * BBOX_NORM_CONSTANT
-#                     ]
-
-#                     page_tokens.append({
-#                         "word": word,
-#                         "raw_bbox": [int(b) for b in raw_bbox],
-#                         "normalized_bbox": [int(b) for b in normalized_bbox]
-#                     })
-
-#             # 2. OCR Fallback
-#             if not page_tokens:
-#                 print(f" (Page {page_num + 1}) No text layer found. Running OCR...")
-#                 page_tokens = ocr_fallback_page(page, page_width, page_height)
-
-#             all_tokens.extend(page_tokens)
-
-#         doc.close()
-#     except Exception as e:
-#         raise RuntimeError(f"Error opening or processing PDF with fitz/OCR: {e}")
-
-#     return all_tokens
-
-
-# extract_tokens_from_pdf = extract_tokens_from_pdf_fitz_with_ocr
-
-
-# def preprocess_and_collate_tokens(all_tokens: List[Dict[str, Any]], word_vocab: Vocab, char_vocab: Vocab,
-#                                   chunk_size: int) -> List[Dict[str, Any]]:
-#     """Chunks the token list, converts to IDs, and prepares batches for inference."""
-#     all_batches = []
-
-#     for i in range(0, len(all_tokens), chunk_size):
-#         chunk = all_tokens[i:i + chunk_size]
-#         if not chunk: continue
-
-#         words = [t["word"] for t in chunk]
-#         bboxes_norm = [t["normalized_bbox"] for t in chunk]
-
-#         # Convert to IDs
-#         word_ids = [word_vocab[w] for w in words]
-
-#         char_ids = []
-#         for w in words:
-#             chs = [char_vocab[ch] for ch in w[:MAX_CHAR_LEN]]
-#             if len(chs) < MAX_CHAR_LEN:
-#                 pad_index = char_vocab.stoi.get(char_vocab.pad_token, 0)
-#                 chs += [pad_index] * (MAX_CHAR_LEN - len(chs))
-#             char_ids.append(chs)
-
-#         # Create padded tensors (using single-sample batches)
-#         word_pad = torch.LongTensor([word_ids]).to(DEVICE)
-#         char_pad = torch.LongTensor([char_ids]).to(DEVICE)
-
-#         # Final normalization to [0, 1] range before feeding to the model
-#         bbox_pad = torch.FloatTensor([bboxes_norm]).to(DEVICE) / BBOX_NORM_CONSTANT
-#         mask = torch.ones(word_pad.size(), dtype=torch.bool).to(DEVICE)
-
-#         all_batches.append({
-#             "words": word_pad,
-#             "chars": char_pad,
-#             "bboxes": bbox_pad,
-#             "mask": mask,
-#             "original_tokens": chunk
-#         })
-
-#     return all_batches
-
-
-# # =========================================================
-# # 3. Model Loading and Caching (Global Variables Defined Here!)
-# # =========================================================
-
-# # Global variables (MODEL, VOCABS) are defined here for use in the wrapper function
-# WORD_VOCAB = None
-# CHAR_VOCAB = None
-# MODEL = None
-
-# try:
-#     WORD_VOCAB, CHAR_VOCAB = load_vocabs(VOCAB_FILE)
-#     MODEL = MCQTagger(len(WORD_VOCAB), len(CHAR_VOCAB), len(LABELS)).to(DEVICE)
-#     MODEL.load_state_dict(torch.load(MODEL_FILE, map_location=DEVICE))
-#     MODEL.eval()
-#     print("✅ Model and Vocabs loaded successfully (Cached).")
-# except Exception as e:
-#     # This prevents the app from crashing if the model files are missing on startup
-#     print(f"❌ Initial Model/Vocab Load Failure: {e}")
-#     print("The Gradio demo will not function until model_CAT.pt and vocabs_CAT.pkl are found.")
-
-
-# # =========================================================
-# # 4. Structuring Logic (Converts BIO to clean JSON)
-# # =========================================================
-
-# def finalize_passage_to_item(item, passage_buffer):
-#     """Adds passage text to the current item and clears the buffer."""
-#     if passage_buffer:
-#         passage_text = re.sub(r'\s{2,}', ' ', ' '.join(passage_buffer)).strip()
-#         if item.get('passage'):
-#             item['passage'] += ' ' + passage_text
-#         else:
-#             item['passage'] = passage_text
-#     passage_buffer.clear()
-#     return item
-
-# def convert_bio_to_structured_json_strict(predictions: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
-#     """
-#     Converts a list of {word, predicted_label} tokens into structured MCQ JSON format.
-#     """
-#     structured_data = []
-#     current_item = None
-#     current_option_key = None
-#     current_passage_buffer = []
-#     current_text_buffer = []
-
-#     first_question_started = False
-#     last_entity_type = None
-
-#     for item in 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)
-
-#         is_passage_label = (label == 'B-PASSAGE' or label == 'I-PASSAGE')
-
-#         # --- BEFORE FIRST QUESTION/METADATA HANDLING ---
-#         if not first_question_started and label != 'B-QUESTION' and not is_passage_label:
-#             continue
-
-#         # --- PASSAGE HANDLING (Before question start) ---
-#         if not first_question_started and is_passage_label:
-#             if label == 'B-PASSAGE' or (label == 'I-PASSAGE' and last_entity_type == 'PASSAGE'):
-#                 current_passage_buffer.append(word)
-#                 last_entity_type = 'PASSAGE'
-#             continue
-
-#         # --- NEW QUESTION START (B-QUESTION) ---
-#         if label == 'B-QUESTION':
-#             # 1. Capture leading text/passage as METADATA
-#             if not first_question_started:
-#                 header_text = ' '.join(current_text_buffer[:-1]).strip()
-#                 if header_text or current_passage_buffer:
-#                     metadata_item = {'type': 'METADATA'}
-#                     metadata_item = 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]
-
-#             # 2. Save previous question block
-#             elif current_item is not None:
-#                 current_item = finalize_passage_to_item(current_item, current_passage_buffer)
-#                 current_item['text'] = ' '.join(current_text_buffer[:-1]).strip()
-#                 structured_data.append(current_item)
-#                 current_text_buffer = [word]
-
-#             # 3. Initialize new question
-#             current_item = {
-#                 'type': 'MCQ',
-#                 'question': word,
-#                 'options': {},
-#                 'answer': '',
-#                 'text': ''
-#             }
-#             current_option_key = None
-#             last_entity_type = 'QUESTION'
-#             continue
-
-#         # --- IF INSIDE A QUESTION BLOCK ---
-#         if current_item is not None:
-
-#             if label.startswith('B-'):
-#                 last_entity_type = entity_type
-
-#                 if entity_type == 'PASSAGE':
-#                     finalize_passage_to_item(current_item, current_passage_buffer)
-#                     current_passage_buffer.append(word)
-#                 elif entity_type == 'OPTION':
-#                     current_option_key = word
-#                     current_item['options'][current_option_key] = word
-#                     current_passage_buffer = []
-#                 elif entity_type == 'ANSWER':
-#                     current_item['answer'] = word
-#                     current_option_key = None
-#                     current_passage_buffer = []
-#                 elif entity_type == 'QUESTION':
-#                     current_item['question'] += f' {word}'
-#                     current_passage_buffer = []
-
-#             elif label.startswith('I-'):
-#                 if entity_type == 'QUESTION' and last_entity_type == 'QUESTION':
-#                     current_item['question'] += f' {word}'
-#                 elif entity_type == 'OPTION' and last_entity_type == 'OPTION' and current_option_key is not None:
-#                     current_item['options'][current_option_key] += f' {word}'
-#                 elif entity_type == 'ANSWER' and last_entity_type == 'ANSWER':
-#                     current_item['answer'] += f' {word}'
-#                 elif entity_type == 'PASSAGE' and last_entity_type == 'PASSAGE':
-#                     current_passage_buffer.append(word)
-
-#             elif label == 'O':
-#                 pass
-
-#     # --- Finalize last item ---
-#     if current_item is not None:
-#         current_item = finalize_passage_to_item(current_item, current_passage_buffer)
-#         current_item['text'] = re.sub(r'\s{2,}', ' ', ' '.join(current_text_buffer)).strip()
-#         structured_data.append(current_item)
-#     elif not structured_data and current_passage_buffer:
-#         # Case: Only passage/metadata was present in the whole document
-#         metadata_item = {'type': 'METADATA'}
-#         metadata_item = finalize_passage_to_item(metadata_item, current_passage_buffer)
-#         metadata_item['text'] = re.sub(r'\s{2,}', ' ', ' '.join(current_text_buffer)).strip()
-#         structured_data.append(metadata_item)
-
-
-#     # --- FINAL CLEANUP ---
-#     for item in structured_data:
-#         # Clean up all text fields for excessive whitespace
-#         item['text'] = re.sub(r'\s{2,}', ' ', item['text']).strip()
-#         if 'passage' in item:
-#             item['passage'] = re.sub(r'\s{2,}', ' ', item['passage']).strip()
-#             if not item['passage']:
-#                 del item['passage']
-#         for field in ['question', 'answer']:
-#             if field in item:
-#                 item[field] = re.sub(r'\s{2,}', ' ', item[field]).strip()
-#         if 'options' in item:
-#             for k, v in item['options'].items():
-#                  item['options'][k] = re.sub(r'\s{2,}', ' ', v).strip()
-
-#     return structured_data
-
-
-# # =========================================================
-# # 5. The Gradio Inference Wrapper Function (Main Entry Point)
-# # =========================================================
-
-# def gradio_inference_wrapper(pdf_file: str) -> Tuple[str, List[Dict[str, Any]]]:
-#     """
-#     Wraps the entire two-stage pipeline: (1) Tagging -> (2) Structuring.
-#     """
-#     # Uses global variables defined in Section 3
-#     if MODEL is None:
-#         return "❌ ERROR: Model failed to load on startup. Check 'model_CAT.pt' and 'vocabs_CAT.pkl'.", []
-
-#     pdf_path = pdf_file
-#     raw_predictions = []
-
-#     try:
-#         # 1. Stage 1: PDF Processing and BIO Tagging
-#         all_tokens = extract_tokens_from_pdf(pdf_path)
-
-#         if not all_tokens:
-#             return "❌ ERROR: No tokens were extracted from the PDF, even after OCR fallback.", []
-
-#         # Uses global variables WORD_VOCAB, CHAR_VOCAB, INFERENCE_CHUNK_SIZE
-#         batches = preprocess_and_collate_tokens(all_tokens, WORD_VOCAB, CHAR_VOCAB, chunk_size=INFERENCE_CHUNK_SIZE)
-
-#         with torch.no_grad():
-#             for batch in batches:
-#                 words, chars, bboxes, mask = (batch[k] for k in ["words", "chars", "bboxes", "mask"])
-#                 preds_batch = MODEL(words, chars, bboxes, mask)
-#                 predictions = preds_batch[0]
-#                 original_tokens = batch["original_tokens"]
-
-#                 for token_data, pred_idx in zip(original_tokens, predictions):
-#                     # Uses global variable IDX2LABEL
-#                     raw_predictions.append({
-#                         "word": token_data["word"],
-#                         "bbox": token_data["raw_bbox"],
-#                         "predicted_label": IDX2LABEL[pred_idx]
-#                     })
-
-#         # 2. Stage 2: Structured JSON Conversion
-#         structured_output = convert_bio_to_structured_json_strict(raw_predictions)
-
-#         mcq_count = len([i for i in structured_output if i.get('type') == 'MCQ'])
-#         status_message = f"✅ Conversion complete. Found {mcq_count} MCQ items and {len(structured_output) - mcq_count} Metadata blocks."
-
-#         return status_message, structured_output
-
-#     except RuntimeError as e:
-#         return f"❌ PDF Processing Error: {e}", []
-#     except Exception as e:
-#         return f"❌ An unexpected processing error occurred: {e}", []
-
-
-# # =========================================================
-# # 6. Define and Launch the Gradio Interface
-# # =========================================================
-
-# if __name__ == "__main__":
-#     title = "MCQ Document Structure Tagger (Bi-LSTM-CRF) - Structured Output"
-#     description = "Upload a PDF document. The system processes it in two stages: 1) BIO-Tagging for structural elements (Question, Option, Answer, Passage) and 2) Converting those tags into a clean, structured JSON list of MCQ items."
-
-#     demo = gr.Interface(
-#         fn=gradio_inference_wrapper,
-#         # Ensure only PDF files are accepted
-#         inputs=gr.File(label="Upload PDF Document"),
-#         outputs=[
-#             gr.Textbox(label="Status Message", interactive=False),
-#             gr.JSON(label="Structured MCQ JSON Output", show_label=True)
-#         ],
-#         title=title,
-#         description=description,
-#         allow_flagging="never",
-#         concurrency_limit=2
-#     )
-
-#     demo.launch(show_error=True)
 
 import os
 import json