run predict
(model_predict, config, args_opt, rank)
| 194 | |
| 195 | |
| 196 | def run_predict(model_predict, config, args_opt, rank): |
| 197 | """run predict""" |
| 198 | from src.generate_finetune import generate_increment |
| 199 | # Define tokenizer |
| 200 | tokenizer = CodeTokenizer(mode='6b') |
| 201 | |
| 202 | # Tokenize input sentence to ids |
| 203 | lang = args_opt.language |
| 204 | data_path = os.path.join(args_opt.code_data, lang, 'test') |
| 205 | dataset = LMDBDataset(data_path) |
| 206 | samples = [] |
| 207 | for i in range(len(dataset)): |
| 208 | prompt, length = dataset[i] |
| 209 | samples.append(prompt[:length]) |
| 210 | generations = [] |
| 211 | batch_size = config.batch_size |
| 212 | verbose = (rank % 8 == 0) |
| 213 | save_path = f'/home/work/sfs/xx/pangu_alpha_code/generation_finetune/code_translation/{lang}/temp_{args_opt.temperature}.txt' # TODO: set as current save path |
| 214 | save_dir = os.path.split(save_path)[0] |
| 215 | if rank == 0: |
| 216 | if not os.path.exists(save_dir): |
| 217 | os.makedirs(save_dir) |
| 218 | if not os.path.exists(save_path): |
| 219 | f = open(save_path, 'w') |
| 220 | f.close() |
| 221 | os.system(f'sudo chmod 777 -R {os.path.split(save_dir)[0]}') |
| 222 | batch = [] |
| 223 | input_length = [] |
| 224 | sample_ids = [] |
| 225 | for i, sample in enumerate(samples): |
| 226 | tokenized_token = sample |
| 227 | input_ids = np.array(tokenized_token).reshape(1, -1) |
| 228 | batch.append(input_ids) |
| 229 | input_length.append(input_ids.shape[1]) |
| 230 | sample_ids.append(i) |
| 231 | if (i + 1) % batch_size == 0: |
| 232 | valid_length = max(input_length) |
| 233 | for j in range(len(batch)): |
| 234 | batch[j] = np.pad(batch[j], ((0, 0), (0, valid_length - input_length[j])), |
| 235 | 'constant', constant_values=(args_opt.end_token, args_opt.end_token)) |
| 236 | input_ids = np.concatenate(batch, axis=0) |
| 237 | t0 = time.perf_counter() |
| 238 | output_ids = generate_increment(model_predict, input_ids, input_length, args_opt, tokenizer, verbose) |
| 239 | t1 = time.perf_counter() |
| 240 | batch, input_length = [], [] |
| 241 | if rank % 8 == 0: |
| 242 | print(f"=== Batch time: {t1 - t0}s") |
| 243 | for k, out in enumerate(output_ids): |
| 244 | print(f"=================== generation {sample_ids[k]} ====================") |
| 245 | print(out, flush=True) |
| 246 | generations.append(out) |
| 247 | if rank == 0: |
| 248 | f = open(save_path, 'a') |
| 249 | f.write(generations[-1]) |
| 250 | if not generations[-1].endswith('\n'): |
| 251 | f.write('\n') |
| 252 | f.close() |
| 253 | sample_ids = [] |
no test coverage detected