feat(transformer-ner) Added class weighting for loss calculation and exposing training params.

medcat-plugins/transformer-ner/src/medcat_transformer_ner/config.py

@@ -29,6 +29,25 @@ class TransformerNER(Ner):
     This helps with low signal words not being detected by the model. This will
     increase computation time, and could reduce precision."""
     learning_rate: float = 1e-5
-    """The learning rate to be used for training the model"""
+    """The learning rate to be used for training the model. Set via
+    TransformerNER.set_learning_rate on a loaded model so the change is propagated 
+    to the optimizer."""
     weight_decay: float = 0.001
-    """The weight decay to be used for training the model"""
+    """The weight decay to be used for training the model. Set via
+    TransformerNER.set_weight_decay on a loaded model so the change is propagated 
+    to the optimizer."""
+    class_weights: dict[str, float] = {
+        "O": 1.0,
+        "B-ENT": 1.0,
+        "I-ENT": 1.0,
+        "E-ENT": 1.0,
+        "S-ENT": 1.0,
+    }
+    """Per-label weights for the token-level cross entropy loss, keyed by
+    label name. Uniform by default. Set via TransformerNER.set_class_weights so the 
+    change is propagated to the loaded model."""
+    crf_loss_weight: float = 1.0
+    """Multiplier on the CRF loss in the total loss:
+    crf_loss_weight * crf_loss + weighted_cross_entropy. Changing this will proioritize 
+    different metrics over the other. Set via TransformerNER.set_crf_loss_weight so the
+    change is propagated to the loaded model."""

medcat-plugins/transformer-ner/src/medcat_transformer_ner/transformer_ner.py

@@ -60,9 +60,61 @@ def _resolve_model_source(path_or_model_name: Union[str, Path]) -> str:
             return str(candidate.resolve())
         return str(path_or_model_name)
 
+    def set_class_weights(self, class_weights: dict[str, float]) -> None:
+        """Set per-label loss weights (labels not given keep their current
+        weight). Updates the config and, if loaded, the model.
+
+        Args: class_weights (dict[str, float]): A dictionary mapping label names to 
+        their corresponding weights. The labels and weight mapping should follow the 
+        BIOES scheme:
+            {
+                "O": not an entity,
+                "B-ENT": the begining token of a multi-token entity,
+                "I-ENT": the intermediate token of a multi-token entity,
+                "E-ENT": the ending token of a multi-token entity,
+                "S-ENT": a single token entity,
+            }
+        """
+        unknown = set(class_weights) - set(self.label2id)
+        if unknown:
+            raise ValueError(f"Unknown labels in class_weights: {sorted(unknown)}")
+        self.cnf_ner.class_weights = {**self.cnf_ner.class_weights, **class_weights}
+        if hasattr(self, "model"):
+            self.model.set_class_weights(self.cnf_ner.class_weights)
+
+    def set_crf_loss_weight(self, crf_loss_weight: float) -> None:
+        """Set the multiplier on the CRF loss. Updates the config and, if
+        loaded, the model."""
+        self.cnf_ner.crf_loss_weight = crf_loss_weight
+        if hasattr(self, "model"):
+            self.model.crf_loss_weight = crf_loss_weight
+
+    def set_learning_rate(self, learning_rate: float) -> None:
+        """Set the learning rate. Updates the config and, if loaded, the
+        optimizer and scheduler."""
+        self.cnf_ner.learning_rate = learning_rate
+        if hasattr(self, "optimizer"):
+            for group in self.optimizer.param_groups:
+                group["lr"] = learning_rate
+                # the scheduler scales from the initial lr, so it must change too
+                group["initial_lr"] = learning_rate
+            self.scheduler.base_lrs = [learning_rate] * len(
+                self.optimizer.param_groups)
+
+    def set_weight_decay(self, weight_decay: float) -> None:
+        """Set the weight decay. Updates the config and, if loaded, the
+        optimizer."""
+        self.cnf_ner.weight_decay = weight_decay
+        if hasattr(self, "optimizer"):
+            for group in self.optimizer.param_groups:
+                group["weight_decay"] = weight_decay
+
     def _get_model_init_kwargs(self) -> dict[str, Any]:
         """Build kwargs passed to ModelForEmbeddingLinking.from_pretrained."""
-        return dict(self._model_init_kwargs)
+        kwargs = dict(self._model_init_kwargs)
+        kwargs.setdefault("class_weights", dict(self.cnf_ner.class_weights))
+        kwargs.setdefault("crf_loss_weight", self.cnf_ner.crf_loss_weight)
+        return kwargs
 
     def load_transformers(self, language_model_name: Union[str, Path]) -> None:
         """Load tokenizer/model from local path or Hugging Face model id."""
@@ -95,9 +147,10 @@ def load_transformers(self, language_model_name: Union[str, Path]) -> None:
             self.model.to(self.device)
             self._loaded_model_source: str = model_source
             self._loaded_model_init_kwargs: dict[str, Any] = model_init_kwargs
-            self.optimizer = torch.optim.AdamW(self.model.parameters(), 
-                                               lr=1e-5, 
-                                               weight_decay=0.001)
+            self.optimizer = torch.optim.AdamW(
+                self.model.parameters(),
+                lr=self.cnf_ner.learning_rate,
+                weight_decay=self.cnf_ner.weight_decay)
             self.scheduler = get_constant_schedule_with_warmup(
                 self.optimizer,
                 num_warmup_steps=20,
@@ -569,9 +622,10 @@ def deserialise_from(
             model_folder,
             device=ner.device,
         )
-        ner.optimizer = torch.optim.AdamW(ner.model.parameters(), 
-                                          lr=1e-5, 
-                                          weight_decay=0.001)
+        ner.optimizer = torch.optim.AdamW(
+            ner.model.parameters(),
+            lr=ner.cnf_ner.learning_rate,
+            weight_decay=ner.cnf_ner.weight_decay)
         ner.scheduler = get_constant_schedule_with_warmup(ner.optimizer, 
                                                           num_warmup_steps=20)
         ner.model.to(ner.device)

medcat-plugins/transformer-ner/src/medcat_transformer_ner/transformer_ner_model.py

@@ -3,7 +3,7 @@
 from typing import Any, Optional, Union
 from torch import Tensor, nn
 from torchcrf import CRF
-from transformers import AutoModelForTokenClassification
+from transformers import AutoModel
 import json
 import logging
 import torch
@@ -15,9 +15,7 @@ class ModelForBinaryNER(nn.Module):
 
     The architecture is: transformer backbone -> linear classifier -> CRF.
     """
