ok

config.json

@@ -1,13 +1,13 @@
 {
-  "_name_or_path": "luxe_ub_sample_private",
+  "_name_or_path": "ubke_ub_sample_private",
   "architectures": [
-    "LuxeForMaskedLM"
+    "UbkeForMaskedLM"
   ],
   "attention_probs_dropout_prob": 0.1,
   "auto_map": {
-    "AutoConfig": "uzabase/UBKE-LUKE--configuration_luxe.LuxeConfig",
-    "AutoModel": "uzabase/UBKE-LUKE--modeling_luxe.LuxeForMaskedLM",
-    "AutoModelForPreTraining": "uzabase/UBKE-LUKE--modeling_luxe.LuxeForMaskedLM"
+    "AutoConfig": "uzabase/UBKE-LUKE--configuration_ubke.UbkeConfig",
+    "AutoModel": "uzabase/UBKE-LUKE--modeling_ubke.UbkeForMaskedLM",
+    "AutoModelForPreTraining": "uzabase/UBKE-LUKE--modeling_ubke.UbkeForMaskedLM"
   },
   "bert_model_name": "cl-tohoku/bert-base-japanese-v3",
   "bos_token_id": null,
@@ -25,7 +25,7 @@
   "intermediate_size": 3072,
   "layer_norm_eps": 1e-12,
   "max_position_embeddings": 512,
-  "model_type": "luxe",
+  "model_type": "ubke",
   "normalize_entity_embeddings": 1,
   "num_attention_heads": 12,
   "num_category_entities": 0,

configuration_luxe.py → configuration_ubke.py

@@ -1,8 +1,8 @@
 from transformers.configuration_utils import PretrainedConfig
 
 
-class LuxeConfig(PretrainedConfig):
-    model_type = "luxe"
+class UbkeConfig(PretrainedConfig):
+    model_type = "ubke"
 
     def __init__(
         self,
@@ -30,7 +30,12 @@ class LuxeConfig(PretrainedConfig):
         eos_token_id=2,
         **kwargs,
     ):
-        super().__init__(pad_token_id=pad_token_id, bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            **kwargs,
+        )
 
         self.vocab_size = vocab_size
         self.entity_vocab_size = entity_vocab_size

modeling_luxe.py → modeling_ubke.py

@@ -11,11 +11,11 @@ from transformers.models.luke.modeling_luke import (
 )
 from transformers.utils import ModelOutput
 
-from .configuration_luxe import LuxeConfig
+from .configuration_ubke import UbkeConfig
 
 
 @dataclass
-class LuxeMaskedLMOutput(ModelOutput):
+class UbkeMaskedLMOutput(ModelOutput):
     loss: Optional[torch.FloatTensor] = None
     mlm_loss: Optional[torch.FloatTensor] = None
     mep_loss: Optional[torch.FloatTensor] = None
@@ -32,8 +32,8 @@ class LuxeMaskedLMOutput(ModelOutput):
     attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
 
 
-class LuxePreTrainedModel(PreTrainedModel):
-    config_class = LuxeConfig
+class UbkePreTrainedModel(PreTrainedModel):
+    config_class = UbkeConfig
     base_model_prefix = "luke"
     supports_gradient_checkpointing = True
     _no_split_modules = ["LukeAttention", "LukeEntityEmbeddings"]
@@ -55,14 +55,14 @@ class LuxePreTrainedModel(PreTrainedModel):
             module.weight.data.fill_(1.0)
 
 
-class LuxeForMaskedLM(LuxePreTrainedModel):
+class UbkeForMaskedLM(UbkePreTrainedModel):
     _tied_weights_keys = [
         "lm_head.decoder.weight",
         "lm_head.decoder.bias",
         "entity_predictions.decoder.weight",
     ]
 
-    def __init__(self, config: LuxeConfig):
+    def __init__(self, config: UbkeConfig):
         super().__init__(config)
 
         self.luke = LukeModel(config)
@@ -114,7 +114,7 @@ class LuxeForMaskedLM(LuxePreTrainedModel):
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
-    ) -> Union[Tuple, LuxeMaskedLMOutput]:
+    ) -> Union[Tuple, UbkeMaskedLMOutput]:
         return_dict = (
             return_dict if return_dict is not None else self.config.use_return_dict
         )
@@ -232,7 +232,7 @@ class LuxeForMaskedLM(LuxePreTrainedModel):
                 if v is not None
             )
 
-        return LuxeMaskedLMOutput(
+        return UbkeMaskedLMOutput(
             loss=loss,
             mlm_loss=mlm_loss,
             mep_loss=mep_loss,