v2_artifacts/_out_surgery_embed.py

"""
ARCHON Embedding Surgery — re-init replaced token rows
=======================================================
For each token slot modified by the tokenizer surgery, this script:
  1. Loads the pretrain checkpoint (ARCHON v3b step ~259K or SFT warmstart)
  2. For each replaced slot, computes a new embedding by averaging the sub-word
     embeddings that tokenizer_v2.1 would assign to the NEW token string.
     This is the standard "mean sub-word init" strategy — better than random
     init and avoids catastrophic drift in neighboring token embeddings.
  3. Clears the lm_head weight rows for replaced slots (same operation, since
     weight tying is assumed unless --no-tie is passed).
  4. Saves a new checkpoint with patched embedding matrix.

Sub-word init rationale:
  - For <|fr|> (was <|teach|>, init artifact): average of tokens that appear
    in the string "fr" under v2.1 → gives a directionally meaningful start.
  - For "aujourd'hui": average of [aujo, urd, 'h, ui] in v2.1 vocab → captures
    the semantic neighborhood. Far better than N(0, 0.02) init.
  - For suffix tokens ("tion", "ment", "eur"): encode the suffix string with
    v2.1, average those embeddings.

Weight tying:
  - If lm_head.weight is tied to embed_tokens.weight (standard practice),
    the same rows must be patched in both. This script handles both cases.
  - If lm_head has a separate weight matrix (--no-tie), only embed_tokens is
    patched and lm_head rows are zeroed for the replaced slots (they will be
    learned quickly in adaptation training).

Usage:
  python _out_surgery_embed.py \
      --ckpt path/to/archon_step259k.pt \
      --tok-old path/to/tokenizer_v21.json \
      --tok-new path/to/tokenizer_v3.json \
      --map path/to/_out_surgery_map.json \
      --out path/to/archon_v3_surgered.pt \
      [--no-tie] \
      [--device cuda] \
      [--dtype bfloat16]
"""

import argparse
import json
import sys
from pathlib import Path
from typing import Optional

import torch


# ---------------------------------------------------------------------------
# Surgery table (same as in _out_surgery_tokenizer.py, hard-coded mirror)
# ---------------------------------------------------------------------------
SURGERY_TABLE: dict[int, str] = {
    5:  "<|fr|>",
    6:  "<|en|>",
    7:  "<|tool_name|>",
    8:  "<|tool_args|>",
    9:  "<|reasoning|>",
    10: "<|reflection|>",
    11: "<|memory|>",
    12: "<|context|>",
    13: "<|plan|>",
    14: "<|step|>",
    15: "<|verdict|>",
    22133: " française",
    24420: "chômage",
    24421: "aujourd'hui",
    24424: "qu'on",
    25336: " qu'il",
    25501: " n'est",
    24167: " n'a",
    23252: " d'une",
    22929: " d'un",
    23675: " l'on",
    23678: " j'ai",
    24695: " m'a",
    23385: " t'es",
    23384: " s'est",
    24054: " qu'est-ce",
    23860: " n'est-ce",
    24917: " j'aime",
    23176: " rejoindre",
    24926: " parce",
    24934: " toujours",
    23850: " encore",
    23251: " jamais",
    23259: " vraiment",
    24052: " maintenant",
    23399: " ailleurs",
    30904: " cependant",
    11602: " également",
    8914:  " notamment",
    15546: "ement",
    24086: "tion",
    8212:  "ique",
    8229:  "ment",
    15231: "eur",
    30173: "ais",
    11028: "ons",
    18722: "ant",
    3315:  "ère",
    8046:  "ité",
    12631: "teur",
    4626:  "aire",
}


def load_tokenizer_vocab(tokenizer_json_path: Path) -> dict[str, int]:
    """Load vocab from HuggingFace tokenizers JSON."""
    with open(tokenizer_json_path, "r", encoding="utf-8") as f:
        tok = json.load(f)
    vocab = tok.get("model", {}).get("vocab", {})
    if not vocab:
        raise ValueError(f"No model.vocab found in {tokenizer_json_path}")
    return vocab


