Hugging Face's logo

Lostboy231
/
cgc

Model card Files
xet
 Community
cgc
File size: 4,561 Bytes
69202a6
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
"""
scripts/compute_density_map.py
------------------------------
Phase 1 + 2: Collect per-head activations and compute capability density map.

Steps:
  1. Load GPT-2 Medium
  2. Load calibration data (WikiText-103-raw-v1, 128 seq Γ— 512 tok)
  3. Run calibration data through model, collect per-head attention outputs
  4. For each of 384 heads: train a TopK SAE, compute β, H, ψ, δ
  5. Save density map to .npz

Runtime: ~15–20 min on T4 GPU (384 SAEs Γ— 5 epochs each)

Usage:
    python scripts/compute_density_map.py --output results/density_map.npz
"""

import argparse
import random
import os
import sys

import numpy as np
import torch

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# ── Reproducibility ───────────────────────────────────────────────────────────
# Note: original experiments were run without fixed seeds.
# This seed ensures approximate reproducibility across runs.
SEED = 42
random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)
torch.cuda.manual_seed_all(SEED)

from transformers import GPT2Model, GPT2Tokenizer
from datasets import load_dataset
from cgc.density import collect_head_activations, compute_density_map


def parse_args():
    p = argparse.ArgumentParser(description="Compute CGC capability density map")
    p.add_argument("--output",       type=str,   default="results/density_map.npz",
                   help="Output path for density map")
    p.add_argument("--n_sequences",  type=int,   default=128,
                   help="Number of calibration sequences (default: 128)")
    p.add_argument("--seq_len",      type=int,   default=512,
                   help="Tokens per sequence (default: 512)")
    p.add_argument("--batch_size",   type=int,   default=8,
                   help="Sequences per forward pass (default: 8)")
    return p.parse_args()


def load_calibration_data(tokenizer, n_sequences, seq_len, device):
    """
    Load WikiText-103-raw-v1 training split as non-overlapping token chunks.
    Same dataset and split used for calibration in the CGC v1 paper.
    """
    print("Loading calibration data (WikiText-103-raw-v1, train split)...")
    dataset   = load_dataset("Salesforce/wikitext",
                             "wikitext-103-raw-v1", split="train")
    full_text = " ".join([
        item["text"].strip()
        for item in dataset
        if len(item["text"].strip()) > 50
    ])
    all_tokens = tokenizer.encode(full_text)

    chunks = []
    for i in range(n_sequences):
        start = i * seq_len
        end   = start + seq_len
        if end > len(all_tokens):
            print(f"Warning: only {len(chunks)} sequences available.")
            break
        chunks.append(torch.tensor(all_tokens[start:end], dtype=torch.long))

    data = torch.stack(chunks).to(device)
    print(f"Calibration data: {data.shape}  "
          f"({data.shape[0] * data.shape[1]:,} tokens total)")
    return data


def main():
    args = parse_args()
    os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Device: {device}")
    if device.type == "cuda":
        print(f"GPU: {torch.cuda.get_device_name(0)}")
        vram = torch.cuda.get_device_properties(0).total_memory / 1e9
        print(f"VRAM: {vram:.1f} GB")

    # Load GPT-2 Medium
    print("\nLoading GPT-2 Medium...")
    tokenizer = GPT2Tokenizer.from_pretrained("gpt2-medium")
    tokenizer.pad_token = tokenizer.eos_token
    model = GPT2Model.from_pretrained("gpt2-medium")
    model = model.to(device)
    model.eval()
    n_params = sum(p.numel() for p in model.parameters())
    print(f"Loaded. Layers={model.config.n_layer}  "
          f"Heads={model.config.n_head}  "
          f"Params={n_params:,}")

    # Calibration data
    cal_data = load_calibration_data(
        tokenizer, args.n_sequences, args.seq_len, device
    )

    # Phase 1: collect per-head activations
    head_acts = collect_head_activations(
        model, cal_data, device, args.batch_size
    )

    # Phase 2: train SAEs and compute density
    density_map, _ = compute_density_map(
        head_acts, device, seq_len=args.seq_len
    )

    # Save
    density_map.save(args.output)
    print(f"\nDensity map saved to: {args.output}")
    print("Next: python scripts/run_compression.py "
          f"--density_map {args.output}")


if __name__ == "__main__":
    main()