同步最新代码

Author: JunnYu

README.md

@@ -104,11 +104,11 @@ max diff : tensor(3.5763e-06)
 base版本
 python compare_model.py
 bert4keras vs pytorch
-mean diff : tensor(4.3925e-07)
-max diff : tensor(7.6294e-06)
+mean diff : tensor(4.3340e-07)
+max diff : tensor(5.7220e-06)
 bert4keras vs tf2.0
-mean diff : tensor(3.4151e-07)
-max diff : tensor(3.8147e-06)
+mean diff : tensor(3.4319e-07)
+max diff : tensor(5.2452e-06)
 ```
 
 

convert_roformer_original_tf_checkpoint_to_pytorch.py

@@ -14,19 +14,18 @@
 # limitations under the License.
 """Convert RoFormer checkpoint."""
 
-
 import argparse
 
 import torch
 
 from roformer import RoFormerConfig, RoFormerForMaskedLM, load_tf_weights_in_roformer
 from transformers.utils import logging
 
-
 logging.set_verbosity_info()
 
 
-def convert_tf_checkpoint_to_pytorch(tf_checkpoint_path, bert_config_file, pytorch_dump_path):
+def convert_tf_checkpoint_to_pytorch(tf_checkpoint_path, bert_config_file,
+                                     pytorch_dump_path):
     # Initialise PyTorch model
     config = RoFormerConfig.from_json_file(bert_config_file)
     print(f"Building PyTorch model from configuration: {config}")
@@ -37,25 +36,34 @@ def convert_tf_checkpoint_to_pytorch(tf_checkpoint_path, bert_config_file, pytor
 
     # Save pytorch-model
     print(f"Save PyTorch model to {pytorch_dump_path}")
-    torch.save(model.state_dict(), pytorch_dump_path, _use_new_zipfile_serialization=False)
+    torch.save(model.state_dict(),
+               pytorch_dump_path,
+               _use_new_zipfile_serialization=False)
 
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     # Required parameters
-    parser.add_argument(
-        "--tf_checkpoint_path", default=None, type=str, required=True, help="Path to the TensorFlow checkpoint path."
-    )
+    parser.add_argument("--tf_checkpoint_path",
+                        default=None,
+                        type=str,
+                        required=True,
+                        help="Path to the TensorFlow checkpoint path.")
     parser.add_argument(
         "--bert_config_file",
         default=None,
         type=str,
         required=True,
-        help="The config json file corresponding to the pre-trained BERT model. \n"
+        help=
+        "The config json file corresponding to the pre-trained BERT model. \n"
         "This specifies the model architecture.",
     )
-    parser.add_argument(
-        "--pytorch_dump_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
-    )
+    parser.add_argument("--pytorch_dump_path",
+                        default=None,
+                        type=str,
+                        required=True,
+                        help="Path to the output PyTorch model.")
     args = parser.parse_args()
-    convert_tf_checkpoint_to_pytorch(args.tf_checkpoint_path, args.bert_config_file, args.pytorch_dump_path)
+    convert_tf_checkpoint_to_pytorch(args.tf_checkpoint_path,
+                                     args.bert_config_file,
+                                     args.pytorch_dump_path)

setup.py

@@ -4,12 +4,12 @@
     name='roformer',
     package_dir={"": "src"},
     packages=find_packages("src"),
-    version='0.0.8',
+    version='0.1.0',
     license='Apache 2.0',
     description='roformer_pytorch',
     author='Jun Yu',
     author_email='[email protected]',
     url='https://github.com/JunnYu/RoFormer_pytorch',
     keywords=['roformer', 'pytorch', 'tf2.0'],
-    install_requires=['transformers>=4.5.0', 'jieba', 'rjieba'],
+    install_requires=['transformers>=4.5.0', 'jieba'],
 )

src/roformer/__init__.py

@@ -19,12 +19,15 @@
 
 from transformers.file_utils import _BaseLazyModule, is_tf_available, is_tokenizers_available, is_torch_available
 
-
 _import_structure = {
-    "configuration_roformer": ["ROFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP", "RoFormerConfig"],
+    "configuration_roformer":
+    ["ROFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP", "RoFormerConfig"],
     "tokenization_roformer": ["RoFormerTokenizer"],
 }
 
+if is_tokenizers_available():
+    _import_structure["tokenization_roformer_fast"] = ["RoFormerTokenizerFast"]
+
 if is_torch_available():
     _import_structure["modeling_roformer"] = [
         "ROFORMER_PRETRAINED_MODEL_ARCHIVE_LIST",
@@ -40,7 +43,6 @@
         "load_tf_weights_in_roformer",
     ]
 
-
 if is_tf_available():
     _import_structure["modeling_tf_roformer"] = [
         "TF_ROFORMER_PRETRAINED_MODEL_ARCHIVE_LIST",
@@ -55,11 +57,13 @@
         "TFRoFormerPreTrainedModel",
     ]
 
-
 if TYPE_CHECKING:
     from .configuration_roformer import ROFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP, RoFormerConfig
     from .tokenization_roformer import RoFormerTokenizer
 
+    if is_tokenizers_available():
+        from .tokenization_roformer_fast import RoFormerTokenizerFast
+
     if is_torch_available():
         from .modeling_roformer import (
             ROFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
@@ -89,7 +93,6 @@
             TFRoFormerPreTrainedModel,
         )
 
-
 else:
     import importlib
     import os

src/roformer/configuration_roformer.py

@@ -17,12 +17,21 @@
 from transformers.configuration_utils import PretrainedConfig
 from transformers.utils import logging
 
-
 logger = logging.get_logger(__name__)
 
 ROFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP = {
-    "junnyu/roformer_chinese_small": "https://huggingface.co/junnyu/roformer_chinese_small/resolve/main/config.json",
-    "junnyu/roformer_chinese_base": "https://huggingface.co/junnyu/roformer_chinese_base/resolve/main/config.json"
+    "junnyu/roformer_chinese_small":
+    "https://huggingface.co/junnyu/roformer_chinese_small/resolve/main/config.json",
+    "junnyu/roformer_chinese_base":
+    "https://huggingface.co/junnyu/roformer_chinese_base/resolve/main/config.json",
+    "junnyu/roformer_chinese_char_small":
+    "https://huggingface.co/junnyu/roformer_chinese_char_small/resolve/main/config.json",
+    "junnyu/roformer_chinese_char_base":
+    "https://huggingface.co/junnyu/roformer_chinese_char_base/resolve/main/config.json",
+    "junnyu/roformer_small_discriminator":
+    "https://huggingface.co/junnyu/roformer_small_discriminator/resolve/main/config.json",
+    "junnyu/roformer_small_generator":
+    "https://huggingface.co/junnyu/roformer_small_generator/resolve/main/config.json",
     # See all RoFormer models at https://huggingface.co/models?filter=roformer
 }
 
@@ -43,7 +52,7 @@ class RoFormerConfig(PretrainedConfig):
             Vocabulary size of the RoFormer model. Defines the number of different tokens that can be represented by
             the :obj:`inputs_ids` passed when calling :class:`~transformers.RoFormerModel` or
             :class:`~transformers.TFRoFormerModel`.
-        embedding_size (:obj:`int`, `optional`, defaults to 768):
+        embedding_size (:obj:`int`, `optional`, defaults to None):
             Dimensionality of the encoder layers and the pooler layer.
         hidden_size (:obj:`int`, `optional`, defaults to 768):
             Dimension of the encoder layers and the pooler layer.
@@ -93,27 +102,25 @@ class RoFormerConfig(PretrainedConfig):
     """
     model_type = "roformer"
 
-    def __init__(
-        self,
-        vocab_size=50000,
-        embedding_size=None,
-        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=1536,
-        type_vocab_size=2,
-        initializer_range=0.02,
-        layer_norm_eps=1e-12,
-        pad_token_id=0,
-        gradient_checkpointing=False,
-        rotary_value=False,
-        use_cache=True,
-        **kwargs
-    ):
+    def __init__(self,
+                 vocab_size=50000,
+                 embedding_size=None,
+                 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=1536,
+                 type_vocab_size=2,
+                 initializer_range=0.02,
+                 layer_norm_eps=1e-12,
+                 pad_token_id=0,
+                 gradient_checkpointing=False,
+                 rotary_value=False,
+                 use_cache=True,
+                 **kwargs):
         super().__init__(pad_token_id=pad_token_id, **kwargs)
 
         self.vocab_size = vocab_size

src/roformer/convert_roformer_original_tf_checkpoint_to_pytorch.py

@@ -14,19 +14,18 @@
 # limitations under the License.
 """Convert RoFormer checkpoint."""
 
-
 import argparse
 
 import torch
 
 from roformer import RoFormerConfig, RoFormerForMaskedLM, load_tf_weights_in_roformer
 from transformers.utils import logging
 
-
 logging.set_verbosity_info()
 
 
-def convert_tf_checkpoint_to_pytorch(tf_checkpoint_path, bert_config_file, pytorch_dump_path):
+def convert_tf_checkpoint_to_pytorch(tf_checkpoint_path, bert_config_file,
+                                     pytorch_dump_path):
     # Initialise PyTorch model
     config = RoFormerConfig.from_json_file(bert_config_file)
     print(f"Building PyTorch model from configuration: {config}")
@@ -37,25 +36,34 @@ def convert_tf_checkpoint_to_pytorch(tf_checkpoint_path, bert_config_file, pytor
 
     # Save pytorch-model
     print(f"Save PyTorch model to {pytorch_dump_path}")
-    torch.save(model.state_dict(), pytorch_dump_path, _use_new_zipfile_serialization=False)
+    torch.save(model.state_dict(),
+               pytorch_dump_path,
+               _use_new_zipfile_serialization=False)
 
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     # Required parameters
-    parser.add_argument(
-        "--tf_checkpoint_path", default=None, type=str, required=True, help="Path to the TensorFlow checkpoint path."
-    )
+    parser.add_argument("--tf_checkpoint_path",
+                        default=None,
+                        type=str,
+                        required=True,
+                        help="Path to the TensorFlow checkpoint path.")
     parser.add_argument(
         "--bert_config_file",
         default=None,
         type=str,
         required=True,
-        help="The config json file corresponding to the pre-trained BERT model. \n"
+        help=
+        "The config json file corresponding to the pre-trained BERT model. \n"
         "This specifies the model architecture.",
     )
-    parser.add_argument(
-        "--pytorch_dump_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
-    )
+    parser.add_argument("--pytorch_dump_path",
+                        default=None,
+                        type=str,
+                        required=True,
+                        help="Path to the output PyTorch model.")
     args = parser.parse_args()
-    convert_tf_checkpoint_to_pytorch(args.tf_checkpoint_path, args.bert_config_file, args.pytorch_dump_path)
+    convert_tf_checkpoint_to_pytorch(args.tf_checkpoint_path,
+                                     args.bert_config_file,
+                                     args.pytorch_dump_path)