#!/usr/bin/env python3
"""원본 Qwen3.8-27B 체크포인트에 CJK 출력 억제를 적용한다 — 텐서 1개만 바꾼다.

가중치를 배포하지 않는 이유는 간단하다. 32개 파일 55.6GB 중 **31개가 원본과 바이트
동일**하고, 달라지는 것은 `lm_head.weight` 하나다. 55GB 를 복제해 받게 하는 것보다
원본을 받아 이 스크립트를 돌리는 편이 모든 면에서 낫다 — 베이스 갱신도 따라간다.

## 무엇을 하나

`lm_head.weight` 의 마스크 대상 행을 은닉상태 평균 방향의 큰 음수 배수로 덮는다.

    W_i := -alpha * mu_h / ||mu_h||^2        (alpha = 200)

`logit_i = h·W_i` 이고 `lm_head` 에 bias 가 없다. 그래서 행을 **0 으로 만들면 안 된다** —
로짓이 0 이 될 뿐이라 다른 후보가 전부 음수인 순간 argmax 가 된다.

`mu_h` 는 가정이 아니라 **측정**한다. 아래 한국어 6문장을 forward 해 최종 은닉상태의
평균을 쓴다. 우리 측정에서 `||mu_h||^2` 는 9845.958 / 9887.059 로 두 번 독립 측정에서
0.4% 안에 재현됐다. 결과 로짓은 마스킹 -217.0 대 최대 16.5 였다.

⛔ "한글 토큰의 W 행 평균"을 mu_h 프록시로 쓰려던 시도는 **기각됐다**
(cos(mu_h, mu_W) = 0.048, 사실상 직교). forward pass 가 필요하다.

## 사용

    hf download Qwen/Qwen3.8-27B --local-dir ./Qwen3.8-27B
    python apply_mask.py --model-dir ./Qwen3.8-27B --mask mask_m2-medium.json

기본은 in-place 패치다. 원본을 남기려면 `--out-dir` 를 준다(바뀐 샤드만 그쪽에 쓰고,
나머지 파일은 직접 복사하거나 심링크한다).

`embed_tokens` 와 토크나이저는 건드리지 않는다 — CJK **입력**은 그대로 읽는다.
"""
from __future__ import annotations

import argparse
import json
import shutil
import sys
from pathlib import Path

# mu_h 측정용 한국어 문장 — 이 목록이 곧 재현 조건이다. 바꾸면 mu_h 가 달라진다.
PROBE_TEXTS = [
    "안녕하세요. 오늘 날씨가 참 좋네요.",
    "이번 분기 실적을 정리해서 보고드리겠습니다.",
    "한국의 사계절은 각각 뚜렷한 특징을 가지고 있습니다.",
    "이 문제를 해결하려면 먼저 원인을 파악해야 합니다.",
    "주말에 뭐 하고 지냈어? 나는 그냥 집에 있었어.",
    "이 기능은 왜 이렇게 동작하나요? 로그를 봐도 이해가 안 됩니다.",
]
MARGIN_MIN = 50.0   # 마스킹 로짓이 최대 로짓보다 이만큼은 아래여야 한다


