from itertools import islice from typing import Dict, List, Optional, Tuple, Union import torch import torch.nn as nn from torchcrf import CRF from transformers import PretrainedConfig, PreTrainedModel from transformers.modeling_outputs import TokenClassifierOutput # Large negative number for masking impossible transitions LARGE_NEGATIVE_NUMBER = -1e9 NUM_PER_LAYER = 16 class MultiHeadCRFConfig(PretrainedConfig): """ Configuration class for Multi-Head CRF models. Args: entity_types: List of entity type names (e.g., ["DRUG", "DISEASE", "SYMPTOM"]) number_of_layers_per_head: Number of dense layers per head before classification crf_reduction: Reduction mode for CRF loss ("mean", "sum", "token_mean", "none") freeze_backbone: Whether to freeze the transformer backbone num_frozen_encoders: Number of encoder layers to freeze (from bottom) classifier_dropout: Dropout rate for classifier heads **kwargs: Additional arguments passed to PretrainedConfig """ model_type = "multihead-crf-tagger" def __init__( self, entity_types: Optional[List[str]] = None, number_of_layers_per_head: int = 1, crf_reduction: str = "mean", freeze_backbone: bool = False, num_frozen_encoders: int = 0, classifier_dropout: float = 0.1, classifier_hidden_layers: Optional[Tuple] = None, class_weights: Optional[List[float]] = None, backbone_model_name: Optional[str] = None, **kwargs, ): self.entity_types = entity_types or [] self.number_of_layers_per_head = number_of_layers_per_head self.crf_reduction = crf_reduction self.freeze_backbone = freeze_backbone self.num_frozen_encoders = num_frozen_encoders self.classifier_dropout = classifier_dropout self.classifier_hidden_layers = classifier_hidden_layers self.class_weights = class_weights self.backbone_model_name = backbone_model_name super().__init__(**kwargs) class MultiHeadCRF(nn.Module): """ Custom CRF implementation with BIO transition masking. This CRF implementation includes: - Proper initialization of transition parameters - Masking of impossible BIO transitions (e.g., O -> I is invalid) - Viterbi decoding for inference Args: num_tags: Number of tags (typically 3 for BIO: O, B, I) batch_first: Whether batch dimension is first """ def __init__(self, num_tags: int, batch_first: bool = True) -> None: if num_tags <= 0: raise ValueError(f"invalid number of tags: {num_tags}") super().__init__() self.num_tags = num_tags self.batch_first = batch_first self.start_transitions = nn.Parameter(torch.empty(num_tags)) self.end_transitions = nn.Parameter(torch.empty(num_tags)) self.transitions = nn.Parameter(torch.empty(num_tags, num_tags)) self.reset_parameters() self.mask_impossible_transitions() def reset_parameters(self) -> None: """Initialize the transition parameters uniformly between -0.1 and 0.1.""" nn.init.uniform_(self.start_transitions, -0.1, 0.1) nn.init.uniform_(self.end_transitions, -0.1, 0.1) nn.init.uniform_(self.transitions, -0.1, 0.1) def mask_impossible_transitions(self) -> None: """ Set impossible BIO transitions to large negative values. For standard BIO tagging with tags [O=0, B=1, I=2]: - Cannot start with I tag - Cannot transition from O to I """ with torch.no_grad(): # Assuming BIO scheme: O=0, B=1, I=2 # Cannot start with I if self.num_tags > 2: self.start_transitions[2] = LARGE_NEGATIVE_NUMBER # Cannot go from O to I self.transitions[0][2] = LARGE_NEGATIVE_NUMBER # If PADDING token exists (index 3+), mask its transitions if self.num_tags > 3: # Cannot start with PADDING self.start_transitions[3] = LARGE_NEGATIVE_NUMBER # Cannot transition to PADDING from valid tags for i in range(3): self.transitions[i][3] = LARGE_NEGATIVE_NUMBER # Cannot transition from PADDING to valid tags for i in range(3): self.transitions[3][i] = LARGE_NEGATIVE_NUMBER def __repr__(self) -> str: return f"{self.__class__.__name__}(num_tags={self.num_tags})" def forward( self, emissions: torch.Tensor, tags: torch.Tensor, mask: Optional[torch.Tensor] = None, reduction: str = "mean", ) -> torch.Tensor: """ Compute the negative log likelihood of a sequence of tags given emission scores. Args: emissions: Emission scores (batch_size, seq_length, num_tags) if batch_first tags: Gold tag sequence (batch_size, seq_length) if batch_first mask: Mask tensor (batch_size, seq_length) if batch_first reduction: Loss reduction mode ("none", "sum", "mean", "token_mean") Returns: Negative log likelihood loss """ self._validate(emissions, tags=tags, mask=mask) if reduction not in ("none", "sum", "mean", "token_mean"): raise ValueError(f"invalid reduction: {reduction}") if mask is None: mask = torch.ones_like(tags, dtype=torch.uint8) # Ensure all tensors are on the same device as emissions device = emissions.device tags = tags.to(device) mask = mask.to(device) if self.batch_first: emissions = emissions.transpose(0, 1) tags = tags.transpose(0, 1) mask = mask.transpose(0, 1) # shape: (batch_size,) numerator = self._compute_score(emissions, tags, mask) # shape: (batch_size,) denominator = self._compute_normalizer(emissions, mask) # shape: (batch_size,) llh = numerator - denominator nllh = -llh if reduction == "none": return nllh if reduction == "sum": return nllh.sum() if reduction == "mean": return nllh.mean() assert reduction == "token_mean" return nllh.sum() / mask.type_as(emissions).sum() def decode( self, emissions: torch.Tensor, mask: Optional[torch.Tensor] = None ) -> List[List[int]]: """ Find the most likely tag sequence using Viterbi algorithm. Args: emissions: Emission scores mask: Mask tensor Returns: List of best tag sequences for each batch """ self._validate(emissions, mask=mask) if mask is None: mask = emissions.new_ones(emissions.shape[:2], dtype=torch.uint8) if self.batch_first: emissions = emissions.transpose(0, 1) mask = mask.transpose(0, 1) return self._viterbi_decode(emissions, mask) def _validate( self, emissions: torch.Tensor, tags: Optional[torch.Tensor] = None, mask: Optional[torch.Tensor] = None, ) -> None: if emissions.dim() != 3: raise ValueError( f"emissions must have dimension of 3, got {emissions.dim()}" ) if emissions.size(2) != self.num_tags: raise ValueError( f"expected last dimension of emissions is {self.num_tags}, " f"got {emissions.size(2)}" ) if tags is not None: if emissions.shape[:2] != tags.shape: raise ValueError( "the first two dimensions of emissions and tags must match, " f"got {tuple(emissions.shape[:2])} and {tuple(tags.shape)}" ) if mask is not None: if emissions.shape[:2] != mask.shape: raise ValueError( "the first two dimensions of emissions and mask must match, " f"got {tuple(emissions.shape[:2])} and {tuple(mask.shape)}" ) no_empty_seq = not self.batch_first and mask[0].all() no_empty_seq_bf = self.batch_first and mask[:, 0].all() if not no_empty_seq and not no_empty_seq_bf: raise ValueError("mask of the first timestep must all be on") def _compute_score( self, emissions: torch.Tensor, tags: torch.Tensor, mask: torch.Tensor ) -> torch.Tensor: # emissions: (seq_length, batch_size, num_tags) # tags: (seq_length, batch_size) # mask: (seq_length, batch_size) assert emissions.dim() == 3 and tags.dim() == 2 assert emissions.shape[:2] == tags.shape assert emissions.size(2) == self.num_tags assert mask.shape == tags.shape assert mask[0].all() # Move all tensors to the same device as emissions device = emissions.device tags = tags.to(device) mask = mask.to(device) seq_length, batch_size = tags.shape mask = mask.type_as(emissions) # Start transition score and first emission # Ensure arange is on the same device as other tensors batch_indices = torch.arange(batch_size, device=device) score = self.start_transitions[tags[0]] score += emissions[0, batch_indices, tags[0]] for i in range(1, seq_length): score += self.transitions[tags[i - 1], tags[i]] * mask[i] score += emissions[i, batch_indices, tags[i]] * mask[i] # End transition score seq_ends = mask.long().sum(dim=0) - 1 last_tags = tags[seq_ends, batch_indices] score += self.end_transitions[last_tags] return score def _compute_normalizer( self, emissions: torch.Tensor, mask: torch.Tensor ) -> torch.Tensor: # emissions: (seq_length, batch_size, num_tags) # mask: (seq_length, batch_size) assert emissions.dim() == 3 and mask.dim() == 2 assert emissions.shape[:2] == mask.shape assert emissions.size(2) == self.num_tags assert mask[0].all() seq_length = emissions.size(0) # Start transition score and first emission score = self.start_transitions + emissions[0] for i in range(1, seq_length): broadcast_score = score.unsqueeze(2) broadcast_emissions = emissions[i].unsqueeze(1) next_score = broadcast_score + self.transitions + broadcast_emissions next_score = torch.logsumexp(next_score, dim=1) score = torch.where(mask[i].unsqueeze(1).bool(), next_score, score) score += self.end_transitions return torch.logsumexp(score, dim=1) def _viterbi_decode( self, emissions: torch.Tensor, mask: torch.Tensor ) -> List[List[int]]: # emissions: (seq_length, batch_size, num_tags) # mask: (seq_length, batch_size) assert emissions.dim() == 3 and mask.dim() == 2 assert emissions.shape[:2] == mask.shape assert emissions.size(2) == self.num_tags assert mask[0].all() seq_length, batch_size = mask.shape # Start transition and first emission score = self.start_transitions + emissions[0] history = [] for i in range(1, seq_length): broadcast_score = score.unsqueeze(2) broadcast_emission = emissions[i].unsqueeze(1) next_score = broadcast_score + self.transitions + broadcast_emission next_score, indices = next_score.max(dim=1) score = torch.where(mask[i].unsqueeze(1).bool(), next_score, score) history.append(indices) score += self.end_transitions # Trace back seq_ends = mask.long().sum(dim=0) - 1 best_tags_list = [] for idx in range(batch_size): _, best_last_tag = score[idx].max(dim=0) best_tags = [best_last_tag.item()] for hist in reversed(history[: seq_ends[idx]]): best_last_tag = hist[idx][best_tags[-1]] best_tags.append(best_last_tag.item()) best_tags.reverse() best_tags_list.append(best_tags) return best_tags_list class TokenClassificationModelCRF(PreTrainedModel): """ Custom token classification model with CRF layer and configurable classifier head. This model can be loaded with trust_remote_code=True for HuggingFace Hub compatibility. """ def __init__( self, config, base_model=None, freeze_backbone=False, classifier_hidden_layers=None, classifier_dropout=0.1, ): super().__init__(config) self.config = config self.num_labels = config.num_labels # If base_model is not provided, load it from config if base_model is None: from transformers import AutoConfig, RobertaForTokenClassification # Use backbone_model_name if available, fallback to name_or_path # This is critical because name_or_path gets overwritten during save/load backbone_name = getattr(config, "backbone_model_name", None) if backbone_name is None: backbone_name = getattr(config, "name_or_path", None) or getattr( config, "_name_or_path", None ) if backbone_name is None: raise ValueError( "config.backbone_model_name (or config.name_or_path) is required to load pretrained backbone" ) # Create a clean config for the backbone backbone_config = AutoConfig.from_pretrained(backbone_name) backbone_config.hidden_dropout_prob = getattr( config, "hidden_dropout_prob", 0.1 ) backbone_config.num_labels = config.num_labels roberta_model = RobertaForTokenClassification.from_pretrained( backbone_name, config=backbone_config ) self.roberta = roberta_model.roberta # Store backbone_model_name in config for future loading if ( not hasattr(config, "backbone_model_name") or config.backbone_model_name is None ): config.backbone_model_name = backbone_name else: if hasattr(base_model, "roberta"): self.roberta = base_model.roberta else: self.roberta = base_model self.lm_output_size = self.roberta.config.hidden_size # Store configuration for saving/loading self.config.freeze_backbone = freeze_backbone self.config.classifier_hidden_layers = classifier_hidden_layers self.config.classifier_dropout = classifier_dropout if freeze_backbone: print("+" * 30, "\n\n", "Freezing backbone...", "+" * 30, "\n\n") for param in self.roberta.parameters(): param.requires_grad = False self.roberta.eval() else: print("+" * 30, "\n\n", "NOT Freezing backbone...", "+" * 30, "\n\n") self.roberta.train(not freeze_backbone) self.dropout = nn.Dropout( config.hidden_dropout_prob if hasattr(config, "hidden_dropout_prob") else 0.1 ) self.crf = CRF(self.num_labels, batch_first=True) self._build_classifier_head(classifier_hidden_layers, classifier_dropout) def _build_classifier_head(self, hidden_layers, dropout_rate): """ Build a flexible classifier head with configurable hidden layers and dropout. Args: hidden_layers: Tuple of integers representing the number of neurons in each hidden layer. None or empty tuple means a simple linear layer. dropout_rate: Dropout probability between layers """ layers = [] input_size = self.lm_output_size # If hidden_layers is None or empty, just create a simple linear layer if not hidden_layers: self.classifier = nn.Sequential( nn.Dropout(dropout_rate), nn.Linear(input_size, self.num_labels) ) return # Build MLP with specified hidden layers for hidden_size in hidden_layers: layers.append(nn.Linear(input_size, hidden_size)) layers.append(nn.ReLU()) layers.append(nn.Dropout(dropout_rate)) input_size = hidden_size # Final classification layer layers.append(nn.Linear(input_size, self.num_labels)) # Create sequential model self.classifier = nn.Sequential(*layers) def forward( self, input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.FloatTensor] = None, token_type_ids: Optional[torch.LongTensor] = None, position_ids: Optional[torch.LongTensor] = None, head_mask: Optional[torch.FloatTensor] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, **kwargs, ) -> Union[Tuple[torch.Tensor], TokenClassifierOutput]: return_dict = ( return_dict if return_dict is not None else self.config.use_return_dict ) outputs = self.roberta( input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids, position_ids=position_ids, head_mask=head_mask, inputs_embeds=inputs_embeds, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, ) sequence_output = self.dropout(outputs.last_hidden_state) logits = self.classifier(sequence_output) # Emissions for CRF loss = None if labels is not None: # CRF calculates the log-likelihood of the correct sequence # We use a negative sign to convert it into a loss # Following ieeta-pt approach: don't pass mask to CRF # All positions have valid labels (O for special/padding tokens) labels_long = labels.long() loss = -self.crf(logits, labels_long, reduction="mean") if not return_dict: output = (logits,) + outputs[2:] return ((loss,) + output) if loss is not None else output return TokenClassifierOutput( loss=loss, logits=logits, hidden_states=outputs.hidden_states, attentions=outputs.attentions, ) @property def device_info(self): return next(self.parameters()).device def get_input_embeddings(self): return self.roberta.get_input_embeddings() def set_input_embeddings(self, value): self.roberta.set_input_embeddings(value) @classmethod def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs): """Override from_pretrained to handle custom model loading""" config = kwargs.pop("config", None) if config is None: from transformers import AutoConfig config = AutoConfig.from_pretrained(pretrained_model_name_or_path, **kwargs) # Extract custom parameters from config if they exist freeze_backbone = getattr(config, "freeze_backbone", False) classifier_hidden_layers = getattr(config, "classifier_hidden_layers", None) classifier_dropout = getattr(config, "classifier_dropout", 0.1) model = cls( config=config, freeze_backbone=freeze_backbone, classifier_hidden_layers=classifier_hidden_layers, classifier_dropout=classifier_dropout, ) # Load state dict if available try: state_dict = torch.load( f"{pretrained_model_name_or_path}/pytorch_model.bin", map_location="cpu" ) model.load_state_dict(state_dict) except: # If loading fails, the model will be initialized with random weights print( "Warning: Could not load pre-trained weights. Using randomly initialized model." ) return model class TokenClassificationModelMultiHeadCRF(PreTrainedModel): """ Multi-Head CRF model for token classification with multiple entity types. Each entity type gets its own classification head and CRF layer, allowing for independent BIO tagging per entity type. This is useful for scenarios where entities can overlap or when different entity types have different transition patterns. Args: config: MultiHeadCRFConfig or compatible config with entity_types base_model: Optional pre-trained RoBERTa model freeze_backbone: Whether to freeze transformer weights """ config_class = MultiHeadCRFConfig base_model_prefix = "roberta" _keys_to_ignore_on_load_unexpected = [r"pooler"] def __init__(self, config, base_model=None, freeze_backbone=None): super().__init__(config) self.config = config # Get entity types from config self.entity_types = getattr(config, "entity_types", []) if not self.entity_types: raise ValueError("entity_types must be provided in config") # Number of labels per head (typically 3 for BIO: O, B, I) + padding self.num_labels = config.num_labels # self.num_labels_with_pad = self.num_labels + 1 # should be self.num_labels? # Configuration parameters self.number_of_layers_per_head = getattr(config, "number_of_layers_per_head", 1) self.crf_reduction = getattr(config, "crf_reduction", "mean") freeze_backbone = ( freeze_backbone if freeze_backbone is not None else getattr(config, "freeze_backbone", False) ) self.num_frozen_encoders = getattr(config, "num_frozen_encoders", 0) classifier_dropout = getattr(config, "classifier_dropout", 0.1) # Initialize the transformer backbone if base_model is None: from transformers import AutoConfig, RobertaModel # Use backbone_model_name if available, fallback to name_or_path backbone_name = getattr(config, "backbone_model_name", None) if backbone_name is None: backbone_name = getattr(config, "name_or_path", None) or getattr( config, "_name_or_path", None ) if backbone_name: # Load pretrained weights backbone_config = AutoConfig.from_pretrained(backbone_name) backbone_config.hidden_dropout_prob = getattr( config, "hidden_dropout_prob", 0.1 ) self.roberta = RobertaModel.from_pretrained( backbone_name, config=backbone_config, add_pooling_layer=False ) # Store backbone_model_name in config for future loading if ( not hasattr(config, "backbone_model_name") or config.backbone_model_name is None ): config.backbone_model_name = backbone_name else: # Fallback: initialize without pretrained weights (not recommended) self.roberta = RobertaModel(config, add_pooling_layer=False) else: if hasattr(base_model, "roberta"): self.roberta = base_model.roberta else: self.roberta = base_model self.hidden_size = config.hidden_size self.dropout = nn.Dropout(getattr(config, "hidden_dropout_prob", 0.1)) # Create heads for each entity type print(f"Creating Multi-Head CRF with entity types: {sorted(self.entity_types)}") for entity_type in self.entity_types: # Dense layers per head for i in range(self.number_of_layers_per_head): setattr( self, f"{entity_type}_dense_{i}", nn.Linear(self.hidden_size, self.hidden_size), ) setattr( self, f"{entity_type}_dense_activation_{i}", nn.GELU(approximate="none"), ) setattr( self, f"{entity_type}_dropout_{i}", nn.Dropout(classifier_dropout) ) # Classifier and CRF per head setattr( self, f"{entity_type}_classifier", nn.Linear(self.hidden_size, self.num_labels), ) setattr( self, f"{entity_type}_crf", MultiHeadCRF(num_tags=self.num_labels, batch_first=True), ) # Handle freezing if freeze_backbone: self._freeze_backbone() def _freeze_backbone(self): """Freeze transformer backbone parameters.""" print("+" * 30, "\n\n", "Freezing backbone...", "+" * 30, "\n\n") # Freeze embeddings for param in self.roberta.embeddings.parameters(): param.requires_grad = False # Optionally freeze some encoder layers if self.num_frozen_encoders > 0: for _, param in islice( self.roberta.encoder.named_parameters(), self.num_frozen_encoders * NUM_PER_LAYER, ): param.requires_grad = False def reset_head_parameters(self): """Reset parameters for all heads (useful after loading pretrained weights).""" for entity_type in self.entity_types: for i in range(self.number_of_layers_per_head): getattr(self, f"{entity_type}_dense_{i}").reset_parameters() getattr(self, f"{entity_type}_classifier").reset_parameters() getattr(self, f"{entity_type}_crf").reset_parameters() getattr(self, f"{entity_type}_crf").mask_impossible_transitions() def forward( self, input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.FloatTensor] = None, token_type_ids: Optional[torch.LongTensor] = None, position_ids: Optional[torch.LongTensor] = None, head_mask: Optional[torch.FloatTensor] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[Dict[str, torch.LongTensor]] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, **kwargs, ): """ Forward pass through the multi-head CRF model. Args: input_ids: Input token IDs attention_mask: Attention mask labels: Dictionary mapping entity types to label tensors e.g., {"DRUG": tensor, "DISEASE": tensor} ... other standard transformer arguments Returns: During training (labels provided): Tuple of (total_loss, logits_dict) where logits_dict maps entity types to logits During inference (no labels): List of prediction tensors, one per entity type (sorted alphabetically) """ return_dict = ( return_dict if return_dict is not None else self.config.use_return_dict ) # Get transformer outputs outputs = self.roberta( input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids, position_ids=position_ids, head_mask=head_mask, inputs_embeds=inputs_embeds, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, ) sequence_output = outputs[0] sequence_output = self.dropout(sequence_output) # (batch, seq_len, hidden) # Compute logits for each head logits = {} for entity_type in self.entity_types: head_output = sequence_output for i in range(self.number_of_layers_per_head): head_output = getattr(self, f"{entity_type}_dense_{i}")(head_output) head_output = getattr(self, f"{entity_type}_dense_activation_{i}")( head_output ) head_output = getattr(self, f"{entity_type}_dropout_{i}")(head_output) logits[entity_type] = getattr(self, f"{entity_type}_classifier")( head_output ) if labels is not None: # Training mode - compute CRF loss for each head # Following ieeta-pt approach: don't pass mask to CRF # All positions have valid labels (O for special/padding tokens) losses = {} for entity_type in self.entity_types: if entity_type in labels: # Ensure labels are on the same device as logits entity_labels = ( labels[entity_type].long().to(logits[entity_type].device) ) crf = getattr(self, f"{entity_type}_crf") # CRF returns negative log likelihood, we want to minimize it losses[entity_type] = crf( logits[entity_type], entity_labels, reduction=self.crf_reduction, ) # Sum losses from all heads total_loss = sum(losses.values()) return total_loss, logits else: # Inference mode - decode each head # Following ieeta-pt approach: don't pass mask, decode all positions predictions = {} for entity_type in self.entity_types: crf = getattr(self, f"{entity_type}_crf") decoded = crf.decode(logits[entity_type]) predictions[entity_type] = torch.tensor(decoded) # Return as list sorted by entity type for consistency return [predictions[ent] for ent in sorted(self.entity_types)] def get_input_embeddings(self): return self.roberta.get_input_embeddings() def set_input_embeddings(self, value): self.roberta.set_input_embeddings(value) @classmethod def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs): """Override from_pretrained to handle custom model loading.""" import json import os config = kwargs.pop("config", None) if config is None: # Load config directly from JSON to get all saved attributes config_file = os.path.join(pretrained_model_name_or_path, "config.json") if os.path.exists(config_file): with open(config_file, "r") as f: config_dict = json.load(f) # Create MultiHeadCRFConfig with all loaded parameters config = MultiHeadCRFConfig(**config_dict) else: from transformers import AutoConfig config = AutoConfig.from_pretrained( pretrained_model_name_or_path, trust_remote_code=kwargs.get("trust_remote_code", True), ) # Ensure config has all required RoBERTa parameters # These are needed to initialize RobertaModel roberta_defaults = { # Core model architecture "layer_norm_eps": 1e-5, "hidden_size": 768, "num_hidden_layers": 12, "num_attention_heads": 12, "intermediate_size": 3072, "hidden_act": "gelu", "hidden_dropout_prob": 0.1, "attention_probs_dropout_prob": 0.1, "max_position_embeddings": 514, "type_vocab_size": 1, "initializer_range": 0.02, "vocab_size": 52000, # Token IDs "pad_token_id": 1, "bos_token_id": 0, "eos_token_id": 2, # Position embeddings "position_embedding_type": "absolute", # Model behavior flags "use_cache": True, "is_decoder": False, "add_cross_attention": False, "chunk_size_feed_forward": 0, "output_hidden_states": False, "output_attentions": False, "torchscript": False, "tie_word_embeddings": True, "return_dict": True, # Gradient checkpointing "gradient_checkpointing": False, # Pruning "pruned_heads": {}, # Problem type (for classification) "problem_type": None, # Embedding layer norm "embedding_size": None, } for key, default_value in roberta_defaults.items(): if not hasattr(config, key) or getattr(config, key) is None: setattr(config, key, default_value) freeze_backbone = getattr(config, "freeze_backbone", False) model = cls(config=config, freeze_backbone=freeze_backbone) # Load state dict if available weight_file = os.path.join(pretrained_model_name_or_path, "pytorch_model.bin") safetensors_file = os.path.join( pretrained_model_name_or_path, "model.safetensors" ) try: if os.path.exists(safetensors_file): from safetensors.torch import load_file state_dict = load_file(safetensors_file) model.load_state_dict(state_dict) elif os.path.exists(weight_file): state_dict = torch.load(weight_file, map_location="cpu") model.load_state_dict(state_dict) else: print( "Warning: No pre-trained weights found. Using randomly initialized model." ) except Exception as e: print(f"Warning: Could not load pre-trained weights: {e}") return model class MultiHeadConfig(PretrainedConfig): """ Configuration class for Multi-Head models (without CRF). Args: entity_types: List of entity type names (e.g., ["DRUG", "DISEASE", "SYMPTOM"]) number_of_layers_per_head: Number of dense layers per head before classification freeze_backbone: Whether to freeze the transformer backbone num_frozen_encoders: Number of encoder layers to freeze (from bottom) classifier_dropout: Dropout rate for classifier heads use_class_weights: Whether to use class weights for loss computation class_weights: Optional dict mapping entity types to weight lists **kwargs: Additional arguments passed to PretrainedConfig """ model_type = "multihead-tagger" def __init__( self, entity_types: Optional[List[str]] = None, number_of_layers_per_head: int = 1, freeze_backbone: bool = False, num_frozen_encoders: int = 0, classifier_dropout: float = 0.1, use_class_weights: bool = False, class_weights: Optional[Dict[str, List[float]]] = None, backbone_model_name: Optional[str] = None, **kwargs, ): self.entity_types = entity_types or [] self.number_of_layers_per_head = number_of_layers_per_head self.freeze_backbone = freeze_backbone self.num_frozen_encoders = num_frozen_encoders self.classifier_dropout = classifier_dropout self.use_class_weights = use_class_weights self.class_weights = class_weights self.backbone_model_name = backbone_model_name super().__init__(**kwargs) class TokenClassificationModelMultiHead(PreTrainedModel): """ Multi-Head model for token classification with multiple entity types (no CRF). Each entity type gets its own classification head, allowing for independent BIO tagging per entity type. This is useful for scenarios where entities can overlap or when different entity types need separate classification. Unlike the CRF variant, this model uses standard CrossEntropyLoss and argmax decoding, which is faster but doesn't enforce valid BIO sequences. Args: config: MultiHeadConfig or compatible config with entity_types base_model: Optional pre-trained RoBERTa model freeze_backbone: Whether to freeze transformer weights """ config_class = MultiHeadConfig base_model_prefix = "roberta" _keys_to_ignore_on_load_unexpected = [r"pooler"] def __init__(self, config, base_model=None, freeze_backbone=None): super().__init__(config) self.config = config # Get entity types from config self.entity_types = getattr(config, "entity_types", []) if not self.entity_types: raise ValueError("entity_types must be provided in config") # Number of labels per head (typically 3 for BIO: O, B, I) self.num_labels = config.num_labels # Configuration parameters self.number_of_layers_per_head = getattr(config, "number_of_layers_per_head", 1) freeze_backbone = ( freeze_backbone if freeze_backbone is not None else getattr(config, "freeze_backbone", False) ) self.num_frozen_encoders = getattr(config, "num_frozen_encoders", 0) classifier_dropout = getattr(config, "classifier_dropout", 0.1) # Class weights for loss computation self.use_class_weights = getattr(config, "use_class_weights", False) self.class_weights = getattr(config, "class_weights", None) # Initialize the transformer backbone if base_model is None: from transformers import AutoConfig, RobertaModel # Use backbone_model_name if available, fallback to name_or_path backbone_name = getattr(config, "backbone_model_name", None) if backbone_name is None: backbone_name = getattr(config, "name_or_path", None) or getattr( config, "_name_or_path", None ) if backbone_name: # Load pretrained weights backbone_config = AutoConfig.from_pretrained(backbone_name) backbone_config.hidden_dropout_prob = getattr( config, "hidden_dropout_prob", 0.1 ) self.roberta = RobertaModel.from_pretrained( backbone_name, config=backbone_config, add_pooling_layer=False ) # Store backbone_model_name in config for future loading if ( not hasattr(config, "backbone_model_name") or config.backbone_model_name is None ): config.backbone_model_name = backbone_name else: # Fallback: initialize without pretrained weights (not recommended) self.roberta = RobertaModel(config, add_pooling_layer=False) else: if hasattr(base_model, "roberta"): self.roberta = base_model.roberta else: self.roberta = base_model self.hidden_size = config.hidden_size self.dropout = nn.Dropout(getattr(config, "hidden_dropout_prob", 0.1)) # Create heads for each entity type print( f"Creating Multi-Head model with entity types: {sorted(self.entity_types)}" ) for entity_type in self.entity_types: # Dense layers per head for i in range(self.number_of_layers_per_head): setattr( self, f"{entity_type}_dense_{i}", nn.Linear(self.hidden_size, self.hidden_size), ) setattr( self, f"{entity_type}_dense_activation_{i}", nn.GELU(approximate="none"), ) setattr( self, f"{entity_type}_dropout_{i}", nn.Dropout(classifier_dropout) ) # Classifier per head (no CRF) setattr( self, f"{entity_type}_classifier", nn.Linear(self.hidden_size, self.num_labels), ) # Set up loss functions per entity type (with optional class weights) self.loss_fns = nn.ModuleDict() for entity_type in self.entity_types: if ( self.use_class_weights and self.class_weights and entity_type in self.class_weights ): weight = torch.tensor( self.class_weights[entity_type], dtype=torch.float ) self.loss_fns[entity_type] = nn.CrossEntropyLoss( weight=weight, ignore_index=-100 ) else: self.loss_fns[entity_type] = nn.CrossEntropyLoss(ignore_index=-100) # Handle freezing if freeze_backbone: self._freeze_backbone() def _freeze_backbone(self): """Freeze transformer backbone parameters.""" print("+" * 30, "\n\n", "Freezing backbone...", "+" * 30, "\n\n") # Freeze embeddings for param in self.roberta.embeddings.parameters(): param.requires_grad = False # Optionally freeze some encoder layers if self.num_frozen_encoders > 0: for _, param in islice( self.roberta.encoder.named_parameters(), self.num_frozen_encoders * NUM_PER_LAYER, ): param.requires_grad = False def reset_head_parameters(self): """Reset parameters for all heads (useful after loading pretrained weights).""" for entity_type in self.entity_types: for i in range(self.number_of_layers_per_head): getattr(self, f"{entity_type}_dense_{i}").reset_parameters() getattr(self, f"{entity_type}_classifier").reset_parameters() def forward( self, input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.FloatTensor] = None, token_type_ids: Optional[torch.LongTensor] = None, position_ids: Optional[torch.LongTensor] = None, head_mask: Optional[torch.FloatTensor] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[Dict[str, torch.LongTensor]] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, **kwargs, ): """ Forward pass through the multi-head model. Args: input_ids: Input token IDs attention_mask: Attention mask labels: Dictionary mapping entity types to label tensors e.g., {"DRUG": tensor, "DISEASE": tensor} ... other standard transformer arguments Returns: During training (labels provided): Tuple of (total_loss, logits_dict) where logits_dict maps entity types to logits During inference (no labels): List of prediction tensors, one per entity type (sorted alphabetically) """ return_dict = ( return_dict if return_dict is not None else self.config.use_return_dict ) # Get transformer outputs outputs = self.roberta( input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids, position_ids=position_ids, head_mask=head_mask, inputs_embeds=inputs_embeds, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, ) sequence_output = outputs[0] sequence_output = self.dropout(sequence_output) # (batch, seq_len, hidden) # Compute logits for each head logits = {} for entity_type in self.entity_types: head_output = sequence_output for i in range(self.number_of_layers_per_head): head_output = getattr(self, f"{entity_type}_dense_{i}")(head_output) head_output = getattr(self, f"{entity_type}_dense_activation_{i}")( head_output ) head_output = getattr(self, f"{entity_type}_dropout_{i}")(head_output) logits[entity_type] = getattr(self, f"{entity_type}_classifier")( head_output ) if labels is not None: # Training mode - compute CrossEntropyLoss for each head losses = {} for entity_type in self.entity_types: if entity_type in labels: entity_labels = ( labels[entity_type].long().to(logits[entity_type].device) ) entity_logits = logits[entity_type] # Reshape for CrossEntropyLoss: (batch * seq_len, num_labels) and (batch * seq_len,) loss_fct = self.loss_fns[entity_type] # Move loss function weights to the same device if needed if hasattr(loss_fct, "weight") and loss_fct.weight is not None: loss_fct.weight = loss_fct.weight.to(entity_logits.device) losses[entity_type] = loss_fct( entity_logits.view(-1, self.num_labels), entity_labels.view(-1), ) # Sum losses from all heads total_loss = sum(losses.values()) return total_loss, logits else: # Inference mode - argmax decoding for each head predictions = {} for entity_type in self.entity_types: # Simple argmax decoding (no CRF constraints) preds = torch.argmax(logits[entity_type], dim=-1) predictions[entity_type] = preds # Return as list sorted by entity type for consistency return [predictions[ent] for ent in sorted(self.entity_types)] def get_input_embeddings(self): return self.roberta.get_input_embeddings() def set_input_embeddings(self, value): self.roberta.set_input_embeddings(value) @classmethod def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs): """Override from_pretrained to handle custom model loading.""" import json import os config = kwargs.pop("config", None) if config is None: # Load config directly from JSON to get all saved attributes config_file = os.path.join(pretrained_model_name_or_path, "config.json") if os.path.exists(config_file): with open(config_file, "r") as f: config_dict = json.load(f) # Create MultiHeadConfig with all loaded parameters config = MultiHeadConfig(**config_dict) else: from transformers import AutoConfig config = AutoConfig.from_pretrained( pretrained_model_name_or_path, trust_remote_code=kwargs.get("trust_remote_code", True), ) # Ensure config has all required RoBERTa parameters roberta_defaults = { "layer_norm_eps": 1e-5, "hidden_size": 768, "num_hidden_layers": 12, "num_attention_heads": 12, "intermediate_size": 3072, "hidden_act": "gelu", "hidden_dropout_prob": 0.1, "attention_probs_dropout_prob": 0.1, "max_position_embeddings": 514, "type_vocab_size": 1, "initializer_range": 0.02, "vocab_size": 52000, "pad_token_id": 1, "bos_token_id": 0, "eos_token_id": 2, "position_embedding_type": "absolute", "use_cache": True, "is_decoder": False, "add_cross_attention": False, "chunk_size_feed_forward": 0, "output_hidden_states": False, "output_attentions": False, "torchscript": False, "tie_word_embeddings": True, "return_dict": True, "gradient_checkpointing": False, "pruned_heads": {}, "problem_type": None, "embedding_size": None, } for key, default_value in roberta_defaults.items(): if not hasattr(config, key) or getattr(config, key) is None: setattr(config, key, default_value) freeze_backbone = getattr(config, "freeze_backbone", False) model = cls(config=config, freeze_backbone=freeze_backbone) # Load state dict if available weight_file = os.path.join(pretrained_model_name_or_path, "pytorch_model.bin") safetensors_file = os.path.join( pretrained_model_name_or_path, "model.safetensors" ) try: if os.path.exists(safetensors_file): from safetensors.torch import load_file state_dict = load_file(safetensors_file) model.load_state_dict(state_dict) elif os.path.exists(weight_file): state_dict = torch.load(weight_file, map_location="cpu") model.load_state_dict(state_dict) else: print( "Warning: No pre-trained weights found. Using randomly initialized model." ) except Exception as e: print(f"Warning: Could not load pre-trained weights: {e}") return model class TokenClassificationModel(PreTrainedModel): """ Custom token classification model with configurable classifier head (no CRF). This model can be loaded with trust_remote_code=True for HuggingFace Hub compatibility. """ def __init__(self, config): super().__init__(config) self.config = config self.num_labels = config.num_labels # Initialize the roberta backbone - load pretrained weights from transformers import AutoConfig, AutoModel # Use backbone_model_name if available, fallback to name_or_path # This is critical because name_or_path gets overwritten during save/load backbone_name = getattr(config, "backbone_model_name", None) if backbone_name is None: backbone_name = getattr(config, "name_or_path", None) or getattr( config, "_name_or_path", None ) if backbone_name is None: raise ValueError( "config.backbone_model_name (or config.name_or_path) is required to load pretrained backbone" ) # Create a clean config for the backbone backbone_config = AutoConfig.from_pretrained(backbone_name) backbone_config.hidden_dropout_prob = getattr( config, "hidden_dropout_prob", 0.1 ) self.roberta = AutoModel.from_pretrained( backbone_name, config=backbone_config, add_pooling_layer=False ) # Store backbone_model_name in config for future loading if ( not hasattr(config, "backbone_model_name") or config.backbone_model_name is None ): config.backbone_model_name = backbone_name self.dropout = nn.Dropout( config.hidden_dropout_prob if hasattr(config, "hidden_dropout_prob") else 0.1 ) # Get classifier configuration classifier_hidden_layers = getattr(config, "classifier_hidden_layers", None) classifier_dropout = getattr(config, "classifier_dropout", 0.1) # Build classifier head if classifier_hidden_layers is not None: # rebuild the MLP head in_size = config.hidden_size layers = [] if classifier_hidden_layers: for h in classifier_hidden_layers: layers += [ nn.Linear(in_size, h), nn.ReLU(), nn.Dropout(classifier_dropout), ] in_size = h layers.append(nn.Linear(in_size, config.num_labels)) self.classifier = nn.Sequential(*layers) else: # Default single linear layer self.classifier = nn.Linear(config.hidden_size, config.num_labels) # Initialize weights self.init_weights() def forward( self, input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.FloatTensor] = None, token_type_ids: Optional[torch.LongTensor] = None, position_ids: Optional[torch.LongTensor] = None, head_mask: Optional[torch.FloatTensor] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, **kwargs, ) -> Union[Tuple[torch.Tensor], TokenClassifierOutput]: return_dict = ( return_dict if return_dict is not None else self.config.use_return_dict ) # Run inputs through the RoBERTa backbone outputs = self.roberta( input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids, position_ids=position_ids, head_mask=head_mask, inputs_embeds=inputs_embeds, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, ) sequence_output = outputs.last_hidden_state sequence_output = self.dropout(sequence_output) logits = self.classifier(sequence_output) loss = None if labels is not None: loss_fct = nn.CrossEntropyLoss() if attention_mask is not None: # Only keep active parts of the sequence active_loss = attention_mask.view(-1) == 1 active_logits = logits.view(-1, self.num_labels)[active_loss] active_labels = labels.view(-1)[active_loss] loss = loss_fct(active_logits, active_labels) else: loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) if not return_dict: output = (logits,) + outputs[2:] return ((loss,) + output) if loss is not None else output return TokenClassifierOutput( loss=loss, logits=logits, hidden_states=outputs.hidden_states, attentions=outputs.attentions, ) def get_input_embeddings(self): return self.roberta.get_input_embeddings() def set_input_embeddings(self, value): self.roberta.set_input_embeddings(value) def load_custom_cardioner_multiclass_model(model_path: str, device: str = "auto"): """ Utility function to easily load a custom CardioNER multiclass model. Args: model_path: Path to the saved model directory device: Device to load model on ("auto", "cpu", "cuda", etc.) Returns: tuple: (model, tokenizer, config) """ # Validate model directory import os import torch from transformers import AutoModelForTokenClassification, AutoTokenizer required_files = ["config.json", "modeling.py", "pytorch_model.bin"] missing_files = [ f for f in required_files if not os.path.exists(os.path.join(model_path, f)) ] if missing_files: raise FileNotFoundError( f"Missing required files in {model_path}: {missing_files}" ) print(f"Loading custom CardioNER multiclass model from: {model_path}") # Load tokenizer tokenizer = AutoTokenizer.from_pretrained(model_path) # Load model with trust_remote_code=True model = AutoModelForTokenClassification.from_pretrained( model_path, trust_remote_code=True ) # Set device if device == "auto": device = "cuda" if torch.cuda.is_available() else "cpu" model = model.to(device) print(f"Model loaded successfully on {device}") print(f"Model type: {type(model).__name__}") print(f"Number of labels: {model.num_labels}") return model, tokenizer, model.config def load_custom_multihead_crf_model(model_path: str, device: str = "auto"): """ Utility function to load a Multi-Head CRF model. Args: model_path: Path to the saved model directory device: Device to load model on ("auto", "cpu", "cuda", etc.) Returns: tuple: (model, tokenizer, config) """ import os from transformers import AutoTokenizer # Validate model directory required_files = ["config.json", "modeling.py"] missing_files = [ f for f in required_files if not os.path.exists(os.path.join(model_path, f)) ] if missing_files: raise FileNotFoundError( f"Missing required files in {model_path}: {missing_files}" ) print(f"Loading Multi-Head CRF model from: {model_path}") # Load tokenizer tokenizer = AutoTokenizer.from_pretrained(model_path) # Load config import json with open(os.path.join(model_path, "config.json"), "r") as f: config_dict = json.load(f) config = MultiHeadCRFConfig(**config_dict) # Load model model = TokenClassificationModelMultiHeadCRF.from_pretrained( model_path, config=config ) # Set device if device == "auto": device = "cuda" if torch.cuda.is_available() else "cpu" model = model.to(device) print(f"Model loaded successfully on {device}") print(f"Model type: {type(model).__name__}") print(f"Entity types: {model.entity_types}") print(f"Number of labels per head: {model.num_labels}") return model, tokenizer, model.config def validate_custom_multiclass_model_directory(model_path: str) -> dict: """ Validate that a model directory contains all necessary files for custom multiclass model loading. Args: model_path: Path to the model directory Returns: dict: Validation results with status and details """ import json import os validation_results = { "valid": True, "errors": [], "warnings": [], "files_found": [], "model_info": {}, } # Required files required_files = { "config.json": "Model configuration", "modeling.py": "Custom model class definition", "pytorch_model.bin": "Model weights", } # Optional files optional_files = { "tokenizer.json": "Tokenizer vocabulary", "tokenizer_config.json": "Tokenizer configuration", "training_args.json": "Training arguments", } # Check required files for filename, description in required_files.items(): filepath = os.path.join(model_path, filename) if os.path.exists(filepath): validation_results["files_found"].append(f"{filename} ({description})") else: validation_results["valid"] = False validation_results["errors"].append( f"Missing required file: {filename} - {description}" ) # Check optional files for filename, description in optional_files.items(): filepath = os.path.join(model_path, filename) if os.path.exists(filepath): validation_results["files_found"].append(f"{filename} ({description})") else: validation_results["warnings"].append( f"Missing optional file: {filename} - {description}" ) # Parse config if available config_path = os.path.join(model_path, "config.json") if os.path.exists(config_path): try: with open(config_path, "r") as f: config = json.load(f) validation_results["model_info"]["num_labels"] = config.get( "num_labels", "Unknown" ) validation_results["model_info"]["model_type"] = config.get( "model_type", "Unknown" ) validation_results["model_info"]["has_auto_map"] = "auto_map" in config validation_results["model_info"]["classifier_hidden_layers"] = config.get( "classifier_hidden_layers", None ) validation_results["model_info"]["freeze_backbone"] = config.get( "freeze_backbone", None ) validation_results["model_info"]["use_crf"] = ( "TokenClassificationModelCRF" in str(config.get("architectures", [])) ) if not config.get("auto_map"): validation_results["warnings"].append( "No auto_map found in config - may not load correctly with trust_remote_code=True" ) except json.JSONDecodeError as e: validation_results["valid"] = False validation_results["errors"].append(f"Invalid config.json: {str(e)}") # Check modeling.py content modeling_path = os.path.join(model_path, "modeling.py") if os.path.exists(modeling_path): try: with open(modeling_path, "r") as f: content = f.read() required_classes = [ "TokenClassificationModel", "TokenClassificationModelCRF", ] missing_classes = [cls for cls in required_classes if cls not in content] if missing_classes: validation_results["valid"] = False validation_results["errors"].append( f"modeling.py missing required classes: {missing_classes}" ) except Exception as e: validation_results["warnings"].append( f"Could not read modeling.py: {str(e)}" ) return validation_results # Register the MultiHeadCRF config for auto loading try: from transformers import AutoConfig AutoConfig.register("multihead-crf-tagger", MultiHeadCRFConfig) except Exception: pass # Config may already be registered def patch_legacy_model( model_path: str, backbone_model_name: str, dry_run: bool = True ) -> bool: """ Patch a legacy saved model by adding backbone_model_name to config.json. Use this to fix models trained before backbone_model_name was added to the config. Args: model_path: Path to the saved model directory backbone_model_name: The original backbone model name used during training (e.g., "CLTL/MedRoBERTa.nl", "GroNLP/bert-base-dutch-cased") dry_run: If True, only print what would be changed without modifying files Returns: bool: True if patch was successful (or would be successful in dry_run mode) Example: >>> # First, do a dry run to see what will change >>> patch_legacy_model("/path/to/saved/model", "CLTL/MedRoBERTa.nl", dry_run=True) >>> # Then apply the patch >>> patch_legacy_model("/path/to/saved/model", "CLTL/MedRoBERTa.nl", dry_run=False) """ import json import os import shutil config_path = os.path.join(model_path, "config.json") if not os.path.exists(config_path): print(f"ERROR: config.json not found at {config_path}") return False # Load existing config with open(config_path, "r") as f: config = json.load(f) # Check if already patched if "backbone_model_name" in config: print(f"Model already has backbone_model_name: {config['backbone_model_name']}") if config["backbone_model_name"] == backbone_model_name: print("No changes needed.") return True else: print(f"WARNING: Existing backbone_model_name differs from provided value!") print(f" Existing: {config['backbone_model_name']}") print(f" Provided: {backbone_model_name}") if dry_run: print("Would update to new value (dry_run=True)") else: print("Updating to new value...") # Add backbone_model_name config["backbone_model_name"] = backbone_model_name if dry_run: print(f"\n[DRY RUN] Would patch {config_path}:") print(f' Adding: backbone_model_name = "{backbone_model_name}"') print("\nTo apply this patch, run with dry_run=False") return True # Create backup backup_path = config_path + ".backup" shutil.copy2(config_path, backup_path) print(f"Created backup at {backup_path}") # Write updated config with open(config_path, "w") as f: json.dump(config, f, indent=2) print(f"Successfully patched {config_path}") print(f' Added: backbone_model_name = "{backbone_model_name}"') return True def patch_multiple_models( model_paths: list, backbone_model_name: str, dry_run: bool = True ) -> dict: """ Patch multiple legacy saved models at once. Args: model_paths: List of paths to saved model directories backbone_model_name: The original backbone model name used during training dry_run: If True, only print what would be changed without modifying files Returns: dict: Results for each model path Example: >>> models = ["/path/to/model1", "/path/to/model2"] >>> patch_multiple_models(models, "CLTL/MedRoBERTa.nl", dry_run=False) """ results = {} for path in model_paths: print(f"\n{'=' * 60}") print(f"Processing: {path}") print("=" * 60) results[path] = patch_legacy_model(path, backbone_model_name, dry_run) # Summary print(f"\n{'=' * 60}") print("SUMMARY") print("=" * 60) success = sum(1 for v in results.values() if v) print( f"Successfully {'would patch' if dry_run else 'patched'}: {success}/{len(model_paths)}" ) return results