File size: 4,322 Bytes
f10f49d | 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 | import os
import gc
import time
import json
import subprocess
import sys
import torch
from safetensors.torch import load_file
from eval_engine import MODEL_DIR, Qwen3FastRunner, decode_in_subprocess
from code_benchmark_problems import CODE_BENCHMARK_PROBLEMS, verify_code_solution
torch.set_num_threads(2)
CODE_TOK_CACHE_PATH = "/home/user/.cache/code_benchmark_tokens.pt"
def ensure_code_tokenized_cache():
if os.path.exists(CODE_TOK_CACHE_PATH):
return
cmd = [
sys.executable,
"-c",
(
"import torch; from transformers import AutoTokenizer; "
"from code_benchmark_problems import CODE_BENCHMARK_PROBLEMS; "
f"tok = AutoTokenizer.from_pretrained('{MODEL_DIR}'); "
"pids = [tok.encode(tok.apply_chat_template([{'role':'user','content':p['prompt']}], tokenize=False, add_generation_prompt=True, enable_thinking=False), add_special_tokens=False) for p in CODE_BENCHMARK_PROBLEMS]; "
f"torch.save({{'prompt_ids_list': pids, 'eos_id': tok.eos_token_id, 'pad_id': tok.pad_token_id}}, '{CODE_TOK_CACHE_PATH}')"
),
]
subprocess.run(cmd, check=True)
def evaluate_code_state_dict(sd, problems=CODE_BENCHMARK_PROBLEMS, max_new_tokens=90, batch_size=7, verbose=True):
ensure_code_tokenized_cache()
tok_data = torch.load(CODE_TOK_CACHE_PATH, weights_only=True)
prompt_ids_list = tok_data["prompt_ids_list"]
eos_id = tok_data["eos_id"]
pad_id = tok_data["pad_id"]
order = sorted(range(len(problems)), key=lambda i: len(prompt_ids_list[i]))
runner = Qwen3FastRunner(sd)
t0 = time.time()
ordered_gen_ids = []
num_batches = (len(problems) + batch_size - 1) // batch_size
for b_idx, start in enumerate(range(0, len(problems), batch_size)):
end = min(start + batch_size, len(problems))
batch_indices = order[start:end]
batch_prompts = [prompt_ids_list[i] for i in batch_indices]
tb0 = time.time()
gen_ids = runner.generate_batch(batch_prompts, eos_id, pad_id, max_new_tokens=max_new_tokens)
ordered_gen_ids.extend(gen_ids)
if verbose:
avg_t = sum(len(g) for g in gen_ids) / len(gen_ids)
print(f" Batch {b_idx+1}/{num_batches} ({end-start} tasks, avg {avg_t:.1f} toks) in {time.time()-tb0:.1f}s", flush=True)
elapsed = time.time() - t0
all_gen_ids = [None] * len(problems)
for orig_idx, g_ids in zip(order, ordered_gen_ids):
all_gen_ids[orig_idx] = g_ids
del runner
gc.collect()
responses = decode_in_subprocess(all_gen_ids)
results = []
passed = 0
for prob, resp, g_ids in zip(problems, responses, all_gen_ids):
ok, detail = verify_code_solution(prob, resp)
if ok:
passed += 1
results.append({
"id": prob["id"],
"category": prob["category"],
"prompt": prob["prompt"],
"passed": ok,
"detail": detail,
"num_tokens": len(g_ids),
"response": resp,
})
summary = {
"total": len(problems),
"passed": passed,
"accuracy": round(100.0 * passed / len(problems), 2),
"avg_tokens": round(sum(r["num_tokens"] for r in results) / len(results), 1),
"elapsed_sec": round(elapsed, 2),
"results": results,
}
return summary
if __name__ == "__main__":
print("Evaluating Baseline Qwen3-0.6B on 28-task Executable Code Benchmark...", flush=True)
sd = load_file(os.path.join(MODEL_DIR, "model.safetensors"), device="cpu")
summary = evaluate_code_state_dict(sd, CODE_BENCHMARK_PROBLEMS, max_new_tokens=90, batch_size=7, verbose=True)
with open("/home/user/code_baseline_results.json", "w") as f:
json.dump(summary, f, indent=2)
print(f"\n==========================================", flush=True)
print(f"CODE BASELINE QWEN3-0.6B: {summary['passed']}/{summary['total']} ({summary['accuracy']}%) | Avg Toks: {summary['avg_tokens']} | {summary['elapsed_sec']}s", flush=True)
print(f"==========================================", flush=True)
for r in summary["results"]:
status = "PASS" if r["passed"] else "FAIL"
print(f"[{status:4s}] {r['id']:8s} ({r['category']:16s}) | Toks: {r['num_tokens']:3d} | {r['detail']}", flush=True)
|