def tokenize_string_with_vocab(text: str, vocab: dict[str, int]) -> list[int]:
    """
    Approximate tokenization: greedy longest-match BPE on the vocab.
    This is used only to find sub-word indices for init — it does not need
    to be byte-perfect. For correctness on the adaptation server, use the
    actual tokenizer; here we just need plausible sub-word neighbors.

    Falls back to character-level if no match found.
    """
    # Convert text to bytes (BPE typically works on byte-level)
    # We'll try character-level greedy for simplicity
    tokens = []
    i = 0
    sorted_vocab = sorted(vocab.keys(), key=len, reverse=True)
    while i < len(text):
        matched = False
        for tok_str in sorted_vocab:
            if text[i:i+len(tok_str)] == tok_str:
                tokens.append(vocab[tok_str])
                i += len(tok_str)
                matched = True
                break
        if not matched:
            # Character not in vocab — use unk or skip
            unk_id = vocab.get("<unk>", vocab.get("[UNK]", 3))
            tokens.append(unk_id)
            i += 1
    return tokens if tokens else [vocab.get("<unk>", 3)]


def find_embed_key(state_dict: dict) -> Optional[str]:
    """Find the embedding weight key in the state dict."""
    candidates = [
        "embed.weight",               # model_v2.py: self.embed = nn.Embedding(...) — FIRST
        "embed_tokens.weight",
        "model.embed_tokens.weight",
        "transformer.wte.weight",
        "embeddings.word_embeddings.weight",
    ]
    for k in candidates:
        if k in state_dict:
            return k
    # fallback: search
    for k in state_dict:
        if "embed" in k.lower() and "weight" in k.lower() and state_dict[k].ndim == 2:
            return k
    return None


def find_lm_head_key(state_dict: dict) -> Optional[str]:
    """Find the lm_head weight key in the state dict."""
    candidates = [
        "lm_head.weight",
        "output.weight",
        "head.weight",
    ]
    for k in candidates:
        if k in state_dict:
            return k
    for k in state_dict:
        if "head" in k.lower() and "weight" in k.lower() and state_dict[k].ndim == 2:
            return k
    return None


def compute_subword_mean_embedding(
    new_token_str: str,
    old_vocab: dict[str, int],
    embed_matrix: torch.Tensor,
    old_idx: int,
) -> torch.Tensor:
    """
    Compute a new embedding for a token by averaging the embeddings of the
    sub-word tokens that tokenizer_v2.1 would produce for new_token_str.

    If the new string tokenizes to a single token that maps to itself (the
    rare case where the new token already exists in old vocab), use that.

    Falls back to the old embedding (perturbed slightly) if tokenization
    returns only the unk or itself — should not happen for our surgery targets.
    """
    sub_ids = tokenize_string_with_vocab(new_token_str, old_vocab)

    # Filter out: the slot being replaced itself (avoid circular init),
    # unk tokens (index 3), pad (0)
    valid_ids = [sid for sid in sub_ids if sid not in {0, 3, old_idx}]

    if not valid_ids:
        # Fallback: use mean of all embeddings + small noise (better than zero)
        print(f"    WARNING: no valid sub-words for '{new_token_str}' (slot {old_idx}), "
              "using corpus mean + noise.")
        mean_embed = embed_matrix.mean(dim=0)
        noise = torch.randn_like(mean_embed) * 0.01
        return (mean_embed + noise).to(embed_matrix.dtype)

    # Average sub-word embeddings
    sub_embeds = embed_matrix[valid_ids]  # [n_subwords, D]
    mean_embed = sub_embeds.mean(dim=0)
    return mean_embed.to(embed_matrix.dtype)


