Create span_ner_model.py

span_ner_model.py

@@ -1,4 +1,3 @@
- 
 """
 Medieval Latin NER - Custom Span-NER Architecture
 =============================================================================
@@ -20,14 +19,14 @@ class Config:
     TEXT_DIM       = 1024
     LABEL_MODEL    = "BAAI/bge-m3"
     LABEL_DIM      = 1024
-
+    
     MAX_SPAN_WIDTH = 80
     WIDTH_EMB_DIM  = 64
     SPAN_HIDDEN    = 512
     ATTENTION_HEADS = 4
-
+    
     MAX_SEQ_LEN    = 512
-    PREDICT_TEMP   = 1.35
+    PREDICT_TEMP   = 1.35 
 
 # ---------------------------------------------------------------------------
 # 2. LABEL DICTIONARY & PROMPTS
@@ -92,16 +91,16 @@ class SpanRepLayer(nn.Module):
         w_emb   = self.width_emb(width)
 
         idx = torch.arange(L, device=seq_out.device).view(1, 1, L)
-        mask = (idx >= spans[:,:,0:1]) & (idx <= spans[:,:,1:2])
+        mask = (idx >= spans[:,:,0:1]) & (idx <= spans[:,:,1:2]) 
 
-        att_logits = self.att_query(seq_out)
+        att_logits = self.att_query(seq_out) 
         att_logits = att_logits.unsqueeze(1).expand(B, S, L, self.num_heads)
 
         mask_expanded = mask.unsqueeze(-1).expand(-1, -1, -1, self.num_heads)
         att_logits = att_logits.masked_fill(~mask_expanded, float('-inf'))
-        att_weights = F.softmax(att_logits, dim=2)
+        att_weights = F.softmax(att_logits, dim=2) 
 
-        h_pool = torch.einsum('bslm,blh->bsmh', att_weights, seq_out)
+        h_pool = torch.einsum('bslm,blh->bsmh', att_weights, seq_out) 
         h_pool = h_pool.reshape(B, S, self.num_heads * H)
 
         return torch.cat([h_start, h_end, h_pool, w_emb], dim=-1)
@@ -112,7 +111,7 @@ class SpanNERModel(nn.Module):
         self.cfg = cfg
         self.text_enc = AutoModel.from_pretrained(cfg.TEXT_MODEL, add_pooling_layer=False)
         self.label_enc = AutoModel.from_pretrained(cfg.LABEL_MODEL)
-
+        
         self.span_layer = SpanRepLayer(cfg.TEXT_DIM, cfg.MAX_SPAN_WIDTH, cfg.WIDTH_EMB_DIM, num_heads=cfg.ATTENTION_HEADS)
 
         self.label_proj = nn.Sequential(
@@ -129,7 +128,7 @@ class SpanNERModel(nn.Module):
             nn.Linear(cfg.SPAN_HIDDEN, cfg.SPAN_HIDDEN)
         )
 
-        self.logit_scale = nn.Parameter(torch.tensor(1.0))
+        self.logit_scale = nn.Parameter(torch.tensor(1.0)) 
         self._raw_label_embs = None
 
     @torch.no_grad()
@@ -220,11 +219,11 @@ class SpanNERModel(nn.Module):
             covered = set(range(ws, we + 1))
             if flat_ner and covered & taken: continue
             if flat_ner: taken |= covered
-
+            
             start_char = tokens_info[ws]["start"]
             end_char = tokens_info[we]["end"]
             text_span = text[start_char:end_char]
-
+            
             result.append({
                 "label": label,
                 "score": round(score, 4),
@@ -235,4 +234,4 @@ class SpanNERModel(nn.Module):
                 "end_word": we
             })
 
-        return result
+        return result