File size: 7,506 Bytes
256a159 |
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 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 |
import argparse
import json
import os
import os.path as osp
import re
import mmengine
from mmengine import Config
from mmengine.utils import mkdir_or_exist
from opencompass.datasets.humanevalx import _clean_up_code
from opencompass.utils import (dataset_abbr_from_cfg, get_infer_output_path,
get_logger, model_abbr_from_cfg)
def parse_args():
parser = argparse.ArgumentParser(
description='Collect Humanevalx dataset predictions.')
parser.add_argument('config', help='Config file path')
parser.add_argument('-r',
'--reuse',
nargs='?',
type=str,
const='latest',
help='Reuse previous outputs & results, and run any '
'missing jobs presented in the config. If its '
'argument is not specified, the latest results in '
'the work_dir will be reused. The argument should '
'also be a specific timestamp, e.g. 20230516_144254'),
args = parser.parse_args()
return args
_LANGUAGE_NAME_DICT = {
'cpp': 'CPP',
'go': 'Go',
'java': 'Java',
'js': 'JavaScript',
'python': 'Python',
'rust': 'Rust',
}
FAILED = 0
SUCCEED = 1
def gpt_python_postprocess(ori_prompt: str, text: str) -> str:
"""Better answer postprocessor for better instruction-aligned models like
GPT."""
if '```' in text:
blocks = re.findall(r'```(.*?)```', text, re.DOTALL)
if len(blocks) == 0:
text = text.split('```')[1] # fall back to default strategy
else:
text = blocks[0] # fetch the first code block
if not text.startswith('\n'): # in case starting with ```python
text = text[max(text.find('\n') + 1, 0):]
match_ori = re.search(r'def(.*?)\(', ori_prompt)
match = re.search(r'def(.*?)\(', text)
if match:
if match.group() == match_ori.group():
text = re.sub('def(.*?)\n', '', text, count=1)
for c_index, c in enumerate(text[:5]):
if c != ' ':
text = ' ' * (4 - c_index) + text
break
text = text.split('\n\n\n')[0]
return text
def wizardcoder_postprocess(text: str) -> str:
"""Postprocess for WizardCoder Models."""
if '```' in text:
blocks = re.findall(r'```(.*?)```', text, re.DOTALL)
if len(blocks) == 0:
text = text.split('```')[1] # fall back to default strategy
else:
text = blocks[0] # fetch the first code block
if not text.startswith('\n'): # in case starting with ```python
text = text[max(text.find('\n') + 1, 0):]
else:
match = re.search(r'Here(.*?)\n', text)
if match:
text = re.sub('Here(.*?)\n', '', text, count=1)
return text
def collect_preds(filename: str):
# in case the prediction is partial
root, ext = osp.splitext(filename)
partial_filename = root + '_0' + ext
# collect all the prediction results
if not osp.exists(osp.realpath(filename)) and not osp.exists(
osp.realpath(partial_filename)):
print(f'No predictions found for {filename}')
return FAILED, None, None
else:
if osp.exists(osp.realpath(filename)):
preds = mmengine.load(filename)
pred_strs = [
preds[str(i)]['prediction'] for i in range(len(preds))
]
ori_prompt_strs = [
preds[str(i)]['origin_prompt'] for i in range(len(preds))
]
else:
filename = partial_filename
pred_strs = []
ori_prompt_strs = []
i = 1
while osp.exists(osp.realpath(filename)):
preds = mmengine.load(filename)
filename = root + f'_{i}' + ext
i += 1
pred_strs += [
preds[str(i)]['prediction'] for i in range(len(preds))
]
ori_prompt_strs += [
preds[str(i)]['origin_prompt'] for i in range(len(preds))
]
return SUCCEED, ori_prompt_strs, pred_strs
def main():
args = parse_args()
# initialize logger
logger = get_logger(log_level='INFO')
cfg = Config.fromfile(args.config)
cfg.setdefault('work_dir', './outputs/default/')
assert args.reuse, 'Please provide the experienment work dir.'
if args.reuse:
if args.reuse == 'latest':
if not os.path.exists(cfg.work_dir) or not os.listdir(
cfg.work_dir):
logger.warning('No previous results to reuse!')
else:
dirs = os.listdir(cfg.work_dir)
dir_time_str = sorted(dirs)[-1]
else:
dir_time_str = args.reuse
logger.info(f'Reusing experiements from {dir_time_str}')
# update "actual" work_dir
cfg['work_dir'] = osp.join(cfg.work_dir, dir_time_str)
for model in cfg.models:
model_abbr = model_abbr_from_cfg(model)
for dataset in cfg.datasets:
dataset_abbr = dataset_abbr_from_cfg(dataset)
filename = get_infer_output_path(
model, dataset, osp.join(cfg.work_dir, 'predictions'))
succeed, ori_prompt_strs, pred_strs = collect_preds(filename)
if not succeed:
continue
# infer the language type
for k, v in _LANGUAGE_NAME_DICT.items():
if k in dataset_abbr:
lang = k
task = v
break
# special postprocess for GPT
if model_abbr in [
'WizardCoder-1B-V1.0',
'WizardCoder-3B-V1.0',
'WizardCoder-15B-V1.0',
'WizardCoder-Python-13B-V1.0',
'WizardCoder-Python-34B-V1.0',
]:
predictions = [{
'task_id': f'{task}/{i}',
'generation': wizardcoder_postprocess(pred),
} for i, pred in enumerate(pred_strs)]
elif 'CodeLlama' not in model_abbr and lang == 'python':
predictions = [{
'task_id':
f'{task}/{i}',
'generation':
gpt_python_postprocess(ori_prompt, pred),
} for i, (ori_prompt,
pred) in enumerate(zip(ori_prompt_strs, pred_strs))]
else:
predictions = [{
'task_id': f'{task}/{i}',
'generation': _clean_up_code(pred, lang),
} for i, pred in enumerate(pred_strs)]
# save processed results if not exists
result_file_path = os.path.join(cfg['work_dir'], 'humanevalx',
model_abbr,
f'humanevalx_{lang}.json')
if osp.exists(result_file_path):
logger.info(
f'File exists for {model_abbr}, skip copy from predictions.' # noqa
)
else:
mkdir_or_exist(osp.split(result_file_path)[0])
with open(result_file_path, 'w') as f:
for pred in predictions:
f.write(json.dumps(pred) + '\n')
if __name__ == '__main__':
main()
|