gemma_inject.py

#!/usr/bin/env python3
"""
gemma_inject.py — FABLE 5 → HPC gate injection for Gemma 4 12B

Adapted from hpc_fable_inject.py for Gemma architecture:
  - hidden_size=3840, intermediate_size=15360 (4× hidden_size)
  - Pauli state: [E/τ | μ/τ | ε/τ | μ/τ] — 4 blocks filling 15360 dims
  - Single model.safetensors file (~23 GB)
  - tie_word_embeddings=True (lm_head == embed_tokens)
"""
import argparse, gc, json, math, os, sys, time
from collections import Counter, defaultdict
from pathlib import Path

import numpy as np
import torch
from safetensors import safe_open
from safetensors.torch import save_file
from transformers import AutoTokenizer


def tokenize_fable(data_path, tokenizer):
    all_ids = []
    with open(data_path) as f:
        for line in f:
            d = json.loads(line)
            text = d.get("text", "")
            for part in text.split("<|im_start|>"):
                if part.startswith("assistant\n"):
                    content = part[len("assistant\n"):].replace("<|im_end|>", "").strip()
                    if content:
                        all_ids.extend(tokenizer.encode(content, add_special_tokens=False))
    return all_ids


def compute_bigram_stats(all_ids, vocab_size, topk):
    cnt = Counter()
    bigram = defaultdict(lambda: Counter())
    for i, tid in enumerate(all_ids):
        cnt[tid] += 1
        if i > 0:
            bigram[all_ids[i - 1]][tid] += 1
    top_tokens = [t for t, _ in cnt.most_common(topk)]
    token_to_idx = {t: i for i, t in enumerate(top_tokens)}
    return cnt, bigram, top_tokens, token_to_idx


def recover_gate_weight(embed, cnt, bigram, top_tokens, token_to_idx, tau):
    D = embed.shape[1]          # 3840
    V = len(top_tokens)

    E_raw = np.array([embed[t] for t in top_tokens]).astype(np.float64)
    mu = np.zeros_like(E_raw)
    eps = np.zeros_like(E_raw)

    for i, ti in enumerate(top_tokens):
        total = sum(bigram[ti].values())
        fi = cnt[ti] or 1
        if total == 0:
            continue
        for tj, c in bigram[ti].items():
            if tj not in token_to_idx:
                continue
            fj = cnt[tj] or 1
            p = c / total
            w = c / math.sqrt(fi * fj)
            c_Z = (1.0 - w) / 2.0
            mu[i] += p * embed[tj]
            eps[i] += c_Z * embed[tj]

    # 4-block Pauli state: [E/τ | μ/τ | ε/τ | μ/τ]
    z = np.hstack([E_raw / tau, mu / tau, eps / tau, mu / tau])
    z = np.clip(z, -30.0, 30.0)

    E_n = (E_raw - E_raw.mean(0, keepdims=True)) / (E_raw.std(0, keepdims=True) + 1e-10)
    reg = 1e-3 * np.eye(D)
    W = np.linalg.solve(E_n.T @ E_n + reg, E_n.T @ z)
    return W.T.astype(np.float32)  # (15360, 3840)


def inject_safetensors(src_path, dst_path, W_hpc_torch, alpha, layer_count):
    """
    Read source (mmap'd), modify gate_proj in RAM, save to new file.
    Unmodified tensors stay mmap'd (zero RAM cost).
    Only 48 gate_proj tensors (~5.6 GB) materialize in RAM.
    """
    with safe_open(str(src_path), framework='pt') as sf:
        keys = list(sf.keys())
        out = {}
        for idx, k in enumerate(keys):
            t = sf.get_tensor(k)
            if 'gate_proj' in k and 'mlp' in k:
                is_gate, layer_idx = _parse_gate_key(k)
                if is_gate and layer_idx < layer_count:
                    t = t.to(torch.float32) + alpha * W_hpc_torch.to(torch.float32)
                    t = t.to(torch.bfloat16).contiguous()
            out[k] = t

            if (idx + 1) % 50 == 0:
                print(f"  [{idx+1}/{len(keys)}] processed", flush=True)
                gc.collect()

    os.makedirs(os.path.dirname(dst_path) or '.', exist_ok=True)
    save_file(out, dst_path, metadata={'format': 'pt'})
    print(f"Wrote {len(keys)} tensors to {dst_path}")


