feedforward flag for HumanEval eval

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Nils
2026-07-14 12:34:57 +02:00
co-authored by Claude Fable 5
parent 12ec01becc
commit d2d372c909
+7 -3
View File
@@ -57,7 +57,8 @@ def run_he_tests(code, item, timeout=10):
@torch.no_grad() @torch.no_grad()
def loop_eval(looper, adapter, tok, items, k, batch=8, max_new=380): def loop_eval(looper, adapter, tok, items, k, batch=8, max_new=380,
feedforward=False):
codes = [] codes = []
for i in range(0, len(items), batch): for i in range(0, len(items), batch):
torch.cuda.empty_cache() torch.cuda.empty_cache()
@@ -67,7 +68,8 @@ def loop_eval(looper, adapter, tok, items, k, batch=8, max_new=380):
add_special_tokens=False).to("cuda") add_special_tokens=False).to("cuda")
gen = looper.generate_frozen_prompt(adapter, tok, enc["input_ids"], k, gen = looper.generate_frozen_prompt(adapter, tok, enc["input_ids"], k,
max_new_tokens=max_new, max_new_tokens=max_new,
attention_mask=enc["attention_mask"]) attention_mask=enc["attention_mask"],
feedforward=feedforward)
for j in range(len(chunk)): for j in range(len(chunk)):
codes.append(extract_code( codes.append(extract_code(
tok.decode(gen[j, enc["input_ids"].shape[1]:], tok.decode(gen[j, enc["input_ids"].shape[1]:],
@@ -83,6 +85,7 @@ def main():
ap.add_argument("--adapter", default=None) ap.add_argument("--adapter", default=None)
ap.add_argument("--tag", default="he") ap.add_argument("--tag", default="he")
ap.add_argument("--ks", default="0,2,4") ap.add_argument("--ks", default="0,2,4")
ap.add_argument("--feedforward", action="store_true")
args = ap.parse_args() args = ap.parse_args()
ks = [int(x) for x in args.ks.split(",")] ks = [int(x) for x in args.ks.split(",")]
@@ -118,7 +121,8 @@ def main():
k0_ok = None k0_ok = None
for k in ks: for k in ks:
t0 = time.time() t0 = time.time()
oks = loop_eval(looper, adapter, tok, items, k) oks = loop_eval(looper, adapter, tok, items, k,
feedforward=args.feedforward)
if k == 0: if k == 0:
k0_ok = oks k0_ok = oks
hard = [i for i, it in enumerate(items) hard = [i for i, it in enumerate(items)