1.修改部分文档中的typo; 2.支持对transformers embedding

Author: yhcc

fastNLP/embeddings/__init__.py

@@ -16,6 +16,9 @@
     "RobertaEmbedding",
     "RobertaWordPieceEncoder",
 
+    "TransformersEmbedding",
+    "TransformersWordPieceEncoder",
+
     "GPT2Embedding",
     "GPT2WordPieceEncoder",
 
@@ -32,6 +35,7 @@
 from .elmo_embedding import ElmoEmbedding
 from .bert_embedding import BertEmbedding, BertWordPieceEncoder
 from .roberta_embedding import RobertaEmbedding, RobertaWordPieceEncoder
+from .transformers_embedding import TransformersEmbedding, TransformersWordPieceEncoder
 from .gpt2_embedding import GPT2WordPieceEncoder, GPT2Embedding
 from .char_embedding import CNNCharEmbedding, LSTMCharEmbedding
 from .stack_embedding import StackEmbedding

fastNLP/embeddings/roberta_embedding.py

@@ -196,7 +196,7 @@ def __init__(self, model_dir_or_name: str, vocab: Vocabulary, layers: str = '-1'
                  include_cls_sep: bool = False, pooled_cls: bool = False, auto_truncate: bool = False, min_freq=2):
         super().__init__()
 
-        self.tokenzier = RobertaTokenizer.from_pretrained(model_dir_or_name)
+        self.tokenizer = RobertaTokenizer.from_pretrained(model_dir_or_name)
         self.encoder = RobertaModel.from_pretrained(model_dir_or_name)
         # 由于RobertaEmbedding中设置了padding_idx为1, 且使用了非常神奇的position计算方式，所以-2
         self._max_position_embeddings = self.encoder.config.max_position_embeddings - 2
