modeling.py

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