(
workdir: PathLike,
model: DecoderBase,
dataset: str,
greedy=False,
n_samples=1,
id_range=None,
version="default",
resume=True,
)
| 36 | |
| 37 | |
| 38 | def codegen( |
| 39 | workdir: PathLike, |
| 40 | model: DecoderBase, |
| 41 | dataset: str, |
| 42 | greedy=False, |
| 43 | n_samples=1, |
| 44 | id_range=None, |
| 45 | version="default", |
| 46 | resume=True, |
| 47 | ): |
| 48 | with Progress( |
| 49 | TextColumn(f"{dataset} •" + "[progress.percentage]{task.percentage:>3.0f}%"), |
| 50 | BarColumn(), |
| 51 | MofNCompleteColumn(), |
| 52 | TextColumn("•"), |
| 53 | TimeElapsedColumn(), |
| 54 | ) as p: |
| 55 | if dataset == "humaneval": |
| 56 | from evalplus.data import get_human_eval_plus |
| 57 | |
| 58 | dataset = get_human_eval_plus(version=version) |
| 59 | elif dataset == "mbpp": |
| 60 | from evalplus.data import get_mbpp_plus |
| 61 | |
| 62 | dataset = get_mbpp_plus(version=version) |
| 63 | |
| 64 | for task_id, task in p.track(dataset.items()): |
| 65 | if id_range is not None: |
| 66 | id_num = int(task_id.split("/")[1]) |
| 67 | low, high = id_range |
| 68 | if id_num < low or id_num >= high: |
| 69 | p.console.print(f"Skipping {task_id} as it is not in {id_range}") |
| 70 | continue |
| 71 | |
| 72 | p_name = task_id.replace("/", "_") |
| 73 | os.makedirs(os.path.join(workdir, p_name), exist_ok=True) |
| 74 | log = f"Codegen: {p_name} @ {model}" |
| 75 | n_existing = 0 |
| 76 | if resume: |
| 77 | # count existing .py files |
| 78 | n_existing = len( |
| 79 | [ |
| 80 | f |
| 81 | for f in os.listdir(os.path.join(workdir, p_name)) |
| 82 | if f.endswith(".py") |
| 83 | ] |
| 84 | ) |
| 85 | if n_existing > 0: |
| 86 | log += f" (resuming from {n_existing})" |
| 87 | |
| 88 | nsamples = n_samples - n_existing |
| 89 | p.console.print(log) |
| 90 | |
| 91 | sidx = n_samples - nsamples |
| 92 | while sidx < n_samples: |
| 93 | outputs = model.codegen( |
| 94 | task["prompt"], |
| 95 | do_sample=not greedy, |
no test coverage detected