import os import gc import time import math import json import torch import torch.nn.functional as F from safetensors.torch import load_file from transformers import AutoTokenizer from eval_engine import MODEL_DIR, rms_norm, rotate_half from code_eval_engine import CODE_TOK_CACHE_PATH, ensure_code_tokenized_cache, evaluate_code_state_dict from code_benchmark_problems import CODE_BENCHMARK_PROBLEMS torch.set_num_threads(2) CACHE_L20_PATH = "/home/user/.cache/code_edit_l20_cache.pt" CALIB_STATES_PATH = "/home/user/.cache/code_calib_states.pt" SURGERY_DELTA_CACHE = "/home/user/.cache/code_surgery_layers_20_27.pt" def get_code_raw_edits(): ensure_code_tokenized_cache() tok = AutoTokenizer.from_pretrained(MODEL_DIR) tok_data = torch.load(CODE_TOK_CACHE_PATH, weights_only=True) prompt_ids_list = tok_data["prompt_ids_list"] def make_diff_edit(prob_idx, prefix, good_cont, bad_cont, weight=1.0): p_ids = prompt_ids_list[prob_idx] ids_pg = tok.encode(prefix + good_cont, add_special_tokens=False) ids_pb = tok.encode(prefix + bad_cont, add_special_tokens=False) diff_idx = 0 while diff_idx < min(len(ids_pg), len(ids_pb)) and ids_pg[diff_idx] == ids_pb[diff_idx]: diff_idx += 1 seq = p_ids + ids_pg[:diff_idx] return (prob_idx, seq, ids_pg[diff_idx], ids_pb[diff_idx], weight) raw_edits = [ # PY_01: merge_intervals make_diff_edit(0, "```python\ndef merge_intervals(intervals: list[list[int]]) -> list[list[int]]:\n ", "intervals.sort()\n merged = []\n for interval in intervals:", "if not intervals:", 1.0), make_diff_edit(0, "```python\ndef merge_intervals(intervals: list[list[int]]) -> list[list[int]]:\n intervals.sort()\n merged = []\n for ", "interval in intervals:\n if not merged or interval[0] > merged[-1][1]:", "s, e in intervals:", 1.0), # PY_04: group_anagrams make_diff_edit(3, "```python\ndef group_anagrams(words: list[str]) -> list[list[str]]:\n ", "groups = {}", "result = []", 1.0), make_diff_edit(3, "```python\ndef group_anagrams(words: list[str]) -> list[list[str]]:\n groups = {}\n for word in words:\n ", "sorted_word = ''.join(sorted(word))", "if word in groups:", 1.15), # PY_05: is_valid_brackets make_diff_edit(4, "```python\ndef is_valid_brackets(s: str) -> bool:\n stack = []\n ", "mapping = {')': '(', ']': '[', '}': '{'}", "for char in s:", 1.0), make_diff_edit(4, "```python\ndef is_valid_brackets(s: str) -> bool:\n stack = []\n mapping = {')': '(', ']': '[', '}': '{'}\n for char in s:\n if char in mapping:\n ", "if not stack or stack.pop() != mapping[char]:\n return False", "top = stack.pop()", 1.15), make_diff_edit(4, "```python\ndef is_valid_brackets(s: str) -> bool:\n stack = []\n mapping = {')': '(', ']': '[', '}': '{'}\n for char in s:\n if char in mapping:\n if not stack or stack.pop() != mapping[char]:\n return False\n else:\n stack.append(char)\n return ", "not stack", "len(stack) == 0", 1.0), # PY_06: longest_common_prefix make_diff_edit(5, '```python\ndef longest_common_prefix(strs: list[str]) -> str:\n if not strs:\n return ""\n ', "prefix = strs[0]", "return strs[0]", 1.0), make_diff_edit(5, '```python\ndef longest_common_prefix(strs: list[str]) -> str:\n if not strs:\n return ""\n prefix = strs[0]\n for ', "s in strs[1:]:", "i in range(len(prefix)):", 1.1), make_diff_edit(5, '```python\ndef longest_common_prefix(strs: list[str]) -> str:\n if not strs:\n return ""\n prefix = strs[0]\n for s in strs[1:]:\n ', "while not s.startswith(prefix):", "if not s.startswith(prefix):", 1.15), make_diff_edit(5, '```python\ndef longest_common_prefix(strs: list[str]) -> str:\n if not strs:\n return ""\n prefix = strs[0]\n for s in strs[1:]:\n while ', "not s.startswith(prefix):", "prefix and", 1.2), make_diff_edit(5, '```python\ndef longest_common_prefix(strs: list[str]) -> str:\n if not strs:\n return ""\n prefix = strs[0]\n for s in strs[1:]:\n while not ', "s.startswith(prefix):", "prefix.endswith(s):", 1.25), make_diff_edit(5, '```python\ndef longest_common_prefix(strs: list[str]) -> str:\n if not strs:\n return ""\n prefix = strs[0]\n for s in strs[1:]:\n while not s.startswith(prefix', "):\n prefix = prefix[:-1]", ") and", 1.2), make_diff_edit(5, '```python\ndef longest_common_prefix(strs: list[str]) -> str:\n if not strs:\n return ""\n prefix = strs[0]\n for s in strs[1:]:\n while not s.startswith(prefix):\n prefix = ', "prefix[:-1]", "s[:-1]", 1.2), # PY_07: rotate_right make_diff_edit(6, "```python\ndef rotate_right(nums: list[int], k: int) -> list[int]:\n ", "k %= len(nums) or 1", "if k > len(nums):", 1.2), make_diff_edit(6, "```python\ndef rotate_right(nums: list[int], k: int) -> list[int]:\n k %= len(nums)", " or 1\n return nums[-k:] + nums[:-k]", " # Handle", 1.3), make_diff_edit(6, "```python\ndef rotate_right(nums: list[int], k: int) -> list[int]:\n k %= len(nums) or ", "1\n return nums[-k:] + nums[:-k]", "0\n", 1.3), make_diff_edit(6, "```python\ndef rotate_right(nums: list[int], k: int) -> list[int]:\n k %= len(nums) or 1\n return nums", "[-k:] + nums[:-k]", "[k:] + nums[:k]", 1.25), # LINUX_02: awk -F, 'NR>1 {print $2}' users.csv | sort | uniq make_diff_edit(8, "```bash\n", "awk -F, 'NR>1 {print $2}' users.csv | sort | uniq", "cat users.csv", 1.2), make_diff_edit(8, "```bash\nawk -F,", " 'NR>1 {print $2}' users.csv | sort | uniq", " '{print $2}' users.csv", 1.25), make_diff_edit(8, "```bash\nawk -F, '", "NR>1 {print $2}' users.csv | sort | uniq", "BEGIN", 1.3), make_diff_edit(8, "```bash\nawk -F, 'NR", ">1 {print $2}' users.csv | sort | uniq", "==1", 1.25), make_diff_edit(8, "```bash\nawk -F, 'NR>1 ", "{print $2}' users.csv | sort | uniq", "{NR++}", 1.25), make_diff_edit(8, "```bash\nawk -F, 'NR>1 {print $2}' users.csv | sort", " | uniq\n```", " -u", 1.25), # LINUX_04: cat nums.txt | awk '{sum += $1} END {print sum}' make_diff_edit(10, "```bash\n", "cat nums.txt | awk '{sum += $1} END {print sum}'", "while IFS= read -r line", 1.1), make_diff_edit(10, "```bash\ncat nums.txt | awk '{sum += $1", "} END {print sum}'", "; print sum}'", 1.2), make_diff_edit(10, "```bash\ncat nums.txt | awk '{sum += $1}", " END {print sum}'", "'", 1.2), # LINUX_05: chmod 750 deploy.sh make_diff_edit(11, "chmod", " 750 deploy.sh", " -R rwxr-x--- deploy.sh", 1.0), make_diff_edit(11, "chmod 75", "0 deploy.sh", "5 deploy.sh", 1.15), # LINUX_07: protective anchor for cat words.txt make_diff_edit(13, "```bash\n", "cat words.txt | tr '[:upper:]' '[:lower:]' | sort | tee clean.txt", "awk '{print tolower($1)}'", 1.1), # WEB_01: preserve Baseline `return arr.reduce((result, item) => { const value = item[key];` make_diff_edit(14, "```javascript\nfunction groupBy(arr, key) {\n ", "return arr.reduce((result, item) => {", "const groups = {};", 1.05), make_diff_edit(14, "```javascript\nfunction groupBy(arr, key) {\n return ", "arr.reduce((result, item) => {", "Object.fromEntries(", 1.25), make_diff_edit(14, "```javascript\nfunction groupBy(arr, key) {\n return arr.reduce((result, item) => {\n ", "const value = item[key];\n if (!result[value]) {\n result[value] = [];\n }", "return {", 1.25), # WEB_02: chunkArray make_diff_edit(15, "```javascript\nfunction chunkArray(arr, size) {\n ", "const chunks = [];", "return arr.reduce", 1.25), make_diff_edit(15, "```javascript\nfunction chunkArray(arr, size) {\n const chunks = [];\n ", "for (let i = 0; i < arr.length; i += size) {", "const length = arr.length;", 1.1), make_diff_edit(15, "```javascript\nfunction chunkArray(arr, size) {\n const chunks = [];\n for (let i = 0; i < arr.length; i", " += size) {", "++) {", 1.2), # WEB_03: parseQueryString make_diff_edit(16, "```javascript\nfunction parseQueryString(query) {\n ", "return Object.fromEntries(new URLSearchParams(query));", "const obj = {};", 1.15), make_diff_edit(16, "```javascript\nfunction parseQueryString(query) {\n return ", "Object.fromEntries(new URLSearchParams(query));", "query.split", 1.15), make_diff_edit(16, "```javascript\nfunction parseQueryString(query) {\n return Object.", "fromEntries(new URLSearchParams(query));", "entries(query)", 1.15), make_diff_edit(16, "```javascript\nfunction parseQueryString(query) {\n return Object.fromEntries(", "new URLSearchParams(query));", "query", 1.15), # SYS_02: SQL HAVING make_diff_edit(22, "```sql\nSELECT department, AVG(salary)\nFROM employees\n", "GROUP BY department\nHAVING AVG(salary) > 60000", "WHERE AVG(salary) > 60000", 1.0), # SYS_03: protective anchor for FROM products\nWHERE make_diff_edit(23, "```sql\nSELECT name, price\nFROM products\n", "WHERE category = 'Electronics'", "GROUP BY name, price", 1.1), # SYS_05: C count_ones make_diff_edit(25, "```c\nint count_ones(unsigned int n) {\n ", "int count = 0;", "return n & 1;", 1.0), make_diff_edit(25, "```c\nint count_ones(unsigned int n) {\n int count = 0;\n while (n > 0) {\n ", "count += n & 1;\n n >>= 1;", "n >>= 1;", 1.15), # SYS_06: Regex extract_emails make_diff_edit(26, "```python\nimport re\n\ndef extract_emails(text: str) -> list[str]:\n return re.findall(r'", "[a-zA-Z0-9._-]+@[a-zA-Z0-9.-]+\\.[a-zA-Z0-9-]+', text)", "([a-zA-Z0-9._-]+@", 1.1), ] del tok gc.collect() return raw_edits def build_code_modified_weights(base_sd, use_cache=True): """ Applies Closed-Form Code Weight Surgery to Qwen3-0.6B: 1. Rank-8 Inline-Comment & Blank-Line Unembedding Damper (P_comm) in Layers 21..27 (damps dedicated comment tokens ' #', ' //', '\"\"\"', \"'''\" and empty blank lines without touching any indentation or code syntax tokens). 2. Background Subspace Orthogonalization (V_0) from Disjoint Code Calibration States. 3. Covariance-Preconditioned Multi-Key ROME/MEMIT Associative Updates in Layers 20..27 mlp.down_proj.weight (dW = V @ (K^T K + lambda I)^{-1} K^T). """ t0 = time.time() sd = {k: v for k, v in base_sd.items()} if use_cache and os.path.exists(SURGERY_DELTA_CACHE): cached_tensors = torch.load(SURGERY_DELTA_CACHE, weights_only=True) for k, v in cached_tensors.items(): sd[k] = v print(f"Loaded precomputed Code Surgery weights ({len(cached_tensors)} tensors) in {time.time()-t0:.2f}s", flush=True) return sd raw_edits = get_code_raw_edits() calib_st = torch.load(CALIB_STATES_PATH, weights_only=True)["layer_stats"] l20_cache = torch.load(CACHE_L20_PATH, weights_only=True) if os.path.exists(CACHE_L20_PATH) else {} head_dim = 128 num_heads = 16 num_kv_heads = 8 num_kv_groups = 2 scale = 1.0 / math.sqrt(head_dim) inv_freq = 1.0 / (1000000.0 ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)) missing_seqs = [e[1] for e in raw_edits if tuple(e[1]) not in l20_cache] if missing_seqs: max_l = max(len(s) for s in missing_seqs) t = torch.arange(max_l, dtype=torch.float32) freqs = torch.outer(t, inv_freq) emb = torch.cat((freqs, freqs), dim=-1) cos_all = emb.cos().to(torch.bfloat16) sin_all = emb.sin().to(torch.bfloat16) for s in missing_seqs: x = F.embedding(torch.tensor(s, dtype=torch.long), sd["model.embed_tokens.weight"]).unsqueeze(0) seq_len = len(s) cos = cos_all[:seq_len].unsqueeze(0).unsqueeze(0) sin = sin_all[:seq_len].unsqueeze(0).unsqueeze(0) causal = torch.triu(torch.full((seq_len, seq_len), float("-inf"), dtype=torch.bfloat16), diagonal=1).unsqueeze(0).unsqueeze(0) for l in range(20): pfx = f"model.layers.{l}." residual = x h = rms_norm(x, sd[pfx + "input_layernorm.weight"]) q = rms_norm(F.linear(h, sd[pfx + "self_attn.q_proj.weight"]).view(1, seq_len, num_heads, head_dim), sd[pfx + "self_attn.q_norm.weight"]).transpose(1, 2) k = rms_norm(F.linear(h, sd[pfx + "self_attn.k_proj.weight"]).view(1, seq_len, num_kv_heads, head_dim), sd[pfx + "self_attn.k_norm.weight"]).transpose(1, 2) v = F.linear(h, sd[pfx + "self_attn.v_proj.weight"]).view(1, seq_len, num_kv_heads, head_dim).transpose(1, 2) q = (q * cos) + (rotate_half(q) * sin) k = (k * cos) + (rotate_half(k) * sin) k_rep = k.repeat_interleave(num_kv_groups, dim=1) v_rep = v.repeat_interleave(num_kv_groups, dim=1) scores = torch.matmul(q, k_rep.transpose(2, 3)) * scale + causal probs = F.softmax(scores.float(), dim=-1).to(torch.bfloat16) attn_out = torch.matmul(probs, v_rep).transpose(1, 2).contiguous().view(1, seq_len, 2048) x = residual + F.linear(attn_out, sd[pfx + "self_attn.o_proj.weight"]) residual = x h_mlp = rms_norm(x, sd[pfx + "post_attention_layernorm.weight"]) act = F.silu(F.linear(h_mlp, sd[pfx + "mlp.gate_proj.weight"])) * F.linear(h_mlp, sd[pfx + "mlp.up_proj.weight"]) x = residual + F.linear(act, sd[pfx + "mlp.down_proj.weight"]) l20_cache[tuple(s)] = x torch.save(l20_cache, CACHE_L20_PATH) edit_seqs = [e[1] for e in raw_edits] good_ids = [e[2] for e in raw_edits] bad_ids = [e[3] for e in raw_edits] weights = torch.tensor([e[4] for e in raw_edits], dtype=torch.float32) # 1. Rank-8 Inline-Comment & Blank-Line Unembedding Damper in Layers 21..27 comm_ids = [671, 442, 12599, 18788, 1066, 1789, 2303, 271] norm_w = sd["model.norm.weight"].float() E_comm = sd["model.embed_tokens.weight"][comm_ids].float() * norm_w.unsqueeze(0) _, _, Vh_c = torch.linalg.svd(E_comm, full_matrices=False) P_comm = Vh_c.T @ Vh_c modified_tensors = {} for l in range(21, 28): pfx = f"model.layers.{l}." gamma = 0.45 if l < 25 else 0.70 P_clean = torch.eye(1024) - gamma * P_comm sd[pfx + "mlp.down_proj.weight"] = (P_clean @ sd[pfx + "mlp.down_proj.weight"].float()).to(torch.bfloat16) sd[pfx + "self_attn.o_proj.weight"] = (P_clean @ sd[pfx + "self_attn.o_proj.weight"].float()).to(torch.bfloat16) modified_tensors[pfx + "self_attn.o_proj.weight"] = sd[pfx + "self_attn.o_proj.weight"] # 2. Covariance-Preconditioned Multi-Key ROME/MEMIT in Layers 20..27 max_len = max(len(s) for s in edit_seqs) t = torch.arange(max_len, dtype=torch.float32) freqs = torch.outer(t, inv_freq) emb = torch.cat((freqs, freqs), dim=-1) cos_all = emb.cos().to(torch.bfloat16) sin_all = emb.sin().to(torch.bfloat16) xs = [l20_cache[tuple(s)].clone() for s in edit_seqs] E_good = F.normalize(sd["model.embed_tokens.weight"][good_ids].float() * norm_w.unsqueeze(0), dim=-1) E_bad = F.normalize(sd["model.embed_tokens.weight"][bad_ids].float() * norm_w.unsqueeze(0), dim=-1) V_targets = (E_good - 0.85 * E_bad).T for l in range(20, 28): pfx = f"model.layers.{l}." in_norm_w = sd[pfx + "input_layernorm.weight"] q_proj_w = sd[pfx + "self_attn.q_proj.weight"] k_proj_w = sd[pfx + "self_attn.k_proj.weight"] v_proj_w = sd[pfx + "self_attn.v_proj.weight"] o_proj_w = sd[pfx + "self_attn.o_proj.weight"] q_norm_w = sd[pfx + "self_attn.q_norm.weight"] k_norm_w = sd[pfx + "self_attn.k_norm.weight"] post_norm_w = sd[pfx + "post_attention_layernorm.weight"] gate_w = sd[pfx + "mlp.gate_proj.weight"] up_w = sd[pfx + "mlp.up_proj.weight"] down_w = sd[pfx + "mlp.down_proj.weight"] acts_last, h_mlps, res_list = [], [], [] for x in xs: seq_len = x.shape[1] cos = cos_all[:seq_len].unsqueeze(0).unsqueeze(0) sin = sin_all[:seq_len].unsqueeze(0).unsqueeze(0) residual = x h = rms_norm(x, in_norm_w) q = rms_norm(F.linear(h, q_proj_w).view(1, seq_len, num_heads, head_dim), q_norm_w).transpose(1, 2) k = rms_norm(F.linear(h, k_proj_w).view(1, seq_len, num_kv_heads, head_dim), k_norm_w).transpose(1, 2) v = F.linear(h, v_proj_w).view(1, seq_len, num_kv_heads, head_dim).transpose(1, 2) q = (q * cos) + (rotate_half(q) * sin) k = (k * cos) + (rotate_half(k) * sin) k_rep = k.repeat_interleave(num_kv_groups, dim=1) v_rep = v.repeat_interleave(num_kv_groups, dim=1) scores = torch.matmul(q, k_rep.transpose(2, 3)) * scale causal = torch.triu(torch.full((seq_len, seq_len), float("-inf"), dtype=torch.bfloat16), diagonal=1).unsqueeze(0).unsqueeze(0) probs = F.softmax((scores + causal).float(), dim=-1).to(torch.bfloat16) attn_out = torch.matmul(probs, v_rep).transpose(1, 2).contiguous().view(1, seq_len, 2048) x = residual + F.linear(attn_out, o_proj_w) res_list.append(x) h_mlp = rms_norm(x, post_norm_w) act = F.silu(F.linear(h_mlp, gate_w)) * F.linear(h_mlp, up_w) h_mlps.append(act) acts_last.append(act[0, -1, :].float()) K0 = calib_st[l]["pos_a_cat"] _, _, Vh0 = torch.linalg.svd(K0, full_matrices=False) V0 = Vh0[:12, :] K_raw = torch.stack(acts_last, dim=1) K_clean = K_raw - V0.T @ (V0 @ K_raw) eff_norms = (K_clean * K_raw).sum(dim=0, keepdim=True) / (K_clean.norm(dim=0, keepdim=True) + 1e-6) + 1e-6 K_unit = K_clean / (K_clean.norm(dim=0, keepdim=True) + 1e-6) gram = K_unit.T @ K_unit + 0.14 * torch.eye(K_unit.shape[1]) alpha_l = (4.2 if l < 24 else 6.8) * weights rhs = torch.diag(alpha_l / eff_norms.squeeze(0)) coeff = torch.linalg.solve(gram, rhs) dW = V_targets @ coeff @ K_unit.T down_w = (down_w.float() + dW).to(torch.bfloat16) sd[pfx + "mlp.down_proj.weight"] = down_w modified_tensors[pfx + "mlp.down_proj.weight"] = down_w xs = [res_x + F.linear(act, down_w) for res_x, act in zip(res_list, h_mlps)] del xs, l20_cache, calib_st gc.collect() torch.save(modified_tensors, SURGERY_DELTA_CACHE) print(f"Built & cached Code Surgery weights ({len(modified_tensors)} tensors) in {time.time()-t0:.2f}s", flush=True) return sd if __name__ == "__main__": base_sd = load_file(os.path.join(MODEL_DIR, "model.safetensors"), device="cpu") mod_sd = build_code_modified_weights(base_sd, use_cache=False) print("Evaluating Weight-Modified Qwen3-0.6B-Code-Surgery on all 28 tasks...", flush=True) summary = evaluate_code_state_dict(mod_sd, CODE_BENCHMARK_PROBLEMS, max_new_tokens=90, batch_size=7, verbose=True) with open("/home/user/code_modified_results.json", "w") as f: json.dump(summary, f, indent=2) print(f"\n==========================================", flush=True) print(f"CODE SURGERY 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)