modeling.py

from itertools import islice
from typing import Dict, List, Optional, Tuple, Union

import torch
import torch.nn as nn
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 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 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
            )

            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",
            ]
            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


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