| import torch | |
| from torch.utils.data import Dataset | |
| class BilingualDataset(Dataset): | |
| def __init__(self, dataset, source_tokenizer, target_tokenizer, source_language, target_language, sequence_length): | |
| super().__init__() | |
| self.dataset = dataset | |
| self.source_tokenizer = source_tokenizer | |
| self.target_tokenizer = target_tokenizer | |
| self.source_language = source_language | |
| self.target_language = target_language | |
| self.sequence_length = sequence_length | |
| self.SOS_token = torch.tensor([target_tokenizer.token_to_id("[SOS]")], dtype=torch.int64) | |
| self.PAD_token = torch.tensor([target_tokenizer.token_to_id("[PAD]")], dtype= torch.int64) | |
| self.EOS_token = torch.tensor([target_tokenizer.token_to_id("[EOS]")], dtype= torch.int64) | |
| def __len__(self): | |
| return len(self.dataset) | |
| def __getitem__(self, index) : | |
| source_target_dataset = self.dataset[index] | |
| source_text = source_target_dataset['translation'][self.source_language] | |
| target_text = source_target_dataset['translation'][self.target_language] | |
| encode_source_tokenizer = self.source_tokenizer.encode(source_text).ids | |
| encode_target_tokenizer = self.target_tokenizer.encode(target_text).ids | |
| encode_source_padding = self.sequence_length - len(encode_source_tokenizer) - 2 | |
| encode_target_padding = self.sequence_length - len(encode_target_tokenizer) - 1 | |
| if encode_source_padding < 0 or encode_target_padding < 0: | |
| raise ValueError("sequence is too long") | |
| encoder_input = torch.cat( | |
| [ | |
| self.SOS_token, | |
| torch.tensor(encode_source_tokenizer, dtype=torch.int64), | |
| self.EOS_token, | |
| torch.tensor([self.PAD_token] * encode_source_padding, dtype=torch.int64) | |
| ] | |
| ) | |
| decoder_input = torch.cat( | |
| [ | |
| self.SOS_token, | |
| torch.tensor(encode_target_tokenizer, dtype=torch.int64), | |
| torch.tensor([self.PAD_token] * encode_target_padding, dtype=torch.int64) | |
| ] | |
| ) | |
| Target = torch.cat( | |
| [ | |
| torch.tensor(encode_target_tokenizer, dtype=torch.int64), | |
| torch.tensor([self.PAD_token] * encode_target_padding, dtype=torch.int64), | |
| self.EOS_token | |
| ] | |
| ) | |
| assert encoder_input.size(0) == self.sequence_length | |
| assert decoder_input.size(0) == self.sequence_length | |
| assert Target.size(0) == self.sequence_length | |
| return { | |
| "encoder_input": encoder_input, | |
| "decoder_input": decoder_input, | |
| "encoder_input_mask": (encoder_input != self.PAD_token).unsqueeze(0).unsqueeze(0).int(), | |
| "decoder_input_mask": (decoder_input != self.PAD_token).unsqueeze(0).int() & casual_mask(decoder_input.size(0)), | |
| "Target": Target, | |
| "source_text": source_text, | |
| "target_text": target_text | |
| } | |
| def casual_mask(size): | |
| mask = torch.triu(torch.ones((1, size, size)), diagonal=1).type(torch.int) | |
| return mask == 0 |