Blame for eval_engine.py

import os
import re
import gc
import time
import json
import ctypes
import subprocess
import sys
import torch
import torch.nn.functional as F
from safetensors.torch import load_file
from benchmark_problems import BENCHMARK_PROBLEMS

torch.set_num_threads(2)
MODEL_DIR = "/home/user/.cache/qwen3-0.6b"
TOK_CACHE_PATH = "/home/user/.cache/benchmark_tokens.pt"


def ensure_tokenized_cache():
    if os.path.exists(TOK_CACHE_PATH):
        return
    # Run tokenization in a separate short-lived subprocess so Rust tokenizer heap never touches main process
    cmd = [
        sys.executable,
        "-c",
        (
            "import torch; from transformers import AutoTokenizer; "
            "from benchmark_problems import BENCHMARK_PROBLEMS; "
            f"tok = AutoTokenizer.from_pretrained('{MODEL_DIR}'); "
            "pids = [tok.encode(tok.apply_chat_template([{'role':'user','content':p['question']+'\\nWrite only the brief calculation (no markdown headers) and put the final number in \\\\boxed{}.'}], tokenize=False, add_generation_prompt=True, enable_thinking=False), add_special_tokens=False) for p in BENCHMARK_PROBLEMS]; "
            f"torch.save({{'prompt_ids_list': pids, 'eos_id': tok.eos_token_id, 'pad_id': tok.pad_token_id}}, '{TOK_CACHE_PATH}')"
        ),
    ]
    subprocess.run(cmd, check=True)


def decode_in_subprocess(all_gen_ids):
    tmp_in = "/home/user/.cache/tmp_gen_ids.pt"
    tmp_out = "/home/user/.cache/tmp_decoded.json"
    torch.save(all_gen_ids, tmp_in)
    cmd = [
        sys.executable,
        "-c",
        (
            "import json, torch; from transformers import AutoTokenizer; "
            f"tok = AutoTokenizer.from_pretrained('{MODEL_DIR}'); "
            f"gids = torch.load('{tmp_in}', weights_only=True); "
            "txts = [tok.decode(g, skip_special_tokens=True).strip() for g in gids]; "
            f"open('{tmp_out}', 'w').write(json.dumps(txts))"
        ),
    ]
    subprocess.run(cmd, check=True)
    with open(tmp_out, "r") as f:
        return json.load(f)


def extract_answer(text):
    boxed = re.findall(r"\\boxed\{([^}]*)\}", text)
    if boxed:
        cand = boxed[-1].strip().replace(",", "").replace("$", "").replace(" ", "")
        m = re.findall(r"-?\d+(?:\.\d+)?", cand)
        if m:
            return m[-1]
    boxed_open = re.findall(r"\\boxed\{(-?\d+(?:\.\d+)?)", text)
    if boxed_open:
        return boxed_open[-1]
    nums = re.findall(r"-?\d+(?:,\d{3})*(?:\.\d+)?", text)
    if nums:
        return nums[-1].replace(",", "")
    return None


def rms_norm(x, w, eps=1e-6):
    dtype = x.dtype
    x_fp32 = x.to(torch.float32)
    var = x_fp32.pow(2).mean(-1, keepdim=True)
    return w * (x_fp32 * torch.rsqrt(var + eps)).to(dtype)


def rotate_half(x):
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)


