增加中文人物关系抽取模型

Author: rui.tao

example/train_people_chinese_bert_softmax.py

@@ -0,0 +1,54 @@
+import sys, json
+import torch
+import os
+import numpy as np
+import opennre
+from opennre import encoder, model, framework
+import argparse
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--mask_entity', action='store_true', help='Mask entity mentions')
+args = parser.parse_args()
+
+# Some basic settings
+root_path = '.'
+sys.path.append(root_path)
+if not os.path.exists('ckpt'):
+    os.mkdir('ckpt')
+ckpt = 'ckpt/people_chinese_bert_softmax.pth.tar'
+
+# Check data
+rel2id = json.load(open(os.path.join(root_path, 'benchmark/people-relation/people-relation_rel2id.json')))
+
+# Define the sentence encoder
+sentence_encoder = opennre.encoder.BERTEncoder(
+    max_length=80, 
+    pretrain_path=os.path.join(root_path, 'pretrain/chinese_wwm_pytorch'),
+    mask_entity=args.mask_entity
+)
+
+# Define the model
+model = opennre.model.SoftmaxNN(sentence_encoder, len(rel2id), rel2id)
+
+# Define the whole training framework
+framework = opennre.framework.SentenceRE(
+    train_path=os.path.join(root_path, 'benchmark/people-relation/people-relation_train.txt'),
+    val_path=os.path.join(root_path, 'benchmark/people-relation/people-relation_val.txt'),
+    test_path=os.path.join(root_path, 'benchmark/people-relation/people-relation_val.txt'),
+    model=model,
+    ckpt=ckpt,
+    batch_size=64, # Modify the batch size w.r.t. your device
+    max_epoch=3,
+    lr=2e-5,
+    opt='adamw'
+)
+
+# Train the model
+framework.train_model()
+
+# Test the model
+framework.load_state_dict(torch.load(ckpt)['state_dict'])
+result = framework.eval_model(framework.test_loader)
+
+# Print the result
+print('Accuracy on test set: {}'.format(result['acc']))

example/train_test_chinese_bert_softmax.py

@@ -0,0 +1,54 @@
+import sys, json
+import torch
+import os
+import numpy as np
+import opennre
+from opennre import encoder, model, framework
+import argparse
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--mask_entity', action='store_true', help='Mask entity mentions')
+args = parser.parse_args()
+
+# Some basic settings
+root_path = '.'
+sys.path.append(root_path)
+if not os.path.exists('ckpt'):
+    os.mkdir('ckpt')
+ckpt = 'ckpt/test_chinese_bert_softmax.pth.tar'
+
+# Check data
+rel2id = json.load(open(os.path.join(root_path, 'benchmark/test_chinese/test_chinese_rel2id.json')))
+
+# Define the sentence encoder
+sentence_encoder = opennre.encoder.BERTEncoder(
+    max_length=80, 
+    pretrain_path=os.path.join(root_path, 'pretrain/chinese_wwm_pytorch'),
+    mask_entity=args.mask_entity
+)
+
+# Define the model
+model = opennre.model.SoftmaxNN(sentence_encoder, len(rel2id), rel2id)
+
+# Define the whole training framework
+framework = opennre.framework.SentenceRE(
+    train_path=os.path.join(root_path, 'benchmark/test_chinese/test_chinese_train.txt'),
+    val_path=os.path.join(root_path, 'benchmark/test_chinese/test_chinese_val.txt'),
+    test_path=os.path.join(root_path, 'benchmark/test_chinese/test_chinese_val.txt'),
+    model=model,
+    ckpt=ckpt,
+    batch_size=64, # Modify the batch size w.r.t. your device
+    max_epoch=7,
+    lr=2e-5,
+    opt='adamw'
+)
+
+# Train the model
+framework.train_model()
+
+# Test the model
+framework.load_state_dict(torch.load(ckpt)['state_dict'])
+result = framework.eval_model(framework.test_loader)
+
+# Print the result
+print('Accuracy on test set: {}'.format(result['acc']))

example/train_wiki80_bert_softmax.py

@@ -15,7 +15,7 @@
 sys.path.append(root_path)
 if not os.path.exists('ckpt'):
     os.mkdir('ckpt')
-ckpt = 'ckpt/wiki80_bert_softmax.pth.tar'
+ckpt = 'ckpt/wiki80_bert_softmax_7epoch.pth.tar'
 
 # Check data
 opennre.download_wiki80(root_path=root_path)
@@ -40,7 +40,7 @@
     model=model,
     ckpt=ckpt,
     batch_size=64, # Modify the batch size w.r.t. your device
-    max_epoch=10,
+    max_epoch=7,
     lr=2e-5,
     opt='adamw'
 )

example/train_wiki80_cnn_softmax.py

@@ -43,7 +43,7 @@
     model=model,
     ckpt=ckpt,
     batch_size=32,
-    max_epoch=100,
+    max_epoch=2,
     lr=0.1,
     weight_decay=1e-5,
     opt='sgd'

