From f65a261333098c5921b630a862672300b871a862 Mon Sep 17 00:00:00 2001 From: Adam Sutton Date: Fri, 16 Oct 2026 20:53:31 +0100 Subject: [PATCH 1/8] fixing tokenization of gold label entities and adding detected name to stats --- medcat/medcat/stats/stats.py | 4 +++- medcat/medcat/tokenizing/spacy_impl/tokens.py | 2 +- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/medcat/medcat/stats/stats.py b/medcat/medcat/stats/stats.py index ae75d4d42..7419dcb9d 100644 --- a/medcat/medcat/stats/stats.py +++ b/medcat/medcat/stats/stats.py @@ -299,7 +299,9 @@ def _extract_predictions( 'start': ent.base.start_char_index, 'end': ent.base.end_char_index, 'cui': ent.cui, - 'text': ent.base.text, + 'text': doc['text'][max(0, ent.base.start_char_index): + ent.base.end_char_index], + 'detected_name': ent.detected_name, 'context': doc['text'][max(0, ent.base.start_char_index - 60): ent.base.end_char_index + 60], 'document_id': doc['id'], diff --git a/medcat/medcat/tokenizing/spacy_impl/tokens.py b/medcat/medcat/tokenizing/spacy_impl/tokens.py index fb6595d2d..aa3e90617 100644 --- a/medcat/medcat/tokenizing/spacy_impl/tokens.py +++ b/medcat/medcat/tokenizing/spacy_impl/tokens.py @@ -236,7 +236,7 @@ def get_tokens(self, start_index: int, end_index: int ) -> list[MutableToken]: char_indices = self._ensure_char_indices() lo = bisect_left(char_indices, start_index) - hi = bisect_right(char_indices, end_index) + hi = bisect_left(char_indices, end_index) return [Token(self._delegate[i]) for i in range(lo, hi)] def set_addon_data(self, path: str, val: Any) -> None: From eb4562e119707fb6e832acab1cc59c1449b208b1 Mon Sep 17 00:00:00 2001 From: Adam Sutton Date: Fri, 16 Oct 2026 21:09:14 +0100 Subject: [PATCH 2/8] mypy! --- medcat/medcat/stats/stats.py | 1 + 1 file changed, 1 insertion(+) diff --git a/medcat/medcat/stats/stats.py b/medcat/medcat/stats/stats.py index 7419dcb9d..e40e96837 100644 --- a/medcat/medcat/stats/stats.py +++ b/medcat/medcat/stats/stats.py @@ -135,6 +135,7 @@ class PredictedAnnotation(TypedDict): end: int cui: str text: str + detected_name: str confidence: float raw: MutableEntity no_tokens: int From 2315bf97498cf0cecea565362e2e36dce193ffc5 Mon Sep 17 00:00:00 2001 From: Adam Sutton Date: Fri, 16 Oct 2026 21:29:19 +0100 Subject: [PATCH 3/8] mypy fix --- medcat/medcat/tokenizing/spacy_impl/tokens.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/medcat/medcat/tokenizing/spacy_impl/tokens.py b/medcat/medcat/tokenizing/spacy_impl/tokens.py index aa3e90617..e7afc0953 100644 --- a/medcat/medcat/tokenizing/spacy_impl/tokens.py +++ b/medcat/medcat/tokenizing/spacy_impl/tokens.py @@ -1,5 +1,5 @@ from typing import Iterator, Union, Optional, overload, cast, Any -from bisect import bisect_left, bisect_right +from bisect import bisect_left import logging from spacy.tokens import Token as SpacyToken From 12d291eb28787d7e27d3be758652c791bdb59349 Mon Sep 17 00:00:00 2001 From: Adam Sutton Date: Mon, 19 Oct 2026 11:35:50 +0100 Subject: [PATCH 4/8] fixed testing --- medcat/tests/utils/test_training_utils.py | 1 + 1 file changed, 1 insertion(+) diff --git a/medcat/tests/utils/test_training_utils.py b/medcat/tests/utils/test_training_utils.py index 1c02e4dcf..289646243 100644 --- a/medcat/tests/utils/test_training_utils.py +++ b/medcat/tests/utils/test_training_utils.py @@ -27,6 +27,7 @@ class _FakeEntity: def __init__(self, start: int, end: int, text: str, cui: str): self.base = _FakeEntityBase(start, end, text) self.cui = cui + self.detected_name = text self.context_similarity = 1.0 self.id = 0 From 2b253910f72c007e475c14807b5d952190ba7390 Mon Sep 17 00:00:00 2001 From: Adam Sutton Date: Mon, 19 Oct 2026 12:21:12 +0100 Subject: [PATCH 5/8] fixed test set up --- medcat/tests/stats/test_stats.py | 1 + 1 file changed, 1 insertion(+) diff --git a/medcat/tests/stats/test_stats.py b/medcat/tests/stats/test_stats.py index 4f42005f1..d46daf778 100644 --- a/medcat/tests/stats/test_stats.py +++ b/medcat/tests/stats/test_stats.py @@ -54,6 +54,7 @@ class DummyEntity: base: SimpleNamespace = field(init=False) def __post_init__(self): + self.detected_name = self.text self.base = SimpleNamespace( text=self.text, start_char_index=self.start_char_index, From 1136d159943145ff2383f671dba8e3badd51c45a Mon Sep 17 00:00:00 2001 From: Adam Sutton Date: Thu, 22 Oct 2026 21:03:29 +0100 Subject: [PATCH 6/8] changed from AutoModelForTokenClassification to AutoModel for consistency --- .../transformer_ner_model.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) 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..228dbb9c1 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 @@ -32,10 +32,11 @@ def __init__( self.num_labels = num_labels self.aux_loss_weight = aux_loss_weight self.id2label = id2label - self.language_model = AutoModelForTokenClassification.from_pretrained( + 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 @@ -43,6 +44,9 @@ def __init__( # 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 + + # Main token classification head: CRF emissions + self.classifier = nn.Linear(hidden_size, self.num_labels) self.start_head = nn.Linear(hidden_size, 1) self.end_head = nn.Linear(hidden_size, 1) @@ -86,9 +90,9 @@ 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] + hidden_states = outputs.last_hidden_state + # Main token-classification logits, used as CRF emissions + emissions = self.classifier(hidden_states) # Linear classifiers for boundary heads start_logits = self.start_head(hidden_states).squeeze(-1) # [B, T] From 4adcf35e50e84143f4ff3f8feddea4ec645483c4 Mon Sep 17 00:00:00 2001 From: Adam Sutton Date: Mon, 26 Oct 2026 21:24:35 +0000 Subject: [PATCH 7/8] changed loss weighting, exposed training params --- .../src/medcat_transformer_ner/config.py | 23 +++- .../medcat_transformer_ner/transformer_ner.py | 68 ++++++++- .../transformer_ner_model.py | 129 +++++++----------- 3 files changed, 133 insertions(+), 87 deletions(-) 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 228dbb9c1..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 @@ -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,53 +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.class_weights_dict = dict(class_weights or {}) + self.crf_loss_weight = crf_loss_weight self.language_model = AutoModel.from_pretrained( embedding_model_name, 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 # Main token classification head: CRF emissions self.classifier = nn.Linear(hidden_size, self.num_labels) - self.start_head = nn.Linear(hidden_size, 1) - self.end_head = nn.Linear(hidden_size, 1) # 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: @@ -93,16 +89,11 @@ def forward(self, **inputs) -> Any: hidden_states = outputs.last_hidden_state # Main token-classification logits, used as CRF emissions emissions = self.classifier(hidden_states) - - # 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] # 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() @@ -129,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 @@ -163,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 @@ -250,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) @@ -278,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 From fba33cb6a12465c02271cd05397c6d3262bf7997 Mon Sep 17 00:00:00 2001 From: Adam Sutton Date: Wed, 28 Oct 2026 13:28:08 +0000 Subject: [PATCH 8/8] removed commented out old call --- .../src/medcat_transformer_ner/transformer_ner.py | 1 - 1 file changed, 1 deletion(-) 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 a135bd752..37071de89 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 @@ -617,7 +617,6 @@ def deserialise_from( f"{model_folder}. Expected both config.json and pytorch_model.bin." ) - # ner.model = AutoModelForTokenClassification.from_pretrained(model_folder) ner.model = ModelForBinaryNER.from_pretrained( model_folder, device=ner.device,