| 56 | |
| 57 | |
| 58 | class CmnDataset: |
| 59 | def __init__(self, path, shuffle=False, batch_size=32, train_ratio=0.8, random_seed=0): |
| 60 | """ |
| 61 | cmn dataset, download from https://www.manythings.org/anki/, contains 29909 Chinese and English translation |
| 62 | pairs, the pair format: English + TAB + Chinese + TAB + Attribution |
| 63 | Args: |
| 64 | path: the path of the dataset |
| 65 | shuffle: shuffle the dataset, default False |
| 66 | batch_size: the size of every batch, default 32 |
| 67 | train_ratio: the proportion of the training set to the total data set, default 0.8 |
| 68 | random_seed: the random seed, used for shuffle operation, default 0 |
| 69 | """ |
| 70 | src_max_len, tgt_max_len, src_sts, tgt_sts = CmnDataset._split_sentences(path) |
| 71 | en_vab, cn_vab = Vocab(src_sts), Vocab(tgt_sts) |
| 72 | src_np, tgt_in_np, tgt_out_np = CmnDataset._encoding_stc(src_sts, tgt_sts, src_max_len, tgt_max_len, |
| 73 | en_vab, cn_vab) |
| 74 | |
| 75 | self.src_max_len, self.tgt_max_len = src_max_len, tgt_max_len |
| 76 | self.en_vab, self.cn_vab = en_vab, cn_vab |
| 77 | self.en_vab_size, self.cn_vab_size = len(en_vab), len(cn_vab) |
| 78 | |
| 79 | self.src_inputs, self.tgt_inputs, self.tgt_outputs = src_np, tgt_in_np, tgt_out_np |
| 80 | |
| 81 | self.shuffle, self.random_seed = shuffle, random_seed |
| 82 | |
| 83 | assert batch_size > 0, "The number of batch_size must be greater than 0" |
| 84 | self.batch_size = batch_size |
| 85 | |
| 86 | assert (0 < train_ratio <= 1.0), "The number of train_ratio must be in (0.0, 1.0]" |
| 87 | self.train_ratio = train_ratio |
| 88 | |
| 89 | self.total_size = len(src_np) |
| 90 | self.train_size = int(self.total_size * train_ratio) |
| 91 | self.test_size = self.total_size - self.train_size |
| 92 | |
| 93 | if shuffle: |
| 94 | index = [i for i in range(self.total_size)] |
| 95 | np.random.seed(self.random_seed) |
| 96 | np.random.shuffle(index) |
| 97 | |
| 98 | self.src_inputs = src_np[index] |
| 99 | self.tgt_inputs = tgt_in_np[index] |
| 100 | self.tgt_outputs = tgt_out_np[index] |
| 101 | |
| 102 | self.train_src_inputs, self.test_src_inputs = self.src_inputs[:self.train_size], self.src_inputs[self.train_size:] |
| 103 | self.train_tgt_inputs, self.test_tgt_inputs = self.tgt_inputs[:self.train_size], self.tgt_inputs[self.train_size:] |
| 104 | self.train_tgt_outputs, self.test_tgt_outputs = self.tgt_outputs[:self.train_size], self.tgt_outputs[self.train_size:] |
| 105 | |
| 106 | @staticmethod |
| 107 | def _split_sentences(path): |
| 108 | en_max_len, cn_max_len = 0, 0 |
| 109 | en_sts, cn_sts = [], [] |
| 110 | with open(path, 'r', encoding='utf-8') as f: |
| 111 | for line in f: |
| 112 | line_split = line.split('\t') |
| 113 | line_split[0] = re.sub(r'[^\w\s\'-]', '', line_split[0]) |
| 114 | line_split[0] = line_split[0].lower() |
| 115 | # [\u4e00-\u9fa5] matching Chinese characters |