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