wehub-resource-sync
520 lines
16 KiB
Python
520 lines
16 KiB
Python
import json
|
|
import pickle
|
|
import random
|
|
|
|
import dgl
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
NODE_TYPE = {"entity": 0, "root": 1, "relation": 2}
|
|
|
|
|
|
def write_txt(batch, seqs, w_file, args):
|
|
# converting the prediction to real text.
|
|
ret = []
|
|
for b, seq in enumerate(seqs):
|
|
txt = []
|
|
for token in seq:
|
|
# copy the entity
|
|
if token >= len(args.text_vocab):
|
|
ent_text = batch["raw_ent_text"][b][
|
|
token - len(args.text_vocab)
|
|
]
|
|
ent_text = filter(lambda x: x != "<PAD>", ent_text)
|
|
txt.extend(ent_text)
|
|
else:
|
|
if int(token) not in [
|
|
args.text_vocab(x) for x in ["<PAD>", "<BOS>", "<EOS>"]
|
|
]:
|
|
txt.append(args.text_vocab(int(token)))
|
|
if int(token) == args.text_vocab("<EOS>"):
|
|
break
|
|
w_file.write(" ".join([str(x) for x in txt]) + "\n")
|
|
ret.append([" ".join([str(x) for x in txt])])
|
|
return ret
|
|
|
|
|
|
def replace_ent(x, ent, V):
|
|
# replace the entity
|
|
mask = x >= V
|
|
if mask.sum() == 0:
|
|
return x
|
|
nz = mask.nonzero()
|
|
fill_ent = ent[nz, x[mask] - V]
|
|
x = x.masked_scatter(mask, fill_ent)
|
|
return x
|
|
|
|
|
|
def len2mask(lens, device):
|
|
max_len = max(lens)
|
|
mask = (
|
|
torch.arange(max_len, device=device)
|
|
.unsqueeze(0)
|
|
.expand(len(lens), max_len)
|
|
)
|
|
mask = mask >= torch.LongTensor(lens).to(mask).unsqueeze(1)
|
|
return mask
|
|
|
|
|
|
def pad(var_len_list, out_type="list", flatten=False):
|
|
if flatten:
|
|
lens = [len(x) for x in var_len_list]
|
|
var_len_list = sum(var_len_list, [])
|
|
max_len = max([len(x) for x in var_len_list])
|
|
if out_type == "list":
|
|
if flatten:
|
|
return [
|
|
x + ["<PAD>"] * (max_len - len(x)) for x in var_len_list
|
|
], lens
|
|
else:
|
|
return [x + ["<PAD>"] * (max_len - len(x)) for x in var_len_list]
|
|
if out_type == "tensor":
|
|
if flatten:
|
|
return (
|
|
torch.stack(
|
|
[
|
|
torch.cat(
|
|
[
|
|
x,
|
|
torch.zeros(
|
|
[max_len - len(x)] + list(x.shape[1:])
|
|
).type_as(x),
|
|
],
|
|
0,
|
|
)
|
|
for x in var_len_list
|
|
],
|
|
0,
|
|
),
|
|
lens,
|
|
)
|
|
else:
|
|
return torch.stack(
|
|
[
|
|
torch.cat(
|
|
[
|
|
x,
|
|
torch.zeros(
|
|
[max_len - len(x)] + list(x.shape[1:])
|
|
).type_as(x),
|
|
],
|
|
0,
|
|
)
|
|
for x in var_len_list
|
|
],
|
|
0,
|
|
)
|
|
|
|
|
|
class Vocab(object):
|
|
def __init__(
|
|
self,
|
|
max_vocab=2**31,
|
|
min_freq=-1,
|
|
sp=["<PAD>", "<BOS>", "<EOS>", "<UNK>"],
|
|
):
|
|
self.i2s = []
|
|
self.s2i = {}
|
|
self.wf = {}
|
|
self.max_vocab, self.min_freq, self.sp = max_vocab, min_freq, sp
|
|
|
|
def __len__(self):
|
|
return len(self.i2s)
|
|
|
|
def __str__(self):
|
|
return "Total " + str(len(self.i2s)) + str(self.i2s[:10])
|
|
|
|
def update(self, token):
|
|
if isinstance(token, list):
|
|
for t in token:
|
|
self.update(t)
|
|
else:
|
|
self.wf[token] = self.wf.get(token, 0) + 1
|
|
|
|
def build(self):
|
|
self.i2s.extend(self.sp)
|
|
sort_kv = sorted(self.wf.items(), key=lambda x: x[1], reverse=True)
|
|
for k, v in sort_kv:
|
|
if (
|
|
len(self.i2s) < self.max_vocab
|
|
and v >= self.min_freq
|
|
and k not in self.sp
|
|
):
|
|
self.i2s.append(k)
|
|
self.s2i.update(list(zip(self.i2s, range(len(self.i2s)))))
|
|
|
|
def __call__(self, x):
|
|
if isinstance(x, int):
|
|
return self.i2s[x]
|
|
else:
|
|
return self.s2i.get(x, self.s2i["<UNK>"])
|
|
|
|
def save(self, fname):
|
|
pass
|
|
|
|
def load(self, fname):
|
|
pass
|
|
|
|
|
|
def at_least(x):
|
|
# handling the illegal data
|
|
if len(x) == 0:
|
|
return ["<UNK>"]
|
|
else:
|
|
return x
|
|
|
|
|
|
class Example(object):
|
|
def __init__(self, title, ent_text, ent_type, rel, text):
|
|
# one object corresponds to a data sample
|
|
self.raw_title = title.split()
|
|
self.raw_ent_text = [at_least(x.split()) for x in ent_text]
|
|
assert min([len(x) for x in self.raw_ent_text]) > 0, str(
|
|
self.raw_ent_text
|
|
)
|
|
self.raw_ent_type = ent_type.split() # <method> .. <>
|
|
self.raw_rel = []
|
|
for r in rel:
|
|
rel_list = r.split()
|
|
for i in range(len(rel_list)):
|
|
if (
|
|
i > 0
|
|
and i < len(rel_list) - 1
|
|
and rel_list[i - 1] == "--"
|
|
and rel_list[i] != rel_list[i].lower()
|
|
and rel_list[i + 1] == "--"
|
|
):
|
|
self.raw_rel.append(
|
|
[
|
|
rel_list[: i - 1],
|
|
rel_list[i - 1] + rel_list[i] + rel_list[i + 1],
|
|
rel_list[i + 2 :],
|
|
]
|
|
)
|
|
break
|
|
self.raw_text = text.split()
|
|
self.graph = self.build_graph()
|
|
|
|
def __str__(self):
|
|
return "\n".join(
|
|
[str(k) + ":\t" + str(v) for k, v in self.__dict__.items()]
|
|
)
|
|
|
|
def __len__(self):
|
|
return len(self.raw_text)
|
|
|
|
@staticmethod
|
|
def from_json(json_data):
|
|
return Example(
|
|
json_data["title"],
|
|
json_data["entities"],
|
|
json_data["types"],
|
|
json_data["relations"],
|
|
json_data["abstract"],
|
|
)
|
|
|
|
def build_graph(self):
|
|
graph = dgl.DGLGraph()
|
|
ent_len = len(self.raw_ent_text)
|
|
rel_len = len(
|
|
self.raw_rel
|
|
) # treat the repeated relation as different nodes, refer to the author's code
|
|
|
|
graph.add_nodes(
|
|
ent_len, {"type": torch.ones(ent_len) * NODE_TYPE["entity"]}
|
|
)
|
|
graph.add_nodes(1, {"type": torch.ones(1) * NODE_TYPE["root"]})
|
|
graph.add_nodes(
|
|
rel_len * 2,
|
|
{"type": torch.ones(rel_len * 2) * NODE_TYPE["relation"]},
|
|
)
|
|
graph.add_edges(ent_len, torch.arange(ent_len))
|
|
graph.add_edges(torch.arange(ent_len), ent_len)
|
|
graph.add_edges(
|
|
torch.arange(ent_len + 1 + rel_len * 2),
|
|
torch.arange(ent_len + 1 + rel_len * 2),
|
|
)
|
|
adj_edges = []
|
|
for i, r in enumerate(self.raw_rel):
|
|
assert len(r) == 3, str(r)
|
|
st, rt, ed = r
|
|
st_ent, ed_ent = self.raw_ent_text.index(
|
|
st
|
|
), self.raw_ent_text.index(ed)
|
|
# according to the edge_softmax operator, we need to reverse the graph
|
|
adj_edges.append([ent_len + 1 + 2 * i, st_ent])
|
|
adj_edges.append([ed_ent, ent_len + 1 + 2 * i])
|
|
adj_edges.append([ent_len + 1 + 2 * i + 1, ed_ent])
|
|
adj_edges.append([st_ent, ent_len + 1 + 2 * i + 1])
|
|
|
|
if len(adj_edges) > 0:
|
|
graph.add_edges(*list(map(list, zip(*adj_edges))))
|
|
return graph
|
|
|
|
def get_tensor(
|
|
self, ent_vocab, rel_vocab, text_vocab, ent_text_vocab, title_vocab
|
|
):
|
|
if hasattr(self, "_cached_tensor"):
|
|
return self._cached_tensor
|
|
else:
|
|
title_data = ["<BOS>"] + self.raw_title + ["<EOS>"]
|
|
title = [title_vocab(x) for x in title_data]
|
|
ent_text = [
|
|
[ent_text_vocab(y) for y in x] for x in self.raw_ent_text
|
|
]
|
|
ent_type = [
|
|
text_vocab(x) for x in self.raw_ent_type
|
|
] # for inference
|
|
rel_data = ["--root--"] + sum(
|
|
[[x[1], x[1] + "_INV"] for x in self.raw_rel], []
|
|
)
|
|
rel = [rel_vocab(x) for x in rel_data]
|
|
|
|
text_data = ["<BOS>"] + self.raw_text + ["<EOS>"]
|
|
text = [text_vocab(x) for x in text_data]
|
|
tgt_text = []
|
|
# the input text and decoding target are different since the consideration of the copy mechanism.
|
|
for i, str1 in enumerate(text_data):
|
|
if str1[0] == "<" and str1[-1] == ">" and "_" in str1:
|
|
a, b = str1[1:-1].split("_")
|
|
text[i] = text_vocab("<" + a + ">")
|
|
tgt_text.append(len(text_vocab) + int(b))
|
|
else:
|
|
tgt_text.append(text[i])
|
|
self._cached_tensor = {
|
|
"title": torch.LongTensor(title),
|
|
"ent_text": [torch.LongTensor(x) for x in ent_text],
|
|
"ent_type": torch.LongTensor(ent_type),
|
|
"rel": torch.LongTensor(rel),
|
|
"text": torch.LongTensor(text[:-1]),
|
|
"tgt_text": torch.LongTensor(tgt_text[1:]),
|
|
"graph": self.graph,
|
|
"raw_ent_text": self.raw_ent_text,
|
|
}
|
|
return self._cached_tensor
|
|
|
|
def update_vocab(
|
|
self, ent_vocab, rel_vocab, text_vocab, ent_text_vocab, title_vocab
|
|
):
|
|
ent_vocab.update(self.raw_ent_type)
|
|
ent_text_vocab.update(self.raw_ent_text)
|
|
title_vocab.update(self.raw_title)
|
|
rel_vocab.update(
|
|
["--root--"]
|
|
+ [x[1] for x in self.raw_rel]
|
|
+ [x[1] + "_INV" for x in self.raw_rel]
|
|
)
|
|
text_vocab.update(self.raw_ent_type)
|
|
text_vocab.update(self.raw_text)
|
|
|
|
|
|
class BucketSampler(torch.utils.data.Sampler):
|
|
def __init__(self, data_source, batch_size=32, bucket=3):
|
|
self.data_source = data_source
|
|
self.bucket = bucket
|
|
self.batch_size = batch_size
|
|
|
|
def __iter__(self):
|
|
# the magic number comes from the author's code
|
|
perm = torch.randperm(len(self.data_source))
|
|
lens = torch.Tensor([len(x) for x in self.data_source])
|
|
lens = lens[perm]
|
|
t1 = []
|
|
t2 = []
|
|
t3 = []
|
|
for i, l in enumerate(lens):
|
|
if l < 100:
|
|
t1.append(perm[i])
|
|
elif l > 100 and l < 220:
|
|
t2.append(perm[i])
|
|
else:
|
|
t3.append(perm[i])
|
|
datas = [t1, t2, t3]
|
|
random.shuffle(datas)
|
|
idxs = sum(datas, [])
|
|
batch = []
|
|
|
|
lens = torch.Tensor([len(x) for x in self.data_source])
|
|
for idx in idxs:
|
|
batch.append(idx)
|
|
mlen = max([0] + [lens[x] for x in batch])
|
|
if (
|
|
(mlen < 100 and len(batch) == 32)
|
|
or (mlen > 100 and mlen < 220 and len(batch) >= 24)
|
|
or (mlen > 220 and len(batch) >= 8)
|
|
or len(batch) == 32
|
|
):
|
|
yield batch
|
|
batch = []
|
|
if len(batch) > 0:
|
|
yield batch
|
|
|
|
def __len__(self):
|
|
return (len(self.data_source) + self.batch_size - 1) // self.batch_size
|
|
|
|
|
|
class GWdataset(torch.utils.data.Dataset):
|
|
def __init__(
|
|
self,
|
|
exs,
|
|
ent_vocab=None,
|
|
rel_vocab=None,
|
|
text_vocab=None,
|
|
ent_text_vocab=None,
|
|
title_vocab=None,
|
|
device=None,
|
|
):
|
|
super(GWdataset, self).__init__()
|
|
self.exs = exs
|
|
(
|
|
self.ent_vocab,
|
|
self.rel_vocab,
|
|
self.text_vocab,
|
|
self.ent_text_vocab,
|
|
self.title_vocab,
|
|
self.device,
|
|
) = (
|
|
ent_vocab,
|
|
rel_vocab,
|
|
text_vocab,
|
|
ent_text_vocab,
|
|
title_vocab,
|
|
device,
|
|
)
|
|
|
|
def __iter__(self):
|
|
return iter(self.exs)
|
|
|
|
def __getitem__(self, index):
|
|
return self.exs[index]
|
|
|
|
def __len__(self):
|
|
return len(self.exs)
|
|
|
|
def batch_fn(self, batch_ex):
|
|
(
|
|
batch_title,
|
|
batch_ent_text,
|
|
batch_ent_type,
|
|
batch_rel,
|
|
batch_text,
|
|
batch_tgt_text,
|
|
batch_graph,
|
|
) = ([], [], [], [], [], [], [])
|
|
batch_raw_ent_text = []
|
|
for ex in batch_ex:
|
|
ex_data = ex.get_tensor(
|
|
self.ent_vocab,
|
|
self.rel_vocab,
|
|
self.text_vocab,
|
|
self.ent_text_vocab,
|
|
self.title_vocab,
|
|
)
|
|
batch_title.append(ex_data["title"])
|
|
batch_ent_text.append(ex_data["ent_text"])
|
|
batch_ent_type.append(ex_data["ent_type"])
|
|
batch_rel.append(ex_data["rel"])
|
|
batch_text.append(ex_data["text"])
|
|
batch_tgt_text.append(ex_data["tgt_text"])
|
|
batch_graph.append(ex_data["graph"])
|
|
batch_raw_ent_text.append(ex_data["raw_ent_text"])
|
|
batch_title = pad(batch_title, out_type="tensor")
|
|
batch_ent_text, ent_len = pad(
|
|
batch_ent_text, out_type="tensor", flatten=True
|
|
)
|
|
batch_ent_type = pad(batch_ent_type, out_type="tensor")
|
|
batch_rel = pad(batch_rel, out_type="tensor")
|
|
batch_text = pad(batch_text, out_type="tensor")
|
|
batch_tgt_text = pad(batch_tgt_text, out_type="tensor")
|
|
batch_graph = dgl.batch(batch_graph)
|
|
batch_graph.to(self.device)
|
|
return {
|
|
"title": batch_title.to(self.device),
|
|
"ent_text": batch_ent_text.to(self.device),
|
|
"ent_len": ent_len,
|
|
"ent_type": batch_ent_type.to(self.device),
|
|
"rel": batch_rel.to(self.device),
|
|
"text": batch_text.to(self.device),
|
|
"tgt_text": batch_tgt_text.to(self.device),
|
|
"graph": batch_graph,
|
|
"raw_ent_text": batch_raw_ent_text,
|
|
}
|
|
|
|
|
|
def get_datasets(
|
|
fnames,
|
|
min_freq=-1,
|
|
sep=";",
|
|
joint_vocab=True,
|
|
device=None,
|
|
save="tmp.pickle",
|
|
):
|
|
# min_freq : not support now since it's very sensitive to the final results, but you can set it via passing min_freq to the Vocab class.
|
|
# sep : not support now
|
|
# joint_vocab : not support now
|
|
ent_vocab = Vocab(sp=["<PAD>", "<UNK>"])
|
|
title_vocab = Vocab(min_freq=5)
|
|
rel_vocab = Vocab(sp=["<PAD>", "<UNK>"])
|
|
text_vocab = Vocab(min_freq=5)
|
|
ent_text_vocab = Vocab(sp=["<PAD>", "<UNK>"])
|
|
datasets = []
|
|
for fname in fnames:
|
|
exs = []
|
|
json_datas = json.loads(open(fname).read())
|
|
for json_data in json_datas:
|
|
# construct one data example
|
|
ex = Example.from_json(json_data)
|
|
if fname == fnames[0]: # only training set
|
|
ex.update_vocab(
|
|
ent_vocab,
|
|
rel_vocab,
|
|
text_vocab,
|
|
ent_text_vocab,
|
|
title_vocab,
|
|
)
|
|
exs.append(ex)
|
|
datasets.append(exs)
|
|
ent_vocab.build()
|
|
rel_vocab.build()
|
|
text_vocab.build()
|
|
ent_text_vocab.build()
|
|
title_vocab.build()
|
|
datasets = [
|
|
GWdataset(
|
|
exs,
|
|
ent_vocab,
|
|
rel_vocab,
|
|
text_vocab,
|
|
ent_text_vocab,
|
|
title_vocab,
|
|
device,
|
|
)
|
|
for exs in datasets
|
|
]
|
|
with open(save, "wb") as f:
|
|
pickle.dump(datasets, f)
|
|
return datasets
|
|
|
|
|
|
if __name__ == "__main__":
|
|
ds = get_datasets(
|
|
[
|
|
"data/unprocessed.val.json",
|
|
"data/unprocessed.val.json",
|
|
"data/unprocessed.test.json",
|
|
]
|
|
)
|
|
print(ds[0].exs[0])
|
|
print(
|
|
ds[0]
|
|
.exs[0]
|
|
.get_tensor(
|
|
ds[0].ent_vocab,
|
|
ds[0].rel_vocab,
|
|
ds[0].text_vocab,
|
|
ds[0].ent_text_vocab,
|
|
ds[0].title_vocab,
|
|
)
|
|
)
|