update

Author: JunnYu

README.md

@@ -17,35 +17,58 @@ https://huggingface.co/junnyu/roformer_chinese_base
 ## 使用
 ```python
 import torch
-from roformer import RoFormerModel, RoFormerTokenizer
+from roformer import RoFormerModel, RoFormerTokenizer, TFRoFormerModel
 tokenizer = RoFormerTokenizer.from_pretrained("junnyu/roformer_chinese_base")
-model = RoFormerModel.from_pretrained("junnyu/roformer_chinese_base")
+pt_model = RoFormerModel.from_pretrained("junnyu/roformer_chinese_base")
+tf_model = TFRoFormerModel.from_pretrained("junnyu/roformer_chinese_base",
+                                           from_pt=True)
 text = "这里基本保留了唐宋遗留下来的坊巷格局和大量明清古建筑，其中各级文保单位29处，被誉为“里坊制度的活化石”“明清建筑博物馆”！"
-inputs = tokenizer(text, return_tensors="pt")
+pt_inputs = tokenizer(text, return_tensors="pt")
+tf_inputs = tokenizer(text, return_tensors="tf")
 with torch.no_grad():
-    outputs = model(**inputs).last_hidden_state
-print(outputs.shape)
+    pt_outputs = pt_model(**pt_inputs).last_hidden_state
+print(pt_outputs.shape)
+tf_outputs = tf_model(**tf_inputs, training=False).last_hidden_state
+print(tf_outputs.shape)
 ```
 ## MLM测试
 ```python
 import torch
-from roformer import RoFormerForMaskedLM, RoFormerTokenizer
+import tensorflow as tf
+from roformer import RoFormerForMaskedLM, RoFormerTokenizer, TFRoFormerForMaskedLM
 text = "今天[MASK]很好，我[MASK]去公园玩。"
 tokenizer = RoFormerTokenizer.from_pretrained("junnyu/roformer_chinese_base")
-model = RoFormerForMaskedLM.from_pretrained("junnyu/roformer_chinese_base")
-inputs = tokenizer(text, return_tensors="pt")
+pt_model = RoFormerForMaskedLM.from_pretrained("junnyu/roformer_chinese_base")
+tf_model = TFRoFormerForMaskedLM.from_pretrained(
+    "junnyu/roformer_chinese_base", from_pt=True)
+pt_inputs = tokenizer(text, return_tensors="pt")
+tf_inputs = tokenizer(text, return_tensors="tf")
+# pytorch
 with torch.no_grad():
-    outputs = model(**inputs).logits[0]
-outputs_sentence = ""
+    pt_outputs = pt_model(**pt_inputs).logits[0]
+pt_outputs_sentence = "pytorch: "
 for i, id in enumerate(tokenizer.encode(text)):
     if id == tokenizer.mask_token_id:
-        tokens = tokenizer.convert_ids_to_tokens(outputs[i].topk(k=5)[1])
-        outputs_sentence += "[" + "||".join(tokens) + "]"
+        tokens = tokenizer.convert_ids_to_tokens(pt_outputs[i].topk(k=5)[1])
+        pt_outputs_sentence += "[" + "||".join(tokens) + "]"
     else:
-        outputs_sentence += "".join(
+        pt_outputs_sentence += "".join(
             tokenizer.convert_ids_to_tokens([id], skip_special_tokens=True))
-print(outputs_sentence)
-# 今天[天气||天||心情||阳光||空气]很好，我[想||要||打算||准备||喜欢]去公园玩。
+print(pt_outputs_sentence)
+# tf
+tf_outputs = tf_model(**tf_inputs, training=False).logits[0]
+tf_outputs_sentence = "tf: "
+for i, id in enumerate(tokenizer.encode(text)):
+    if id == tokenizer.mask_token_id:
+        tokens = tokenizer.convert_ids_to_tokens(
+            tf.math.top_k(tf_outputs[i], k=5)[1])
+        tf_outputs_sentence += "[" + "||".join(tokens) + "]"
+    else:
+        tf_outputs_sentence += "".join(
+            tokenizer.convert_ids_to_tokens([id], skip_special_tokens=True))
+print(tf_outputs_sentence)
+# pytorch: 今天[天气||天||心情||阳光||空气]很好，我[想||要||打算||准备||喜欢]去公园玩。
+# tf:      今天[天气||天||心情||阳光||空气]很好，我[想||要||打算||准备||喜欢]去公园玩。
 ```
 
 ## 手动权重转换

src/roformer/__init__.py

@@ -15,6 +15,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from typing import TYPE_CHECKING
 
 from transformers.file_utils import (
     _BaseLazyModule,
@@ -61,21 +62,58 @@
         "TFRoFormerModel",
         "TFRoFormerPreTrainedModel",
     ]
-import importlib
-import os
-import sys
+if TYPE_CHECKING:
+    from .configuration_roformer import ROFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP, RoFormerConfig
+    from .tokenization_roformer import CustomBasicTokenizer, RoFormerTokenizer
 
+    if is_torch_available():
+        from .modeling_roformer import (
+            ROFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
+            RoFormerForMaskedLM,
+            RoFormerForMultipleChoice,
+            RoFormerForNextSentencePrediction,
+            RoFormerForPreTraining,
+            RoFormerForQuestionAnswering,
+            RoFormerForSequenceClassification,
+            RoFormerForTokenClassification,
+            RoFormerLayer,
+            RoFormerLMHeadModel,
+            RoFormerModel,
+            RoFormerPreTrainedModel,
+            load_tf_weights_in_roformer,
+        )
 
-class _LazyModule(_BaseLazyModule):
-    """
-    Module class that surfaces all objects but only performs associated imports when the objects are requested.
-    """
+    if is_tf_available():
+        from .modeling_tf_roformer import (
+            TF_ROFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
+            TFRoFormerEmbeddings,
+            TFRoFormerForMaskedLM,
+            TFRoFormerForMultipleChoice,
+            TFRoFormerForNextSentencePrediction,
+            TFRoFormerForPreTraining,
+            TFRoFormerForQuestionAnswering,
+            TFRoFormerForSequenceClassification,
+            TFRoFormerForTokenClassification,
+            TFRoFormerLMHeadModel,
+            TFRoFormerMainLayer,
+            TFRoFormerModel,
+            TFRoFormerPreTrainedModel,
+        )
 
-    __file__ = globals()["__file__"]
-    __path__ = [os.path.dirname(__file__)]
+else:
+    import importlib
+    import os
+    import sys
 
-    def _get_module(self, module_name: str):
-        return importlib.import_module("." + module_name, self.__name__)
+    class _LazyModule(_BaseLazyModule):
+        """
+        Module class that surfaces all objects but only performs associated imports when the objects are requested.
+        """
 
+        __file__ = globals()["__file__"]
+        __path__ = [os.path.dirname(__file__)]
 
-sys.modules[__name__] = _LazyModule(__name__, _import_structure)
+        def _get_module(self, module_name: str):
+            return importlib.import_module("." + module_name, self.__name__)
+
+    sys.modules[__name__] = _LazyModule(__name__, _import_structure)