| 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 = -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 |
|
|