class Qwen3FastRunner:
    """
    Pure PyTorch implementation of Qwen3ForCausalLM (exact bit-for-bit parity with HuggingFace transformers)
    that avoids loading the 160MB transformers library in the inference process.
    """
    def __init__(self, sd):
        self.num_layers = 28
        self.num_heads = 16
        self.num_kv_heads = 8
        self.head_dim = 128
        self.inv_freq = 1.0 / (1000000.0 ** (torch.arange(0, 128, 2, dtype=torch.float32) / 128.0))
        self.embed = sd["model.embed_tokens.weight"]
        if sd.get("_untied_lm_head", False):
            self.lm_head = sd["lm_head.weight"]
        else:
            self.lm_head = self.embed
        self.norm = sd["model.norm.weight"]
        self.layers = []
        for i in range(self.num_layers):
            p = f"model.layers.{i}."
            self.layers.append({
                "in_norm": sd[p + "input_layernorm.weight"],
                "q_proj": sd[p + "self_attn.q_proj.weight"],
                "k_proj": sd[p + "self_attn.k_proj.weight"],
                "v_proj": sd[p + "self_attn.v_proj.weight"],
                "o_proj": sd[p + "self_attn.o_proj.weight"],
                "q_norm": sd[p + "self_attn.q_norm.weight"],
                "k_norm": sd[p + "self_attn.k_norm.weight"],
                "post_norm": sd[p + "post_attention_layernorm.weight"],
                "gate_proj": sd[p + "mlp.gate_proj.weight"],
                "up_proj": sd[p + "mlp.up_proj.weight"],
                "down_proj": sd[p + "mlp.down_proj.weight"],
            })

    @torch.no_grad()
    def forward_step(self, input_ids, position_ids, kv_cache, attn_mask):
        B, S = input_ids.shape
        h = F.embedding(input_ids, self.embed)
        freqs = position_ids[..., None].float() * self.inv_freq
        emb = torch.cat((freqs, freqs), dim=-1)
        cos = emb.cos().to(h.dtype).unsqueeze(1)
        sin = emb.sin().to(h.dtype).unsqueeze(1)

        for i, lyr in enumerate(self.layers):
            residual = h
            x = rms_norm(h, lyr["in_norm"])
            q = F.linear(x, lyr["q_proj"]).view(B, S, self.num_heads, self.head_dim)
            k = F.linear(x, lyr["k_proj"]).view(B, S, self.num_kv_heads, self.head_dim)
            v = F.linear(x, lyr["v_proj"]).view(B, S, self.num_kv_heads, self.head_dim).transpose(1, 2)

            q = rms_norm(q, lyr["q_norm"]).transpose(1, 2)
            k = rms_norm(k, lyr["k_norm"]).transpose(1, 2)

            q = (q * cos) + (rotate_half(q) * sin)
            k = (k * cos) + (rotate_half(k) * sin)

            if kv_cache[i] is not None:
                pk, pv = kv_cache[i]
                k = torch.cat([pk, k], dim=2)
                v = torch.cat([pv, v], dim=2)
            kv_cache[i] = (k, v)

            attn_out = F.scaled_dot_product_attention(
                q, k, v, attn_mask=attn_mask, enable_gqa=True
            )
            attn_out = attn_out.transpose(1, 2).contiguous().view(B, S, -1)
            h = residual + F.linear(attn_out, lyr["o_proj"])

            residual = h
            x = rms_norm(h, lyr["post_norm"])
            mlp_out = F.silu(F.linear(x, lyr["gate_proj"])) * F.linear(x, lyr["up_proj"])
            h = residual + F.linear(mlp_out, lyr["down_proj"])

        h_last = rms_norm(h[:, -1:, :], self.norm)
        logits = F.linear(h_last, self.lm_head)
        return logits[:, 0, :]

    @torch.no_grad()
    def generate_batch(self, batch_prompt_ids, eos_id, pad_id, max_new_tokens=95):
        B = len(batch_prompt_ids)
        max_len = max(len(x) for x in batch_prompt_ids)
        input_ids = torch.full((B, max_len), pad_id, dtype=torch.long)
        valid_mask = torch.zeros((B, max_len), dtype=torch.bool)
        for i, ids in enumerate(batch_prompt_ids):
            L = len(ids)
            input_ids[i, max_len - L :] = torch.tensor(ids, dtype=torch.long)
            valid_mask[i, max_len - L :] = True

        pos_ids = (valid_mask.long().cumsum(dim=-1) - 1).clamp(min=0)
        causal = torch.tril(torch.ones((max_len, max_len), dtype=torch.bool))
        attn_mask = causal.unsqueeze(0).unsqueeze(0) & valid_mask[:, None, None, :]

        kv_cache = [None] * self.num_layers
        active_indices = list(range(B))
        gen_tokens = [[] for _ in range(B)]

        cur_ids = input_ids
        cur_pos = pos_ids
        cur_attn_mask = attn_mask

        for step in range(max_new_tokens):
            if not active_indices:
                break
            logits = self.forward_step(cur_ids, cur_pos, kv_cache, cur_attn_mask)
            next_tokens = logits.argmax(dim=-1)

            still_active_local = []
            new_active_indices = []
            for local_i, orig_i in enumerate(active_indices):
                tok = int(next_tokens[local_i].item())
                if tok == eos_id:
                    continue
                gen_tokens[orig_i].append(tok)
                # Early stop once "}" after "\boxed{" is generated and newline/period follows
                # Wait: NO early stop heuristics! Let the model naturally emit eos_id so decoding is 100% standard greedy!
                still_active_local.append(local_i)
                new_active_indices.append(orig_i)

            if not still_active_local:
                break

            if len(still_active_local) < len(active_indices):
                idx = torch.tensor(still_active_local, dtype=torch.long)
                next_tokens = next_tokens.index_select(0, idx)
                valid_mask = valid_mask.index_select(0, idx)
                cur_pos = cur_pos.index_select(0, idx)
                for l_idx in range(self.num_layers):
                    pk, pv = kv_cache[l_idx]
                    kv_cache[l_idx] = (pk.index_select(0, idx), pv.index_select(0, idx))
                active_indices = new_active_indices

            cur_ids = next_tokens.unsqueeze(1)
            cur_pos = cur_pos[:, -1:] + 1
            valid_mask = torch.cat([valid_mask, torch.ones((valid_mask.shape[0], 1), dtype=torch.bool)], dim=1)
            cur_attn_mask = valid_mask[:, None, None, :]

        return gen_tokens