@@ -233,14 +233,14 @@ def __init__(self, model_dir_or_name: str, vocab: Vocabulary, layers: str = '-1'
                 word = '<unk>'
             elif vocab.word_count[word]<min_freq:
                 word = '<unk>'
-            word_pieces = self.tokenzier.tokenize(word)
-            word_pieces = self.tokenzier.convert_tokens_to_ids(word_pieces)
+            word_pieces = self.tokenizer.tokenize(word)
+            word_pieces = self.tokenizer.convert_tokens_to_ids(word_pieces)
             word_to_wordpieces.append(word_pieces)
             word_pieces_lengths.append(len(word_pieces))
-        self._cls_index = self.tokenzier.encoder['<s>']
-        self._sep_index = self.tokenzier.encoder['</s>']
+        self._cls_index = self.tokenizer.encoder['<s>']
+        self._sep_index = self.tokenizer.encoder['</s>']
         self._word_pad_index = vocab.padding_idx
-        self._wordpiece_pad_index = self.tokenzier.encoder['<pad>']  # 需要用于生成word_piece
+        self._wordpiece_pad_index = self.tokenizer.encoder['<pad>']  # 需要用于生成word_piece
         self.word_to_wordpieces = np.array(word_to_wordpieces)
         self.register_buffer('word_pieces_lengths', torch.LongTensor(word_pieces_lengths))
         logger.debug("Successfully generate word pieces.")
@@ -352,20 +352,19 @@ def forward(self, words):
         return outputs
 
     def save(self, folder):
-        self.tokenzier.save_pretrained(folder)
+        self.tokenizer.save_pretrained(folder)
         self.encoder.save_pretrained(folder)
 
 
 class RobertaWordPieceEncoder(nn.Module):
     r"""
-    读取bert模型，读取之后调用index_dataset方法在dataset中生成word_pieces这一列。
+    读取roberta模型，读取之后调用index_dataset方法在dataset中生成word_pieces这一列。
 
     RobertaWordPieceEncoder可以支持自动下载权重，当前支持的模型:
         en: roberta-base
         en-large: roberta-large
 
     """
-
     def __init__(self, model_dir_or_name: str = 'en', layers: str = '-1', pooled_cls: bool = False,
                  word_dropout=0, dropout=0, requires_grad: bool = True, **kwargs):
         r"""
@@ -417,11 +416,10 @@ def index_datasets(self, *datasets, field_name, add_cls_sep=True, add_prefix_spa
 
     def forward(self, word_pieces, token_type_ids=None):
         r"""
-        计算words的bert embedding表示。传入的words中应该自行包含[CLS]与[SEP]的tag。
+        计算words的bert embedding表示。传入的words中应该自行包含<s>与</s>>的tag。
 
         :param words: batch_size x max_len
-        :param token_type_ids: batch_size x max_len, 用于区分前一句和后一句话. 如果不传入，则自动生成(大部分情况，都不需要输入),
-            第一个[SEP]及之前为0, 第二个[SEP]及到第一个[SEP]之间为1; 第三个[SEP]及到第二个[SEP]之间为0，依次往后推。
+        :param token_type_ids: batch_size x max_len, 用于区分前一句和后一句话. 如果不传入，则自动生成(大部分情况，都不需要输入)。
         :return: torch.FloatTensor. batch_size x max_len x (768*len(self.layers))
         """
         word_pieces = self.drop_word(word_pieces)
@@ -484,7 +482,7 @@ class _WordPieceRobertaModel(nn.Module):
     def __init__(self, model_dir_or_name: str, layers: str = '-1', pooled_cls: bool=False):
         super().__init__()
 
-        self.tokenzier = RobertaTokenizer.from_pretrained(model_dir_or_name)
+        self.tokenizer = RobertaTokenizer.from_pretrained(model_dir_or_name)
         self.encoder = RobertaModel.from_pretrained(model_dir_or_name)
         #  检查encoder_layer_number是否合理
         encoder_layer_number = len(self.encoder.encoder.layer)
@@ -504,25 +502,25 @@ def __init__(self, model_dir_or_name: str, layers: str = '-1', pooled_cls: bool=
                 assert layer <= encoder_layer_number, f"The layer index:{layer} is out of scope for " \
                     f"a RoBERTa model with {encoder_layer_number} layers."
 
-        self._cls_index = self.tokenzier.encoder['<s>']
-        self._sep_index = self.tokenzier.encoder['</s>']
-        self._wordpiece_pad_index = self.tokenzier.encoder['<pad>']  # 需要用于生成word_piece
-        self._wordpiece_unknown_index = self.tokenzier.encoder['<unk>']
+        self._cls_index = self.tokenizer.encoder['<s>']
+        self._sep_index = self.tokenizer.encoder['</s>']
+        self._wordpiece_pad_index = self.tokenizer.encoder['<pad>']  # 需要用于生成word_piece
+        self._wordpiece_unknown_index = self.tokenizer.encoder['<unk>']
         self.pooled_cls = pooled_cls
 
     def index_datasets(self, *datasets, field_name, add_cls_sep=True, add_prefix_space=True):
         r"""
-        使用bert的tokenizer新生成word_pieces列加入到datasets中，并将他们设置为input。如果首尾不是
-            [CLS]与[SEP]会在首尾额外加入[CLS]与[SEP], 且将word_pieces这一列的pad value设置为了bert的pad value。
+        使用roberta的tokenizer新生成word_pieces列加入到datasets中，并将他们设置为input。如果首尾不是
+            <s>与</s>会在首尾额外加入<s>与</s>, 且将word_pieces这一列的pad value设置为了bert的pad value。
 
         :param datasets: DataSet对象
-        :param field_name: 基于哪一列index
+        :param field_name: 基于哪一列index, 这一列一般是raw_string
         :param bool add_cls_sep: 是否在句首句尾添加cls和sep的index
         :param bool add_prefix_space: 是否在句子开头添加空格，预训练时RoBERTa该值为True
         :return:
         """
 
-        encode_func = partial(self.tokenzier.encode, add_special_tokens=add_cls_sep, add_prefix_space=add_prefix_space)
+        encode_func = partial(self.tokenizer.encode, add_special_tokens=add_cls_sep, add_prefix_space=add_prefix_space)
 
         for index, dataset in enumerate(datasets):
             try:
@@ -555,5 +553,5 @@ def forward(self, word_pieces):
         return outputs
 
     def save(self, folder):
-        self.tokenzier.save_pretrained(folder)
+        self.tokenizer.save_pretrained(folder)
         self.encoder.save_pretrained(folder)