Update app.py

app.py

@@ -554,7 +554,6 @@
 #     demo.launch(show_error=True)
 
 
-
 import os
 import json
 import pickle
@@ -571,8 +570,7 @@ import sys
 from types import ModuleType
 
 # --- 1. CRITICAL: MOCK THE TRAINING MODULE ---
-# Your model was saved with references to a 'train_model' module.
-# We create a fake module and inject our local classes into it so torch.load works.
+# We create a fake module to satisfy pickle/torch.load
 train_mod = ModuleType("train_model")
 sys.modules["train_model"] = train_mod
 
@@ -580,9 +578,14 @@ sys.modules["train_model"] = train_mod
 try:
     from torch_crf import CRF
 except ImportError:
-    from TorchCRF import CRF
+    try:
+        from TorchCRF import CRF
+    except ImportError:
+        # Fallback if libraries are missing (prevents crash, but model won't load)
+        class CRF(nn.Module):
+            def __init__(self, *args, **kwargs): super().__init__()
 
-# --- 3. CONFIG (Matching Training Script) ---
+# --- 3. CONFIG ---
 MODEL_FILE = "model_enhanced.pt"
 VOCAB_FILE = "vocabs_enhanced.pkl"
 DEVICE = torch.device("cpu") 
@@ -605,7 +608,19 @@ LABELS = [
 ]
 IDX2LABEL = {i: l for i, l in enumerate(LABELS)}
 
-# --- 4. MODEL ARCHITECTURE (Exact Match to Training Script) ---
+# --- 4. CLASSES (Re-defined to match training) ---
+
+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 __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)):
@@ -681,12 +696,14 @@ class MCQTagger(nn.Module):
         emissions = self.ff(torch.cat([lstm_out, attn_out], dim=-1))
         return self.crf.viterbi_decode(emissions, mask=mask)
 
-# Link classes to the fake module
+# --- 5. CRITICAL FIX: LINK CLASSES TO FAKE MODULE ---
+# This tells pickle: "When you look for 'train_model.Vocab', look here instead."
+train_mod.Vocab = Vocab
 train_mod.MCQTagger = MCQTagger
 train_mod.CharCNNEncoder = CharCNNEncoder
+train_mod.SpatialAttention = SpatialAttention
 
-# --- 5. FEATURE EXTRACTION HELPERS (From Training Script) ---
-
+# --- 6. FEATURE EXTRACTORS ---
 def extract_spatial_features(tokens, idx):
     curr = tokens[idx]
     f = []
@@ -732,22 +749,23 @@ def extract_context_features(tokens, idx, window=3):
     f.extend([dq, dopt])
     return f
 
-# --- 6. INFERENCE PIPELINE ---
-
+# --- 7. INFERENCE WRAPPER ---
 def gradio_inference(pdf_file):
     if not os.path.exists(MODEL_FILE) or not os.path.exists(VOCAB_FILE):
         return "❌ Missing model/vocab files.", []
     
     try:
+        # Load Vocab
         with open(VOCAB_FILE, "rb") as f:
             word_vocab, char_vocab = pickle.load(f)
         
+        # Load Model
         model = MCQTagger(len(word_vocab), len(char_vocab), len(LABELS)).to(DEVICE)
         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()
 
-        # Extract tokens
+        # Parse PDF
         doc = fitz.open(pdf_file.name)
         all_tokens = []
         for page in doc:
@@ -756,16 +774,28 @@ def gradio_inference(pdf_file):
                 all_tokens.append({'text': text, 'x0': x0, 'y0': y0, 'x1': x1, 'y1': y1})
         doc.close()
 
-        # Generate features
+        if not all_tokens: return "❌ No text found.", []
+
+        # Feature Extraction
         for i in range(len(all_tokens)):
             all_tokens[i]['spatial_features'] = extract_spatial_features(all_tokens, i)
             all_tokens[i]['context_features'] = extract_context_features(all_tokens, i)
 
+        # Predict
         results = []
         for i in range(0, len(all_tokens), INFERENCE_CHUNK_SIZE):
             chunk = all_tokens[i : i + INFERENCE_CHUNK_SIZE]
+            
+            # Prepare Inputs
             w_ids = torch.LongTensor([[word_vocab[t['text']] for t in chunk]]).to(DEVICE)
-            c_ids = torch.LongTensor([[([char_vocab[c] for c in t['text'][:MAX_CHAR_LEN]] + [0]*MAX_CHAR_LEN)[:MAX_CHAR_LEN] for t in chunk]]).to(DEVICE)
+            
+            c_ids_list = []
+            for t in chunk:
+                chars = [char_vocab[c] for c in t['text'][: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)
@@ -776,17 +806,18 @@ def gradio_inference(pdf_file):
                 for t, p in zip(chunk, preds):
                     results.append({"word": t['text'], "label": IDX2LABEL[p]})
 
-        return "✅ Processing Complete", results
+        return "✅ Success", results
+
     except Exception as e:
+        import traceback
         return f"❌ Error: {str(e)}", []
 
-# --- 7. INTERFACE ---
-
+# --- 8. UI ---
 demo = gr.Interface(
     fn=gradio_inference,
-    inputs=gr.File(label="Upload MCQ PDF"),
+    inputs=gr.File(label="Upload PDF"),
     outputs=[gr.Textbox(label="Status"), gr.JSON(label="Predictions")],
-    title="Enhanced MCQ Tagger (Spatial Attention + BiLSTM-CRF)"
+    title="MCQ Enhanced Tagger"
 )
 
 if __name__ == "__main__":