-
Notifications
You must be signed in to change notification settings - Fork 0
/
filter_correct.py
62 lines (53 loc) · 2.21 KB
/
filter_correct.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
import sys
from opennre import encoder
try:
sys.path.append('..')
from server.adversary_utils import *
from config import *
except ImportError:
pass
logging.basicConfig(level=logging.INFO)
# Silent unimportant log messages
for logger_name in ['transformers.configuration_utils',
'transformers.modeling_utils',
'transformers.tokenization_utils_base']:
logging.getLogger(logger_name).setLevel(logging.WARNING)
parser = argparse.ArgumentParser()
parser.add_argument('--dataset', '-d', required=True,
choices=['tacred', 'wiki80'])
parser.add_argument('--model', '-m', required=True,
choices=['bert', 'bertentity'])
parser.add_argument('--max_seq_len', '-l', type=int, default=128,
help='Maximum sequence length of bert model')
args = parser.parse_args()
# Initialize model
rel2id = json.load(open(RELATION_PATH_DICT[args.dataset], 'r'))
id2rel = {v: k for k, v in rel2id.items()}
device = torch.device('cuda:0')
sentence_encoder = {'bert': encoder.BERTEncoder,
'bertentity': encoder.BERTEntityEncoder}[args.model](
max_length=args.max_seq_len,
pretrain_path=PRETRAIN_PATH,
mask_entity=False
)
model = SoftmaxNN(sentence_encoder, len(rel2id), rel2id)
model.load_state_dict(torch.load(MODEL_PATH_DICT[args.dataset][args.model],
map_location=device)['state_dict'])
model.to(device)
model.eval()
model = REClassifier(model, rel2id, id2rel, device)
# Create directory for files containing filtered samples
target_dir = os.path.join(DATA_ROOT, args.dataset, args.model)
os.makedirs(target_dir, exist_ok=True)
for data_type in DATA_TYPES[args.dataset]:
input_file = os.path.join(os.path.join(
DATA_ROOT, args.dataset, data_type + '.txt'))
output_file = os.path.join(target_dir, data_type + '_correct.index')
# Filter out correct samples
correct_samples = []
samples = [json.loads(line) for line in open(input_file, 'r').readlines()]
with open(output_file, 'w') as f:
for idx, sample in enumerate(tqdm(samples)):
label, _ = model.infer(sample)
if label == sample['relation']:
f.write(str(idx) + '\n')