-    # for mypy checking
-    label_is_start: Tensor
-    label_is_end: Tensor
+    class_weights: Tensor
 
     def __init__(
         self,
@@ -26,49 +24,51 @@ def __init__(
         num_labels: int = 5,
         top_n_layers_to_unfreeze: int = -1,
         device: Optional[Union[str, torch.device]] = None,
-        aux_loss_weight: float = 0.5,
+        class_weights: Optional[dict[str, float]] = None,
+        crf_loss_weight: float = 1.0,
     ) -> None:
         super().__init__()
         self.num_labels = num_labels
-        self.aux_loss_weight = aux_loss_weight
         self.id2label = id2label
-        self.language_model = AutoModelForTokenClassification.from_pretrained(
+        self.class_weights_dict = dict(class_weights or {})
+        self.crf_loss_weight = crf_loss_weight
+        self.language_model = AutoModel.from_pretrained(
             embedding_model_name,
-            num_labels=self.num_labels,
+            trust_remote_code=True,  # generally harmless for standard models
         )
-        # Make sure hidden states are available for the auxiliary heads.
-        self.language_model.config.output_hidden_states = True
+
         self.base_model_name = self.language_model.config.name_or_path
-
-        # For the auxiliary start/end position heads, we use the hidden states 
-        # from the last layer of the transformer.
         hidden_size = self.language_model.config.hidden_size
-        self.start_head = nn.Linear(hidden_size, 1)
-        self.end_head = nn.Linear(hidden_size, 1)
+
+        # Main token classification head: CRF emissions
+        self.classifier = nn.Linear(hidden_size, self.num_labels)
 
         # For state transitions
         self.crf = CRF(num_tags=self.num_labels, batch_first=True)
 
-        # Build boundary lookup tables from BIOES label names.
-        # This is future proof in the sense that "B" and "E" would still be the same
-        # for multiple different types of entities.
-        start_flags = []
-        end_flags = []
-        for i in range(self.num_labels):
-            label = self.id2label[i]
-            prefix = label.split("-", 1)[0]
-
-            start_flags.append(1.0 if prefix == "B" else 0.0)
-            end_flags.append(1.0 if prefix == "E" else 0.0)
-
-        self.register_buffer("label_is_start", 
-                             torch.tensor(start_flags, dtype=torch.float))
-        self.register_buffer("label_is_end", 
-                             torch.tensor(end_flags, dtype=torch.float))
+        # non-persistent so that existing checkpoints still load
+        self.register_buffer("class_weights",
+                             self._build_class_weights_tensor(),
+                             persistent=False)
 
         target_device = self._resolve_device(device)
         self.to(target_device)
 
+    def _build_class_weights_tensor(self) -> Tensor:
+        """cross_entropy needs a 1-D tensor indexed by label id; labels
+        without a given weight default to 1.0."""
+        return torch.tensor(
+            [self.class_weights_dict.get(self.id2label[i], 1.0)
+             for i in range(self.num_labels)],
+            dtype=torch.float)
+
+    def set_class_weights(self, class_weights: dict[str, float]) -> None:
+        """Update the class weights (by label name) on an already built model.
+        If all weights are 0 the weighted cross entropy is skipped (CRF loss
+        only)."""
+        self.class_weights_dict = dict(class_weights)
+        self.class_weights = self._build_class_weights_tensor().to(self.device)
+
     @staticmethod
     def _resolve_device(device: Optional[Union[str, torch.device]]) -> torch.device:
         if device is None:
@@ -86,19 +86,14 @@ def forward(self, **inputs) -> Any:
         outputs = self.language_model(**inputs, 
                                       return_dict=True, 
                                       output_hidden_states=True)
-        emissions = outputs.logits
-        # the last layer's hidden states for the start/end heads
-        hidden_states = outputs.hidden_states[-1]
-
-        # Linear classifiers for boundary heads
-        start_logits = self.start_head(hidden_states).squeeze(-1)  # [B, T]
-        end_logits = self.end_head(hidden_states).squeeze(-1)      # [B, T]
+        hidden_states = outputs.last_hidden_state
+        # Main token-classification logits, used as CRF emissions
+        emissions = self.classifier(hidden_states)
 
         # CRF can't handle -100 labels so this handles it
         loss = None
         crf_loss = None
-        start_loss = None
-        end_loss = None
+        weighted_ce_loss = None
 
         valid_mask = attention_mask.bool()
 
@@ -125,17 +120,21 @@ def forward(self, **inputs) -> Any:
                 mask=crf_mask,
                 reduction="token_mean",
             )
-
-            # Auxiliary start/end targets
-            start_targets, end_targets = self._build_boundary_targets(labels)
-
-            # Use the standard attention mask, but exclude -100 positions
-            aux_mask = valid_mask & (labels != -100)
 
-            start_loss = self._masked_bce_loss(start_logits, start_targets, aux_mask)
-            end_loss = self._masked_bce_loss(end_logits, end_targets, aux_mask)
-
-            loss = crf_loss + self.aux_loss_weight * (start_loss + end_loss)
+            loss = self.crf_loss_weight * crf_loss
+
+            # The CRF NLL is sequence-level and cannot be weighted per class,
+            # so class weights are applied via a token-level cross entropy
+            # on the emissions, added to the (scaled) CRF loss.
+            use_ce = self.class_weights.sum() > 0
+            if use_ce:
+                weighted_ce_loss = nn.functional.cross_entropy(
+                    emissions.reshape(-1, self.num_labels),
+                    labels.masked_fill(~valid_mask, -100).reshape(-1),
+                    weight=self.class_weights,
+                    ignore_index=-100,
+                )
+                loss = loss + weighted_ce_loss
 
         decoded_sequences = None
         decoded_tensor = None
@@ -159,44 +158,12 @@ def forward(self, **inputs) -> Any:
         return SimpleNamespace(
             loss=loss,
             crf_loss=crf_loss,
-            start_loss=start_loss,
-            end_loss=end_loss,
+            weighted_ce_loss=weighted_ce_loss,
             logits=emissions,
-            start_logits=start_logits,
-            end_logits=end_logits,
             predictions=decoded_tensor,
             decoded_sequences=decoded_sequences,
         )
 
-    def _masked_bce_loss(self, logits: Tensor, targets: Tensor, mask: Tensor) -> Tensor:
-        """
-        Normal BCE doesn't handle masking (i.e. handling [-100]), 
-        so this implements a masked version.
-        """
-        loss_fn = nn.BCEWithLogitsLoss(reduction="none")
-        loss = loss_fn(logits, targets.float())
-        loss = loss * mask.float()
-
-        denom = mask.float().sum().clamp_min(1.0)
-        return loss.sum() / denom
-
-    def _build_boundary_targets(self, labels: Tensor) -> tuple[Tensor, Tensor]:
-        """
-        Convert BIOES token labels into binary start/end targets.
-        - start = 1 for B-*, (and maybe S-*)
-        - end   = 1 for E-*, (and maybe S-*)
-        """
-        safe_labels = labels.clone()
-        safe_labels[safe_labels == -100] = 0
-
-        start_targets = self.label_is_start[safe_labels].to(labels.device)
-        end_targets = self.label_is_end[safe_labels].to(labels.device)
-
-        start_targets = start_targets.masked_fill(labels == -100, 0.0)
-        end_targets = end_targets.masked_fill(labels == -100, 0.0)
-
-        return start_targets, end_targets
-
     def _freeze_all_parameters(self) -> None:
         for param in self.language_model.parameters():
             param.requires_grad = False
@@ -246,7 +213,8 @@ def save_pretrained(self, save_directory: Union[str, Path]) -> None:
             "embedding_model_name": self.base_model_name,
             "num_labels": self.num_labels,
             "id2label": self.id2label,
-            "aux_loss_weight": self.aux_loss_weight,
+            "class_weights": self.class_weights_dict,
+            "crf_loss_weight": self.crf_loss_weight,
         }
         with open(save_path / "config.json", "w", encoding="utf-8") as f:
             json.dump(config, f, indent=2)
@@ -274,9 +242,18 @@ def from_pretrained(
                 for key, value in config["id2label"].items():
                     id2label[int(key)] = value
                 config["id2label"] = id2label
+            # the auxiliary start/end loss was removed; tolerate old checkpoints
+            config.pop("aux_loss_weight", None)
+            config.pop("class_weighted_loss_coef", None)
             config.update(kwargs)
             model = cls(device=target_device, **config)
             state_dict = torch.load(weights_path, map_location="cpu")
+            legacy_prefixes = ("start_head.", "end_head.")
+            legacy_buffers = ("label_is_start", "label_is_end")
+            state_dict = {
+                k: v for k, v in state_dict.items()
+                if not k.startswith(legacy_prefixes) and k not in legacy_buffers
+            }
             model.load_state_dict(state_dict)
             model.to(target_device)
             return model

medcat/medcat/stats/stats.py

@@ -121,7 +121,6 @@ class Metrics(BaseModel):
     overall: OverallMetrics
     per_cui: dict[str, CUIMetrics] = Field(default_factory=dict)
 
-
 class ModeStats(BaseModel):
     """Accumulated state and calculated metrics for one evaluation mode."""