opennre/framework/data_loader.py

@@ -140,7 +140,9 @@ def __init__(self, path, rel2id, tokenizer, entpair_as_bag=False, bag_size=None,
             self.bag_name = []
             self.facts = {}
             for idx, item in enumerate(self.data):
+               
                 fact = (item['h']['id'], item['t']['id'], item['relation'])
+
                 if item['relation'] != 'NA':
                     self.facts[fact] = 1
                 if entpair_as_bag:

opennre/pretrain.py

@@ -7,14 +7,14 @@
 import json
 import numpy as np
 
-default_root_path = os.path.join(os.getenv('HOME'), '.opennre')
+default_root_path = os.path.join(os.getenv('openNRE'), '.')
 
 def check_root(root_path=default_root_path):
     if not os.path.exists(root_path):
         os.mkdir(root_path)
         os.mkdir(os.path.join(root_path, 'benchmark'))
         os.mkdir(os.path.join(root_path, 'pretrain'))
-        os.mkdir(os.path.join(root_path, 'pretrain/nre'))
+        os.mkdir(os.path.join(root_path, 'ckpt'))
 
 def download_wiki80(root_path=default_root_path):
     check_root()
@@ -49,14 +49,18 @@ def download_bert_base_uncased(root_path=default_root_path):
         os.system('wget -P ' + os.path.join(root_path, 'pretrain/bert-base-uncased') + ' http://193.112.16.83:8080/opennre/pretrain/bert-base-uncased/vocab.txt')
 
 def download_pretrain(model_name, root_path=default_root_path):
-    ckpt = os.path.join(root_path, 'pretrain/nre/' + model_name + '.pth.tar')
+    ckpt = os.path.join(root_path, 'ckpt/' + model_name + '.pth.tar')
     if not os.path.exists(ckpt):
-        os.system('wget -P ' + os.path.join(root_path, 'pretrain/nre')  + ' http://193.112.16.83:8080/opennre/pretrain/nre/' + model_name + '.pth.tar')
+        print("*"*20)
+        print("下载ckpt")
+        os.system('wget -P ' + os.path.join(root_path, 'ckpt/')  + ' http://193.112.16.83:8080/opennre/ckpt/' + model_name + '.pth.tar')
 
 def get_model(model_name, root_path=default_root_path):
     check_root()
-    ckpt = os.path.join(root_path, 'pretrain/nre/' + model_name + '.pth.tar')
+    ckpt = os.path.join(root_path, 'ckpt/' + model_name + '.pth.tar')
+    
     if model_name == 'wiki80_cnn_softmax':
+        print("*"*20+"taorui")
         download_pretrain(model_name)
         download_glove()
         download_wiki80()
@@ -86,5 +90,38 @@ def get_model(model_name, root_path=default_root_path):
         m = model.SoftmaxNN(sentence_encoder, len(rel2id), rel2id)
         m.load_state_dict(torch.load(ckpt)['state_dict'])
         return m
+    elif model_name == 'test_chinese_bert_softmax':
+        download_pretrain(model_name)
+        download_bert_base_uncased()
+        download_wiki80()
+        rel2id = json.load(open(os.path.join(root_path, 'benchmark/test_chinese/test_chinese_rel2id.json')))
+        sentence_encoder = encoder.BERTEncoder(
+            max_length=80, pretrain_path=os.path.join(root_path, 'pretrain/chinese_wwm_pytorch'))
+        m = model.SoftmaxNN(sentence_encoder, len(rel2id), rel2id)
+        m.load_state_dict(torch.load(ckpt)['state_dict'])
+        return m
+
+    elif model_name == 'people_chinese_bert_softmax':
+        download_pretrain(model_name)
+        download_bert_base_uncased()
+        download_wiki80()
+        rel2id = json.load(open(os.path.join(root_path, 'benchmark/people-relation/people-relation_rel2id.json')))
+        sentence_encoder = encoder.BERTEncoder(
+            max_length=80, pretrain_path=os.path.join(root_path, 'pretrain/chinese_wwm_pytorch'))
+        m = model.SoftmaxNN(sentence_encoder, len(rel2id), rel2id)
+        m.load_state_dict(torch.load(ckpt)['state_dict'])
+        return m 
+
+    elif model_name == 'people_delunknown_chinese_bert_softmax':
+        download_pretrain(model_name)
+        download_bert_base_uncased()
+        download_wiki80()
+        rel2id = json.load(open(os.path.join(root_path, 'benchmark/people-relation-delunknow/people-relation_rel2id.json')))
+        sentence_encoder = encoder.BERTEncoder(
+            max_length=80, pretrain_path=os.path.join(root_path, 'pretrain/chinese_wwm_pytorch'))
+        m = model.SoftmaxNN(sentence_encoder, len(rel2id), rel2id)
+        m.load_state_dict(torch.load(ckpt)['state_dict'])
+        return m
+    
     else:
         raise NotImplementedError