-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathdataset.py
92 lines (72 loc) · 3.11 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
import torch
import torch.nn as nn
from torch.utils.data import Dataset
class BilingualDataset(Dataset):
def __init__(self, ds, src_tokenizer, tgt_tokenizer, src_lang, tgt_lang, seq_len):
super().__init__()
self.src_tokenizer = src_tokenizer
self.tgt_tokenizer = tgt_tokenizer
self.src_lang = src_lang
self.tgt_lang = tgt_lang
self.seq_len = seq_len
self.ds = ds
self.sos_token = torch.tensor([tgt_tokenizer.token_to_id("[SOS]")], dtype=torch.int64)
self.eos_token = torch.tensor([tgt_tokenizer.token_to_id("[EOS]")], dtype=torch.int64)
self.pad_token = torch.tensor([tgt_tokenizer.token_to_id("[PAD]")], dtype=torch.int64)
def __len__(self):
return len(self.ds)
def __getitem__(self, idx):
src_tgt_pair = self.ds[idx]
src_text = src_tgt_pair['translation'][self.src_lang]
tgt_text = src_tgt_pair['translation'][self.tgt_lang]
#print("BBBBBIIIIII")
#print(tgt_text)
enc_input_tokens = self.src_tokenizer.encode(src_text).ids
dec_input_tokens = self.tgt_tokenizer.encode(tgt_text).ids
#print(len(enc_input_tokens))
#print(len(dec_input_tokens))
enc_num_padding_tokens = self.seq_len-len(enc_input_tokens)-2
dec_num_padding_tokens = self.seq_len-len(dec_input_tokens)-1
#print(enc_num_padding_tokens,dec_num_padding_tokens,self.seq_len)
if enc_num_padding_tokens < 0 or dec_num_padding_tokens < 0:
raise ValueError("Sentence too long")
encoder_input = torch.cat(
[
self.sos_token,
torch.tensor(enc_input_tokens, dtype=torch.int64),
self.eos_token,
torch.tensor([self.pad_token]*enc_num_padding_tokens, dtype=torch.int64)
],
dim=0
)
decoder_input = torch.cat(
[
self.sos_token,
torch.tensor(dec_input_tokens, dtype=torch.int64),
torch.tensor([self.pad_token]*dec_num_padding_tokens, dtype=torch.int64)
],
dim=0
)
label = torch.cat(
[
torch.tensor(dec_input_tokens, dtype=torch.int64),
self.eos_token,
torch.tensor([self.pad_token]*dec_num_padding_tokens, dtype=torch.int64)
],
dim=0
)
assert encoder_input.size(0) == self.seq_len
assert decoder_input.size(0) == self.seq_len
assert label.size(0) == self.seq_len
return {
"encoder_input" : encoder_input,
"decoder_input" : decoder_input,
"encoder_mask" : (encoder_input != self.pad_token).unsqueeze(0).unsqueeze(0).int(),
"decoder_mask" : (decoder_input != self.pad_token).unsqueeze(0).int() & causal_mask(decoder_input.size(0)),
"label":label,
"src_text": src_text,
"tgt_text":tgt_text
}
def causal_mask(size):
mask = torch.triu(torch.ones((1,size,size)),diagonal=1).type(torch.int)
return mask == 0