def find_shard(model_dir: Path, key: str = "lm_head.weight") -> str:
    """어느 safetensors 파일에 lm_head 가 들어 있는지 인덱스에서 찾는다."""
    idx = model_dir / "model.safetensors.index.json"
    if idx.exists():
        wm = json.loads(idx.read_text())["weight_map"]
        if key not in wm:
            sys.exit(f"⛔ index 에 {key} 가 없다. tied embedding 모델일 수 있다.")
        return wm[key]
    single = model_dir / "model.safetensors"
    if single.exists():
        return single.name
    sys.exit(f"⛔ {model_dir} 에서 safetensors 를 못 찾았다.")


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--model-dir", required=True, help="원본 체크포인트 로컬 경로")
    ap.add_argument("--mask", default="mask_m2-medium.json")
    ap.add_argument("--alpha", type=float, default=200.0)
    ap.add_argument("--out-dir", default=None, help="생략하면 in-place 패치")
    ap.add_argument("--dry-run", action="store_true", help="측정·검증만 하고 쓰지 않는다")
    a = ap.parse_args()

    import torch
    from transformers import AutoTokenizer, AutoModelForCausalLM
    from safetensors.torch import load_file, save_file

    model_dir = Path(a.model_dir).resolve()
    mask = json.loads(Path(a.mask).read_text())
    ids = mask["target_token_ids"]
    print(f"[1/4] 마스크 {mask['name']} · 대상 토큰 {len(ids):,}개 · alpha={a.alpha}", flush=True)
    if mask.get("model") and mask["model"] not in str(model_dir):
        print(f"      ⚠️ 마스크는 {mask['model']} 기준이다. 토큰 id 는 토크나이저에 종속된다 —"
              f" 다른 모델이면 build_masks.py 로 다시 만들어라.", flush=True)

    # ── mu_h: 측정한다 (가정하지 않는다) ────────────────────────────────
    tok = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        model_dir, dtype=torch.bfloat16, device_map="cpu", trust_remote_code=True).eval()
    hs = []
    with torch.no_grad():
        for t in PROBE_TEXTS:
            o = model(**tok(t, return_tensors="pt"), output_hidden_states=True)
            hs.append(o.hidden_states[-1][0].float().mean(0))
    mu = torch.stack(hs).mean(0)
    mu_n2 = float((mu * mu).sum())
    repl32 = (-a.alpha / mu_n2) * mu
    print(f"[2/4] mu_h 측정 · ||mu_h||^2={mu_n2:.3f} (표본 {len(PROBE_TEXTS)}문장, 참고값 9887.059)",
          flush=True)

    # ⛔ 쓰기 전에 확인한다 — 이 mu 로 덮으면 실제로 로짓이 내려가는가
    with torch.no_grad():
        enc = tok("한국의 수도는", return_tensors="pt")
        h = model(**enc, output_hidden_states=True).hidden_states[-1][0, -1].float()
        pred = float(h @ repl32)
        top = float(model(**enc).logits[0, -1].float().max())
    print(f"[2/4] 예상 마스킹 로짓 {pred:.1f} · 현재 최대 로짓 {top:.1f} · 마진 {top - pred:.1f}",
          flush=True)
    if pred > top - MARGIN_MIN:
        print("⛔ 마진 부족 — alpha 를 올려라", file=sys.stderr)
        return 2
    del model
    import gc; gc.collect()

    # ── 샤드 하나만 수술 ────────────────────────────────────────────────
    shard = find_shard(model_dir)
    src = model_dir / shard
    t = load_file(str(src))
    W = t["lm_head.weight"]
    print(f"[3/4] {shard} · lm_head.weight {tuple(W.shape)} {W.dtype} · 이 샤드 텐서 {len(t)}개",
          flush=True)
    before = W[ids[0]].float().norm().item()
    W[torch.tensor(ids, dtype=torch.long)] = repl32.to(W.dtype).unsqueeze(0).expand(len(ids), -1)
    print(f"[3/4] 대상 행 노름 {before:.3f} -> {W[ids[0]].float().norm().item():.3f}", flush=True)

    if a.dry_run:
        print("[4/4] --dry-run — 아무것도 쓰지 않았다", flush=True)
        return 0

    dst_dir = Path(a.out_dir).resolve() if a.out_dir else model_dir
    dst_dir.mkdir(parents=True, exist_ok=True)
    dst = dst_dir / shard
    t["lm_head.weight"] = W
    tmp = dst.with_suffix(dst.suffix + ".tmp")
    save_file(t, str(tmp), metadata={"format": "pt"})

    # 재로드 검증 — 의도한 텐서 외에 아무것도 안 바뀌었는지
    chk = load_file(str(tmp))
    diff = [k for k in chk if not torch.equal(chk[k], t[k])]
    assert set(chk) == set(t) and not diff, f"⛔ 재로드 불일치: {diff}"
    tmp.replace(dst)
    print(f"[4/4] 재로드 검증 OK — 텐서 {len(chk)}개, 의도 외 변경 0개 -> {dst}", flush=True)

    if a.out_dir and dst_dir != model_dir:
        print(f"      ⚠️ 나머지 파일은 원본에서 복사/심링크해야 한다: {model_dir}/*", flush=True)

    (dst_dir / "PRUNE_INFO.json").write_text(json.dumps({
        "mask": mask["name"], "n_masked": len(ids), "alpha": a.alpha,
        "mu_h_norm2": mu_n2, "expected_masked_logit": pred, "top_logit": top,
        "probe_texts": PROBE_TEXTS,
        "method": "lm_head 행을 -alpha*mu_h/||mu_h||^2 로 대체. 샤드 1개만 변경.",
        "embed_tokens": "untouched — CJK 입력은 그대로 읽는다",
        "tokenizer": "unchanged — vocab_size 동일",
    }, ensure_ascii=False, indent=2), encoding="utf-8")
    print("PRUNE_OK", flush=True)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())