from torch import nn
from transformers import (
    ModernBertConfig,
    ModernBertForMaskedLM,
    ModernBertModel,
    ModernBertPreTrainedModel,
)
from transformers.models.modernbert.modeling_modernbert import (
    ModernBertEncoderLayer,
    ModernBertPredictionHead,
)


class ModernBertConfigCustom(ModernBertConfig):
    model_type = "modern_bert"

    def __init__(self, repetition_factor=1, **kwargs):
        super().__init__(**kwargs)
        self.repetition_factor = repetition_factor

    def standardize_rope_params(self):
        pass


class ModernBertModelCustom(ModernBertModel):
    config_class = ModernBertConfigCustom

    def __init__(self, config):
        super().__init__(config)
        r = getattr(config, "repetition_factor", 1)
        n = config.num_hidden_layers
        self.layers = nn.ModuleList(
            [ModernBertEncoderLayer(config, layer_idx=i // r) for i in range(n * r)]
        )


class ModernBertForMaskedLMCustom(ModernBertForMaskedLM):
    config_class = ModernBertConfigCustom
    _tied_weights_keys = {"decoder.weight": "model.embeddings.tok_embeddings.weight"}

    def __init__(self, config):
        ModernBertPreTrainedModel.__init__(self, config)
        self.model = ModernBertModelCustom(config)
        self.head = ModernBertPredictionHead(config)
        self.decoder = nn.Linear(
            config.hidden_size, config.vocab_size, bias=config.decoder_bias
        )
        self.sparse_prediction = config.sparse_prediction
        self.sparse_pred_ignore_index = config.sparse_pred_ignore_index
        self.post_init()