def patch_checkpoint(
    ckpt_path: Path,
    old_vocab: dict[str, int],
    surgery_table: dict[int, str],
    out_path: Path,
    no_tie: bool,
    device: str,
    dtype: torch.dtype,
) -> None:
    print(f"Loading checkpoint: {ckpt_path}")
    ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)

    # HF Trainer saves {"model": state_dict, ...} or just state_dict
    if isinstance(ckpt, dict) and "model" in ckpt and isinstance(ckpt["model"], dict):
        state_dict = ckpt["model"]
        ckpt_wrapper = ckpt
    elif isinstance(ckpt, dict) and any("weight" in k for k in ckpt):
        state_dict = ckpt
        ckpt_wrapper = None
    else:
        raise ValueError(
            "Unrecognized checkpoint format. Expected dict with 'model' key "
            "or direct state_dict."
        )

    embed_key = find_embed_key(state_dict)
    lm_head_key = find_lm_head_key(state_dict)

    if embed_key is None:
        raise ValueError("Could not find embedding matrix in checkpoint. "
                         "Check key naming conventions.")

    print(f"  Embedding key: {embed_key}")
    print(f"  LM head key:   {lm_head_key if lm_head_key else 'not found / tied'}")

    embed_matrix = state_dict[embed_key].float()  # work in fp32 for precision
    vocab_size, hidden_dim = embed_matrix.shape
    print(f"  Embedding shape: {embed_matrix.shape}")

    # Check tied weights
    tied = False
    if lm_head_key and lm_head_key != embed_key:
        lm_head_matrix = state_dict[lm_head_key].float()
        if lm_head_matrix.data_ptr() == embed_matrix.data_ptr():
            tied = True
            print("  Weight tying: detected (shared tensor).")
        elif lm_head_matrix.shape == embed_matrix.shape:
            # Not the same tensor object but same shape — treat as tied unless --no-tie
            if not no_tie:
                tied = True
                print("  Weight tying: assumed (same shape, --no-tie not set).")
            else:
                print("  Weight tying: disabled via --no-tie.")
        else:
            print(f"  Weight tying: shapes differ ({lm_head_matrix.shape} vs {embed_matrix.shape}), treating as separate.")
    elif not lm_head_key or lm_head_key == embed_key:
        tied = True
        print("  Weight tying: yes (same key or no lm_head found).")

    # Validate surgery slots are within vocab size
    for idx in surgery_table:
        if idx >= vocab_size:
            raise ValueError(
                f"Surgery slot {idx} >= vocab_size {vocab_size}. "
                "Check that this checkpoint matches tokenizer_v2.1."
            )

    # Apply surgery
    print(f"\nApplying {len(surgery_table)} embedding patches...")
    patched = 0
    for slot_idx, new_tok_str in surgery_table.items():
        old_row = embed_matrix[slot_idx].clone()
        new_row = compute_subword_mean_embedding(new_tok_str, old_vocab, embed_matrix, slot_idx)

        # Sanity check: new row should not be all zeros or identical to old
        if torch.allclose(new_row, old_row, atol=1e-6):
            print(f"  WARN slot {slot_idx} ('{new_tok_str}'): new embedding identical to old. "
                  "This may indicate the new token tokenizes to itself in old vocab.")

        embed_matrix[slot_idx] = new_row
        patched += 1
        if patched % 10 == 0:
            print(f"  Patched {patched}/{len(surgery_table)} slots...")

    print(f"  Done: {patched} embedding rows updated.")

    # Cast back to target dtype
    embed_matrix = embed_matrix.to(dtype)
    state_dict[embed_key] = embed_matrix

    # Patch lm_head if not tied (or tied but different tensor object)
    if lm_head_key and lm_head_key != embed_key:
        if tied:
            # Update lm_head to same data
            state_dict[lm_head_key] = embed_matrix
            print("  lm_head updated (tied, shared update).")
        else:
            # --no-tie mode: zero out the replaced rows in lm_head
            # They will be re-learned quickly in adaptation
            lm_head_matrix = state_dict[lm_head_key].float()
            for slot_idx in surgery_table:
                lm_head_matrix[slot_idx].zero_()
            state_dict[lm_head_key] = lm_head_matrix.to(dtype)
            print("  lm_head rows zeroed for surgery slots (--no-tie mode).")

    # Re-pack checkpoint
    if ckpt_wrapper is not None:
        ckpt_wrapper["model"] = state_dict
        ckpt_wrapper["surgery_applied"] = True
        ckpt_wrapper["surgery_slots"] = list(surgery_table.keys())
        ckpt_wrapper["surgery_new_tokens"] = list(surgery_table.values())
        save_obj = ckpt_wrapper
    else:
        state_dict["_surgery_applied"] = torch.tensor(True)
        save_obj = state_dict

    out_path.parent.mkdir(parents=True, exist_ok=True)
    print(f"\nSaving surgered checkpoint → {out_path}")
    torch.save(save_obj, out_path)
    print(f"  Saved. Size: {out_path.stat().st_size / 1e9:.2f} GB")

    # Write embedding diff stats
    stats_path = out_path.parent / "surgery_embed_stats.json"
    stats = {
        "ckpt_src": str(ckpt_path),
        "ckpt_dst": str(out_path),
        "vocab_size": vocab_size,
        "hidden_dim": hidden_dim,
        "slots_patched": patched,
        "dtype": str(dtype),
        "tied_weights": tied,
    }
    with open(stats_path, "w") as f:
        json.dump(stats, f, indent=2)
    print(f"  Stats → {stats_path}")