def _parse_gate_key(key):
    parts = key.split('.')
    for i, p in enumerate(parts):
        if p == 'layers' and i + 1 < len(parts):
            try:
                return True, int(parts[i + 1])
            except ValueError:
                pass
    return False, -1


def main():
    parser = argparse.ArgumentParser(description="Gemma FABLE → HPC gate injector")
    parser.add_argument("--model", default="./gemma-4-12B-it", help="Gemma model directory")
    parser.add_argument("--output", default="./gemma-4-12B-it-hpc", help="Output directory")
    parser.add_argument("--data", default="/tmp/fable5_sft.jsonl", help="FABLE 5 JSONL path")
    parser.add_argument("--alpha", type=float, default=0.3, help="Injection strength")
    parser.add_argument("--tau", type=float, default=0.003, help="Temperature")
    parser.add_argument("--topk", type=int, default=30000, help="Top-k tokens")
    args = parser.parse_args()

    t0 = time.time()
    model_dir = Path(args.model)
    src_safetensors = model_dir / "model.safetensors"
    dst_safetensors = Path(args.output) / "model.safetensors"

    # ── 1. Tokenizer & embedding ──
    print("[1/6] Loading tokenizer & embedding...")
    tokenizer = AutoTokenizer.from_pretrained(str(model_dir), trust_remote_code=True)
    embed = None
    with safe_open(str(src_safetensors), framework='pt') as sf:
        for k in sf.keys():
            if 'embed_tokens' in k:
                embed = sf.get_tensor(k).float().numpy()
                break
    assert embed is not None, "embed_tokens not found"
    print(f"  Embedding: {embed.shape}")

    # ── 2. Tokenize FABLE ──
    print(f"[2/6] Tokenizing {args.data}...")
    all_ids = tokenize_fable(args.data, tokenizer)
    print(f"  {len(all_ids)} tokens")

    # ── 3. Bigram stats ──
    print(f"[3/6] Bigrams (top-{args.topk})...")
    cnt, bigram, top_tokens, token_to_idx = compute_bigram_stats(all_ids, tokenizer.vocab_size, args.topk)
    print(f"  {len(top_tokens)} tokens")

    # ── 4. Recover gate weight ──
    print(f"[4/6] Recovering gate weight (τ={args.tau})...")
    W_hpc = recover_gate_weight(embed, cnt, bigram, top_tokens, token_to_idx, args.tau)
    target_std = 0.012
    scale = target_std / W_hpc.std()
    W_hpc *= scale
    print(f"  W_hpc: {W_hpc.shape}, std={W_hpc.std():.6f}")

    # ── 5. Get layer count ──
    print(f"[5/6] Scanning source file...")
    layer_count = 0
    with safe_open(str(src_safetensors), framework='pt') as sf:
        for k in sf.keys():
            is_gate, layer = _parse_gate_key(k)
            if is_gate:
                layer_count = max(layer_count, layer + 1)
    print(f"  {layer_count} layers detected")

    W_hpc_t = torch.from_numpy(W_hpc).to(torch.bfloat16)

    # ── 6. Inject ──
    print(f"[6/6] Injecting (α={args.alpha}) → {args.output}...")
    os.makedirs(args.output, exist_ok=True)
    inject_safetensors(str(src_safetensors), str(dst_safetensors),
                       W_hpc_t, args.alpha, layer_count)

    import shutil
    for fn in ["config.json", "tokenizer.json", "tokenizer_config.json",
               "generation_config.json", "chat_template.jinja", "processor_config.json"]:
        src = model_dir / fn
        if src.exists():
            shutil.copy2(src, Path(args.output) / fn)

    elapsed = time.time() - t0
    print(f"Done in {elapsed:.0f}s. Output: {args.output}/")


if __name__ == "__main__":
    main()