| import argparse |
| import ast |
| import json |
| import platform |
| import re |
| import time |
| from collections import Counter |
| from importlib.metadata import version |
| from common import ROOT, inside, rows, save_json, append, stable_seed, permit_compute |
| from evaluate import model_fingerprint,recover_generations |
| from inference_queue import completed_outputs |
| from resources import snapshot |
| from code_sandbox import check_code |
|
|
|
|
| def extract_code(text): |
| text = text.split('</think>')[-1].strip() |
| blocks = re.findall(r'```(?:python|py)?\s*\n(.*?)```',text,flags=re.S|re.I) |
| if blocks: |
| return blocks[0].strip() |
| if text.startswith('```'): |
| return '\n'.join(text.split('\n')[1:]).removesuffix('```').strip() |
| return text |
|
|
|
|
| def public_prelude(row): |
| if row['source'] != 'humaneval': |
| return row.get('prelude','') |
| tree = ast.parse(row['problem']) |
| return '\n'.join(ast.unparse(node) for node in tree.body if isinstance(node,(ast.Import,ast.ImportFrom))) |
|
|
|
|
| def main(): |
| p = argparse.ArgumentParser() |
| p.add_argument('--model',default=str(ROOT/'models/original')) |
| p.add_argument('--data',required=True) |
| p.add_argument('--out',required=True) |
| p.add_argument('--config',default=str(ROOT/'configs/coding_protocol.json')) |
| p.add_argument('--limit',type=int) |
| p.add_argument('--samples',type=int) |
| args = p.parse_args() |
| permit_compute() |
| from transformers import AutoTokenizer |
| from vllm import LLM,SamplingParams |
| model,data,out = inside(args.model),inside(args.data),inside(args.out) |
| cfg = json.loads(inside(args.config).read_text()) |
| records = rows(data) |
| if args.limit: |
| records = records[:args.limit] |
| samples = args.samples or cfg['samples'] |
| out.mkdir(parents=True,exist_ok=True) |
| job = {'model':str(model),'fingerprint':model_fingerprint(model),'data':str(data), |
| 'data_sha256':__import__('hashlib').sha256(data.read_bytes()).hexdigest(),'protocol':cfg, |
| 'samples':samples,'max_tokens':cfg['max_tokens'],'set':'coding','limit':args.limit, |
| 'runtime_versions':{x:version(x) for x in ('vllm','torch','transformers')},'platform':platform.platform(), |
| 'seed_scheme':'sha256(base_seed:coding/item_id:sample)31'} |
| if (out/'job.json').exists() and json.loads((out/'job.json').read_text()) != job: |
| raise ValueError('Existing job differs') |
| save_json(out/'job.json',job) |
| completed = recover_generations(out/'generations.jsonl') |
| done = {(r['id'],r['sample']) for r in completed} |
| if len(done)!=len(completed): |
| raise ValueError('Duplicate generations') |
| pending = [(row,sample) for row in records for sample in range(samples) if (row['id'],sample) not in done] |
| if pending: |
| tok = AutoTokenizer.from_pretrained(model,local_files_only=True) |
| engine = LLM(model=str(model),dtype='bfloat16',seed=cfg['seed'],max_model_len=cfg['max_model_len'], |
| gpu_memory_utilization=cfg['gpu_memory_utilization'],kv_cache_memory_bytes=cfg['kv_cache_memory_bytes'], |
| max_num_seqs=cfg['max_num_seqs'],max_num_batched_tokens=cfg['max_num_batched_tokens'], |
| enable_prefix_caching=False,enforce_eager=False,limit_mm_per_prompt={'image':0,'video':0}) |
| append(out/'resources.jsonl',snapshot('code_engine_initialized')) |
| prompts,params = [],[] |
| for row,sample in pending: |
| prompt = tok.apply_chat_template([{'role':'user','content':cfg['instruction']+'\n\n'+row['problem']}], |
| tokenize=False,add_generation_prompt=True,enable_thinking=cfg['enable_thinking']) |
| prompts.append(prompt) |
| params.append(SamplingParams(n=1,seed=stable_seed('coding/'+row['id'],sample,cfg['seed']), |
| temperature=cfg['temperature'],top_p=cfg['top_p'],top_k=cfg['top_k'], |
| presence_penalty=cfg['presence_penalty'],max_tokens=cfg['max_tokens'])) |
| start = time.time() |
| for index,output in completed_outputs(engine,prompts,params,stream=True): |
| row,sample = pending[index] |
| generation = output.outputs[0] |
| if generation.finish_reason not in ('stop','length'): |
| raise RuntimeError('Incomplete code generation') |
| code = public_prelude(row)+'\n'+extract_code(generation.text) |
| verdict = check_code(code,row['entry_point'],row['test_code']) |
| result = {'id':row['id'],'source':row['source'],'sample':sample, |
| 'seed':stable_seed('coding/'+row['id'],sample,cfg['seed']),'correct':verdict['passed'], |
| 'verdict':verdict,'code':code,'text':generation.text,'tokens':len(generation.token_ids), |
| 'prompt':prompts[index], |
| 'prompt_token_ids':tok.encode(prompts[index],add_special_tokens=False), |
| 'token_ids':list(generation.token_ids),'finish_reason':generation.finish_reason, |
| 'queue_elapsed_seconds':time.time()-start} |
| append(out/'generations.jsonl',result) |
| completed.append(result) |
| print(json.dumps({k:result[k] for k in ('id','sample','correct','tokens','finish_reason')},ensure_ascii=False),flush=True) |
| append(out/'resources.jsonl',snapshot('code_queue_finished')) |
| n = len(completed) |
| if not n: |
| raise ValueError('No coding results') |
| by_source = {} |
| for source in sorted({r['source'] for r in completed}): |
| subset = [r for r in completed if r['source']==source] |
| by_source[source] = {'correct':sum(r['correct'] for r in subset),'generations':len(subset), |
| 'accuracy':sum(r['correct'] for r in subset)/len(subset)} |
| by_sample = {} |
| for sample in range(samples): |
| subset = [r for r in completed if r['sample']==sample] |
| by_sample[str(sample)] = {'correct':sum(r['correct'] for r in subset),'generations':len(subset)} |
| task_results = {} |
| for result in completed: |
| task_results.setdefault(result['id'],[]).append(result['correct']) |
| save_json(out/'summary.json',{'generations':n,'planned_generations':len(records)*samples, |
| 'complete':n==len(records)*samples, |
| 'problems':len(task_results),'samples_per_problem':samples, |
| 'correct':sum(r['correct'] for r in completed),'accuracy':sum(r['correct'] for r in completed)/n, |
| 'mean_tokens':sum(r['tokens'] for r in completed)/n,'by_source':by_source, |
| 'by_sample':by_sample, |
| 'tasks_correct_all_samples':sum(len(v)==samples and all(v) for v in task_results.values()), |
| 'verdict_counts':dict(Counter('passed' if r['correct'] else r['verdict'].get('error_type','unknown') for r in completed)), |
| 'guard_rejections':sum(not r['correct'] and r['verdict'].get('error','').startswith('unsupported_') for r in completed), |
| 'cut_share':sum(r['finish_reason']=='length' for r in completed)/n, |
| 'scope':'Compact fixed HumanEval/MBPP subsets, custom chat prompt and restricted stdlib sandbox; not official leaderboard scores.'}) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
