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): """ 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