-
Notifications
You must be signed in to change notification settings - Fork 3
/
dataset.py
100 lines (78 loc) · 2.88 KB
/
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
import torch
from torch.utils import data
from transformers import AutoTokenizer
# map lm name to huggingface's pre-trained model names
lm_mp = {'roberta': 'roberta-base',
'distilbert': 'distilbert-base-uncased'}
def get_tokenizer(lm):
if lm in lm_mp:
return AutoTokenizer.from_pretrained(lm_mp[lm])
else:
return AutoTokenizer.from_pretrained(lm)
class PretrainDataset(data.Dataset):
"""EM dataset"""
def __init__(self,
path,
max_len=256,
size=None,
lm='roberta'):
self.tokenizer = get_tokenizer(lm)
self.pairs = []
self.labels = []
self.max_len = max_len
self.size = size
if isinstance(path, list):
lines = path
else:
lines = open(path)
for line in lines:
s1, s2, label = line.strip().split('\t')
self.pairs.append((s1, s2))
self.labels.append(int(label))
self.pairs = self.pairs[:size]
self.labels = self.labels[:size]
def __len__(self):
"""Return the size of the dataset."""
return len(self.pairs)
def __getitem__(self, idx):
"""Return a tokenized item of the dataset.
Args:
idx (int): the index of the item
Returns:
List of int: token ID's of the two entities
List of int: token ID's of the two entities augmented (if da is set)
int: the label of the pair (0: unmatch, 1: match)
"""
left = self.pairs[idx][0]
right = self.pairs[idx][1]
# left + right
x = self.tokenizer.encode(text=left,
text_pair=right,
max_length=self.max_len,
truncation=True)
return x, self.labels[idx]
@staticmethod
def pad(batch):
"""Merge a list of dataset items into a train/test batch
Args:
batch (list of tuple): a list of dataset items
Returns:
LongTensor: x1 of shape (batch_size, seq_len)
LongTensor: x2 of shape (batch_size, seq_len).
Elements of x1 and x2 are padded to the same length
LongTensor: a batch of labels, (batch_size,)
"""
if len(batch[0]) == 3:
x1, x2, y = zip(*batch)
maxlen = max([len(x) for x in x1+x2])
x1 = [xi + [0]*(maxlen - len(xi)) for xi in x1]
x2 = [xi + [0]*(maxlen - len(xi)) for xi in x2]
return torch.LongTensor(x1), \
torch.LongTensor(x2), \
torch.LongTensor(y)
else:
x12, y = zip(*batch)
maxlen = max([len(x) for x in x12])
x12 = [xi + [0]*(maxlen - len(xi)) for xi in x12]
return torch.LongTensor(x12), \
torch.LongTensor(y)