File size: 6,771 Bytes
580cb69 | 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 | #!/usr/bin/env python3
"""Generate EvalPlus-compatible HumanEval(+)/MBPP(+) samples."""
from __future__ import annotations
import argparse
import importlib
import json
import re
import sys
from pathlib import Path
from typing import Any
def arguments() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--dataset", choices=("humaneval", "mbpp"), required=True)
parser.add_argument("--checkpoint", type=Path, required=True)
parser.add_argument("--tokenizer", type=Path, required=True)
parser.add_argument("--model-path", type=Path, required=True)
parser.add_argument("--model-module", required=True)
parser.add_argument("--model-class", default="DnaPeerV21")
parser.add_argument("--adapter-module", help="Optional module whose load(args) returns an object with generate(prompt, max_new_tokens)")
parser.add_argument("--out", type=Path, required=True)
parser.add_argument("--device", default="cuda")
parser.add_argument("--max-new-tokens", type=int, default=512)
parser.add_argument("--limit", type=int, help="Smoke testing only; never report limited runs as full benchmark")
return parser.parse_args()
def strip_fence(text: str) -> str:
text = text.strip()
if "```" not in text:
return text
parts = text.split("```")
fenced = parts[1] if len(parts) >= 3 else parts[-1]
if fenced.lstrip().startswith("python"):
fenced = fenced.lstrip()[6:]
return fenced.strip()
class BuiltinRecurrentAdapter:
"""Adapter for v23/v24-style DnaPeer recurrent blocks.
v25 model changes should use --adapter-module instead of editing evaluator code.
"""
def __init__(self, args: argparse.Namespace):
import torch
from tokenizers import Tokenizer, decoders
sys.path.insert(0, str(args.model_path))
module = importlib.import_module(args.model_module)
model_type = getattr(module, args.model_class)
checkpoint = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
self.model = model_type(**checkpoint["config"]).to(args.device)
self.model.load_state_dict({key.replace("_orig_mod.", ""): value for key, value in checkpoint["model"].items()}, strict=True)
self.model.eval()
self.tokenizer = Tokenizer.from_file(str(args.tokenizer))
self.tokenizer.decoder = decoders.ByteLevel()
self.bos = self.tokenizer.token_to_id("<s>")
self.eos = self.tokenizer.token_to_id("</s>")
self.user = self.tokenizer.token_to_id("<|user|>")
self.assistant = self.tokenizer.token_to_id("<|assistant|>")
self.device, self.torch = args.device, torch
if None in (self.bos, self.eos, self.user, self.assistant):
raise ValueError("tokenizer must contain <s>, </s>, <|user|>, and <|assistant|>")
def generate(self, prompt: str, max_new_tokens: int) -> str:
torch = self.torch
model = self.model
states = [torch.zeros(1, model.d, device=self.device) for _ in model.blocks]
def step(token: int):
x = model.embed(torch.tensor([token], device=self.device))
for index, block in enumerate(model.blocks):
key, value, receptance, gate = block.proj(block.n1(x)).chunk(4, -1)
gate = torch.sigmoid(gate + block.decay)
states[index] = gate * states[index] + (1 - gate) * torch.tanh(key)
x = x + block.o(torch.sigmoid(receptance) * states[index] * torch.sigmoid(value))
x = x + block.ffn(block.n2(x))
memory, _, _ = model.route(x)
return model.norm(x + memory)
prefix = [self.bos, self.user] + self.tokenizer.encode("\n" + prompt).ids + [self.eos, self.assistant]
with torch.no_grad():
feature = None
for token in prefix:
feature = step(token)
generated = []
for _ in range(max_new_tokens):
token = int(torch.nn.functional.linear(feature, model.embed.weight)[0].argmax())
if token == self.eos:
break
generated.append(token)
text = self.tokenizer.decode(generated)
markers = ("\n<|user|>", "\nif __name__ ==", "\n# End")
hits = [text.find(marker) for marker in markers if marker in text]
if hits:
generated = self.tokenizer.encode(text[: min(hits)]).ids
break
feature = step(token)
return self.tokenizer.decode(generated)
def main() -> None:
args = arguments()
try:
from evalplus.data import get_human_eval_plus, get_mbpp_plus
except ImportError as exc:
raise SystemExit("install pinned EvalPlus first: pip install 'evalplus==0.3.1'") from exc
problems: dict[str, dict[str, Any]] = (get_human_eval_plus() if args.dataset == "humaneval" else get_mbpp_plus())
if args.adapter_module:
sys.path.insert(0, str(args.model_path))
adapter = importlib.import_module(args.adapter_module).load(args)
else:
adapter = BuiltinRecurrentAdapter(args)
args.out.parent.mkdir(parents=True, exist_ok=True)
selected = list(problems.items())[: args.limit]
with args.out.open("w", encoding="utf-8", buffering=1) as output:
for index, (task_id, problem) in enumerate(selected, 1):
prompt = problem["prompt"]
instruction = "Complete this Python program. Return only valid Python code, without Markdown fences.\n\n" + prompt
completion = strip_fence(adapter.generate(instruction, args.max_new_tokens))
entry_point = problem.get("entry_point")
definition = rf"\bdef\s+{re.escape(entry_point)}\s*\(" if entry_point else None
if definition and re.search(definition, completion):
solution = completion
elif definition and re.search(definition, prompt):
solution = prompt + completion
else:
solution = completion
output.write(json.dumps({"task_id": task_id, "solution": solution}) + "\n")
print(f"GENERATED {index}/{len(selected)} {task_id}", flush=True)
metadata = {"dataset": args.dataset, "samples": len(selected), "full_dataset": args.limit is None,
"generation": "greedy", "max_new_tokens": args.max_new_tokens,
"checkpoint": str(args.checkpoint), "model_module": args.model_module,
"adapter_module": args.adapter_module}
args.out.with_suffix(args.out.suffix + ".meta.json").write_text(json.dumps(metadata, indent=2) + "\n", encoding="utf-8")
if __name__ == "__main__":
main()
|