Upload generate_and_eval_v6.py with huggingface_hub

generate_and_eval_v6.py

@@ -0,0 +1,291 @@
+#!/usr/bin/env python3
+import argparse
+import json
+import random
+from pathlib import Path
+from typing import Iterable
+
+import numpy as np
+import torch
+
+from collect_experts import (
+    _load_llm_with_fallback,
+    _load_tokenizer_with_fallback,
+    collect_router_scores,
+)
+from v6_model import EncoderOnlyModel
+
+
+def _set_seed(seed: int) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+def _default_device() -> str:
+    return "cuda" if torch.cuda.is_available() else "cpu"
+
+
+def _load_ckpt(path: str) -> dict:
+    return torch.load(path, map_location="cpu")
+
+
+def _load_state_dict(path: str) -> dict:
+    ckpt = _load_ckpt(path)
+    state = ckpt.get("model", ckpt)
+    if any(key.startswith("_orig_mod.") for key in state.keys()):
+        state = {key.replace("_orig_mod.", ""): value for key, value in state.items()}
+    return state
+
+
+def _load_ckpt_config(path: str) -> dict:
+    ckpt = _load_ckpt(path)
+    cfg = ckpt.get("config", None)
+    return cfg if isinstance(cfg, dict) else {}
+
+
+def non_overlapping_chunks(
+    token_ids: list[int],
+    expert_topk_idx: torch.Tensor,
+    seq_len: int,
+    pad_id: int,
+) -> Iterable[tuple[list[int], torch.Tensor, list[bool]]]:
+    n_tokens = len(token_ids)
+    if n_tokens == 0:
+        return
+
+    start = 0
+    while start < n_tokens:
+        end = min(start + seq_len, n_tokens)
+        chunk_len = end - start
+        chunk_tokens = token_ids[start:end]
+        chunk_experts = expert_topk_idx[start:end]
+
+        if chunk_len < seq_len:
+            chunk_tokens = chunk_tokens + [pad_id] * (seq_len - chunk_len)
+            if chunk_len > 0:
+                pad_row = chunk_experts[-1].unsqueeze(0)
+            else:
+                pad_row = torch.zeros_like(expert_topk_idx[:1])
+            pad_block = pad_row.expand(seq_len - chunk_len, -1, -1)
+            chunk_experts = torch.cat([chunk_experts, pad_block], dim=0)
+
+        attention_mask = [True] * chunk_len + [False] * (seq_len - chunk_len)
+        yield chunk_tokens, chunk_experts, attention_mask
+        start += seq_len
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="Decode text from GPT-OSS expert selections with the V6 inverter."
+    )
+    parser.add_argument("--checkpoint", default="inverter_v6.pt")
+    parser.add_argument("--text-file", default="text.txt")
+    parser.add_argument("--text", default=None)
+
+    parser.add_argument("--model", default="openai/gpt-oss-20b")
+    parser.add_argument("--model-revision", default=None)
+    parser.add_argument(
+        "--attn-impl",
+        choices=["auto", "flash_attention_2", "sdpa", "eager"],
+        default="auto",
+    )
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--limit-tokens", type=int, default=None)
+
+    parser.add_argument("--router-topk", type=int, default=4)
+    parser.add_argument("--router-chunk-size", type=int, default=2048)
+    parser.add_argument("--router-min-chunk-size", type=int, default=128)
+
+    parser.add_argument("--seq-len", type=int, default=256)
+    parser.add_argument("--batch-size", type=int, default=8)
+    parser.add_argument("--eval-topk", default="1,5,10")
+
+    parser.add_argument("--use-ckpt-config", action="store_true", default=True)
+    parser.add_argument("--no-use-ckpt-config", action="store_false", dest="use_ckpt_config")
+    parser.add_argument("--layers", type=int, default=24)
+    parser.add_argument("--d-model", type=int, default=768)
+    parser.add_argument("--n-head", type=int, default=12)
+    parser.add_argument("--d-ff", type=int, default=2048)
+    parser.add_argument("--n-layer", type=int, default=6)
+    parser.add_argument("--layer-hidden", type=int, default=64)
+    parser.add_argument("--layer-proj", type=int, default=64)
+    parser.add_argument("--dropout", type=float, default=0.1)
+    parser.add_argument("--logit-softcap", type=float, default=0.0)
+    parser.add_argument("--layer-gating", action="store_true", default=False)
+    parser.add_argument("--position-type", choices=["auto", "learned", "rope"], default="auto")
+    parser.add_argument("--rope-theta", type=float, default=10000.0)
+    parser.add_argument("--qk-norm", action="store_true", default=True)
+    parser.add_argument("--no-qk-norm", action="store_false", dest="qk_norm")
+    parser.add_argument("--qk-norm-eps", type=float, default=1e-5)
+
+    parser.add_argument("--out", default="text_eval_v6.json")
+    parser.add_argument("--decoded-out", default="decoded_top1_v6.txt")
+    args = parser.parse_args()
+
+    device = _default_device()
+    if device == "cuda":
+        try:
+            torch.backends.cuda.matmul.fp32_precision = "tf32"
+            torch.backends.cudnn.conv.fp32_precision = "tf32"
+        except AttributeError:
+            torch.backends.cuda.matmul.allow_tf32 = True
+            torch.backends.cudnn.allow_tf32 = True
+        torch.set_float32_matmul_precision("high")
+
+    _set_seed(args.seed)
+
+    ckpt_cfg = _load_ckpt_config(args.checkpoint)
+    state_dict = _load_state_dict(args.checkpoint)
+    ckpt_has_gate = bool(ckpt_cfg.get("layer_gating", False)) or ("encoder_in.layer_gate" in state_dict)
+    if ckpt_has_gate and not args.layer_gating:
+        args.layer_gating = True
+
+    if args.use_ckpt_config and ckpt_cfg:
+        mapping = {
+            "seq_len": "seq_len",
+            "layers": "layers",
+            "d_model": "d_model",
+            "n_head": "n_head",
+            "d_ff": "d_ff",
+            "n_layer": "n_layer",
+            "layer_hidden": "layer_hidden",
+            "layer_proj": "layer_proj",
+            "dropout": "dropout",
+            "logit_softcap": "logit_softcap",
+            "rope_theta": "rope_theta",
+            "qk_norm_eps": "qk_norm_eps",
+        }
+        for ckpt_key, arg_key in mapping.items():
+            if ckpt_key in ckpt_cfg:
+                setattr(args, arg_key, ckpt_cfg[ckpt_key])
+        if "position_type" in ckpt_cfg and args.position_type == "auto":
+            args.position_type = ckpt_cfg["position_type"]
+        if "qk_norm" in ckpt_cfg:
+            args.qk_norm = bool(ckpt_cfg["qk_norm"])
+
+    if args.position_type == "auto":
+        args.position_type = "learned" if "pos_emb.weight" in state_dict else "rope"
+
+    text = args.text
+    if text is None:
+        text = Path(args.text_file).read_text(encoding="utf-8")
+
+    tokenizer = _load_tokenizer_with_fallback(args.model, args.model_revision)
+    if tokenizer.pad_token_id is None:
+        tokenizer.pad_token_id = tokenizer.eos_token_id
+
+    attn_impl = args.attn_impl
+    if attn_impl == "auto":
+        attn_impl = "flash_attention_2" if device != "cpu" else "eager"
+    llm = _load_llm_with_fallback(args.model, args.model_revision, device, attn_impl)
+
+    inverter = EncoderOnlyModel(
+        vocab_size=len(tokenizer),
+        num_experts=32,
+        num_layers=int(args.layers),
+        topk=int(args.router_topk),
+        d_model=int(args.d_model),
+        n_head=int(args.n_head),
+        d_ff=int(args.d_ff),
+        n_layer=int(args.n_layer),
+        dropout=float(args.dropout),
+        max_len=int(args.seq_len),
+        layer_gating=bool(args.layer_gating),
+        logit_softcap=float(args.logit_softcap),
+        layer_hidden=int(args.layer_hidden),
+        layer_proj=int(args.layer_proj),
+        position_type=str(args.position_type),
+        rope_theta=float(args.rope_theta),
+        qk_norm=bool(args.qk_norm),
+        qk_norm_eps=float(args.qk_norm_eps),
+    ).to(device)
+    inverter.load_state_dict(state_dict, strict=True)
+    inverter.eval()
+
+    input_ids = tokenizer(text, return_tensors="pt", add_special_tokens=False)["input_ids"][0].tolist()
+    if args.limit_tokens is not None:
+        input_ids = input_ids[: args.limit_tokens]
+    if not input_ids:
+        raise ValueError("No tokens found in the provided text.")
+
+    input_ids_cpu = torch.tensor([input_ids], dtype=torch.long, device="cpu")
+    topk_scores = collect_router_scores(
+        model=llm,
+        input_ids_cpu=input_ids_cpu,
+        chunk_size=max(1, int(args.router_chunk_size)),
+        min_chunk_size=max(1, int(args.router_min_chunk_size)),
+        topk=int(args.router_topk),
+        save_dtype=torch.float16,
+        pin_memory=(device == "cuda"),
+    )
+    topk_idx_cpu = topk_scores["topk_idx"][:, : int(args.layers), :]
+
+    eval_topk = sorted({int(item) for item in args.eval_topk.split(",") if item.strip() and int(item) > 0})
+    correct = {k: 0 for k in eval_topk}
+    total = 0
+    predicted_ids: list[int] = []
+
+    def run_batch(batch_tokens, batch_experts, batch_attn) -> None:
+        nonlocal total
+        target_ids = torch.tensor(batch_tokens, dtype=torch.long, device=device)
+        expert_idx = torch.stack(batch_experts, dim=0).to(device=device, dtype=torch.long)
+        attention_mask = torch.tensor(batch_attn, dtype=torch.bool, device=device)
+
+        with torch.autocast(device_type=device, dtype=torch.bfloat16, enabled=(device == "cuda")):
+            logits = inverter(expert_idx, attention_mask)
+
+        top1 = torch.argmax(logits, dim=-1)
+        for row_idx in range(top1.shape[0]):
+            valid_len = int(attention_mask[row_idx].sum().item())
+            predicted_ids.extend(top1[row_idx, :valid_len].tolist())
+
+        for k in eval_topk:
+            topk_pred = torch.topk(logits, k=k, dim=-1).indices
+            match = (topk_pred == target_ids.unsqueeze(-1)).any(dim=-1)
+            match = match & attention_mask
+            correct[k] += int(match.sum().item())
+
+        total += int(attention_mask.sum().item())
+
+    batch_tokens = []
+    batch_experts = []
+    batch_attn = []
+    for chunk_tokens, chunk_experts, attn_mask in non_overlapping_chunks(
+        token_ids=input_ids,
+        expert_topk_idx=topk_idx_cpu,
+        seq_len=int(args.seq_len),
+        pad_id=int(tokenizer.pad_token_id),
+    ):
+        batch_tokens.append(chunk_tokens)
+        batch_experts.append(chunk_experts)
+        batch_attn.append(attn_mask)
+        if len(batch_tokens) >= int(args.batch_size):
+            run_batch(batch_tokens, batch_experts, batch_attn)
+            batch_tokens, batch_experts, batch_attn = [], [], []
+
+    if batch_tokens:
+        run_batch(batch_tokens, batch_experts, batch_attn)
+
+    decoded_text = tokenizer.decode(
+        predicted_ids,
+        skip_special_tokens=False,
+        clean_up_tokenization_spaces=False,
+    )
+    Path(args.decoded_out).write_text(decoded_text, encoding="utf-8")
+
+    result = {
+        "text_file": args.text_file if args.text is None else None,
+        "tokens": total,
+        "accuracy": {str(k): correct[k] / total for k in eval_topk},
+        "checkpoint": args.checkpoint,
+        "decoded_out": args.decoded_out,
+    }
+    Path(args.out).write_text(json.dumps(result, indent=2), encoding="utf-8")
+    print(json.dumps(result, indent=2))
+
+
+if __name__ == "__main__":
+    main()