run predict
(model_predict, config, args_opt, rank)
| 193 | |
| 194 | |
| 195 | def run_predict(model_predict, config, args_opt, rank): |
| 196 | """run predict""" |
| 197 | from src.generate_humaneval import generate_increment |
| 198 | # Define tokenizer |
| 199 | tokenizer = CodeTokenizer(mode='6b') |
| 200 | |
| 201 | # Tokenize input sentence to ids |
| 202 | humaneval_path = '/home/work/sfs/xx/human_eval_x/data/humaneval_cpp.jsonl' # TODO: set as current humaneval path |
| 203 | humaneval = open(humaneval_path, 'r').readlines() |
| 204 | humaneval = [json.loads(task) for task in humaneval if len(task) != 0] |
| 205 | samples = [task['prompt'] for task in humaneval] |
| 206 | generations = [] |
| 207 | batch_size = config.batch_size |
| 208 | verbose = (rank % 8 == 0) |
| 209 | part = int(args_opt.part) |
| 210 | gen_times = 12 # TODO: set as generation times of current task |
| 211 | print(f"gen times: {gen_times}, part: {part}") |
| 212 | save_path = f'/home/work/sfs/xx/pangu_alpha_code/generation_humanevalx/cpp/temp_{args_opt.temperature}/samples_{args_opt.load_ckpt_epoch}_part_{part}.jsonl' # TODO: set as current save path |
| 213 | if rank == 0 and not os.path.exists(save_path): |
| 214 | os.makedirs(os.path.split(save_path)[0], exist_ok=True) |
| 215 | f = open(save_path, 'w') |
| 216 | f.close() |
| 217 | os.system(f'sudo chmod 777 {save_path}') |
| 218 | for i, sample in enumerate(samples): |
| 219 | tag = "// language: C++\n" |
| 220 | sample = tag + sample |
| 221 | if rank % 8 == 0: |
| 222 | print(f"=================== prompt {i} ====================") |
| 223 | print(sample, flush=True) |
| 224 | for j in range((gen_times + batch_size - 1) // batch_size): |
| 225 | tokenized_token = tokenizer.encode_code(sample) |
| 226 | input_ids = np.array(tokenized_token).reshape(1, -1).repeat(batch_size, axis=0) |
| 227 | # Call inference |
| 228 | mindspore.set_seed(j + 8 * part) |
| 229 | generate_func = generate_increment |
| 230 | t0 = time.perf_counter() |
| 231 | output_ids = generate_func(model_predict, input_ids, args_opt, tokenizer, verbose) |
| 232 | t1 = time.perf_counter() |
| 233 | if rank % 8 == 0: |
| 234 | print(f"=== Batch time: {t1 - t0}s") |
| 235 | for k, out in enumerate(output_ids): |
| 236 | print(f"=================== generation {j * batch_size + k} ====================") |
| 237 | print(out, flush=True) |
| 238 | generations.append(json.dumps({'task_id': humaneval[i]['task_id'], 'completion': out})) |
| 239 | if rank == 0: |
| 240 | f = open(save_path, 'a') |
| 241 | f.write(generations[-1] + '\n') |
| 242 | f.close() |
| 243 | |
| 244 | |
| 245 | def main(): |
no test coverage detected