Update working_yolo_pipeline.py

working_yolo_pipeline.py

@@ -190,12 +190,20 @@ class MCQTagger(nn.Module):
         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)
+        
+        # --- FIX: ADD THIS LINE TO MATCH SAVED MODEL ---
+        self.layer_norm = nn.LayerNorm(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)
@@ -203,17 +211,24 @@ class MCQTagger(nn.Module):
         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)
+        
+        # Note: Even if layer_norm isn't explicitly used in the forward pass logic here,
+        # it must be defined in __init__ to satisfy the strict state_dict loading.
+        
         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