| 125 | |
| 126 | @staticmethod |
| 127 | def _encoding_stc(src_tokens, tgt_tokens, src_max_len, tgt_max_len, src_vocab, tgt_vocab): |
| 128 | src_list = [] |
| 129 | for line in src_tokens: |
| 130 | if len(line) > src_max_len: |
| 131 | line = line[:src_max_len] |
| 132 | lst = src_vocab[line + ['<pad>'] * (src_max_len + 1 - len(line))] |
| 133 | src_list.append(lst) |
| 134 | tgt_in_list, tgt_out_list = [], [] |
| 135 | for line in tgt_tokens: |
| 136 | if len(line) > tgt_max_len: |
| 137 | line = line[:tgt_max_len] |
| 138 | in_lst = tgt_vocab[['<bos>'] + line + ['<pad>'] * (tgt_max_len - len(line))] |
| 139 | out_lst = tgt_vocab[line + ['<eos>'] + ['<pad>'] * (tgt_max_len - len(line))] |
| 140 | tgt_in_list.append(in_lst) |
| 141 | tgt_out_list.append(out_lst) |
| 142 | src_np = np.asarray(src_list, dtype=np.int32) |
| 143 | tgt_in_np = np.asarray(tgt_in_list, dtype=np.int32) |
| 144 | tgt_out_np = np.asarray(tgt_out_list, dtype=np.int32) |
| 145 | return src_np, tgt_in_np, tgt_out_np |
| 146 | |
| 147 | def get_batch_data(self, batch, mode='train'): |
| 148 | assert (mode == 'train' or mode == 'test'), "The mode must be 'train' or 'test'." |