from transformers import BertConfig 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 try: from transformers.models.eurobert.modeling_eurobert import EuroBertModel except Exception: try: from transformers import EuroBertModel except Exception: EuroBertModel = None print("COULD NOT IMPORT EUROBERT MODEL") # Large negative number for masking impossible transitions LARGE_NEGATIVE_NUMBER = -1e9 NUM_PER_LAYER = 16 def _build_backbone_from_config(config): """ Build a backbone model structure from config only. Never call from_pretrained() here; outer model loading will restore weights. """ from transformers import AutoConfig, AutoModel backbone_name = getattr(config, "backbone_model_name", None) if backbone_name is None: backbone_name = 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 backbone" ) backbone_config = AutoConfig.from_pretrained( backbone_name, trust_remote_code=True, ) if hasattr(config, "hidden_dropout_prob"): backbone_config.hidden_dropout_prob = getattr( config, "hidden_dropout_prob", 0.1 ) if hasattr(config, "num_labels"): backbone_config.num_labels = getattr(config, "num_labels") if "eurobert" in backbone_name.lower() and EuroBertModel is not None: backbone = EuroBertModel(backbone_config) else: backbone = AutoModel.from_config( backbone_config, trust_remote_code=True, ) if getattr(config, "backbone_model_name", None) is None: config.backbone_model_name = backbone_name return backbone, backbone_name class MultiHeadCRFConfig(PretrainedConfig): """ Configuration class for Multi-Head CRF models. """ 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. """ 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: 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: with torch.no_grad(): if self.num_tags > 2: self.start_transitions[2] = LARGE_NEGATIVE_NUMBER self.transitions[0][2] = LARGE_NEGATIVE_NUMBER if self.num_tags > 3: self.start_transitions[3] = LARGE_NEGATIVE_NUMBER for i in range(3): self.transitions[i][3] = LARGE_NEGATIVE_NUMBER 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: 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) 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) numerator = self._compute_score(emissions, tags, mask) denominator = self._compute_normalizer(emissions, mask) llh = numerator - denominator nllh = -llh if reduction == "none": return nllh if reduction == "sum": return nllh.sum() if reduction == "mean": return nllh.mean() return nllh.sum() / mask.type_as(emissions).sum() def decode( self, emissions: torch.Tensor, mask: Optional[torch.Tensor] = None ) -> List[List[int]]: 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 and 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: 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() device = emissions.device tags = tags.to(device) mask = mask.to(device) seq_length, batch_size = tags.shape mask = mask.type_as(emissions) 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] 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: 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) 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]]: 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 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 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. """ 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 None: self.roberta, backbone_name = _build_backbone_from_config(config) else: if hasattr(base_model, "roberta"): self.roberta = base_model.roberta else: self.roberta = base_model backbone_name = ( getattr(getattr(self.roberta, "config", None), "_name_or_path", None) or getattr(config, "backbone_model_name", None) or getattr(config, "_name_or_path", None) ) if getattr(config, "backbone_model_name", None) is None: config.backbone_model_name = backbone_name self.lm_output_size = self.roberta.config.hidden_size 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(True) self.dropout = nn.Dropout(getattr(config, "hidden_dropout_prob", 0.1)) self.crf = CRF(self.num_labels, batch_first=True) self._build_classifier_head(classifier_hidden_layers, classifier_dropout) self.post_init() def _build_classifier_head(self, hidden_layers, dropout_rate): layers = [] input_size = self.lm_output_size if not hidden_layers: self.classifier = nn.Sequential( nn.Dropout(dropout_rate), nn.Linear(input_size, self.num_labels) ) return 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 layers.append(nn.Linear(input_size, self.num_labels)) 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 ) try: 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, ) except TypeError: outputs = self.roberta( input_ids, attention_mask=attention_mask, position_ids=position_ids, 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) loss = None if labels is not None: labels_long = labels.long() if attention_mask is not None: mask = attention_mask.bool() loss = -self.crf(logits, labels_long, mask=mask, reduction="mean") else: if not getattr(self, "_warned_no_attention_mask", False): print( "WARNING: attention_mask is None; CRF loss will include padding tokens." ) self._warned_no_attention_mask = True 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): config = kwargs.pop("config", None) if config is None: from transformers import AutoConfig config = AutoConfig.from_pretrained(pretrained_model_name_or_path, **kwargs) 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, ) try: state_dict = torch.load( f"{pretrained_model_name_or_path}/pytorch_model.bin", map_location="cpu" ) model.load_state_dict(state_dict) except Exception: 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. """ 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 self.entity_types = getattr(config, "entity_types", []) if not self.entity_types: raise ValueError("entity_types must be provided in config") self.num_labels = config.num_labels 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) if base_model is None: self.roberta, backbone_name = _build_backbone_from_config(config) else: if hasattr(base_model, "roberta"): self.roberta = base_model.roberta else: self.roberta = base_model backbone_name = ( getattr(getattr(self.roberta, "config", None), "_name_or_path", None) or getattr(config, "backbone_model_name", None) or getattr(config, "_name_or_path", None) ) if getattr(config, "backbone_model_name", None) is None: config.backbone_model_name = backbone_name self.hidden_size = self.roberta.config.hidden_size self.dropout = nn.Dropout(getattr(config, "hidden_dropout_prob", 0.1)) print(f"Creating Multi-Head CRF with entity types: {sorted(self.entity_types)}") for entity_type in self.entity_types: 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) ) 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), ) if freeze_backbone: self._freeze_backbone() self.post_init() def _freeze_backbone(self): print("+" * 30, "\n\n", "Freezing backbone...", "+" * 30, "\n\n") for param in self.roberta.embeddings.parameters(): param.requires_grad = False 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): 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, ): return_dict = ( return_dict if return_dict is not None else self.config.use_return_dict ) try: 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, ) except TypeError: outputs = self.roberta( input_ids, attention_mask=attention_mask, position_ids=position_ids, 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) 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: losses = {} mask = attention_mask.bool() if attention_mask is not None else None for entity_type in self.entity_types: if entity_type in labels: entity_labels = ( labels[entity_type].long().to(logits[entity_type].device) ) crf = getattr(self, f"{entity_type}_crf") if mask is not None: losses[entity_type] = crf( logits[entity_type], entity_labels, mask=mask, reduction=self.crf_reduction, ) else: if not getattr(self, "_warned_no_attention_mask", False): print( "WARNING: attention_mask is None; CRF loss will include padding tokens." ) self._warned_no_attention_mask = True losses[entity_type] = crf( logits[entity_type], entity_labels, reduction=self.crf_reduction, ) total_loss = sum(losses.values()) return total_loss, logits predictions = {} mask = attention_mask.bool() if attention_mask is not None else None for entity_type in self.entity_types: crf = getattr(self, f"{entity_type}_crf") if mask is not None: decoded = crf.decode(logits[entity_type], mask=mask) else: decoded = crf.decode(logits[entity_type]) predictions[entity_type] = torch.tensor(decoded) 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): import json import os config = kwargs.pop("config", None) if config is None: 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) 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), ) 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) 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). """ 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). """ 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 self.entity_types = getattr(config, "entity_types", []) if not self.entity_types: raise ValueError("entity_types must be provided in config") self.num_labels = config.num_labels 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) self.use_class_weights = getattr(config, "use_class_weights", False) self.class_weights = getattr(config, "class_weights", None) if base_model is None: self.roberta, backbone_name = _build_backbone_from_config(config) else: if hasattr(base_model, "roberta"): self.roberta = base_model.roberta else: self.roberta = base_model backbone_name = ( getattr(getattr(self.roberta, "config", None), "_name_or_path", None) or getattr(config, "backbone_model_name", None) or getattr(config, "_name_or_path", None) ) if getattr(config, "backbone_model_name", None) is None: config.backbone_model_name = backbone_name self.hidden_size = self.roberta.config.hidden_size self.dropout = nn.Dropout(getattr(config, "hidden_dropout_prob", 0.1)) print( f"Creating Multi-Head model with entity types: {sorted(self.entity_types)}" ) for entity_type in self.entity_types: 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) ) setattr( self, f"{entity_type}_classifier", nn.Linear(self.hidden_size, self.num_labels), ) 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) if freeze_backbone: self._freeze_backbone() self.post_init() def _freeze_backbone(self): print("+" * 30, "\n\n", "Freezing backbone...", "+" * 30, "\n\n") for param in self.roberta.embeddings.parameters(): param.requires_grad = False 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): 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, ): return_dict = ( return_dict if return_dict is not None else self.config.use_return_dict ) try: 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, ) except TypeError: outputs = self.roberta( input_ids, attention_mask=attention_mask, position_ids=position_ids, 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) 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: 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] loss_fct = self.loss_fns[entity_type] 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), ) total_loss = sum(losses.values()) return total_loss, logits predictions = {} for entity_type in self.entity_types: preds = torch.argmax(logits[entity_type], dim=-1) predictions[entity_type] = preds 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): import json import os config = kwargs.pop("config", None) if config is None: 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) 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), ) 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) 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): config_class = BertConfig """ Custom token classification model with configurable classifier head (no CRF). """ def __init__(self, config, base_model=None): super().__init__(config) self.config = config self.num_labels = config.num_labels if base_model is None: self.roberta, backbone_name = _build_backbone_from_config(config) else: if hasattr(base_model, "roberta"): self.roberta = base_model.roberta else: self.roberta = base_model backbone_name = ( getattr(getattr(self.roberta, "config", None), "_name_or_path", None) or getattr(config, "backbone_model_name", None) or getattr(config, "_name_or_path", None) ) if getattr(config, "backbone_model_name", None) is None: config.backbone_model_name = backbone_name self.dropout = nn.Dropout(getattr(config, "hidden_dropout_prob", 0.1)) classifier_hidden_layers = getattr(config, "classifier_hidden_layers", None) classifier_dropout = getattr(config, "classifier_dropout", 0.1) if classifier_hidden_layers is not None: in_size = self.roberta.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: self.classifier = nn.Linear( self.roberta.config.hidden_size, config.num_labels ) if isinstance(self.classifier, nn.Sequential): for module in self.classifier: if isinstance(module, nn.Linear): self._init_weights(module) elif isinstance(self.classifier, nn.Linear): self._init_weights(self.classifier) self.post_init() 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 ) try: 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, ) except TypeError: outputs = self.roberta( input_ids, attention_mask=attention_mask, position_ids=position_ids, 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: 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"): 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}") tokenizer = AutoTokenizer.from_pretrained(model_path) model = AutoModelForTokenClassification.from_pretrained( model_path, trust_remote_code=True, use_safetensors=True, ) 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"): import json import os from transformers import AutoTokenizer 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}") tokenizer = AutoTokenizer.from_pretrained(model_path) with open(os.path.join(model_path, "config.json"), "r") as f: config_dict = json.load(f) config = MultiHeadCRFConfig(**config_dict) model = TokenClassificationModelMultiHeadCRF.from_pretrained( model_path, config=config ) 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: import json import os validation_results = { "valid": True, "errors": [], "warnings": [], "files_found": [], "model_info": {}, } required_files = { "config.json": "Model configuration", "modeling.py": "Custom model class definition", "pytorch_model.bin": "Model weights", } optional_files = { "tokenizer.json": "Tokenizer vocabulary", "tokenizer_config.json": "Tokenizer configuration", "training_args.json": "Training arguments", } 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}" ) 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}" ) 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)}") 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 try: from transformers import AutoConfig AutoConfig.register("multihead-crf-tagger", MultiHeadCRFConfig) except Exception: pass def patch_legacy_model( model_path: str, backbone_model_name: str, dry_run: bool = True ) -> bool: 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 with open(config_path, "r") as f: config = json.load(f) 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...") 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 backup_path = config_path + ".backup" shutil.copy2(config_path, backup_path) print(f"Created backup at {backup_path}") 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: 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) 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