run predict
(model_predict, config, args_opt, rank)
| 172 | |
| 173 | |
| 174 | def run_predict(model_predict, config, args_opt, rank): |
| 175 | """run predict""" |
| 176 | # Define tokenizer |
| 177 | tokenizer = CodeTokenizer(mode='6b') |
| 178 | |
| 179 | # Tokenize input sentence to ids |
| 180 | batch_size = config.batch_size |
| 181 | input_ids = np.array( |
| 182 | [8189, 11059, 198, 29584, 25, 198, 11377, 1398, 28186, 1391, 198, 50268, 11377, 9037, 25131, 468, 26125, 36, |
| 183 | 3639, 7, 600, 21737, 997, 82, 11, 493, 11387, 8, 1391, 198, 50272, 1640, 357, 600, 1312, 796, 657, 26, 1312, |
| 184 | 1279, 997, 82, 13, 13664, 532, 352, 26, 1312, 29577, 1391, 198, 50276, 1640, 357, 600, 474, 796, 1312, 1343, |
| 185 | 352, 26, 474, 1279, 997, 82, 13, 13664, 26, 474, 29577, 1391, 198, 50280, 361, 357, 37372, 13, 8937, 7, 77, |
| 186 | 5700, 58, 72, 60, 532, 997, 82, 58, 73, 12962, 1279, 11387, 8, 1391, 198, 50284, 7783, 2081, 26, 198, 50280, |
| 187 | 92, 198, 50276, 92, 198, 50272, 92, 198, 50272, 7783, 3991, 26, 198, 50268, 92, 198, 92, 198, 5247, 25, 198], |
| 188 | dtype=np.int32) |
| 189 | valid_length = input_ids.shape[0] |
| 190 | input_ids = np.concatenate((input_ids, np.ones(2048 - valid_length, dtype=np.int32) * 50256)) |
| 191 | attention_mask = np.tril(np.ones((2048, 2048))) |
| 192 | attention_mask[valid_length:] = 0 |
| 193 | input_ids = input_ids.reshape(1, -1).repeat(config.batch_size, axis=0) |
| 194 | current_index = valid_length - 1 if valid_length - 1 > 0 else 0 |
| 195 | init = Tensor([False], mstype.bool_) |
| 196 | model_predict.predict_network.add_flags_recursive(is_first_iteration=True) |
| 197 | batch_valid_length = Tensor(np.array([current_index for _ in range(batch_size)]), mstype.int32) |
| 198 | output_logits = model_predict.predict(Tensor(input_ids, mstype.int32), |
| 199 | init, batch_valid_length) |
| 200 | output = output_logits.asnumpy() |
| 201 | if rank == 0: |
| 202 | np.save("/home/work/sfs/xx/pangu_alpha_code/output_6_7375_8.13.npy", output) # TODO: set as current save path |
| 203 | os.system( |
| 204 | "chmod 777 /home/work/sfs/xx/pangu_alpha_code/output_6_7375_8.13.npy") # TODO: set as current save path |
| 205 | print("== Output shape: ", output.shape) |
| 206 | |
| 207 | |
| 208 | def main(): |