| import argparse |
| import json |
| import os |
| import random |
| import sys |
| import tempfile |
| import subprocess |
| import torch |
| import torch.nn.functional as F |
| import pyarrow.parquet as pq |
| from tqdm import tqdm |
| from transformers import AutoTokenizer, AutoModel |
|
|
| MASK_ID = 126336 |
|
|
| def add_gumbel_noise(logits, temperature): |
| if temperature == 0: |
| return logits |
| logits = logits.to(torch.float64) |
| noise = torch.rand_like(logits, dtype=torch.float64) |
| gumbel_noise = (-torch.log(noise)) ** temperature |
| return logits.exp() / gumbel_noise |
|
|
| def get_num_transfer_tokens(mask_index, steps): |
| mask_num = mask_index.sum(dim=1, keepdim=True) |
| base = mask_num // steps |
| remainder = mask_num % steps |
| num_transfer_tokens = torch.zeros(mask_num.size(0), steps, device=mask_index.device, dtype=torch.int64) + base |
| for i in range(mask_num.size(0)): |
| num_transfer_tokens[i, :remainder[i]] += 1 |
| return num_transfer_tokens |
|
|
| @torch.no_grad() |
| def generate_trajectory(model, prompt_ids, attention_mask=None, steps=16, gen_length=256, block_length=32, temperature=0.0, remasking="low_confidence"): |
| x = torch.full((prompt_ids.shape[0], prompt_ids.shape[1] + gen_length), MASK_ID, dtype=torch.long).to(model.device) |
| x[:, :prompt_ids.shape[1]] = prompt_ids.clone() |
| if attention_mask is not None: |
| attention_mask = torch.cat( |
| [attention_mask, torch.ones((prompt_ids.shape[0], gen_length), dtype=attention_mask.dtype, device=model.device)], |
| dim=-1, |
| ) |
| prompt_index = x != MASK_ID |
| assert gen_length % block_length == 0 |
| num_blocks = gen_length // block_length |
| assert steps % num_blocks == 0 |
| steps = steps // num_blocks |
| trajectory = {} |
| step_idx = 0 |
| trajectory[f"step{step_idx}"] = x[:, prompt_ids.shape[1]:].clone() |
| step_idx += 1 |
| for num_block in range(num_blocks): |
| block_mask_index = (x[:, prompt_ids.shape[1] + num_block * block_length: prompt_ids.shape[1] + (num_block + 1) * block_length] == MASK_ID) |
| num_transfer_tokens = get_num_transfer_tokens(block_mask_index, steps) |
| for i in range(steps): |
| mask_index = x == MASK_ID |
| logits = model(x, attention_mask=attention_mask).logits |
| logits_with_noise = add_gumbel_noise(logits, temperature=temperature) |
| x0 = torch.argmax(logits_with_noise, dim=-1) |
| if remasking == "low_confidence": |
| p = F.softmax(logits, dim=-1) |
| x0_p = torch.squeeze(torch.gather(p, dim=-1, index=torch.unsqueeze(x0, -1)), -1) |
| elif remasking == "random": |
| x0_p = torch.rand((x0.shape[0], x0.shape[1]), device=x0.device) |
| else: |
| raise NotImplementedError(remasking) |
| x0_p[:, prompt_ids.shape[1] + (num_block + 1) * block_length :] = -torch.inf |
| x0 = torch.where(mask_index, x0, x) |
| confidence = torch.where(mask_index, x0_p, -torch.inf) |
| transfer_index = torch.zeros_like(x0, dtype=torch.bool, device=x0.device) |
| for j in range(confidence.shape[0]): |
| k = int(num_transfer_tokens[j, i].item()) |
| if k > 0: |
| _, select_index = torch.topk(confidence[j], k=k) |
| transfer_index[j, select_index] = True |
| x[transfer_index] = x0[transfer_index] |
| trajectory[f"step{step_idx}"] = x[:, prompt_ids.shape[1]:].clone() |
| step_idx += 1 |
| return x, trajectory |
|
|
| def build_prompt(tokenizer, text): |
| messages = [{"role": "user", "content": text}] |
| return tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False) |
|
|
| def extract_code_from_text(text): |
| l = text.split("```") |
| if len(l) >= 3: |
| body = l[1] |
| if body.strip().lower().startswith("python"): |
| body = body.split("\n", 1)[1] if "\n" in body else "" |
| return body |
| return text |
|
|
| def run_tests(solution_code, test_code, timeout=20): |
| tmp = tempfile.mkdtemp(prefix="kodcode_") |
| try: |
| sol_path = os.path.join(tmp, "solution.py") |
| tst_path = os.path.join(tmp, "test_content.py") |
| runner_path = os.path.join(tmp, "test_runner.py") |
| with open(sol_path, "w", encoding="utf-8") as f: |
| f.write(solution_code) |
| with open(tst_path, "w", encoding="utf-8") as f: |
| f.write(test_code) |
| runner_src = ( |
| "import importlib.util,sys\n" |
| "sys.path.insert(0,'.')\n" |
| "spec=importlib.util.spec_from_file_location('solution','solution.py')\n" |
| "m=importlib.util.module_from_spec(spec)\n" |
| "spec.loader.exec_module(m)\n" |
| "g={'__name__':'__main__'}\n" |
| "exec(open('test_content.py','r',encoding='utf-8').read(),g)\n" |
| "fails=0\n" |
| "for k,v in list(g.items()):\n" |
| " if callable(v) and k.startswith('test_'):\n" |
| " try:\n" |
| " v()\n" |
| " except Exception:\n" |
| " fails+=1\n" |
| "import sys\n" |
| "sys.exit(1 if fails>0 else 0)\n" |
| ) |
| with open(runner_path, "w", encoding="utf-8") as f: |
| f.write(runner_src) |
| p = subprocess.run([sys.executable, runner_path], cwd=tmp, stdout=subprocess.PIPE, stderr=subprocess.PIPE, timeout=timeout) |
| return p.returncode == 0 |
| except Exception: |
| return False |
| finally: |
| try: |
| for fn in os.listdir(tmp): |
| fp = os.path.join(tmp, fn) |
| if os.path.isfile(fp): |
| os.remove(fp) |
| os.rmdir(tmp) |
| except Exception: |
| pass |
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--model_name", type=str, default="Model Path") |
| parser.add_argument("--parquet_path", type=str, default="Dataset Path") |
| parser.add_argument("--output_path", type=str, required=True) |
| parser.add_argument("--limit", type=int, default=0) |
| parser.add_argument("--num_samples", type=int, default=1) |
| parser.add_argument("--steps", type=int, default=256) |
| parser.add_argument("--max_new_tokens", type=int, default=256) |
| parser.add_argument("--block_length", type=int, default=32) |
| parser.add_argument("--temperature", type=float, default=0.0) |
| parser.add_argument("--seed", type=int, default=42) |
| parser.add_argument("--target_count", type=int, default=500) |
| args = parser.parse_args() |
| random.seed(args.seed) |
| torch.manual_seed(args.seed) |
| model = AutoModel.from_pretrained(args.model_name, trust_remote_code=True, torch_dtype=torch.bfloat16, device_map="auto").eval() |
| tokenizer = AutoTokenizer.from_pretrained(args.model_name, trust_remote_code=True) |
| if tokenizer.padding_side != "left": |
| tokenizer.padding_side = "left" |
| os.makedirs(os.path.dirname(args.output_path), exist_ok=True) |
| existing = set() |
| if os.path.exists(args.output_path): |
| with open(args.output_path, "r", encoding="utf-8") as f: |
| for line in f: |
| line = line.strip() |
| if not line: |
| continue |
| try: |
| obj = json.loads(line) |
| key = (int(obj.get("row_id", -1)), int(obj.get("sample_id", -1))) |
| existing.add(key) |
| except Exception: |
| continue |
| out_f = open(args.output_path, "a", encoding="utf-8") |
| pf = pq.ParquetFile(args.parquet_path) |
| total_rows = 0 |
| for i in range(pf.num_row_groups): |
| rg = pf.read_row_group(i) |
| total_rows += rg.num_rows |
| target = args.target_count |
| collected = 0 |
| total_attempts = 0 |
| processed = 0 |
| progress = tqdm(total=total_rows if args.limit <= 0 else min(args.limit, total_rows), desc="processing") |
| for i in range(pf.num_row_groups): |
| rg = pf.read_row_group(i) |
| tbl = rg.to_pydict() |
| n = rg.num_rows |
| for r in range(n): |
| if args.limit > 0 and processed >= args.limit: |
| break |
| processed += 1 |
| progress.update(1) |
| row = {k: tbl[k][r] for k in tbl.keys()} |
| row_id = int(row.get("metadata", {}).get("row_id")) if isinstance(row.get("metadata"), dict) and "row_id" in row["metadata"] else processed |
| for sample_id in range(args.num_samples): |
| if (row_id, sample_id) in existing: |
| continue |
| question = str(row.get("question", "")) |
| solution_ref = str(row.get("solution", "")) |
| test_code = str(row.get("test", "")) |
| teacher_text = question + "\nReference Answer: " + solution_ref + "\nAfter understanding the reference solution, please try to solve this problem using your own approach below and just OUTPUT YOUR SOLUTION CODE, DON'T EXPLAIN IT:" |
| prompt_text = build_prompt(tokenizer, teacher_text) |
| encoded = tokenizer([prompt_text], add_special_tokens=False, padding=True, return_tensors="pt") |
| input_ids = encoded["input_ids"].to(model.device) |
| attention_mask = encoded["attention_mask"].to(model.device) |
| final_x, trajectory = generate_trajectory( |
| model, |
| input_ids, |
| attention_mask=attention_mask, |
| steps=args.steps, |
| gen_length=args.max_new_tokens, |
| block_length=args.block_length, |
| temperature=args.temperature, |
| remasking="low_confidence", |
| ) |
| output_text = tokenizer.batch_decode(final_x[:, input_ids.shape[1] :], skip_special_tokens=True)[0] |
| code_text = extract_code_from_text(output_text) |
| passed = run_tests(code_text, test_code, timeout=30) |
| total_attempts += 1 |
| success_rate = collected / total_attempts * 100 if total_attempts > 0 else 0.0 |
| if not passed: |
| print(f"[idx={row_id} sample={sample_id}] x Test failed, skipped " |
| f"(success rate: {collected}/{total_attempts} = {success_rate:.2f}%)") |
| continue |
| print(f"[idx={row_id} sample={sample_id}] v Test passed, collecting trajectory " |
| f"(success rate: {collected + 1}/{total_attempts} = {(collected + 1) / total_attempts * 100:.2f}%)") |
| traj_dict = {k: v.squeeze(0).tolist() for k, v in trajectory.items()} |
| rec = { |
| "row_id": int(row_id), |
| "sample_id": int(sample_id), |
| "question": question, |
| "solution_ref": solution_ref, |
| "teacher_text": teacher_text, |
| "trajectory": traj_dict, |
| "output_text": output_text, |
| "solution_code": code_text, |
| } |
| out_f.write(json.dumps(rec, ensure_ascii=False) + "\n") |
| out_f.flush() |
| collected += 1 |
| if collected >= target: |
| break |
| if collected >= target: |
| break |
| if collected >= target: |
| break |
| progress.close() |
| out_f.close() |
|
|
| |
| final_rate = collected / total_attempts * 100 if total_attempts > 0 else 0.0 |
| summary = { |
| "summary": "collection_stats", |
| "collected": collected, |
| "total_attempts": total_attempts, |
| "success_rate": round(final_rate, 4), |
| "processed_rows": processed, |
| "target_count": target, |
| } |
| with open(args.output_path, "a", encoding="utf-8") as f: |
| f.write(json.dumps(summary, ensure_ascii=False) + "\n") |
|
|
| print(f"[Done] Collected {collected} trajectories / Total attempts {total_attempts} / " |
| f"Success rate {final_rate:.2f}% / Processed rows {processed}") |
| print(f"[Done] Output: {args.output_path}") |
|
|
| if __name__ == "__main__": |
| main() |
|
|
