fastNLP/envs models embeddings 文档

fastNLP/core/controllers/trainer.py

@@ -305,7 +305,7 @@ class Trainer(TrainerEventTrigger):
           :class:`~fastNLP.core.drivers.paddle_driver.PaddleSingleDriver`；
         * *fairscale_kwargs* -- ``FairScaleDriver`` 所需的其它参数，详见 :class:`~fastNLP.core.drivers.torch_driver.FairScaleDriver`；
         * *deepspeed_kwargs* -- ``DeepSpeedDriver`` 所需的其它参数，详见 :class:`~fastNLP.core.drivers.torch_driver.DeepSpeedDriver`；
-        * *torch_kwargs* -- ``OneflowDriver`` 所需的其它参数，详见 :class:`~fastNLP.core.drivers.oneflow_driver.OneflowSingleDriver` 和
+        * *oneflow_kwargs* -- ``OneflowDriver`` 所需的其它参数，详见 :class:`~fastNLP.core.drivers.oneflow_driver.OneflowSingleDriver` 和
           :class:`~fastNLP.core.drivers.oneflow_driver.OneflowDDPDriver`；
         * *data_device* -- 一个具体的 driver 实例中，有 ``model_device`` 和 ``data_device``，前者表示模型所在的设备，后者表示
          当 ``model_device`` 为 None 时应当将数据迁移到哪个设备；

fastNLP/embeddings/torch/char_embedding.py

