File size: 2,857 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 | import ast
import gzip
import hashlib
import json
import subprocess
from pathlib import Path
from download_model import ROOT, fetch
provenance = json.loads((ROOT/'data/provenance.json').read_text())
revision = next(r['revision'] for r in provenance if r.get('repo')=='google-research/google-research')
mbpp_path = ROOT/'sources/mbpp/sanitized-mbpp.json'
fetch(f'https://raw.githubusercontent.com/google-research/google-research/{revision}/mbpp/sanitized-mbpp.json',mbpp_path)
records = []
with gzip.open(ROOT/'sources/human-eval/data/HumanEval.jsonl.gz','rt') as f:
for row in map(json.loads,f):
records.append({'id':'humaneval:'+row['task_id'],'source':'humaneval','problem':row['prompt'],
'entry_point':row['entry_point'],'test_code':row['test'],
'reference_code':row['prompt']+row['canonical_solution'],'prelude':''})
skipped = []
for row in json.loads(mbpp_path.read_text()):
tree = ast.parse(row['code'])
functions = {node.name:node for node in tree.body if isinstance(node,ast.FunctionDef)}
entry = None
for test in row['test_list']:
entry = next((node.func.id for node in ast.walk(ast.parse(test)) if isinstance(node,ast.Call) and isinstance(node.func,ast.Name) and node.func.id in functions),None)
if entry:
break
if entry is None:
skipped.append(row['task_id'])
continue
node = functions[entry]
signature = 'def '+entry+'('+ast.unparse(node.args)+')'+(' -> '+ast.unparse(node.returns) if node.returns else '')+':'
imports = '\n'.join(row.get('test_imports',[]))
tests = '\n'.join(' '+line for test in row['test_list'] for line in test.split('\n'))
records.append({'id':'mbpp:'+str(row['task_id']),'source':'mbpp','problem':row['prompt']+'\nRequired signature:\n'+signature,
'entry_point':entry,'test_code':imports+'\ndef check(candidate):\n '+entry+' = candidate\n'+tests,
'reference_code':imports+'\n'+row['code'],'prelude':imports})
out = ROOT/'data/coding'
out.mkdir(parents=True,exist_ok=True)
(out/'all_tasks.jsonl').write_text(''.join(json.dumps(r,ensure_ascii=False)+'\n' for r in records))
human_commit = subprocess.check_output(['git','rev-parse','HEAD'],cwd=ROOT/'sources/human-eval',text=True).strip()
(out/'provenance.json').write_text(json.dumps({'humaneval':{'repo':'openai/human-eval','revision':human_commit},
'mbpp':{'repo':'google-research/google-research','revision':revision,
'file_sha256':hashlib.sha256(mbpp_path.read_bytes()).hexdigest()},
'rows':len(records),'unresolved_mbpp_entries':skipped},indent=2))
print('Prepared',len(records),'coding tasks; splits follow canonical sandbox validation.')
|