Update modeling_jargon.py

modeling_jargon.py

@@ -1,27 +1,38 @@
+from transformers import AutoModel
 from huggingface_hub import PyTorchModelHubMixin
-from transformers import PreTrainedModel, AutoConfig, AutoModel
 from torchcrf import CRF
 import torch.nn as nn
-import torch
 
 class CRFTokenClassificationModel(nn.Module, PyTorchModelHubMixin):
     def __init__(self, config):
         super().__init__()
-        self.transformer = AutoModel.from_config(config)
-        self.dropout = nn.Dropout(config.hidden_dropout_prob)
-        self.classifier = nn.Linear(config.hidden_size, config.num_labels)
-        self.crf = CRF(config.num_labels, batch_first=True)
+        # Load base transformer model
+        self.transformer = AutoModel.from_pretrained(config["pretrained_model_name"])
+        
+        # Classification layers
+        self.dropout = nn.Dropout(config["hidden_dropout_prob"])
+        self.classifier = nn.Linear(config["hidden_size"], config["num_labels"])
+        
+        # CRF layer for sequence labeling
+        self.crf = CRF(config["num_labels"], batch_first=True)
+        
+        # Label mappings
+        self.id2label = {v: k for k, v in config["label_map"].items()}
 
     def forward(self, input_ids, attention_mask=None, labels=None, **kwargs):
+        # Get transformer outputs
         outputs = self.transformer(input_ids=input_ids, attention_mask=attention_mask)
         sequence_output = self.dropout(outputs.last_hidden_state)
         logits = self.classifier(sequence_output)
 
+        # Calculate loss if labels provided (training mode)
         if labels is not None:
             loss = -self.crf(logits, labels, mask=attention_mask.bool(), reduction='mean')
             return {"loss": loss, "logits": logits}
         
+        # Return logits only (inference mode)
         return {"logits": logits}
 
     def decode(self, logits, mask):
+        # Use CRF to decode best sequence
         return self.crf.decode(logits, mask.bool())