| 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'." |
| 149 | total_size = self.train_size |
| 150 | if mode == 'test': |
| 151 | total_size = self.test_size |
| 152 | |
| 153 | max_batch = total_size // self.batch_size |
| 154 | if total_size % self.batch_size > 0: |
| 155 | max_batch += 1 |
| 156 | assert batch < max_batch, "The batch number is out of bounds." |
| 157 | |
| 158 | low = batch * self.batch_size |
| 159 | if (batch + 1) * self.batch_size < total_size: |
| 160 | high = (batch + 1) * self.batch_size |
| 161 | else: |
| 162 | high = total_size |
| 163 | if mode == 'train': |
| 164 | if high-low != self.batch_size: |
| 165 | return (np.concatenate((self.train_src_inputs[low:high], self.train_src_inputs[:self.batch_size-high+low]), axis=0), |
| 166 | np.concatenate((self.train_tgt_inputs[low:high], self.train_tgt_inputs[:self.batch_size-high+low]), axis=0), |
| 167 | np.concatenate((self.train_tgt_outputs[low:high], self.train_tgt_outputs[:self.batch_size-high+low]), axis=0)) |
| 168 | else: |
| 169 | return self.train_src_inputs[low:high], self.train_tgt_inputs[low:high], self.train_tgt_outputs[low:high] |
| 170 | else: |
| 171 | if high-low != self.batch_size: |
| 172 | return (np.concatenate((self.test_src_inputs[low:high], self.test_src_inputs[:self.batch_size-high+low]), axis=0), |
| 173 | np.concatenate((self.test_tgt_inputs[low:high], self.test_tgt_inputs[:self.batch_size-high+low]), axis=0), |
| 174 | np.concatenate((self.test_tgt_outputs[low:high], self.test_tgt_outputs[:self.batch_size-high+low]), axis=0)) |
| 175 | return self.test_src_inputs[low:high], self.test_tgt_inputs[low:high], self.test_tgt_outputs[low:high] |
| 176 | |
| 177 | def __len__(self): |
| 178 | return self.src_inputs.shape[0] |