"""BaichuanCharRM numerical sensitivity (CharacterEval matrix V07a), run by the operator of the reward-model endpoint.

Why: the anchor (V07) reproduced the paper's Table 4 within one step, not within 0.005. All 150 pairwise orders hold and
the cell differences have mean +0.0003, but 17 of 75 cells differ by 0.005-0.012. Read as noise on single scores, the
differences correspond to a root-mean-square of 0.121 per score (the RMS over the 60 metric cells of diff * sqrt(n)).
This check measures how far single scores move when only the numerics change - the weights, the inputs and the official
code stay the same - so that we can tell whether ordinary numerical variation (another GPU, another torch or cuBLAS
version) is large enough to explain the anchor.

Inputs: 1,200 reward-model inputs, the first 20 records of each of the 12 metrics in each of the five frozen result files
of github.com/morecry/CharacterEval at c3d44a6 (results/generation_trans_*.jsonl), built by the official format_input
(run_char_rm.py, reproduced below). Each input is scored three times, batch size 1, the last 4,096 tokens kept:

  bf16          run_char_rm.py as published: bfloat16 weights, PyTorch defaults
  bf16-strict   the same, with torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
                (cuBLAS accumulates bfloat16 products differently; a kernel-level change of the kind another GPU or
                library version makes)
  fp32          float32 weights, TF32 disabled (the precision ceiling)

Criteria, fixed on 2026-10-28 before any result. For each of bf16-strict and fp32, let r be the root-mean-square of
(variant - bf16) over the 1,200 inputs:
  sufficient   r >= 0.06 (half the anchor's 0.121): numerical variation is large enough to explain the anchor
  ruled out    r <  0.02 (it would move a 308-record cell by about 0.001): numerical variation cannot explain it
  partial      otherwise
The bf16 scores are also compared, input by input, with the platform's scores of the same inputs (on our side).

    python3 charactereval_rm_numerics_check.py --model-dir /path/to/BaichuanCharRM --repo /path/to/CharacterEval \\
        --device cuda:0 [--out charactereval_rm_numerics_check.json]

Needs torch, transformers and one GPU with ~60 GB free (fp32). Takes about 30-45 minutes on one A800. The output holds
scores, token counts and the environment only - no text from the dataset - and should be sent back unchanged.
Reads files; changes nothing.
"""
import argparse
import datetime
import importlib.util
import json
import math
import os
import platform
import socket
import subprocess
import sys
import time

FROZEN = ["chatglm3_6b", "baichuan_7b", "xverse_7b", "internlm_7b", "qwen_7b"]
PER_METRIC = 20
MAX_SEQ_LENGTH = 4096  # run_char_rm.py
COMMIT = "c3d44a6"
ANCHOR_RMS, SUFFICIENT, RULED_OUT = 0.121, 0.06, 0.02


def format_input(example, character_profile):
    # run_char_rm.py at c3d44a6, lines 15-18 (character_profile passed in instead of read as a global)
    input_text = "<RoleInfo>\n\n" \
        + str(character_profile[example['role']]) + "\n\n<Context>\n\n" + example['context'] + "\n\n<Response>\n\n" + example['model_output'] + "\n\n<Dimension>\n\n" + example["metric_zh"]
    return input_text


def select(repo):
    with open(os.path.join(repo, "data", "character_profiles.json"), "r") as f:
        profiles = json.load(f)
    items = []
    for name in FROZEN:
        with open(os.path.join(repo, "results", f"generation_trans_{name}.jsonl"), mode="r") as f:
            records = json.load(f)
        taken = {}
        for i, r in enumerate(records):
            if taken.get(r["metric_en"], 0) < PER_METRIC:
                taken[r["metric_en"]] = taken.get(r["metric_en"], 0) + 1
                items.append({"file": name, "row": i, "metric_en": r["metric_en"], "text": format_input(r, profiles)})
    return items


def load_classes(model_dir):
    def load(name):
        spec = importlib.util.spec_from_file_location(f"BaichuanCharRM.{name}", os.path.join(model_dir, f"{name}.py"))
        mod = importlib.util.module_from_spec(spec)
        sys.modules[spec.name] = mod
        spec.loader.exec_module(mod)
        return mod

    pkg = type(sys)("BaichuanCharRM")
    pkg.__path__ = [model_dir]
    sys.modules["BaichuanCharRM"] = pkg
    load("configuration_baichuan")
    return load("modeling_baichuan").BaichuanCharRM, load("tokenization_baichuan").BaichuanTokenizer


