from transformers.models.bert import BertConfig
from transformers.models.xlm_roberta import XLMRobertaConfig


def _init_function(
    self,
    entity_vocab_size: int | None = 10000,
    entity_embedding_size: int = 768,
    entity_fusion_method: str = "multihead_attention",
    use_entity_position_embeddings: bool = True,
    entity_fusion_activation: str = "softmax",
    num_entity_fusion_attention_heads: int = 12,
    similarity_function: str = "dot",
    similarity_temperature: float = 1.0,
    *args,
    **kwargs,
):
    self.entity_vocab_size = entity_vocab_size
    self.entity_embedding_size = entity_embedding_size
    self.entity_fusion_method = entity_fusion_method
    self.use_entity_position_embeddings = use_entity_position_embeddings
    self.entity_fusion_activation = entity_fusion_activation
    self.num_entity_fusion_attention_heads = num_entity_fusion_attention_heads
    self.similarity_function = similarity_function
    self.similarity_temperature = similarity_temperature

    super(self.__class__, self).__init__(*args, **kwargs)


class KPRConfigForBert(BertConfig):
    __init__ = _init_function
    model_type = "kpr-bert"


class KPRConfigForXLMRoberta(XLMRobertaConfig):
    __init__ = _init_function
    model_type = "kpr-xlm-roberta"