dmlc--dgl
Guo Qipeng
cb2bc69b55
In the bucket sampler, the length array should be recovered after random picking samples.
324 行
13 KiB
Python
324 行
13 KiB
Python
import torch
|
|
import dgl
|
|
import numpy as np
|
|
import json
|
|
import pickle
|
|
import random
|
|
|
|
|
|
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))
|
|
|