def evaluate_state_dict(sd, problems=BENCHMARK_PROBLEMS, max_new_tokens=95, batch_size=6, verbose=True):
    ensure_tokenized_cache()
    tok_data = torch.load(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"]

    # Sort indices by prompt length to minimize padding waste
    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} probs, avg {avg_t:.1f} toks) in {time.time()-tb0:.1f}s", flush=True)
    elapsed = time.time() - t0

    # Restore original problem order
    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 = []
    correct = 0
    for prob, resp, g_ids in zip(problems, responses, all_gen_ids):
        pred = extract_answer(resp)
        is_ok = (pred == str(prob["answer"]))
        if is_ok:
            correct += 1
        results.append({
            "id": prob["id"],
            "category": prob["category"],
            "question": prob["question"],
            "expected": str(prob["answer"]),
            "predicted": pred,
            "correct": is_ok,
            "num_tokens": len(g_ids),
            "response": resp,
        })

    summary = {
        "total": len(problems),
        "correct": correct,
        "accuracy": round(100.0 * correct / len(problems), 2),
        "elapsed_sec": round(elapsed, 2),
        "results": results,
    }
    return summary


if __name__ == "__main__":
    print("Evaluating Baseline Qwen3-0.6B on 30-problem Math Benchmark...", flush=True)
    sd = load_file(os.path.join(MODEL_DIR, "model.safetensors"), device="cpu")
    summary = evaluate_state_dict(sd, BENCHMARK_PROBLEMS, max_new_tokens=95, batch_size=6, verbose=True)
    with open("/home/user/baseline_results.json", "w") as f:
        json.dump(summary, f, indent=2)
    print(f"\n==========================================", flush=True)
    print(f"BASELINE QWEN3-0.6B: {summary['correct']}/{summary['total']} ({summary['accuracy']}%) in {summary['elapsed_sec']}s", flush=True)
    print(f"==========================================", flush=True)
    for r in summary["results"]:
        status = "OK" if r["correct"] else "FAIL"
        print(f"[{status:4s}] {r['id']:8s} ({r['category']:24s}) | Exp: {r['expected']:>6s} | Pred: {str(r['predicted']):>6s} | Toks: {r['num_tokens']:3d}", flush=True)