@@ -1,6 +1,6 @@
 r"""
-该文件中主要包含的是character的Embedding，包括基于CNN与LSTM的character Embedding。与其它Embedding一样，这里的Embedding输入也是
-词的index而不需要使用词语中的char的index来获取表达。
+该文件中主要包含的是 character 的 Embedding ，包括基于 CNN 与 LSTM 的 character Embedding。与其它 Embedding 一样，这里的 Embedding 输入也是
+词的 index 而不需要使用词语中的 char 的 index 来获取表达。
 """
 
 __all__ = [
@@ -30,8 +30,8 @@ from ...core.vocabulary import Vocabulary
 
 class CNNCharEmbedding(TokenEmbedding):
     r"""
-    使用 ``CNN`` 生成 ``character embedding``。``CNN`` 的结构为, char_embed(x) -> Dropout(x) -> CNN(x) -> activation(x) -> pool -> fc -> Dropout.
-    不同的 ``kernel`` 大小的 ``fitler`` 结果是拼起来然后通过一层``fully connected layer,`` 然后输出``word``的表示。
+    使用 ``CNN`` 生成 ``character embedding``。``CNN`` 的结构为：char_embed(x) -> Dropout(x) -> CNN(x) -> activation(x) -> pool -> fc -> Dropout.
+    不同的 ``kernel`` 大小的 ``fitler`` 结果是拼起来然后通过一层 **全连接层** 然后输出 ``word`` 的表示。
 
     Example::
 
@@ -43,32 +43,33 @@ class CNNCharEmbedding(TokenEmbedding):
         >>> words = torch.LongTensor([[vocab.to_index(word) for word in "The whether is good .".split()]])
         >>> outputs = embed(words)
         >>> outputs.size()
-        # torch.Size([1, 5，50])
+        torch.Size([1, 5，50])
 
+    :param vocab: 词表
+    :param embed_size: 该 :class:`CNNCharEmbedding` 的输出维度大小。
+    :param char_emb_size: character 的 embed 的维度。character 是从 ``vocab`` 中生成的。
+    :param word_dropout: 按照一定概率随机将 word 设置为 ``unk_index`` ，这样可以使得 ``<UNK>`` 这个 token 得到足够的训练，
+        且会对网络有一定的 regularize 作用。
+    :param dropout: 以多大的概率 drop 分布式表示与 char embedding 的输出。
+    :param filter_nums: filter 的数量。长度需要和 ``kernel_sizes`` 一致。
+    :param kernel_sizes: kernel 的大小。
+    :param pool_method: character 的表示在合成一个表示时所使用的 pool 池化方法，支持 ``['avg', 'max']`` 。
+    :param activation: CNN 之后使用的激活方法，支持 ``['relu', 'sigmoid', 'tanh']`` 或者自定义函数。
+    :param min_char_freq: character 的最少出现次数。
+    :param pre_train_char_embed: 可以有两种方式调用预训练好的 :class:`CNNCharEmbedding` ：
+    
+            1. 传入 embedding 文件夹（文件夹下应该只有一个以 **.txt** 作为后缀的文件）或文件路径；
+            2. 传入 embedding 的名称，第二种情况将自动查看缓存中是否存在该模型，没有的话将自动下载；
+            3. 如果输入为 ``None`` 则使用 ``embedding_dim`` 的维度随机初始化一个 embedding；
+    :param requires_grad: 是否更新权重
+    :param include_word_start_end: 是否在每个 word 开始的 character 前和结束的 character 增加特殊标示符号
     """
     
     def __init__(self, vocab: Vocabulary, embed_size: int = 50, char_emb_size: int = 50, word_dropout: float = 0,
                  dropout: float = 0, filter_nums: List[int] = (40, 30, 20), kernel_sizes: List[int] = (5, 3, 1),
                  pool_method: str = 'max', activation='relu', min_char_freq: int = 2, pre_train_char_embed: str = None,
                  requires_grad:bool=True, include_word_start_end:bool=True):
-        r"""
-        
-        :param vocab: 词表
-        :param embed_size: 该CNNCharEmbedding的输出维度大小，默认值为50.
-        :param char_emb_size: character的embed的维度。character是从vocab中生成的。默认值为50.
-        :param word_dropout: 以多大的概率将一个词替换为unk。这样既可以训练unk也是一定的regularize。
-        :param dropout: 以多大的概率drop分布式表示与char embedding的输出。
-        :param filter_nums: filter的数量. 长度需要和kernels一致。默认值为[40, 30, 20].
-        :param kernel_sizes: kernel的大小. 默认值为[5, 3, 1].
-        :param pool_method: character的表示在合成一个表示时所使用的pool方法，支持'avg', 'max'.
-        :param activation: CNN之后使用的激活方法，支持'relu', 'sigmoid', 'tanh' 或者自定义函数.
-        :param min_char_freq: character的最少出现次数。默认值为2.
-        :param pre_train_char_embed: 可以有两种方式调用预训练好的character embedding：第一种是传入embedding文件夹
-            (文件夹下应该只有一个以.txt作为后缀的文件)或文件路径；第二种是传入embedding的名称，第二种情况将自动查看缓存中是否存在该模型，
-            没有的话将自动下载。如果输入为None则使用embedding_dim的维度随机初始化一个embedding.
-        :param requires_grad: 是否更新权重
-        :param include_word_start_end: 是否在每个word开始的character前和结束的character增加特殊标示符号；
-        """
+
         super(CNNCharEmbedding, self).__init__(vocab, word_dropout=word_dropout, dropout=dropout)
         
         for kernel in kernel_sizes:
@@ -128,10 +129,10 @@ class CNNCharEmbedding(TokenEmbedding):
 
     def forward(self, words):
         r"""
-        输入words的index后，生成对应的words的表示。
+        输入 ``words`` 的 index 后，生成对应的 ``words`` 的表示。
 
-        :param words: [batch_size, max_len]
-        :return: [batch_size, max_len, embed_size]
+        :param words: 形状为 ``[batch_size, max_len]``
+        :return: 形状为 ``[batch_size, max_len, embed_size]`` 的结果
         """
         words = self.drop_word(words)
         batch_size, max_len = words.size()
@@ -161,7 +162,7 @@ class CNNCharEmbedding(TokenEmbedding):
 
 class LSTMCharEmbedding(TokenEmbedding):
     r"""
-    使用 ``LSTM`` 的方式对 ``character`` 进行 ``encode``. embed(x) -> Dropout(x) -> LSTM(x) -> activation(x) -> pool -> Dropout
+    使用 ``LSTM`` 的方式对 ``character`` 进行 ``encode``。结构为：embed(x) -> Dropout(x) -> LSTM(x) -> activation(x) -> pool -> Dropout 。
 
     Example::
 
@@ -175,30 +176,32 @@ class LSTMCharEmbedding(TokenEmbedding):
         >>> outputs.size()
         >>> # torch.Size([1, 5，50])
 
+    :param vocab: 词表
+    :param embed_size: :class:`LSTMCharEmbedding` 的输出维度。
+    :param char_emb_size: character 的 embedding 的维度。
+    :param word_dropout: 按照一定概率随机将 word 设置为 ``unk_index`` ，这样可以使得 ``<UNK>`` 这个 token 得到足够的训练，
+        且会对网络有一定的 regularize 作用。
+    :param dropout: 以多大的概率 drop 分布式表示与 char embedding 的输出。
+    :param hidden_size: ``LSTM`` 的中间 hidden 的大小，如果为 ``bidirectional`` 的，hidden 会除二。
+    :param pool_method: character 的表示在合成一个表示时所使用的 pool 池化方法，支持 ``['avg', 'max']`` 。
+    :param activation: LSTM 之后使用的激活方法，支持 ``['relu', 'sigmoid', 'tanh']`` 或者自定义函数。
+    :param min_char_freq: character 的最少出现次数。
+    :param bidirectional: 是否使用双向的 LSTM 进行 encode。
+    :param pre_train_char_embed: 可以有两种方式调用预训练好的 :class:`LSTMCharEmbedding` ：
+    
+            1. 传入 embedding 文件夹（文件夹下应该只有一个以 **.txt** 作为后缀的文件）或文件路径；
+            2. 传入 embedding 的名称，第二种情况将自动查看缓存中是否存在该模型，
+               没有的话将自动下载；
+            3. 如果输入为 ``None`` 则使用 ``embedding_dim`` 的维度随机初始化一个 embedding；
+    :param requires_grad: 是否更新权重
+    :param include_word_start_end: 是否在每个 word 开始的 character 前和结束的 character 增加特殊标示符号
     """
     
     def __init__(self, vocab: Vocabulary, embed_size: int = 50, char_emb_size: int = 50, word_dropout: float = 0,
                  dropout: float = 0, hidden_size=50, pool_method: str = 'max', activation='relu',
                  min_char_freq: int = 2, bidirectional=True, pre_train_char_embed: str = None,
                  requires_grad:bool=True, include_word_start_end:bool=True):
-        r"""
-        
-        :param vocab: 词表
-        :param embed_size: LSTMCharEmbedding的输出维度。默认值为50.
-        :param char_emb_size: character的embedding的维度。默认值为50.
-        :param word_dropout: 以多大的概率将一个词替换为unk。这样既可以训练unk也是一定的regularize。
-        :param dropout: 以多大概率drop character embedding的输出以及最终的word的输出。
-        :param hidden_size: LSTM的中间hidden的大小，如果为bidirectional的，hidden会除二，默认为50.
-        :param pool_method: 支持'max', 'avg'。
-        :param activation: 激活函数，支持'relu', 'sigmoid', 'tanh', 或者自定义函数.
-        :param min_char_freq: character的最小出现次数。默认值为2.
-        :param bidirectional: 是否使用双向的LSTM进行encode。默认值为True。
-        :param pre_train_char_embed: 可以有两种方式调用预训练好的character embedding：第一种是传入embedding文件夹
-            (文件夹下应该只有一个以.txt作为后缀的文件)或文件路径；第二种是传入embedding的名称，第二种情况将自动查看缓存中是否存在该模型，
-            没有的话将自动下载。如果输入为None则使用embedding_dim的维度随机初始化一个embedding.
-        :param requires_grad: 是否更新权重
-        :param include_word_start_end: 是否在每个word开始的character前和结束的character增加特殊标示符号；
-        """
+
         super(LSTMCharEmbedding, self).__init__(vocab, word_dropout=word_dropout, dropout=dropout)
         
         assert hidden_size % 2 == 0, "Only even kernel is allowed."
@@ -256,10 +259,10 @@ class LSTMCharEmbedding(TokenEmbedding):
     
     def forward(self, words):
         r"""
-        输入words的index后，生成对应的words的表示。
+        输入 ``words`` 的 index 后，生成对应的 ``words`` 的表示。
 
-        :param words: [batch_size, max_len]
-        :return: [batch_size, max_len, embed_size]
+        :param words: 形状为 ``[batch_size, max_len]``
+        :return: 形状为 ``[batch_size, max_len, embed_size]`` 的结果
         """
         words = self.drop_word(words)
         batch_size, max_len = words.size()

fastNLP/embeddings/torch/embedding.py

@@ -1,5 +1,5 @@
 r"""
-该模块中的Embedding主要用于随机初始化的embedding(更推荐使用 :class:`fastNLP.embeddings.StaticEmbedding` )，或按照预训练权重初始化Embedding。
+该模块中的 :class:`Embedding` 主要用于随机初始化的 embedding （更推荐使用 :class:`fastNLP.embeddings.torch.StaticEmbedding` ），或按照预训练权重初始化 Embedding。
 
 """
 
@@ -36,19 +36,20 @@ class Embedding(Module):
         >>> init_embed = np.zeros((2000, 100))
         >>> embed = Embedding(init_embed)  # 使用numpy.ndarray的值作为初始化值初始化一个Embedding
 
+    :param init_embed: 支持传入 Embedding 的大小。支持以下类型：
+
+            1. 传入 tuple(int, int)，第一个 int 为 ``vocab_size``, 第二个 int  ``为embed_dim``；
+            2. 传入 :class:`Tensor`, :class:`Embedding`, :class:`numpy.ndarray` 等则直接使用该值初始化 Embedding；
+    
+    :param word_dropout: 按照一定概率随机将 word 设置为 ``unk_index`` ，这样可以使得 ``<UNK>`` 这个 token 得到足够的训练，
+        且会对网络有一定的 regularize 作用。设置该值时，必须同时设置 ``unk_index``。
+    :param dropout: 对 Embedding 的输出的 dropout。
+    :param unk_index: drop word 时替换为的 index。**fastNLP** 的 :class:`fastNLP.Vocabulary`` 的 ``unk_index`` 默认为 1。
     """
     
     def __init__(self, init_embed:Union[Tuple[int,int],'torch.FloatTensor','nn.Embedding',np.ndarray],
                  word_dropout:float=0, dropout:float=0.0, unk_index:int=None):
-        r"""
-        
-        :param init_embed: 支持传入Embedding的大小(传入tuple(int, int),
-            第一个int为vocab_zie, 第二个int为embed_dim); 或传入Tensor, Embedding, numpy.ndarray等则直接使用该值初始化Embedding;
-        :param word_dropout: 按照一定概率随机将word设置为unk_index，这样可以使得unk这个token得到足够的训练, 且会对网络有
-            一定的regularize的作用。设置该值时，必须同时设置unk_index
-        :param dropout: 对Embedding的输出的dropout。
-        :param unk_index: drop word时替换为的index。fastNLP的Vocabulary的unk_index默认为1。
-        """
+
         super(Embedding, self).__init__()
         
         self.embed = get_embeddings(init_embed)
@@ -69,10 +70,10 @@ class Embedding(Module):
         self.unk_index = unk_index
         self.word_dropout = word_dropout
     
-    def forward(self, words):
+    def forward(self, words: "torch.LongTensor") -> "torch.Tensor":
         r"""
-        :param torch.LongTensor words: [batch, seq_len]
-        :return: torch.Tensor : [batch, seq_len, embed_dim]
+        :param words: 形状为 ``[batch, seq_len]``
+        :return: 形状为 ``[batch, seq_len, embed_dim]`` 的张量
         """
         if self.word_dropout > 0 and self.training:
             mask = torch.ones_like(words).float() * self.word_dropout
@@ -102,7 +103,11 @@ class Embedding(Module):
     @property
     def requires_grad(self):
         r"""
-        Embedding的参数是否允许优化。True: 所有参数运行优化; False: 所有参数不允许优化; None: 部分允许优化、部分不允许
+        Embedding 的参数是否允许优化：
+        
+            - ``True`` -- 所有参数运行优化
+            - ``False`` -- 所有参数不允许优化
+            - ``None`` -- 部分允许优化、部分不允许
         :return:
         """
         if not isinstance(self.embed, TokenEmbedding):

fastNLP/embeddings/torch/stack_embedding.py

@@ -20,7 +20,7 @@ from .utils import _check_vocab_has_same_index
 
 class StackEmbedding(TokenEmbedding):
     r"""
-    支持将多个embedding集合成一个embedding。
+    支持将多个 embedding 集合成一个 embedding。
 
     Example::
 
@@ -31,16 +31,16 @@ class StackEmbedding(TokenEmbedding):
         >>> embed_2 = StaticEmbedding(vocab, model_dir_or_name='en-word2vec-300', requires_grad=True)
         >>> embed = StackEmbedding([embed_1, embed_2])
 
+    :param embeds: 一个由若干个 :class:`~fastNLP.embeddings.torch.embedding.TokenEmbedding` 组成的 :class:`list` ，要求
+        每一个 ``TokenEmbedding`` 的词表都保持一致
+    :param word_dropout: 按照一定概率随机将 word 设置为 ``unk_index`` ，这样可以使得 ``<UNK>`` 这个 token 得到足够的训练，
+        且会对网络有一定的 regularize 作用。不同 embedidng 会在相同的位置被设置为 ``<UNK>`` 。 如果这里设置了 dropout，则
+        组成的 embedding 就不要再设置 dropout 了。
+    :param dropout: 以多大的概率对 embedding 的表示进行 Dropout。0.1 即随机将 10% 的值置为 0。
     """
     
     def __init__(self, embeds: List[TokenEmbedding], word_dropout=0, dropout=0):
-        r"""
-        
-        :param embeds: 一个由若干个TokenEmbedding组成的list，要求每一个TokenEmbedding的词表都保持一致
-        :param word_dropout: 以多大的概率将一个词替换为unk。这样既可以训练unk也是一定的regularize。不同embedidng会在相同的位置
-            被设置为unknown。如果这里设置了dropout，则组成的embedding就不要再设置dropout了。
-        :param dropout: 以多大的概率对embedding的表示进行Dropout。0.1即随机将10%的值置为0。
-        """
+
         vocabs = []
         for embed in embeds:
             if hasattr(embed, 'get_word_vocab'):
@@ -59,9 +59,10 @@ class StackEmbedding(TokenEmbedding):
     
     def append(self, embed: TokenEmbedding):
         r"""
-        添加一个embedding到结尾。
+        添加一个 embedding 到结尾。
+
         :param embed:
-        :return:
+        :return: 自身
         """
         assert isinstance(embed, TokenEmbedding)
         _check_vocab_has_same_index(self.get_word_vocab(), embed.get_word_vocab())
@@ -71,8 +72,9 @@ class StackEmbedding(TokenEmbedding):
     
     def pop(self):
         r"""
-        弹出最后一个embed
-        :return:
+        弹出最后一个 embedding
+
+        :return: 被弹出的 embedding
         """
         embed = self.embeds.pop()
         self._embed_size -= embed.embed_size
@@ -81,17 +83,16 @@ class StackEmbedding(TokenEmbedding):
     @property
     def embed_size(self):
         r"""
-        该Embedding输出的vector的最后一维的维度。
-        :return:
+        该 Embedding 输出的 vector 的最后一维的维度。
         """
         return self._embed_size
     
     def forward(self, words):
         r"""
-        得到多个embedding的结果，并把结果按照顺序concat起来。
+        得到多个 embedding 的结果，并把结果按照顺序连接起来。
 
-        :param words: batch_size x max_len
-        :return: 返回的shape和当前这个stack embedding中embedding的组成有关
+        :param words: 形状为 ``[batch_size, max_len]``
+        :return: 形状和当前这个 :class:`StackEmbedding` 中 embedding 的组成有关
         """
         outputs = []
         words = self.drop_word(words)

fastNLP/embeddings/torch/static_embedding.py

@@ -10,7 +10,7 @@ import os
 from collections import defaultdict
 from copy import deepcopy
 import json
-from typing import Union
+from typing import Callable, Union
 
 import numpy as np
 
@@ -34,29 +34,27 @@ STATIC_EMBED_FILENAME = 'static.txt'
 
 class StaticEmbedding(TokenEmbedding):
     r"""
-    StaticEmbedding组件. 给定预训练embedding的名称或路径，根据vocab从embedding中抽取相应的数据(只会将出现在vocab中的词抽取出来，
-    如果没有找到，则会随机初始化一个值(但如果该word是被标记为no_create_entry的话，则不会单独创建一个值，而是会被指向unk的index))。
-    当前支持自动下载的预训练vector有:
-    
-    .. code::
-    
-        en: 实际为en-glove-840b-300d(常用)
-        en-glove-6b-50d: glove官方的50d向量
-        en-glove-6b-100d: glove官方的100d向量
-        en-glove-6b-200d: glove官方的200d向量
-        en-glove-6b-300d: glove官方的300d向量
-        en-glove-42b-300d: glove官方使用42B数据训练版本
-        en-glove-840b-300d:
-        en-glove-twitter-27b-25d:
-        en-glove-twitter-27b-50d:
-        en-glove-twitter-27b-100d:
-        en-glove-twitter-27b-200d:
-        en-word2vec-300d: word2vec官方发布的300d向量
-        en-fasttext-crawl: fasttext官方发布的300d英文预训练
-        cn-char-fastnlp-100d: fastNLP训练的100d的character embedding
-        cn-bi-fastnlp-100d: fastNLP训练的100d的bigram embedding
-        cn-tri-fastnlp-100d: fastNLP训练的100d的trigram embedding
-        cn-fasttext: fasttext官方发布的300d中文预训练embedding
+    ``StaticEmbedding`` 组件。给定预训练 embedding 的名称或路径，根据 ``vocab`` 从 embedding 中抽取相应的数据（只会将出现在 ``vocab`` 中的词抽取出来，
+    如果没有找到，则会随机初始化一个值；但如果该 word 是被标记为 ``no_create_entry`` 的话，则不会单独创建一个值，而是被指向 ``<UNK>`` 的 index）。
+    当前支持自动下载的预训练 vector 有:
+
+        - ``en`` -- 实际为 ``en-glove-840b-300d`` （常用）
+        - ``en-glove-6b-50d`` -- **glove** 官方的 50d 向量
+        - ``en-glove-6b-100d`` -- **glove** 官方的 100d 向量
+        - ``en-glove-6b-200d`` -- **glove** 官方的 200d 向量
+        - ``en-glove-6b-300d`` -- **glove** 官方的 300d 向量
+        - ``en-glove-42b-300d`` -- **glove** 官方使用 42B 数据训练版本
+        - ``en-glove-840b-300d``
+        - ``en-glove-twitter-27b-25d``
+        - ``en-glove-twitter-27b-50d``
+        - ``en-glove-twitter-27b-100d``
+        - ``en-glove-twitter-27b-200d``
+        - ``en-word2vec-300d`` -- **word2vec** 官方发布的 300d 向量
+        - ``en-fasttext-crawl`` -- **fasttext** 官方发布的 300d 英文预训练
+        - ``cn-char-fastnlp-100d`` -- **fastNLP** 训练的 100d 的 character embedding
+        - ``cn-bi-fastnlp-100d`` -- **fastNLP** 训练的 100d 的 bigram embedding
+        - ``cn-tri-fastnlp-100d`` -- **fastNLP** 训练的 100d 的 trigram embedding
+        - ``cn-fasttext`` -- **fasttext** 官方发布的 300d 中文预训练 embedding
 
     Example::
         
@@ -77,34 +75,34 @@ class StaticEmbedding(TokenEmbedding):
                      [ 0.5773,  0.7251, -0.3104,  0.0777,  0.4849],
                      [ 0.5773,  0.7251, -0.3104,  0.0777,  0.4849]]],
                    grad_fn=<EmbeddingBackward>)  # 每种word的输出是一致的。
+        
+    :param vocab: 词表。``StaticEmbedding`` 只会加载包含在词表中的词的词向量，在预训练向量中没找到的使用随机初始化
+    :param model_dir_or_name: 可以有两种方式调用预训练好的 :class:`StaticEmbedding` ：
+    
+            1. 传入 embedding 文件夹（文件夹下应该只有一个以 **.txt** 作为后缀的文件）或文件路径；
+            2. 传入 embedding 的名称，第二种情况将自动查看缓存中是否存在该模型，没有的话将自动下载;
+            3. 如果输入为 ``None`` 则使用 ``embedding_dim`` 的维度随机初始化一个 embedding；
+    :param embedding_dim: 随机初始化的 embedding 的维度，当该值为大于 0 的值时，将忽略 ``model_dir_or_name`` 。
+    :param requires_grad: 是否需要梯度。
+    :param init_method: 如何初始化没有找到的值。可以使用 :mod:`torch.nn.init` 中的各种方法，传入的方法应该接受一个 tensor，并
+        inplace 地修改其值。
+    :param lower: 是否将 ``vocab`` 中的词语小写后再和预训练的词表进行匹配。如果你的词表中包含大写的词语，或者就是需要单独
+        为大写的词语开辟一个 vector 表示，则将 ``lower`` 设置为 ``False``。
+    :param dropout: 以多大的概率对 embedding 的表示进行 Dropout。0.1 即随机将 10% 的值置为 0。
+    :param word_dropout: 按照一定概率随机将 word 设置为 ``unk_index`` ，这样可以使得 ``<UNK>`` 这个 token 得到足够的训练，
+        且会对网络有一定的 regularize 作用。
+    :param normalize: 是否对 vector 进行 ``normalize`` ，使得每个 vector 的 norm 为 1。
+    :param min_freq: Vocabulary 词频数小于这个数量的 word 将被指向 ``<UNK>``。
+    :kwargs:
+        * *only_train_min_freq* (*bool*) -- 仅对 train 中的词语使用 ``min_freq`` 筛选
+        * *only_norm_found_vector* (*bool*) -- 默认为 ``False``，是否仅对在预训练中找到的词语使用 ``normalize``
+        * *only_use_pretrain_word* (*bool*) -- 默认为 ``False``，仅使用出现在 pretrain 词表中的词，如果该词没有在预训练的词表中出现
+          则为 ``<UNK>`` 。如果 embedding 不需要更新建议设置为 ``True`` 。
 
     """
     
     def __init__(self, vocab: Vocabulary, model_dir_or_name: Union[str, None] = 'en', embedding_dim=-1, requires_grad: bool = True,
-                 init_method=None, lower=False, dropout=0, word_dropout=0, normalize=False, min_freq=1, **kwargs):
-        r"""
-        
-        :param Vocabulary vocab: 词表. StaticEmbedding只会加载包含在词表中的词的词向量，在预训练向量中没找到的使用随机初始化
-        :param model_dir_or_name: 可以有两种方式调用预训练好的static embedding：第一种是传入embedding文件夹(文件夹下应该只有一个
-            以.txt作为后缀的文件)或文件路径；第二种是传入embedding的名称，第二种情况将自动查看缓存中是否存在该模型，没有的话将自动下载。
-            如果输入为None则使用embedding_dim的维度随机初始化一个embedding。
-        :param embedding_dim: 随机初始化的embedding的维度，当该值为大于0的值时，将忽略model_dir_or_name。
-        :param requires_grad: 是否需要gradient. 默认为True
-        :param callable init_method: 如何初始化没有找到的值。可以使用torch.nn.init.*中各种方法, 传入的方法应该接受一个tensor，并
-            inplace地修改其值。
-        :param lower: 是否将vocab中的词语小写后再和预训练的词表进行匹配。如果你的词表中包含大写的词语，或者就是需要单独
-            为大写的词语开辟一个vector表示，则将lower设置为False。
-        :param dropout: 以多大的概率对embedding的表示进行Dropout。0.1即随机将10%的值置为0。
-        :param word_dropout: 以多大的概率将一个词替换为unk。这样既可以训练unk也是一定的regularize。
-        :param normalize: 是否对vector进行normalize，使得每个vector的norm为1。
-        :param min_freq: Vocabulary词频数小于这个数量的word将被指向unk。
-        :param kwargs:
-            * only_train_min_freq * (*bool*) -- 仅对 train 中的词语使用 ``min_freq`` 筛选;
-            * only_norm_found_vector * (*bool*) -- 默认为False, 是否仅对在预训练中找到的词语使用normalize;
-            * only_use_pretrain_word * (*bool*) -- 默认为False, 仅使用出现在pretrain词表中的词，如果该词没有在预训练的词表中出现
-             则为unk。如果embedding不需要更新建议设置为True。
-
-        """
+                 init_method: Callable = None, lower=False, dropout=0, word_dropout=0, normalize=False, min_freq=1, **kwargs):
         super(StaticEmbedding, self).__init__(vocab, word_dropout=word_dropout, dropout=dropout)
         if embedding_dim > 0:
             if model_dir_or_name:
@@ -327,12 +325,12 @@ class StaticEmbedding(TokenEmbedding):
 
             return vectors
     
-    def forward(self, words):
+    def forward(self, words: "torch.LongTensor") -> "torch.FloatTensor":
         r"""
-        传入words的index
+        传入 ``words`` 的 index
 
-        :param words: torch.LongTensor, [batch_size, max_len]
-        :return: torch.FloatTensor, [batch_size, max_len, embed_size]
+        :param words: 形状为 ``[batch, seq_len]``
+        :return: 形状为 ``[batch, seq_len, embed_dim]`` 的张量
         """
         if hasattr(self, 'words_to_words'):
             words = self.words_to_words[words]
@@ -341,14 +339,16 @@ class StaticEmbedding(TokenEmbedding):
         words = self.dropout(words)
         return words
 
-    def save(self, folder):
+    def save(self, folder: str):
         """
-        将embedding存储到folder下，之后可以通过使用load方法读取
+        将 embedding 存储到 ``folder`` 下，之后可以通过使用 :meth:`load` 方法读取
+
+        :param folder: 会在该 ``folder`` 下生成三个文件：
 
-        :param str folder: 会在该folder下生成三个文件, vocab.txt, static_embed_hyper.txt, static_embed_hyper.json.
-            其中vocab.txt可以用Vocabulary通过load读取; embedding.txt按照word2vec的方式存储，以空格的方式隔开元素,
-            第一行只有两个元素，剩下的行首先是word然后是各个维度的值; static_embed_hyper.json是StaticEmbedding的超参数
-        :return:
+                - ``vocab.txt``，可以通过 :meth:`fastNLP.core.Vocabulary.load` 读取；
+                - ``embedding.txt`` 按照 *word2vec* 的方式存储，以空格的方式隔开元素，第一行只有两个元素，剩下的行首先是
+                  word 然后是各个维度的值；
+                - ``static_embed_hyper.json``，:class:`StaticEmbedding` 的超参数；
         """
         os.makedirs(folder, exist_ok=True)
 
@@ -391,11 +391,11 @@ class StaticEmbedding(TokenEmbedding):
         logger.debug(f"StaticEmbedding has been saved to {folder}.")
 
     @classmethod
-    def load(cls, folder):
+    def load(cls, folder: str):
         """
 
-        :param str folder: 该folder下应该有以下三个文件vocab.txt, static_embed.txt, static_hyper.json
-        :return:
+        :param folder: 该 ``folder`` 下应该有以下三个文件 ``vocab.txt``, ``static_embed.txt``, ``static_hyper.json``
+        :return: 加载后的 embedding
         """
         for name in [VOCAB_FILENAME, STATIC_EMBED_FILENAME, STATIC_HYPER_FILENAME]:
             assert os.path.exists(os.path.join(folder, name)), f"{name} not found in {folder}."

fastNLP/embeddings/torch/utils.py

@@ -36,15 +36,17 @@ def _construct_char_vocab_from_vocab(vocab: Vocabulary, min_freq: int = 1, inclu
 
 def get_embeddings(init_embed, padding_idx=None):
     r"""
-    根据输入的init_embed返回Embedding对象。如果输入是tuple, 则随机初始化一个nn.Embedding; 如果输入是numpy.ndarray, 则按照ndarray
-    的值将nn.Embedding初始化; 如果输入是torch.Tensor, 则按该值初始化nn.Embedding; 如果输入是fastNLP中的embedding将不做处理
-    返回原对象。
-
-    :param init_embed: 可以是 tuple:(num_embedings, embedding_dim), 即embedding的大小和每个词的维度;也可以传入
-        nn.Embedding 对象, 此时就以传入的对象作为embedding; 传入np.ndarray也行，将使用传入的ndarray作为作为Embedding初始化;
-        传入torch.Tensor, 将使用传入的值作为Embedding初始化。
-    :param padding_idx: 当传入tuple时，padding_idx有效
-    :return nn.Embedding:  embeddings
+    根据输入的 ``init_embed`` 返回 ``Embedding`` 对象。
+
+    :param init_embed: 支持以下几种输入类型：
+
+            - ``tuple(num_embedings, embedding_dim)``，即 embedding 的大小和每个词的维度，此时将随机初始化一个 :class:`torch.nn.Embedding` 实例；
+            - :class:`torch.nn.Embedding` 或 **fastNLP** 的 ``Embedding`` 对象，此时就以传入的对象作为 embedding；
+            - :class:`numpy.ndarray` ，将使用传入的 ndarray 作为 Embedding 初始化；
+            - :class:`torch.Tensor`，此时将使用传入的值作为 Embedding 初始化；
+
+    :param padding_idx: 当传入 :class:`tuple` 时，``padding_idx`` 有效
+    :return:
     """
     if isinstance(init_embed, tuple):
         res = nn.Embedding(
@@ -64,14 +66,14 @@ def get_embeddings(init_embed, padding_idx=None):
     return res
 
 
-def get_sinusoid_encoding_table(n_position, d_hid, padding_idx=None):
+def get_sinusoid_encoding_table(n_position: int, d_hid: int, padding_idx=None) -> "torch.FloatTensor":
     """
-    sinusoid的embedding，其中position的表示中，偶数维(0,2,4,...)是sin, 奇数(1,3,5...)是cos
+    sinusoid 的 embedding，其中 ``position`` 的表示中，偶数维 ``(0,2,4,...)`` 是 sin，奇数 ``(1,3,5...)`` 是 cos。
 
-    :param int n_position: 一共多少个position
+    :param int n_position: 一共多少个 position
     :param int d_hid: 多少维度，需要为偶数
     :param padding_idx:
-    :return: torch.FloatTensor, shape为n_position x d_hid
+    :return: 形状为 ``[n_position, d_hid]`` 的张量
     """
 
     def cal_angle(position, hid_idx):

fastNLP/envs/distributed.py

@@ -18,9 +18,9 @@ from fastNLP.envs.env import FASTNLP_GLOBAL_RANK, FASTNLP_NO_SYNC
 
 def is_cur_env_distributed() -> bool:
     """
-    单卡模式该函数一定返回 False；
-    注意进程 0 在多卡的训练模式下前后的值是不一样的，例如在开启多卡的 driver 之前，在进程 0 上的该函数返回 False；但是在开启后，在进程 0 上
-    的该函数返回的值是 True；多卡模式下除了进程 0 外的其它进程返回的值一定是 True；
+    判断当前是否处于分布式的环境下。单卡模式该函数一定返回 ``False``；
+    注意进程 0 在多卡的训练模式下前后的值是不一样的，例如在开启多卡的 driver 之前，在进程 0 上的该函数返回 ``False`` ；但是在开启后，在进程 0 上
+    的该函数返回的值是 ``True`` ；多卡模式下除了进程 0 外的其它进程返回的值一定是 ``True`` 。
     """
     return FASTNLP_GLOBAL_RANK in os.environ
 
@@ -50,8 +50,8 @@ def rank_zero_call(fn: Callable):
             return a+b
         rank_zero_call(add)(1, 2)
 
-    同时，该函数还会设置 FASTNLP_NO_SYNC 为 2，在这个环境下，所有的 fastNLP 内置的 barrier 接口，gather/broadcast 操作都没有任何
-    意义。
+    同时，该函数还会设置环境变量 ``FASTNLP_NO_SYNC`` 为 **2** ，在这个环境下，所有的 **fastNLP** 内置的 :meth:`barrier` 接口和 ``gather`` / ``broadcast``
+    操作都没有任何意义。
 
     :param fn: 需要包裹的可执行的函数。
     :return:
@@ -66,13 +66,12 @@ def rank_zero_call(fn: Callable):
 
 
 @contextmanager
-def fastnlp_no_sync_context(level=2):
+def fastnlp_no_sync_context(level: int = 2):
     """
-    用于让 fastNLP 的 barrier 以及 gather/broadcast等操作等同于只有 1 卡的多卡程序。如果为 1 表示 fastNLP 里的barrier 操作失效；
-    如果为 2 表示 barrier 与 gather/broadcast 都失效。
+    用于让 **fastNLP** 的 :meth:`barrier` 以及 ``gather`` / ``broadcast`` 等操作等同于只有 1 卡的多卡程序。如果为 1 表示 **fastNLP** 里的
+    :meth:`barrier` 操作失效；如果为 2 表示 :meth:`barrier` 与 ``gather`` / ``broadcast`` 都失效。
 
-    :param int level: 可选 [0, 1, 2]
-    :return:
+    :param level: 可选 ``[0, 1, 2]``
     """
     old_level = os.environ.get(FASTNLP_NO_SYNC, None)
     os.environ[FASTNLP_NO_SYNC] = f'{level}'
@@ -86,15 +85,13 @@ def fastnlp_no_sync_context(level=2):
 @contextmanager
 def all_rank_call_context():
     """
-    在多卡模式下，该环境内，会暂时地将 FASTNLP_GLOBAL_RANK 设置为 "0"，使得 rank_zero_call 函数失效，使得每个进程都会运行该函数。
+    在多卡模式下，该环境内，会暂时地将 ``FASTNLP_GLOBAL_RAN``K 设置为 **"0"** ，使得 :func:`rank_zero_call` 函数失效，使得每个进程都会运行该函数。
 
     使用方式::
 
         with all_rank_call_context():
             do_something  # all rank will do
 
-    :param fn:
-    :return:
     """
     old_fastnlp_global_rank = os.environ[FASTNLP_GLOBAL_RANK] if FASTNLP_GLOBAL_RANK in os.environ else None
     os.environ[FASTNLP_GLOBAL_RANK] = '0'
@@ -109,12 +106,11 @@ def all_rank_call_context():
 
 def rank_zero_rm(path: Optional[Union[str, Path]]):
     """
-    这个是因为在分布式文件系统中可能会发生错误，rank0下发删除成功后就运行走了，但实际的删除需要rank0的机器发送到远程文件系统再去执行，这个时候
-    在rank0那里，确实已经删除成功了，但是在远程文件系统那里这个操作还没完成，rank1读取的时候还是读取到存在这个文件；
-    该函数会保证所有进程都检测到 path 删除之后才退出，请保证不同进程上 path 是完全一样的，否则会陷入死锁状态。
+    仅在 rank 0 下删除文件的函数。普通的删除文件操作在分布式文件系统中可能会发生错误，rank 0 下发删除成功后就运行走了，但实际的删除需要 rank 0 的机器
+    发送到远程文件系统再去执行，这个时候在 rank 0 已经删除成功了，但是在远程文件系统那里这个操作还没完成，rank 1 读取的时候还是读取到存在这个文件；
+    该函数会保证所有进程都检测到 ``path`` 删除之后才退出，请保证不同进程上 ``path`` 是完全一样的，否则会陷入死锁状态。
 
     :param path:
-    :return:
     """
     if int(os.environ.get(FASTNLP_GLOBAL_RANK, 0)) == 0:
         if path is None:

fastNLP/envs/set_backend.py

@@ -10,6 +10,7 @@ from fastNLP.envs.utils import _module_available, get_gpu_count
 
 SUPPORT_BACKENDS = ['torch', 'paddle', 'jittor', 'oneflow']
 
+__all__ = []
 
 def _set_backend():
     """
@@ -152,11 +153,13 @@ def set_env(global_seed=None):
 def dump_fastnlp_backend(default:bool = False, backend=None):
     """
     将 fastNLP 的设置写入到 ~/.fastNLP/envs/ 文件夹下，
-        若 default 为 True，则保存的文件为 ~/.fastNLP/envs/default.json 。
-        如 default 为 False，则保存的文件为 ~/.fastNLP/envs/{CONDA_DEFAULT_ENV}.json ，当CONDA_DEFAULT_ENV这个环境变量不存在时
+
+        - 若 default 为 True，则保存的文件为 ~/.fastNLP/envs/default.json 。
+        - 如 default 为 False，则保存的文件为 ~/.fastNLP/envs/{CONDA_DEFAULT_ENV}.json ，当CONDA_DEFAULT_ENV这个环境变量不存在时
         ，报错。
+
     当 fastNLP 被 import 时，会默认尝试从 ~/.fastNLP/envs/{CONDA_DEFAULT_ENV}.json 读取配置文件，如果文件不存在，则尝试从
-     ~/.fastNLP/envs/default.json （如果有）读取环境变量。不过这些变量的优先级低于代码运行时的环境变量注入。
+    ~/.fastNLP/envs/default.json （如果有）读取环境变量。不过这些变量的优先级低于代码运行时的环境变量注入。
 
     会保存的环境变量为 FASTNLP_BACKEND 。
 

fastNLP/envs/set_env_on_import.py

@@ -5,6 +5,7 @@ import sys
 from .env import *
 import datetime
 
+__all__ = []
 
 def remove_local_rank_in_argv():
     """

fastNLP/models/torch/biaffine_parser.py

@@ -1,5 +1,5 @@
 r"""
-Biaffine Dependency Parser 的 Pytorch 实现.
+**Biaffine Dependency Parser** 的 Pytorch 实现.
 """
 __all__ = [
     "BiaffineParser",
@@ -125,7 +125,7 @@ def _find_cycle(vertices, edges):
 
 class GraphParser(nn.Module):
     r"""
-    基于图的parser base class, 支持贪婪解码和最大生成树解码
+    基于图的 parser base class，支持 **贪婪解码** 和 **最大生成树解码**
     """
     
     def __init__(self):
@@ -134,12 +134,12 @@ class GraphParser(nn.Module):
     @staticmethod
     def greedy_decoder(arc_matrix, mask=None):
         r"""
-        贪心解码方式, 输入图, 输出贪心解码的parsing结果, 不保证合法的构成树
+        贪心解码方式，输入图，输出贪心解码的 parsing 结果，不保证合法地构成树。
 
-        :param arc_matrix: [batch, seq_len, seq_len] 输入图矩阵
-        :param mask: [batch, seq_len] 输入图的padding mask, 有内容的部分为 1, 否则为 0.
-            若为 ``None`` 时, 默认为全1向量. Default: ``None``
-        :return heads: [batch, seq_len] 每个元素在树中对应的head(parent)预测结果
+        :param arc_matrix: 输入图矩阵，形状为 ``[batch, seq_len, seq_len]``。
+        :param mask: 输入图的padding mask，形状为 ``[batch, seq_len]`` ， 有内容的部分为 **1** ， 否则为 **0** 。
+            若为 ``None`` ，则默认为全1向量。
+        :return: 每个元素在树中对应的 ``head(parent)`` 预测结果，形状为 ``[batch, seq_len]``。
         """
         _, seq_len, _ = arc_matrix.shape
         matrix = arc_matrix + torch.diag(arc_matrix.new(seq_len).fill_(-np.inf))
@@ -153,12 +153,12 @@ class GraphParser(nn.Module):
     @staticmethod
     def mst_decoder(arc_matrix, mask=None):
         r"""
-        用最大生成树算法, 计算parsing结果, 保证输出合法的树结构
+        用最大生成树算法，计算 parsing 结果，保证输出合法的树结构
 
-        :param arc_matrix: [batch, seq_len, seq_len] 输入图矩阵
-        :param mask: [batch, seq_len] 输入图的padding mask, 有内容的部分为 1, 否则为 0.
-            若为 ``None`` 时, 默认为全1向量. Default: ``None``
-        :return heads: [batch, seq_len] 每个元素在树中对应的head(parent)预测结果
+        :param arc_matrix: 输入图矩阵，形状为 ``[batch, seq_len, seq_len]``。
+        :param mask: 输入图的padding mask，形状为 ``[batch, seq_len]`` ， 有内容的部分为 **1** ， 否则为 **0** 。
+            若为 ``None`` ，则默认为全1向量。
+        :return: 每个元素在树中对应的 ``head(parent)`` 预测结果，形状为 ``[batch, seq_len]``。
         """
         batch_size, seq_len, _ = arc_matrix.shape
         matrix = arc_matrix.clone()
@@ -238,9 +238,27 @@ class LabelBilinear(nn.Module):
 
 class BiaffineParser(GraphParser):
     r"""
-    Biaffine Dependency Parser 实现.
-    论文参考 `Deep Biaffine Attention for Neural Dependency Parsing (Dozat and Manning, 2016) <https://arxiv.org/abs/1611.01734>`_ .
-
+    **Biaffine Dependency Parser** 实现。
+    论文参考 `Deep Biaffine Attention for Neural Dependency Parsing (Dozat and Manning, 2016) <https://arxiv.org/abs/1611.01734>`_ 。
+
+    :param embed: 单词词典，支持以下几种输入类型：
+
+            - ``tuple(num_embedings, embedding_dim)``，即 embedding 的大小和每个词的维度，此时将随机初始化一个 :class:`torch.nn.Embedding` 实例；
+            - :class:`torch.nn.Embedding` 或 **fastNLP** 的 ``Embedding`` 对象，此时就以传入的对象作为 embedding；
+            - :class:`numpy.ndarray` ，将使用传入的 ndarray 作为 Embedding 初始化；
+            - :class:`torch.Tensor`，此时将使用传入的值作为 Embedding 初始化；
+
+    :param pos_vocab_size: part-of-speech 词典大小
+    :param pos_emb_dim: part-of-speech 向量维度
+    :param num_label: 边的类别个数
+    :param rnn_layers: rnn encoder 的层数
+    :param rnn_hidden_size: rnn encoder 的隐状态维度
+    :param arc_mlp_size: 边预测的 MLP 维度
+    :param label_mlp_size: 类别预测的 MLP 维度
+    :param dropout: dropout 概率
+    :param encoder: encoder 类别，可选 ``['lstm', 'var-lstm', 'transformer']``。
+    :param use_greedy_infer: 是否在 inference 时使用 :meth:`贪心算法 <GraphParser.greedy_decoder>` ，若为 ``False`` ，
+        将使用更加精确但相对缓慢的 :meth:`MST算法 <GraphParser.mst_decoder>` 。
     """
     
     def __init__(self,
@@ -255,23 +273,6 @@ class BiaffineParser(GraphParser):
                  dropout=0.3,
                  encoder='lstm',
                  use_greedy_infer=False):
-        r"""
-        
-        :param embed: 单词词典, 可以是 tuple, 包括(num_embedings, embedding_dim), 即
-            embedding的大小和每个词的维度. 也可以传入 nn.Embedding 对象,
-            此时就以传入的对象作为embedding
-        :param pos_vocab_size: part-of-speech 词典大小
-        :param pos_emb_dim: part-of-speech 向量维度
-        :param num_label: 边的类别个数
-        :param rnn_layers: rnn encoder的层数
-        :param rnn_hidden_size: rnn encoder 的隐状态维度
-        :param arc_mlp_size: 边预测的MLP维度
-        :param label_mlp_size: 类别预测的MLP维度
-        :param dropout: dropout概率.
-        :param encoder: encoder类别, 可选 ('lstm', 'var-lstm', 'transformer'). Default: lstm
-        :param use_greedy_infer: 是否在inference时使用贪心算法.
-            若 ``False`` , 使用更加精确但相对缓慢的MST算法. Default: ``False``
-        """
         super(BiaffineParser, self).__init__()
         rnn_out_size = 2 * rnn_hidden_size
         word_hid_dim = pos_hid_dim = rnn_hidden_size
@@ -336,20 +337,19 @@ class BiaffineParser(GraphParser):
                     nn.init.normal_(p, 0, 0.1)
     
     def forward(self, words1, words2, seq_len, target1=None):
-        r"""模型forward阶段
+        r"""
+        模型 forward 阶段
 
-        :param words1: [batch_size, seq_len] 输入word序列
-        :param words2: [batch_size, seq_len] 输入pos序列
-        :param seq_len: [batch_size, seq_len] 输入序列长度
-        :param target1: [batch_size, seq_len] 输入真实标注的heads, 仅在训练阶段有效,
-            用于训练label分类器. 若为 ``None`` , 使用预测的heads输入到label分类器
-            Default: ``None``
-        :return dict: parsing
-                结果::
+        :param words1: 输入 word 序列，形状为 ``[batch_size, seq_len]``
+        :param words2: 输入 pos 序列，形状为 ``[batch_size, seq_len]`` 
+        :param seq_len: 输入序列长度，形状为 ``[batch_size, seq_len]`` 
+        :param target1: 输入真实标注的 heads ，形状为 ``[batch_size, seq_len]`` ，仅在训练阶段有效，
+            用于训练 label 分类器. 若为 ``None`` ，则使用预测的 heads 输入到 label 分类器。
+        :return: 类型为字典的 parsing 结果，各个键的含义为：
 
-                    pred1: [batch_size, seq_len, seq_len] 边预测logits
-                    pred2: [batch_size, seq_len, num_label] label预测logits
-                    pred3: [batch_size, seq_len] heads的预测结果, 在 ``target1=None`` 时预测
+                * ``pred1`` -- **边** 预测 logits，形状为 ``[batch_size, seq_len, seq_len]``；
+                * ``pred2`` -- **label** 预测 logits，形状为 ``[batch_size, seq_len, num_label]`` ；
+                * ``pred3`` -- **heads** 的预测结果，形状为 ``[batch_size, seq_len]`` ，在 ``target1=None`` 时预测；
 
         """
         # prepare embeddings
@@ -416,6 +416,17 @@ class BiaffineParser(GraphParser):
         return res_dict
 
     def train_step(self, words1, words2, seq_len, target1, target2):
+        """
+        模型的训练接口。
+
+        :param words1: 输入 word 序列，形状为 ``[batch_size, seq_len]``
+        :param words2: 输入 pos 序列，形状为 ``[batch_size, seq_len]`` 
+        :param target1: 输入真实标注的 heads ，形状为 ``[batch_size, seq_len]`` ，仅在训练阶段有效，
+            用于训练 label 分类器. 若为 ``None`` ，则使用预测的 heads 输入到 label 分类器。
+        :param target2: 真实类别的标注，形状为 ``[batch_size, seq_len]``
+        :param seq_len: 输入序列长度，形状为 ``[batch_size, seq_len]``
+        :return: 类型为字典的结果，仅包含一个键 ``loss``，表示当次训练的 loss
+        """
         res = self(words1, words2, seq_len, target1)
         arc_pred = res['pred1']
         label_pred = res['pred2']
@@ -425,14 +436,14 @@ class BiaffineParser(GraphParser):
     @staticmethod
     def loss(pred1, pred2, target1, target2, seq_len):
         r"""
-        计算parser的loss
-
-        :param pred1: [batch_size, seq_len, seq_len] 边预测logits
-        :param pred2: [batch_size, seq_len, num_label] label预测logits
-        :param target1: [batch_size, seq_len] 真实边的标注
-        :param target2: [batch_size, seq_len] 真实类别的标注
-        :param seq_len: [batch_size, seq_len] 真实目标的长度
-        :return loss: scalar
+        计算 parser 的 loss
+
+        :param pred1: 边预测 logits，形状为 ``[batch_size, seq_len, seq_len]``
+        :param pred2: **label** 预测 logits，形状为 ``[batch_size, seq_len, num_label]``
+        :param target1: 真实边的标注，形状为 ``[batch_size, seq_len]``
+        :param target2: 真实类别的标注，形状为 ``[batch_size, seq_len]``
+        :param seq_len: 真实目标的长度，形状为 ``[batch_size, seq_len]``
+        :return: 计算出的 loss。
         """
         
         batch_size, length, _ = pred1.shape
@@ -456,14 +467,13 @@ class BiaffineParser(GraphParser):
     def evaluate_step(self, words1, words2, seq_len):
         r"""模型预测API
 
-        :param words1: [batch_size, seq_len] 输入word序列
-        :param words2: [batch_size, seq_len] 输入pos序列
-        :param seq_len: [batch_size, seq_len] 输入序列长度
-        :return dict: parsing
-                结果::
+        :param words1: 输入 word 序列，形状为 ``[batch_size, seq_len]``
+        :param words2: 输入 pos 序列，形状为 ``[batch_size, seq_len]`` 
+        :param seq_len: 输入序列长度，形状为 ``[batch_size, seq_len]``
+        :return: 字典类型的 parsing 结果，各个键的含义为：
 
-                    pred1: [batch_size, seq_len] heads的预测结果
-                    pred2: [batch_size, seq_len, num_label] label预测logits
+                * ``pred1`` -- **heads** 的预测结果，形状为 ``[batch_size, seq_len]``；
+                * ``pred2`` -- **label** 预测 logits，形状为 ``[batch_size, seq_len, num_label]`` ；
 
         """
         res = self(words1, words2, seq_len)

fastNLP/models/torch/cnn_text_classification.py

@@ -7,6 +7,7 @@ __all__ = [
     "CNNText"
 ]
 
+from typing import Union, Tuple
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -18,24 +19,27 @@ from ...modules.torch import encoder
 
 class CNNText(torch.nn.Module):
     r"""
-    使用CNN进行文本分类的模型
-    'Yoon Kim. 2014. Convolution Neural Networks for Sentence Classification.'
-    
-    """
+    使用 **CNN** 进行文本分类的模型。
+    论文参考 `Yoon Kim. 2014. Convolution Neural Networks for Sentence Classification <https://arxiv.org/abs/1408.5882>`_ 。
+
+    :param embed: 单词词典，支持以下几种输入类型：
 
+            - ``tuple(num_embedings, embedding_dim)``，即 embedding 的大小和每个词的维度，此时将随机初始化一个 :class:`torch.nn.Embedding` 实例；
+            - :class:`torch.nn.Embedding` 或 **fastNLP** 的 ``Embedding`` 对象，此时就以传入的对象作为 embedding；
+            - :class:`numpy.ndarray` ，将使用传入的 ndarray 作为 Embedding 初始化；
+            - :class:`torch.Tensor`，此时将使用传入的值作为 Embedding 初始化；
+
+    :param num_classes: 一共有多少类
+    :param kernel_nums: 输出 channel 的 kernel 数目。
+        如果为 :class:`list` 或 :class:`tuple`，则需要与 ``kernel_sizes`` 的大小保持一致。
+    :param kernel_sizes: 输出 channel 的 kernel 大小。
+    :param dropout: Dropout 的大小
+    """
     def __init__(self, embed,
-                 num_classes,
-                 kernel_nums=(30, 40, 50),
-                 kernel_sizes=(1, 3, 5),
-                 dropout=0.5):
-        r"""
-        
-        :param tuple(int,int),torch.FloatTensor,nn.Embedding,numpy.ndarray embed: Embedding的大小(传入tuple(int, int),
-            第一个int为vocab_zie, 第二个int为embed_dim); 如果为Tensor, Embedding, ndarray等则直接使用该值初始化Embedding
-        :param int num_classes: 一共有多少类
-        :param int,tuple(int) kernel_sizes: 输出channel的kernel大小。
-        :param float dropout: Dropout的大小
-        """
+                 num_classes: int,
+                 kernel_nums: Union[int, Tuple[int]] = (30, 40, 50),
+                 kernel_sizes: Union[int, Tuple[int]] = (1, 3, 5),
+                 dropout: float = 0.5):
         super(CNNText, self).__init__()
 
         # no support for pre-trained embedding currently
@@ -47,14 +51,12 @@ class CNNText(torch.nn.Module):
         self.dropout = nn.Dropout(dropout)
         self.fc = nn.Linear(sum(kernel_nums), num_classes)
 
-    def forward(self, words, seq_len=None):
+    def forward(self, words: "torch.LongTensor", seq_len: "torch.LongTensor"=None):
         r"""
 
-        :param torch.LongTensor words: [batch_size, seq_len]，句子中word的index
-        :param torch.LongTensor seq_len:  [batch,] 每个句子的长度
-        :param target: 每个 sample 的目标值。
-
-        :return output:
+        :param words: 句子中 word 的 index，形状为 ``[batch_size, seq_len]``
+        :param seq_len:  每个句子的长度，形状为 ``[batch,]`` 
+        :return: 前向传播的结果，为仅包含一个键 ``pred`` 的字典
         """
         x = self.embed(words)  # [N,L] -> [N,L,C]
         if seq_len is not None:
@@ -70,22 +72,22 @@ class CNNText(torch.nn.Module):
     def train_step(self, words, target, seq_len=None):
         """
 
-        :param words:
-        :param target:
-        :param seq_len:
-        :return:
+        :param words: 句子中 word 的 index，形状为 ``[batch_size, seq_len]``
+        :param target: 每个 sample 的目标值
+        :param seq_len: 每个句子的长度，形状为 ``[batch,]`` 
+        :return: 类型为字典的结果，仅包含一个键 ``loss``，表示当次训练的 loss
         """
         res = self(words, seq_len)
         x = res['pred']
         loss = F.cross_entropy(x, target)
         return {'loss': loss}
 
-    def evaluate_step(self, words, seq_len=None):
+    def evaluate_step(self, words: "torch.LongTensor", seq_len: "torch.LongTensor"=None):
         r"""
-        :param torch.LongTensor words: [batch_size, seq_len]，句子中word的index
-        :param torch.LongTensor seq_len:  [batch,] 每个句子的长度
 
-        :return predict: dict of torch.LongTensor, [batch_size, ]
+        :param words: 句子中 word 的 index，形状为 ``[batch_size, seq_len]``
+        :param seq_len:  每个句子的长度，形状为 ``[batch_size,]`` 
+        :return: 预测结果，仅包含一个键 ``pred``，值为形状为 ``[batch_size,]`` 的 :class:`torch.LongTensor`
         """
         output = self(words, seq_len)
         _, predict = output['pred'].max(dim=1)

fastNLP/models/torch/seq2seq_generator.py

@@ -1,5 +1,3 @@
-r"""undocumented"""
-
 import torch
 from torch import nn
 import torch.nn.functional as F
@@ -13,30 +11,27 @@ __all__ = ['SequenceGeneratorModel']
 
 class SequenceGeneratorModel(nn.Module):
     """
-    通过使用本模型封装seq2seq_model使得其既可以用于训练也可以用于生成。训练的时候，本模型的forward函数会被调用，生成的时候本模型的predict
-        函数会被调用。
+    通过使用本模型封装 seq2seq_model 使得其既可以用于训练也可以用于生成。训练的时候，本模型的 :meth:`forward` 函数会被调用，
+    生成的时候本模型的 :meth:`predict` 函数会被调用。
 
+    :param seq2seq_model: 序列到序列模型
+    :param bos_token_id: 句子开头的 token id
+    :param eos_token_id: 句子结束的 token id
+    :param max_length: 生成句子的最大长度, 每句话的 decode 长度为 ``max_length + max_len_a * src_len``
+    :param max_len_a: 每句话的 decode 长度为 ``max_length + max_len_a*src_len``。如果不为 0，需要保证 State 中包含 encoder_mask
+    :param num_beams: **beam search** 的大小
+    :param do_sample: 是否通过采样的方式生成
+    :param temperature: 只有在 do_sample 为 ``True`` 才有意义
+    :param top_k: 只从 ``top_k`` 中采样
+    :param top_p: 只从 ``top_p`` 的 token 中采样（ **nucleus sampling** ）
+    :param repetition_penalty: 多大程度上惩罚重复的 token
+    :param length_penalty: 对长度的惩罚，**小于 1** 鼓励长句，**大于 1** 鼓励短句
+    :param pad_token_id: 当某句话生成结束之后，之后生成的内容用 ``pad_token_id`` 补充
     """
 
-    def __init__(self, seq2seq_model: Seq2SeqModel, bos_token_id, eos_token_id=None, max_length=30, max_len_a=0.0,
-                 num_beams=1, do_sample=True, temperature=1.0, top_k=50, top_p=1.0,
-                 repetition_penalty=1, length_penalty=1.0, pad_token_id=0):
-        """
-
-        :param Seq2SeqModel seq2seq_model: 序列到序列模型
-        :param int,None bos_token_id: 句子开头的token id
-        :param int,None eos_token_id: 句子结束的token id
-        :param int max_length: 生成句子的最大长度, 每句话的decode长度为max_length + max_len_a*src_len
-        :param float max_len_a: 每句话的decode长度为max_length + max_len_a*src_len。 如果不为0，需要保证State中包含encoder_mask
-        :param int num_beams: beam search的大小
-        :param bool do_sample: 是否通过采样的方式生成
-        :param float temperature: 只有在do_sample为True才有意义
-        :param int top_k: 只从top_k中采样
-        :param float top_p: 只从top_p的token中采样，nucles sample
-        :param float repetition_penalty: 多大程度上惩罚重复的token
-        :param float length_penalty: 对长度的惩罚，小于1鼓励长句，大于1鼓励短剧
-        :param int pad_token_id: 当某句话生成结束之后，之后生成的内容用pad_token_id补充
-        """
+    def __init__(self, seq2seq_model: Seq2SeqModel, bos_token_id: int=None, eos_token_id: int=None, max_length: int=30,
+                 max_len_a: float=0.0, num_beams: int=1, do_sample: bool=True, temperature: float=1.0, top_k: int=50,
+                 top_p: float=1.0, repetition_penalty: float=1, length_penalty: float=1.0, pad_token_id: int=0):
         super().__init__()
         self.seq2seq_model = seq2seq_model
         self.generator = SequenceGenerator(seq2seq_model.decoder, max_length=max_length, max_len_a=max_len_a,
@@ -47,19 +42,28 @@ class SequenceGeneratorModel(nn.Module):
                                            repetition_penalty=repetition_penalty, length_penalty=length_penalty,
                                            pad_token_id=pad_token_id)
 
-    def forward(self, src_tokens, tgt_tokens, src_seq_len=None, tgt_seq_len=None):
+    def forward(self, src_tokens: "torch.LongTensor", tgt_tokens: "torch.LongTensor",
+                src_seq_len: "torch.LongTensor"=None, tgt_seq_len: "torch.LongTensor"=None):
         """
-        透传调用seq2seq_model的forward。
+        调用 seq2seq_model 的 :meth:`forward` 。
 
-        :param torch.LongTensor src_tokens: bsz x max_len
-        :param torch.LongTensor tgt_tokens: bsz x max_len'
-        :param torch.LongTensor src_seq_len: bsz
-        :param torch.LongTensor tgt_seq_len: bsz
-        :return:
+        :param src_tokens: source 的 token，形状为 ``[batch_size, max_len]``
+        :param tgt_tokens: target 的 token，形状为 ``[batch_size, max_len]``
+        :param src_seq_len: source的长度，形状为 ``[batch_size,]``
+        :param tgt_seq_len: target的长度，形状为 ``[batch_size,]``
+        :return: 字典 ``{'pred': torch.Tensor}``, 其中 ``pred`` 的形状为 ``[batch_size, max_len, vocab_size]``
         """
         return self.seq2seq_model(src_tokens, tgt_tokens, src_seq_len, tgt_seq_len)
 
-    def train_step(self, src_tokens, tgt_tokens, src_seq_len=None, tgt_seq_len=None):
+    def train_step(self, src_tokens: "torch.LongTensor", tgt_tokens: "torch.LongTensor",
+                    src_seq_len: "torch.LongTensor"=None, tgt_seq_len: "torch.LongTensor"=None):
+        """
+        :param src_tokens: source 的 token，形状为 ``[batch_size, max_len]``
+        :param tgt_tokens: target 的 token，形状为 ``[batch_size, max_len]``
+        :param src_seq_len: source的长度，形状为 ``[batch_size,]``
+        :param tgt_seq_len: target的长度，形状为 ``[batch_size,]``
+        :return: 字典 ``{'loss': torch.Tensor}``
+        """
         res = self(src_tokens, tgt_tokens, src_seq_len, tgt_seq_len)
         pred = res['pred']
         if tgt_seq_len is not None:
@@ -68,13 +72,13 @@ class SequenceGeneratorModel(nn.Module):
         loss = F.cross_entropy(pred[:, :-1].transpose(1, 2), tgt_tokens[:, 1:])
         return {'loss': loss}
 
-    def evaluate_step(self, src_tokens, src_seq_len=None):
+    def evaluate_step(self, src_tokens: "torch.LongTensor", src_seq_len: "torch.LongTensor"=None):
         """
-        给定source的内容，输出generate的内容。
+        给定 source 的内容，输出 generate 的内容。
 
-        :param torch.LongTensor src_tokens: bsz x max_len
-        :param torch.LongTensor src_seq_len: bsz
-        :return:
+        :param src_tokens: source 的 token，形状为 ``[batch_size, max_len]``
+        :param src_seq_len: source的长度，形状为 ``[batch_size,]``
+        :return: 字典 ``{'pred': torch.Tensor}`` ，表示生成结果
         """
         state = self.seq2seq_model.prepare_state(src_tokens, src_seq_len)
         result = self.generator.generate(state)

fastNLP/models/torch/seq2seq_model.py

@@ -18,31 +18,33 @@ __all__ = ['Seq2SeqModel', 'TransformerSeq2SeqModel', 'LSTMSeq2SeqModel']
 
 
 class Seq2SeqModel(nn.Module):
+    """
+    可以用于在 :class:`~fastNLP.core.controllers.Trainer` 中训练的 **Seq2Seq模型** 。正常情况下，继承了该函数之后，只需要
+    实现 classmethod ``build_model`` 即可。如果需要使用该模型进行生成，需要把该模型输入到 :class:`~fastNLP.models.torch.SequenceGeneratorModel`
+    中。在本模型中， :meth:`forward` 会把 encoder 后的结果传入到 decoder 中，并将 decoder 的输出 output 出来。
+
+    :param encoder: :class:`~fastNLP.modules.torch.encoder.Seq2SeqEncoder` 对象，需要实现对应的 :meth:`forward` 函数，接受两个参数，第一个为
+        ``[batch_size, max_len]`` 的 source tokens, 第二个为 ``[batch_size,]`` 的 source 的长度；需要返回两个 tensor： 
+        
+            - ``encoder_outputs`` : ``[batch_size, max_len, hidden_size]``
+            - ``encoder_mask`` :  ``[batch_size, max_len]``，为 **0** 的地方为 pad。
+        如果encoder的输出或者输入有变化，可以重载本模型的 :meth:`prepare_state` 函数或者 :meth:`forward` 函数。
+    :param decoder: :class:`~fastNLP.modules.torch.decoder.Seq2SeqEncoder` 对象，需要实现 :meth:`init_state` 函数，需要接受两个参数，分别为
+        上述的 ``encoder_outputs`` 和 ``encoder_mask``。若decoder需要更多输入，请重载当前模型的 :meth:`prepare_state` 或 :meth:`forward` 函数。
+    """
     def __init__(self, encoder: Seq2SeqEncoder, decoder: Seq2SeqDecoder):
-        """
-        可以用于在Trainer中训练的Seq2Seq模型。正常情况下，继承了该函数之后，只需要实现classmethod build_model即可。如果需要使用该模型
-            进行生成，需要把该模型输入到 :class:`~fastNLP.models.SequenceGeneratorModel` 中。在本模型中，forward()会把encoder后的
-            结果传入到decoder中，并将decoder的输出output出来。
-
-        :param encoder: Seq2SeqEncoder 对象，需要实现对应的forward()函数，接受两个参数，第一个为bsz x max_len的source tokens, 第二个为
-            bsz的source的长度；需要返回两个tensor: encoder_outputs: bsz x max_len x hidden_size, encoder_mask: bsz x max_len
-            为1的地方需要被attend。如果encoder的输出或者输入有变化，可以重载本模型的prepare_state()函数或者forward()函数
-        :param decoder: Seq2SeqDecoder 对象，需要实现init_state()函数，输出为两个参数，第一个为bsz x max_len x hidden_size是
-            encoder的输出; 第二个为bsz x max_len，为encoder输出的mask，为0的地方为pad。若decoder需要更多输入，请重载当前模型的
-            prepare_state()或forward()函数
-        """
         super().__init__()
         self.encoder = encoder
         self.decoder = decoder
 
-    def forward(self, src_tokens, tgt_tokens, src_seq_len=None, tgt_seq_len=None):
+    def forward(self, src_tokens: "torch.LongTensor", tgt_tokens: "torch.LongTensor",
+                src_seq_len: "torch.LongTensor"=None, tgt_seq_len: "torch.LongTensor"=None):
         """
-
-        :param torch.LongTensor src_tokens: source的token
-        :param torch.LongTensor tgt_tokens: target的token
-        :param torch.LongTensor src_seq_len: src的长度
-        :param torch.LongTensor tgt_seq_len: target的长度，默认用不上
-        :return: {'pred': torch.Tensor}, 其中pred的shape为bsz x max_len x vocab_size
+        :param src_tokens: source 的 token，形状为 ``[batch_size, max_len]``
+        :param tgt_tokens: target 的 token，形状为 ``[batch_size, max_len]``
+        :param src_seq_len: source的长度，形状为 ``[batch_size,]``
+        :param tgt_seq_len: target的长度，形状为 ``[batch_size,]``
+        :return: 字典 ``{'pred': torch.Tensor}``, 其中 ``pred`` 的形状为 ``[batch_size, max_len, vocab_size]``
         """
         state = self.prepare_state(src_tokens, src_seq_len)
         decoder_output = self.decoder(tgt_tokens, state)
@@ -53,7 +55,15 @@ class Seq2SeqModel(nn.Module):
         else:
             raise TypeError(f"Unsupported return type from Decoder:{type(self.decoder)}")
 
-    def train_step(self, src_tokens, tgt_tokens, src_seq_len=None, tgt_seq_len=None):
+    def train_step(self, src_tokens: "torch.LongTensor", tgt_tokens: "torch.LongTensor",
+                    src_seq_len: "torch.LongTensor"=None, tgt_seq_len: "torch.LongTensor"=None):
+        """
+        :param src_tokens: source 的 token，形状为 ``[batch_size, max_len]``
+        :param tgt_tokens: target 的 token，形状为 ``[batch_size, max_len]``
+        :param src_seq_len: source的长度，形状为 ``[batch_size,]``
+        :param tgt_seq_len: target的长度，形状为 ``[batch_size,]``
+        :return: 字典 ``{'loss': torch.Tensor}``
+        """
         res = self(src_tokens, tgt_tokens, src_seq_len, tgt_seq_len)
         pred = res['pred']
         if tgt_seq_len is not None:
@@ -62,13 +72,13 @@ class Seq2SeqModel(nn.Module):
         loss = F.cross_entropy(pred[:, :-1].transpose(1, 2), tgt_tokens[:, 1:])
         return {'loss': loss}
 
-    def prepare_state(self, src_tokens, src_seq_len=None):
+    def prepare_state(self, src_tokens: "torch.LongTensor", src_seq_len: "torch.LongTensor"=None):
         """
-        调用encoder获取state，会把encoder的encoder_output, encoder_mask直接传入到decoder.init_state中初始化一个state
+        调用 encoder 获取 state，会把 encoder 的 ``encoder_output``, ``encoder_mask`` 直接传入到 :meth:`decoder.init_state` 中初始化一个 state
 
-        :param src_tokens:
-        :param src_seq_len:
-        :return:
+        :param src_tokens: source 的 token，形状为 ``[batch_size, max_len]``
+        :param src_seq_len: source的长度，形状为 ``[batch_size,]``
+        :return: decode 初始化的 state
         """
         encoder_output, encoder_mask = self.encoder(src_tokens, src_seq_len)
         state = self.decoder.init_state(encoder_output, encoder_mask)
@@ -77,43 +87,54 @@ class Seq2SeqModel(nn.Module):
     @classmethod
     def build_model(cls, *args, **kwargs):
         """
-        需要实现本方法来进行Seq2SeqModel的初始化
+        需要实现本方法来进行 :class:`Seq2SeqModel` 的初始化
 
         :return:
         """
-        raise NotImplemented
+        raise NotImplementedError("A `Seq2SeqModel` must implement its own classmethod `build_model()`.")
 
 
 class TransformerSeq2SeqModel(Seq2SeqModel):
     """
-    Encoder为TransformerSeq2SeqEncoder, decoder为TransformerSeq2SeqDecoder，通过build_model方法初始化
-
+    Encoder 为 :class:`~fastNLP.modules.torch.encoder.TransformerSeq2SeqEncoder` ，decoder 为 
+    :class:`~fastNLP.modules.torch.decoder.TransformerSeq2SeqDecoder` 的 :class:`Seq2SeqModel` ，
+    通过 :meth:`build_model` 方法初始化。
     """
-
     def __init__(self, encoder, decoder):
         super().__init__(encoder, decoder)
 
     @classmethod
     def build_model(cls, src_embed, tgt_embed=None,
-                    pos_embed='sin', max_position=1024, num_layers=6, d_model=512, n_head=8, dim_ff=2048, dropout=0.1,
-                    bind_encoder_decoder_embed=False,
-                    bind_decoder_input_output_embed=True):
+                    pos_embed: str='sin', max_position: int=1024, num_layers: int=6, d_model: int=512,
+                    n_head: int=8, dim_ff: int=2048, dropout: float=0.1,
+                    bind_encoder_decoder_embed: bool=False,
+                    bind_decoder_input_output_embed: bool=True):
         """
-        初始化一个TransformerSeq2SeqModel
-
-        :param nn.Module, StaticEmbedding, Tuple[int, int] src_embed: source的embedding
-        :param nn.Module, StaticEmbedding, Tuple[int, int] tgt_embed: target的embedding，如果bind_encoder_decoder_embed为
-            True，则不要输入该值
-        :param str pos_embed: 支持sin, learned两种
-        :param int max_position: 最大支持长度
-        :param int num_layers: encoder和decoder的层数
-        :param int d_model: encoder和decoder输入输出的大小
-        :param int n_head: encoder和decoder的head的数量
-        :param int dim_ff: encoder和decoder中FFN中间映射的维度
-        :param float dropout: Attention和FFN dropout的大小
-        :param bool bind_encoder_decoder_embed: 是否对encoder和decoder使用相同的embedding
-        :param bool bind_decoder_input_output_embed: decoder的输出embedding是否与其输入embedding是一样的权重
-        :return: TransformerSeq2SeqModel
+        初始化一个 :class:`TransformerSeq2SeqModel` 。
+
+        :param src_embed: source 的 embedding，支持以下几种输入类型：
+
+                - ``tuple(num_embedings, embedding_dim)``，即 embedding 的大小和每个词的维度，此时将随机初始化一个 :class:`torch.nn.Embedding` 实例；
+                - :class:`torch.nn.Embedding` 或 **fastNLP** 的 ``Embedding`` 对象，此时就以传入的对象作为 embedding；
+                - :class:`numpy.ndarray` ，将使用传入的 ndarray 作为 Embedding 初始化；
+                - :class:`torch.Tensor`，此时将使用传入的值作为 Embedding 初始化；
+        :param tgt_embed: target 的 embedding，如果 ``bind_encoder_decoder_embed``
+            为 ``True`` ，则不要输入该值。支持以下几种输入类型：
+
+                - ``tuple(num_embedings, embedding_dim)``，即 embedding 的大小和每个词的维度，此时将随机初始化一个 :class:`torch.nn.Embedding` 实例；
+                - :class:`torch.nn.Embedding` 或 **fastNLP** 的 ``Embedding`` 对象，此时就以传入的对象作为 embedding；
+                - :class:`numpy.ndarray` ，将使用传入的 ndarray 作为 Embedding 初始化；
+                - :class:`torch.Tensor`，此时将使用传入的值作为 Embedding 初始化；
+        :param pos_embed: 支持 ``['sin', 'learned']`` 两种
+        :param max_position: 最大支持长度
+        :param num_layers: ``encoder`` 和 ``decoder`` 的层数
+        :param d_model: ``encoder`` 和 ``decoder`` 输入输出的大小
+        :param n_head: ``encoder`` 和 ``decoder`` 的 head 的数量
+        :param dim_ff: ``encoder`` 和 ``decoder`` 中 FFN 中间映射的维度
+        :param dropout: Attention 和 FFN dropout的大小
+        :param bind_encoder_decoder_embed: 是否对 ``encoder`` 和 ``decoder`` 使用相同的 embedding
+        :param bind_decoder_input_output_embed: ``decoder`` 的输出 embedding 是否与其输入 embedding 是一样的权重
+        :return: :class:`TransformerSeq2SeqModel` 模型
         """
         if bind_encoder_decoder_embed and tgt_embed is not None:
             raise RuntimeError("If you set `bind_encoder_decoder_embed=True`, please do not provide `tgt_embed`.")
@@ -152,7 +173,8 @@ class TransformerSeq2SeqModel(Seq2SeqModel):
 
 class LSTMSeq2SeqModel(Seq2SeqModel):
     """
-    使用LSTMSeq2SeqEncoder和LSTMSeq2SeqDecoder的model
+    使用 :class:`~fastNLP.modules.torch.encoder.LSTMSeq2SeqEncoder` 和 :class:`~fastNLP.modules.torch.decoder.LSTMSeq2SeqDecoder` 的
+    :class:`Seq2SeqModel`，通过 :meth:`build_model` 方法初始化。
 
     """
     def __init__(self, encoder, decoder):
@@ -160,22 +182,32 @@ class LSTMSeq2SeqModel(Seq2SeqModel):
 
     @classmethod
     def build_model(cls, src_embed, tgt_embed=None,
-                    num_layers = 3, hidden_size = 400, dropout = 0.3, bidirectional=True,
-                    attention=True, bind_encoder_decoder_embed=False,
-                    bind_decoder_input_output_embed=True):
+                    num_layers: int = 3, hidden_size: int = 400, dropout: float = 0.3, bidirectional: bool=True,
+                    attention: bool=True, bind_encoder_decoder_embed: bool=False,
+                    bind_decoder_input_output_embed: bool=True):
         """
 
-        :param nn.Module, StaticEmbedding, Tuple[int, int] src_embed: source的embedding
-        :param nn.Module, StaticEmbedding, Tuple[int, int] tgt_embed: target的embedding，如果bind_encoder_decoder_embed为
-            True，则不要输入该值
-        :param int num_layers: Encoder和Decoder的层数
-        :param int hidden_size: encoder和decoder的隐藏层大小
-        :param float dropout: 每层之间的Dropout的大小
-        :param bool bidirectional: encoder是否使用双向LSTM
-        :param bool attention: decoder是否使用attention attend encoder在所有时刻的状态
-        :param bool bind_encoder_decoder_embed: 是否对encoder和decoder使用相同的embedding
-        :param bool bind_decoder_input_output_embed: decoder的输出embedding是否与其输入embedding是一样的权重
-        :return: LSTMSeq2SeqModel
+        :param src_embed: source 的 embedding，支持以下几种输入类型：
+
+                - ``tuple(num_embedings, embedding_dim)``，即 embedding 的大小和每个词的维度，此时将随机初始化一个 :class:`torch.nn.Embedding` 实例；
+                - :class:`torch.nn.Embedding` 或 **fastNLP** 的 ``Embedding`` 对象，此时就以传入的对象作为 embedding；
+                - :class:`numpy.ndarray` ，将使用传入的 ndarray 作为 Embedding 初始化；
+                - :class:`torch.Tensor`，此时将使用传入的值作为 Embedding 初始化；
+        :param tgt_embed: target 的 embedding，如果 ``bind_encoder_decoder_embed``
+            为 ``True`` ，则不要输入该值，支持以下几种输入类型：
+
+                - ``tuple(num_embedings, embedding_dim)``，即 embedding 的大小和每个词的维度，此时将随机初始化一个 :class:`torch.nn.Embedding` 实例；
+                - :class:`torch.nn.Embedding` 或 **fastNLP** 的 ``Embedding`` 对象，此时就以传入的对象作为 embedding；
+                - :class:`numpy.ndarray` ，将使用传入的 ndarray 作为 Embedding 初始化；
+                - :class:`torch.Tensor`，此时将使用传入的值作为 Embedding 初始化；
+        :param num_layers: ``encoder`` 和 ``decoder`` 的层数
+        :param hidden_size: ``encoder`` 和 ``decoder`` 的隐藏层大小
+        :param dropout: 每层之间的 Dropout 的大小
+        :param bidirectional: ``encoder`` 是否使用 **双向LSTM**
+        :param attention: 是否在 ``decoder`` 中使用 attention 来添加 ``encoder`` 在所有时刻的状态
+        :param bind_encoder_decoder_embed: 是否对 ``encoder`` 和 ``decoder`` 使用相同的 embedding
+        :param bind_decoder_input_output_embed: ``decoder`` 的输出 embedding 是否与其输入 embedding 是一样的权重
+        :return: :class:`LSTMSeq2SeqModel` 模型
         """
         if bind_encoder_decoder_embed and tgt_embed is not None:
             raise RuntimeError("If you set `bind_encoder_decoder_embed=True`, please do not provide `tgt_embed`.")

fastNLP/models/torch/sequence_labeling.py

@@ -1,5 +1,5 @@
 r"""
-本模块实现了几种序列标注模型
+本模块实现了几种序列标注模型。
 """
 __all__ = [
     "SeqLabeling",
@@ -21,20 +21,24 @@ from ...modules.torch.decoder.crf import allowed_transitions
 
 class BiLSTMCRF(nn.Module):
     r"""
-    结构为embedding + BiLSTM + FC + Dropout + CRF.
+    结构为 ``Embedding`` + :class:`BiLSTM <fastNLP.modules.torch.encoder.LSTM>` + ``FC`` + ``Dropout`` +
+    :class:`CRF <fastNLP.modules.torch.decoder.ConditionalRandomField>` 。
 
+    :param embed: 支持以下几种输入类型：
+
+        - ``tuple(num_embedings, embedding_dim)``，即 embedding 的大小和每个词的维度，此时将随机初始化一个 :class:`torch.nn.Embedding` 实例；
+        - :class:`torch.nn.Embedding` 或 **fastNLP** 的 ``Embedding`` 对象，此时就以传入的对象作为 embedding；
+        - :class:`numpy.ndarray` ，将使用传入的 ndarray 作为 Embedding 初始化；
+        - :class:`torch.Tensor`，此时将使用传入的值作为 Embedding 初始化；
+
+    :param num_classes: 一共多少个类
+    :param num_layers: **BiLSTM** 的层数
+    :param hidden_size: **BiLSTM** 的 ``hidden_size``，实际 hidden size 为该值的 **两倍** （前向、后向）
+    :param dropout: dropout 的概率，0 为不 dropout
+    :param target_vocab: :class:`~fastNLP.core.Vocabulary` 对象，target 与 index 的对应关系。如果传入该值，将自动避免非法的解码序列。
     """
     def __init__(self, embed, num_classes, num_layers=1, hidden_size=100, dropout=0.5,
                   target_vocab=None):
-        r"""
-        
-        :param embed: 支持(1)fastNLP的各种Embedding, (2) tuple, 指明num_embedding, dimension, 如(1000, 100)
-        :param num_classes: 一共多少个类
-        :param num_layers: BiLSTM的层数
-        :param hidden_size: BiLSTM的hidden_size，实际hidden size为该值的两倍(前向、后向)
-        :param dropout: dropout的概率，0为不dropout
-        :param target_vocab: Vocabulary对象，target与index的对应关系。如果传入该值，将自动避免非法的解码序列。
-        """
         super().__init__()
         self.embed = get_embeddings(embed)
 
@@ -55,7 +59,13 @@ class BiLSTMCRF(nn.Module):
 
         self.crf = ConditionalRandomField(num_classes, include_start_end_trans=True, allowed_transitions=trans)
 
-    def forward(self, words, seq_len=None, target=None):
+    def forward(self, words: "torch.LongTensor", target: "torch.LongTensor"=None, seq_len: "torch.LongTensor"=None):
+        """
+        :param words: 句子中 word 的 index，形状为 ``[batch_size, seq_len]``
+        :param target: 每个 sample 的目标值
+        :param seq_len: 每个句子的长度，形状为 ``[batch,]``
+        :return: 如果 ``target`` 为 ``None``，则返回预测结果 ``{'pred': torch.Tensor}``，否则返回 loss ``{'loss': torch.Tensor}``
+        """
         words = self.embed(words)
         feats, _ = self.lstm(words, seq_len=seq_len)
         feats = self.fc(feats)
@@ -69,28 +79,40 @@ class BiLSTMCRF(nn.Module):
             loss = self.crf(logits, target, mask).mean()
             return {'loss':loss}
 
-    def train_step(self, words, seq_len, target):
+    def train_step(self, words: "torch.LongTensor", target: "torch.LongTensor", seq_len: "torch.LongTensor"):
+        """
+        :param words: 句子中 word 的 index，形状为 ``[batch_size, seq_len]``
+        :param target: 每个 sample 的目标值
+        :param seq_len: 每个句子的长度，形状为 ``[batch,]``
+        :return: 如果 ``target`` 为 ``None``，则返回预测结果 ``{'pred': torch.Tensor}``，否则返回 loss ``{'loss': torch.Tensor}``
+        """
         return self(words, seq_len, target)
 
-    def evaluate_step(self, words, seq_len):
+    def evaluate_step(self, words: "torch.LongTensor", seq_len: "torch.LongTensor"):
+        """
+        :param words: 句子中 word 的 index，形状为 ``[batch_size, seq_len]``
+        :param seq_len: 每个句子的长度，形状为 ``[batch,]``
+        :return: 预测结果 ``{'pred': torch.Tensor}``
+        """
         return self(words, seq_len)
 
 
 class SeqLabeling(nn.Module):
     r"""
-    一个基础的Sequence labeling的模型。
-    用于做sequence labeling的基础类。结构包含一层Embedding，一层LSTM(单向，一层)，一层FC，以及一层CRF。
-    
+    一个基础的 Sequence labeling 的模型。
+    用于做 sequence labeling 的基础类。结构包含一层 ``Embedding`` ，一层 :class:`~fastNLP.modules.torch.encoder.LSTM` (单向，一层)，
+    一层全连接层，以及一层 :class:`CRF <fastNLP.modules.torch.decoder.ConditionalRandomField>` 。
+
+    :param embed: 支持以下几种输入类型：
+
+            - ``tuple(num_embedings, embedding_dim)``，即 embedding 的大小和每个词的维度，此时将随机初始化一个 :class:`torch.nn.Embedding` 实例；
+            - :class:`torch.nn.Embedding` 或 **fastNLP** 的 ``Embedding`` 对象，此时就以传入的对象作为 embedding；
+            - :class:`numpy.ndarray` ，将使用传入的 ndarray 作为 Embedding 初始化；
+            - :class:`torch.Tensor`，此时将使用传入的值作为 Embedding 初始化；
+    :param hidden_size: :class:`fastNLP.modules.torch.encoder.LSTM` 隐藏层的大小
+    :param num_classes: 一共有多少类
     """
-    
-    def __init__(self, embed, hidden_size, num_classes):
-        r"""
-        
-        :param tuple(int,int),torch.FloatTensor,nn.Embedding,numpy.ndarray embed: Embedding的大小(传入tuple(int, int),
-            第一个int为vocab_zie, 第二个int为embed_dim); 如果为Tensor, embedding, ndarray等则直接使用该值初始化Embedding
-        :param int hidden_size: LSTM隐藏层的大小
-        :param int num_classes: 一共有多少类
-        """
+    def __init__(self, embed, hidden_size: int, num_classes: int):
         super(SeqLabeling, self).__init__()
         
         self.embedding = get_embeddings(embed)
@@ -98,11 +120,11 @@ class SeqLabeling(nn.Module):
         self.fc = nn.Linear(hidden_size, num_classes)
         self.crf = decoder.ConditionalRandomField(num_classes)
 
-    def forward(self, words, seq_len):
+    def forward(self, words: "torch.LongTensor", seq_len: "torch.LongTensor"):
         r"""
-        :param torch.LongTensor words: [batch_size, max_len]，序列的index
-        :param torch.LongTensor seq_len: [batch_size,], 这个序列的长度
-        :return
+        :param words: 句子中 word 的 index，形状为 ``[batch_size, seq_len]``
+        :param seq_len: 每个句子的长度，形状为 ``[batch,]``
+        :return: 预测结果 ``{'pred': torch.Tensor}``
         """
         x = self.embedding(words)
         # [batch_size, max_len, word_emb_dim]
@@ -112,19 +134,23 @@ class SeqLabeling(nn.Module):
         return {'pred': x}
         # [batch_size, max_len, num_classes]
 
-    def train_step(self, words, seq_len, target):
+    def train_step(self, words, target, seq_len):
+        """
+        :param words: 句子中 word 的 index，形状为 ``[batch_size, seq_len]``
+        :param target: 每个 sample 的目标值
+        :param seq_len: 每个句子的长度，形状为 ``[batch,]``
+        :return: 如果 ``target`` 为 ``None``，则返回预测结果 ``{'pred': torch.Tensor}``，否则返回 loss ``{'loss': torch.Tensor}``
+        """
         res = self(words, seq_len)
         pred = res['pred']
         mask = seq_len_to_mask(seq_len, max_len=target.size(1))
         return {'loss': self._internal_loss(pred, target, mask)}
 
     def evaluate_step(self, words, seq_len):
-        r"""
-        用于在预测时使用
-
-        :param torch.LongTensor words: [batch_size, max_len]
-        :param torch.LongTensor seq_len: [batch_size,]
-        :return: {'pred': xx}, [batch_size, max_len]
+        """
+        :param words: 句子中 word 的 index，形状为 ``[batch_size, seq_len]``
+        :param seq_len: 每个句子的长度，形状为 ``[batch,]``
+        :return: 预测结果 ``{'pred': torch.Tensor}``
         """
         mask = seq_len_to_mask(seq_len, max_len=words.size(1))
 
@@ -158,22 +184,25 @@ class SeqLabeling(nn.Module):
 
 class AdvSeqLabel(nn.Module):
     r"""
-    更复杂的Sequence Labelling模型。结构为Embedding, LayerNorm, 双向LSTM(两层)，FC，LayerNorm，DropOut，FC，CRF。
+    更复杂的 Sequence Labelling 模型。结构为 ``Embedding``, ``LayerNorm``, :class:`BiLSTM <fastNLP.modules.torch.encoder.LSTM>` （两层），
+    ``FC``，``LayerNorm``，``Dropout``，``FC``，:class:`CRF <fastNLP.modules.torch.decoder.ConditionalRandomField>`。
+
+    :param embed: 支持以下几种输入类型：
+
+        - ``tuple(num_embedings, embedding_dim)``，即 embedding 的大小和每个词的维度，此时将随机初始化一个 :class:`torch.nn.Embedding` 实例；
+        - :class:`torch.nn.Embedding` 或 **fastNLP** 的 ``Embedding`` 对象，此时就以传入的对象作为 embedding；
+        - :class:`numpy.ndarray` ，将使用传入的 ndarray 作为 Embedding 初始化；
+        - :class:`torch.Tensor`，此时将使用传入的值作为 Embedding 初始化；
+    :param hidden_size: :class:`~fastNLP.modules.torch.LSTM` 的隐藏层大小
+    :param num_classes: 有多少个类
+    :param dropout: :class:`~fastNLP.modules.torch.LSTM` 中以及 DropOut 层的 drop 概率
+    :param id2words: tag id 转为其 tag word 的表。用于在 CRF 解码时防止解出非法的顺序，比如 **'BMES'** 这个标签规范中，``'S'``
+        不能出现在 ``'B'`` 之后。这里也支持类似于 ``'B-NN'``，即 ``'-'`` 前为标签类型的指示，后面为具体的 tag 的情况。这里不但会保证
+        ``'B-NN'`` 后面不为 ``'S-NN'`` 还会保证 ``'B-NN'`` 后面不会出现 ``'M-xx'`` （任何非 ``'M-NN'`` 和 ``'E-NN'`` 的情况）。
+    :param encoding_type: 支持 ``["BIO", "BMES", "BEMSO"]`` ，只有在 ``id2words`` 不为 ``None`` 的情况有用。
     """
     
-    def __init__(self, embed, hidden_size, num_classes, dropout=0.3, id2words=None, encoding_type='bmes'):
-        r"""
-        
-        :param tuple(int,int),torch.FloatTensor,nn.Embedding,numpy.ndarray embed: Embedding的大小(传入tuple(int, int),
-            第一个int为vocab_zie, 第二个int为embed_dim); 如果为Tensor, Embedding, ndarray等则直接使用该值初始化Embedding
-        :param int hidden_size: LSTM的隐层大小
-        :param int num_classes: 有多少个类
-        :param float dropout: LSTM中以及DropOut层的drop概率
-        :param dict id2words: tag id转为其tag word的表。用于在CRF解码时防止解出非法的顺序，比如'BMES'这个标签规范中，'S'
-            不能出现在'B'之后。这里也支持类似与'B-NN'，即'-'前为标签类型的指示，后面为具体的tag的情况。这里不但会保证
-            'B-NN'后面不为'S-NN'还会保证'B-NN'后面不会出现'M-xx'(任何非'M-NN'和'E-NN'的情况。)
-        :param str encoding_type: 支持"BIO", "BMES", "BEMSO", 只有在id2words不为None的情况有用。
-        """
+    def __init__(self, embed, hidden_size: int, num_classes: int, dropout: float=0.3, id2words: dict=None, encoding_type: str='bmes'):
         super().__init__()
         
         self.Embedding = get_embeddings(embed)
@@ -217,13 +246,12 @@ class AdvSeqLabel(nn.Module):
         total_loss = self.Crf(x, y, mask)
         return torch.mean(total_loss)
     
-    def forward(self, words, seq_len, target=None):
-        r"""
-        :param torch.LongTensor words: [batch_size, mex_len]
-        :param torch.LongTensor seq_len:[batch_size, ]
-        :param torch.LongTensor target: [batch_size, max_len]
-        :return y: If truth is None, return list of [decode path(list)]. Used in testing and predicting.
-                   If truth is not None, return loss, a scalar. Used in training.
+    def forward(self, words: "torch.LongTensor", target: "torch.LongTensor"=None, seq_len: "torch.LongTensor"=None):
+        """
+        :param words: 句子中 word 的 index，形状为 ``[batch_size, seq_len]``
+        :param target: 每个 sample 的目标值
+        :param seq_len: 每个句子的长度，形状为 ``[batch,]``
+        :return: 如果 ``target`` 为 ``None``，则返回预测结果 ``{'pred': torch.Tensor}``，否则返回 loss ``{'loss': torch.Tensor}``
         """
         
         words = words.long()
@@ -251,21 +279,19 @@ class AdvSeqLabel(nn.Module):
         else:
             return {"pred": self._decode(x, mask)}
     
-    def train_step(self, words, seq_len, target):
-        r"""
-        
-        :param torch.LongTensor words: [batch_size, mex_len]
-        :param torch.LongTensor seq_len: [batch_size, ]
-        :param torch.LongTensor target: [batch_size, max_len], 目标
-        :return torch.Tensor: a scalar loss
+    def train_step(self, words: "torch.LongTensor", target: "torch.LongTensor", seq_len: "torch.LongTensor"):
+        """
+        :param words: 句子中 word 的 index，形状为 ``[batch_size, seq_len]``
+        :param target: 每个 sample 的目标值
+        :param seq_len: 每个句子的长度，形状为 ``[batch,]``
+        :return: 如果 ``target`` 为 ``None``，则返回预测结果 ``{'pred': torch.Tensor}``，否则返回 loss ``{'loss': torch.Tensor}``
         """
         return self(words, seq_len, target)
     
-    def evaluate_step(self, words, seq_len):
-        r"""
-        
-        :param torch.LongTensor words: [batch_size, mex_len]
-        :param torch.LongTensor seq_len: [batch_size, ]
-        :return torch.LongTensor: [batch_size, max_len]
+    def evaluate_step(self, words: "torch.LongTensor", seq_len: "torch.LongTensor"):
+        """
+        :param words: 句子中 word 的 index，形状为 ``[batch_size, seq_len]``
+        :param seq_len: 每个句子的长度，形状为 ``[batch,]``
+        :return: 预测结果 ``{'pred': torch.Tensor}``
         """
         return self(words, seq_len)