diff --git a/medcat-plugins/transformer-ner/src/medcat_transformer_ner/config.py b/medcat-plugins/transformer-ner/src/medcat_transformer_ner/config.py index d9c7cf2a0..940d0b80a 100644 --- a/medcat-plugins/transformer-ner/src/medcat_transformer_ner/config.py +++ b/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""" \ No newline at end of file + """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.""" \ No newline at end of file diff --git a/medcat-plugins/transformer-ner/src/medcat_transformer_ner/transformer_ner.py b/medcat-plugins/transformer-ner/src/medcat_transformer_ner/transformer_ner.py index 7d98b6b0a..a135bd752 100644 --- a/medcat-plugins/transformer-ner/src/medcat_transformer_ner/transformer_ner.py +++ b/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) diff --git a/medcat-plugins/transformer-ner/src/medcat_transformer_ner/transformer_ner_model.py b/medcat-plugins/transformer-ner/src/medcat_transformer_ner/transformer_ner_model.py index 3c904c3cb..a521e0ef7 100644 --- a/medcat-plugins/transformer-ner/src/medcat_transformer_ner/transformer_ner_model.py +++ b/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 diff --git a/medcat/medcat/stats/stats.py b/medcat/medcat/stats/stats.py index bb7f83914..448c11bad 100644 --- a/medcat/medcat/stats/stats.py +++ b/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."""