OpenNRE源码结构解读

OpenNRE源码结构解读

THUNLP OpenNRE开源模型代码解读，以BagRE为例。

structure

层次：framework(BagRE) - model(BagAttention) - encoder(PCNN)

• framework(BagRE)

  • dataLoader = BagRELoader

    • datasest = BagREDataset
      • tokenizer = model.sentence_encoder.tokenize（传入的model的成员的函数对象）
      • 其getitem函数为
        • 见下

  • model = BagAttention（传入的参数）

    • sentence_encoder = PCNNEncoder（在实例化时传入的参数）
      • 继承自父类BaseEncoder，其中包含
        • tokenizer = WordTokenizer(vocab=self.token2id)
          • token2id (dictionary of token->idx mapping)
        • word ＆ pos1,2 embeddings，
      • 其tokenize函数为
        • Sentence -> Token: 利用BaseEncoder的tokenizer.tokenize进行分词（得到切分好的List of wordpieces & pos）
        • Token -> index: 利用tokenizer.convert_tokens_to_ids得到indexed_tokens（通过utils.convert_by_vocab）
        • Position -> index | Mask | Padding: pos1, pos2, mask
        • return indexed_tokens, pos1, pos2, mask
      • 其forward函数为
        • arg：按batch的上述返回值，尺寸为(B, L)，L为sentence length
        • ret：按batch输出向量表示，尺寸为(B, EMBED)

  • num classes, id2rel mapping，用于模型内部的卷积、线性层

  • forward函数：

    • 先通过view展平为nsum（一个bag有n个sentence，批量处理B个bag对应nsum个sentence，每个bag中的sentence利用scope确定）
    • 然后encode为H=EMBED（隐藏层尺寸即词嵌入dim）
    • 再经过attention得到score，经过softmax得到logits

    • def forward(self, label, scope, token, pos1, pos2, mask=None, train=True, bag_size=0):
          """
          Args:
              label: (B), label of the bag
              scope: (B), scope for each bag
              token: (nsum, L), index of tokens
              pos1: (nsum, L), relative position to head entity
              pos2: (nsum, L), relative position to tail entity
              mask: (nsum, L), used for piece-wise CNN
          Return:
              logits, (B, N)
          """
          
          # 1. get representation, size (nsum, H)
          if bag_size