def score_all(model_dir, items, device, variant):
    import torch
    BaichuanCharRM, BaichuanTokenizer = load_classes(model_dir)
    torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = variant != "bf16-strict"
    torch.backends.cuda.matmul.allow_tf32 = False
    torch.backends.cudnn.allow_tf32 = False
    dtype = torch.float32 if variant == "fp32" else torch.bfloat16
    # --- run_char_rm.py, lines 23-39 (model path, device, dtype and input list as above) ---
    tokenizer = BaichuanTokenizer.from_pretrained(model_dir)
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "left"
    base_model = BaichuanCharRM.from_pretrained(model_dir, torch_dtype=dtype).to(device)
    scores, tokens, started = [], [], time.time()
    for k, item in enumerate(items):
        input_ids = tokenizer.encode(text=item["text"], add_special_tokens=False) + [tokenizer.eos_token_id]
        tokens.append(len(input_ids))
        if len(input_ids) > MAX_SEQ_LENGTH:
            input_ids = input_ids[-MAX_SEQ_LENGTH:]
        input_ids = torch.tensor(input_ids).unsqueeze(0).to(device)
        with torch.no_grad():
            score = base_model(input_ids=input_ids)[1].item() * 4 + 1
        scores.append(score)
        if (k + 1) % 100 == 0:
            print(f"  {variant}: {k + 1}/{len(items)} ({time.time() - started:.0f} s)", flush=True)
    del base_model
    torch.cuda.empty_cache()
    return scores, tokens, round(time.time() - started, 1)


def compare(base, other):
    finite = [(a, b) for a, b in zip(base, other) if math.isfinite(a) and math.isfinite(b)]
    d = [b - a for a, b in finite]
    rms = math.sqrt(sum(x * x for x in d) / len(d))
    verdict = "sufficient" if rms >= SUFFICIENT else "ruled out" if rms < RULED_OUT else "partial"
    return {"n": len(d), "non_finite": len(base) - len(finite), "rms": rms, "max_abs": max(map(abs, d)),
            "mean_signed": sum(d) / len(d), "changed": sum(x != 0 for x in d), "verdict": verdict}


def environment(device):
    import torch
    env = {"torch": torch.__version__, "cuda": torch.version.cuda, "device": torch.cuda.get_device_name(device)}
    try:
        import transformers
        env["transformers"] = transformers.__version__
    except ImportError:
        pass
    return env


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model-dir", required=True)
    ap.add_argument("--repo", required=True, help="github.com/morecry/CharacterEval checked out at c3d44a6")
    ap.add_argument("--device", default="cuda:0")
    ap.add_argument("--out", default="charactereval_rm_numerics_check.json")
    args = ap.parse_args()
    try:
        head = subprocess.run(["git", "-C", args.repo, "rev-parse", "HEAD"], capture_output=True, text=True).stdout.strip()
    except OSError:
        head = ""
    items = select(args.repo)
    print(f"{len(items)} inputs from {len(FROZEN)} files; repo HEAD {head[:7] or 'unknown'}", flush=True)
    result = {"script": os.path.basename(__file__), "time_utc": datetime.datetime.now(datetime.timezone.utc).isoformat(timespec="seconds"),
              "host": socket.gethostname(), "python": platform.python_version(), "repo_head": head,
              "criteria": {"sufficient": f"rms >= {SUFFICIENT}", "ruled_out": f"rms < {RULED_OUT}", "anchor_rms": ANCHOR_RMS},
              "environment": environment(args.device),
              "inputs": [{"file": x["file"], "row": x["row"], "metric_en": x["metric_en"]} for x in items], "variants": {}}
    for variant in ["bf16", "bf16-strict", "fp32"]:
        scores, tokens, seconds = score_all(args.model_dir, items, args.device, variant)
        result["variants"][variant] = {"scores": scores, "seconds": seconds}
        result["tokens"] = tokens
    base = result["variants"]["bf16"]["scores"]
    result["comparison"] = {v: compare(base, result["variants"][v]["scores"]) for v in ["bf16-strict", "fp32"]}
    with open(args.out, "w") as f:
        json.dump(result, f, indent=1)
    for v, c in result["comparison"].items():
        print(f"{v} vs bf16: rms {c['rms']:.4f}, max |diff| {c['max_abs']:.4f}, changed {c['changed']}/{c['n']}, "
              f"non-finite {c['non_finite']} -> {c['verdict']}")
    print(f"written: {args.out}")


if __name__ == "__main__":
    main()