File size: 7,415 Bytes
fc32096 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 | 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()
|