def main() -> None:
    parser = argparse.ArgumentParser(description="ARCHON embedding surgery v2.1 → v3")
    parser.add_argument("--ckpt", required=True, type=Path,
                        help="Input checkpoint path (pretrain or SFT warmstart)")
    parser.add_argument("--tok-old", required=True, type=Path,
                        help="tokenizer_v21.json (old vocab for sub-word init)")
    parser.add_argument("--tok-new", required=True, type=Path,
                        help="tokenizer_v3.json (new vocab, for validation)")
    parser.add_argument("--map", type=Path, default=None,
                        help="Optional: _out_surgery_map.json (overrides hard-coded table)")
    parser.add_argument("--out", required=True, type=Path,
                        help="Output checkpoint path")
    parser.add_argument("--no-tie", action="store_true",
                        help="Treat embed_tokens and lm_head as separate matrices")
    parser.add_argument("--device", default="cpu",
                        help="Device for tensor ops (cpu or cuda). Default: cpu.")
    parser.add_argument("--dtype", default="bfloat16",
                        choices=["float32", "float16", "bfloat16"],
                        help="Output dtype for embedding matrix. Default: bfloat16.")
    args = parser.parse_args()

    dtype_map = {
        "float32": torch.float32,
        "float16": torch.float16,
        "bfloat16": torch.bfloat16,
    }
    dtype = dtype_map[args.dtype]

    table = SURGERY_TABLE
    if args.map:
        with open(args.map, "r", encoding="utf-8") as f:
            surgery_map = json.load(f)
        replacements = surgery_map.get("replacements", {})
        table = {}
        for k, v in replacements.items():
            if not k.isdigit():
                continue
            idx = int(k)
            new_tok = v.get("new_revised", v.get("new", ""))
            if new_tok:
                table[idx] = new_tok
        print(f"Loaded {len(table)} surgery entries from {args.map}")

    old_vocab = load_tokenizer_vocab(args.tok_old)
    print(f"Old vocab size: {len(old_vocab)}")

    # Validate new tokenizer exists and has same size
    new_vocab = load_tokenizer_vocab(args.tok_new)
    if len(new_vocab) != len(old_vocab):
        raise ValueError(
            f"Vocab size mismatch: old={len(old_vocab)}, new={len(new_vocab)}. "
            "Tokenizer surgery must preserve vocab size."
        )
    print(f"New vocab size: {len(new_vocab)} (matches old — OK)")

    patch_checkpoint(
        ckpt_path=args.ckpt,
        old_vocab=old_vocab,
        surgery_table=table,
        out_path=args.out,
        no_tie=args.no_tie,
        device=args.device,
        dtype=dtype,
    )


if __name__ == "__main__":
    main()