This repository has been archived by the owner on Feb 21, 2023. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 1
/
prepare_dataset.py
145 lines (123 loc) · 5.16 KB
/
prepare_dataset.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
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
from pathlib import Path
import json
from itertools import zip_longest
import torch
from scripts.utils.dataset_process_utils import parse_arguments, build_vocab
from scripts.utils.dataset_process_utils import split_train_data, encode_data
from scripts.utils.utils import seed_everything
from scripts.dataset.process_task8_files import process_file
if __name__ == "__main__":
# parse command
args = parse_arguments()
dataset_dir = args.dataset_dir # semeval task8 directory to text files
save_dir = Path(args.save_dir) # directory to save output files
check = args.check
# set some pathes
train_path = Path(dataset_dir) / args.train_path # path to train txt file
test_path = Path(dataset_dir) / args.test_path # path to test txt file
config_path = args.config_path # path to configuration file
# load config files
with open(config_path, "r") as json_file:
config = json.load(json_file)
# seed
seed = config["seed"]
seed_everything(seed)
config = config["dataset"]
stanza_models_dir = config["stanza_models_dir"]
# read deaset files and extract data
data = {}
for split_path, split in zip([train_path, test_path], ["train", "test"]):
data[split] = process_file(split_path, split, stanza_models_dir,
config["tags"], check)
# split training data into validation and training splits
# get indices for tain and val split based of the distribution of relation
# label + relation direction
print("splitting training data...")
ids = list(range(len(data["train"][0])))
y = [
r + d if r != "Other" else r for r, d in list(zip(*data["train"][-2:]))
]
train_id, val_id, _, _ = split_train_data(ids, y, config["val_size"], seed)
# split training data based on the ids
train_data, val_data = [], []
for t_idx, v_idx in zip_longest(train_id, val_id):
train_data.append(
[data["train"][i][t_idx] for i in range(len(data["train"]))])
if v_idx is not None:
val_data.append(
[data["train"][i][v_idx] for i in range(len(data["train"]))])
data["train"] = list(zip(*train_data))
data["val"] = list(zip(*val_data))
print("splitting finished.")
print("saving files...")
# create vocab
_, tokens, ents, pos, dep, _, _, _, _, rels, rels_dir = data["train"]
tokens_vocab = build_vocab(list(tokens),
special_tokens=["<unk>", "<pad>"],
min_freq=config["min_freq"])
ents_vocab = build_vocab(list(ents), special_tokens=["<pad>"])
pos_vocab = build_vocab(list(pos), special_tokens=["<unk>", "<pad>"])
dep_vocab = build_vocab(list(dep), special_tokens=["<unk>", "<pad>"])
rels_vocab = build_vocab([list(rels)])
rels_dir_vocab = build_vocab([list(rels_dir)])
tokens_vocab.set_default_index(tokens_vocab["<unk>"])
pos_vocab.set_default_index(pos_vocab["<unk>"])
dep_vocab.set_default_index(dep_vocab["<unk>"])
# encode, convert to torch, save
for split in ["train", "val", "test"]:
data2save = {}
tokens = list(data[split][1])
ents = list(data[split][2])
pos = list(data[split][3])
dep = list(data[split][4])
rels = list(data[split][-2])
rels_dir = list(data[split][-1])
tokens_encoded = encode_data(tokens, tokens_vocab)
ents_encoded = encode_data(ents, ents_vocab)
pos_encoded = encode_data(pos, pos_vocab)
dep_encoded = encode_data(dep, dep_vocab)
rels_encoded = encode_data([rels], rels_vocab)[0]
rels_dir_encoded = encode_data([rels_dir], rels_dir_vocab)[0]
# save data
idx = torch.LongTensor([int(i) for i in data[split][0]])
trees = list(data[split][5:7])
length = torch.LongTensor(data[split][7])
data2save["idx"] = idx
data2save["tokens"] = tokens_encoded
data2save["ents"] = ents_encoded
data2save["pos"] = pos_encoded
data2save["dep"] = dep_encoded
data2save["length"] = length
data2save["dep_tree"] = trees[0]
data2save["shortest_path"] = trees[1]
data2save["rels"] = rels_encoded
data2save["rels_dir"] = rels_dir_encoded
# write to desk
file_name = f"{split}.pt"
save_path = str(save_dir / file_name)
torch.save(data2save, save_path)
raw_file_name = f"raw_{split}.pt"
raw_save_path = str(save_dir / raw_file_name)
torch.save(data[split], raw_save_path)
# save vocab
file_name = "vocabs.pt"
save_path = str(save_dir / file_name)
vocabs = {
"tokens": tokens_vocab,
"ents": ents_vocab,
"pos": pos_vocab,
"dep": dep_vocab,
"rels": rels_vocab,
"rel_dir": rels_dir_vocab
}
torch.save(vocabs, save_path)
# write answer keys
answer_keys = ""
for lid, label, ldir in zip(idx, rels, rels_dir):
answer_dir = ldir if label != "Other" else ""
answer_keys += f"{lid}\t{label}{answer_dir}\n"
file_name = "answer_keys.txt"
save_path = str(save_dir / file_name)
with open(save_path, "w") as file:
file.write(answer_keys)
print("done.")