Find features in a language model with sparse autoencoders

flyte create secret anthropic-api-key
git clone https://github.com/unionai/unionai-examples
cd unionai-examples/v2/tutorials/sparse_autoencoders
python find_features.py            # 3 layers x 2 widths, 8M tokens
python find_features.py --smoke    # 1 layer x 1 width, 0.5M tokens

image = flyte.Image.from_debian_base(name="sae-features").with_pip_packages(
    "torch==2.9.1", "transformers==4.57.3", "datasets==4.4.2", "numpy<3", "anthropic==2.6.0",
)

# One GPU per task: the harvest, and every SAE trainer. sae.py is imported inside the tasks.
gpu_env = flyte.TaskEnvironment(
    name="sae-gpu",
    image=image,
    resources=flyte.Resources(cpu=6, memory="40Gi", gpu="L4:1", disk="80Gi"),
    cache="auto",
    include=["sae.py"],
)

# Autointerp is API calls, so CPU only. Only this environment holds the API key.
interp_env = flyte.TaskEnvironment(
    name="sae-autointerp",
    image=image,
    resources=flyte.Resources(cpu=2, memory="4Gi"),
    secrets=[flyte.Secret(key=ANTHROPIC_SECRET, as_env_var="ANTHROPIC_API_KEY")],
    cache="auto",
)

driver_env = flyte.TaskEnvironment(
    name="sae-driver",
    image=image,
    resources=flyte.Resources(cpu=2, memory="8Gi"),
    # The driver declares label_features when it calls it, so it needs the same secret name.
    env_vars={"ANTHROPIC_SECRET": ANTHROPIC_SECRET},
    depends_on=[gpu_env, interp_env],
)

@gpu_env.task
async def harvest_activations(
    layers: list[int],
    n_tokens: int = 8_000_000,
    ctx_len: int = 128,
    dataset: str = "HuggingFaceFW/fineweb:sample-10BT",
    shard_seqs: int = 4096,
) -> Harvest:
    """Run Pythia over FineWeb and save the residual stream after each layer in ``layers``.

    Each context is the first ``ctx_len`` tokens of one document. Activations are float16 shards,
    one directory per layer, so a trainer downloads only the layer it needs.
    """
    import numpy as np
    import torch
    from datasets import load_dataset
    from transformers import AutoModelForCausalLM, AutoTokenizer

    tok = AutoTokenizer.from_pretrained(MODEL)
    model = AutoModelForCausalLM.from_pretrained(MODEL).to(DEVICE).eval()
    name, _, config = dataset.partition(":")
    stream = load_dataset(name, name=config or None, split="train", streaming=True)

    n_seqs = n_tokens // ctx_len
    rows: list[list[int]] = []
    texts: list[str] = []
    for doc in stream:
        texts.append(doc["text"])
        if len(texts) == 1000:
            enc = tok(texts, truncation=True, max_length=ctx_len)["input_ids"]
            rows += [ids for ids in enc if len(ids) == ctx_len]
            texts = []
            if len(rows) >= n_seqs:
                break
    tokens = np.array(rows[:n_seqs], dtype=np.int32)
    print(f"{len(tokens):,} contexts of {ctx_len} tokens from {dataset}")

    captured: dict[int, torch.Tensor] = {}
    for layer in layers:
        model.gpt_neox.layers[layer].register_forward_hook(
            lambda _m, _i, out, layer=layer: captured.__setitem__(layer, out[0] if isinstance(out, tuple) else out))

    root = Path(tempfile.mkdtemp())
    for layer in layers:
        (root / f"layer_{layer}").mkdir()
    d_model = model.config.hidden_size
    with torch.no_grad():
        for shard, s in enumerate(range(0, len(tokens), shard_seqs)):
            chunk = tokens[s:s + shard_seqs]
            out = {layer: np.empty((len(chunk), ctx_len, d_model), dtype=np.float16) for layer in layers}
            for b in range(0, len(chunk), 256):
                model(input_ids=torch.as_tensor(chunk[b:b + 256], dtype=torch.long, device=DEVICE))
                for layer in layers:
                    out[layer][b:b + 256] = captured[layer].clamp(-6e4, 6e4).half().cpu().numpy()
            for layer in layers:
                np.save(root / f"layer_{layer}" / f"shard_{shard:03d}.npy", out[layer])
    np.save(root / "tokens.npy", tokens)
    print(f"saved {len(layers)} layers x {len(tokens) * ctx_len:,} tokens")

    acts = await asyncio.gather(*[Dir.from_local(str(root / f"layer_{layer}")) for layer in layers])
    return Harvest(model=MODEL, dataset=dataset, layers=layers, ctx_len=ctx_len, d_model=d_model,
                   n_seqs=len(tokens), acts=list(acts), tokens=await File.from_local(str(root / "tokens.npy")))

class TopKSAE(nn.Module):
    def __init__(self, d_model: int, width: int, k: int):
        super().__init__()
        self.d_model, self.width, self.k = d_model, width, k
        w_dec = torch.randn(width, d_model)
        w_dec /= w_dec.norm(dim=1, keepdim=True)
        self.W_dec = nn.Parameter(w_dec)  # one unit-norm direction per feature
        self.W_enc = nn.Parameter(w_dec.T.clone())  # tied at init, trained separately
        self.b_enc = nn.Parameter(torch.zeros(width))
        self.b_dec = nn.Parameter(torch.zeros(d_model))

    def pre_acts(self, x: torch.Tensor) -> torch.Tensor:
        return (x - self.b_dec) @ self.W_enc + self.b_enc

    def encode(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """The k active features per token: their indices, their values, and every pre-activation."""
        pre = self.pre_acts(x)
        vals, idx = pre.topk(self.k, dim=-1)
        return idx, F.relu(vals), pre

    def decode(self, idx: torch.Tensor, vals: torch.Tensor) -> torch.Tensor:
        # A sum of k decoder rows per token, without materializing the width-sized code.
        return F.embedding_bag(idx, self.W_dec, per_sample_weights=vals, mode="sum") + self.b_dec

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        idx, vals, _ = self.encode(x)
        return self.decode(idx, vals)

    @torch.no_grad()
    def normalize_decoder(self) -> None:
        self.W_dec.data /= self.W_dec.data.norm(dim=1, keepdim=True).clamp_min(1e-8)

def train(
    acts: torch.Tensor,  # [sequences, positions, d_model] float16, on the GPU
    width: int,
    ks: list[int],
    steps: int,
    batch_size: int = 4096,
    lr: float | None = None,
    seed: int = 0,
    log_every: int = 250,
    on_log=None,
) -> tuple[list[TopKSAE], float, list[dict]]:
    """Train one TopK SAE per k on the same batches. Position 0 is skipped (see ``token_index``)."""
    torch.manual_seed(seed)
    device = acts.device
    n_seq, ctx, d = acts.shape
    scale = activation_scale(acts)
    # Gao et al. scale the learning rate as 1/sqrt(width).
    lr = lr or 2e-4 * math.sqrt(16384 / width)
    saes = [TopKSAE(d, width, k).to(device) for k in ks]
    sample = gather(acts, token_index(n_seq, ctx, 32_768, device), scale)
    for sae in saes:
        sae.b_dec.data = sample.median(dim=0).values  # a robust center for the decoder bias
    opts = [torch.optim.Adam(s.parameters(), lr=lr, betas=(0.9, 0.999)) for s in saes]
    decay_from = int(steps * 0.8)
    sched = lambda step: min(1.0, (step + 1) / 100) * min(1.0, (steps - step) / max(steps - decay_from, 1))
    last_fired = [torch.zeros(width, dtype=torch.long, device=device) for _ in saes]
    dead_after = max(1, 1_000_000 // batch_size)  # a feature is dead after 1M tokens without firing
    curve: list[dict] = []
    for step in range(steps):
        x = gather(acts, token_index(n_seq, ctx, batch_size, device), scale)
        point = {"step": step + 1}
        for i, (sae, opt) in enumerate(zip(saes, opts)):
            for g in opt.param_groups:
                g["lr"] = lr * sched(step)
            idx, vals, pre = sae.encode(x)
            x_hat = sae.decode(idx, vals)
            err = x - x_hat
            nmse = err.pow(2).sum(-1).mean() / (x - x.mean(0)).pow(2).sum(-1).mean()
            dead = (step - last_fired[i]) > dead_after
            loss = nmse + aux_loss(sae, pre, err.detach(), dead) / 32
            opt.zero_grad(set_to_none=True)
            loss.backward()
            opt.step()
            sae.normalize_decoder()
            last_fired[i][idx[vals > 0]] = step
            point[f"k{sae.k}"] = nmse.item()
            point[f"dead_k{sae.k}"] = float(dead.float().mean())
        if step % log_every == 0 or step == steps - 1:
            curve.append(point)
            if on_log:
                on_log(point)
    return saes, scale, curve

@gpu_env.task(report=True, cache=flyte.Cache("auto", salt=SAE_VERSION))
async def train_sae(
    harvest: Harvest,
    layer: int,
    width: int,
    ks: list[int],
    steps: int = 8000,
    batch_size: int = 4096,
    label_k: int = 32,
    n_features: int = 40,
    n_top: int = 12,
    n_sampled: int = 8,
    held_out: float = 0.05,
    seed: int = 0,
) -> Trained:
    """Train one TopK SAE per k on one layer, evaluate each on held-out text, and sample features to label."""
    import sae as S
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer

    torch.backends.cuda.matmul.allow_tf32 = True
    acts, tokens = await load_layer(harvest, layer)
    n_train = int(len(acts) * (1 - held_out))
    print(f"layer {layer}: {n_train:,} training and {len(acts) - n_train:,} held-out contexts, width {width}, k {ks}")

    started = time.time()
    curve_so_far: list[dict] = []

    def on_log(point: dict) -> None:  # stream the training curves into this task's report
        curve_so_far.append(point)
        flyte.report.replace(reports.training(layer, width, ks, steps, curve_so_far, time.time() - started))
        flyte.report.flush()

    # Training runs in a worker thread, so the report can stream from the callback.
    saes, scale, curve = await asyncio.to_thread(S.train, acts[:n_train], width, ks, steps, batch_size,
                                                 seed=seed, on_log=on_log)
    train_s = time.time() - started
    print(f"trained {len(ks)} SAEs for {steps} steps in {train_s / 60:.1f} min")

    # Evaluate on held-out contexts: reconstruction, sparsity, and what the model loses with the SAE in place.
    model = AutoModelForCausalLM.from_pretrained(harvest.model).to(DEVICE).eval()
    eval_tokens = torch.as_tensor(tokens[n_train:][:512], dtype=torch.long, device=DEVICE)
    ce = S.spliced_losses(model, eval_tokens, layer, saes, scale)
    stats = []
    for sae in saes:
        m = S.reconstruction_metrics(sae, acts[n_train:], scale)
        st = SAEStats(layer=layer, width=width, k=sae.k, fve=m["fve"], l0=m["l0"], dead_frac=m["dead_frac"],
                      ce_clean=ce["clean"], ce_sae=ce[f"k{sae.k}"], ce_zero=ce["zero"],
                      loss_recovered=S.loss_recovered(ce["clean"], ce[f"k{sae.k}"], ce["zero"]),
                      density_hist=S.density_histogram(m["density"]))
        print(f"  k={sae.k}: FVE {st.fve:.3f}, L0 {st.l0:.1f}, dead {st.dead_frac:.1%}, "
              f"loss {st.ce_clean:.3f} -> {st.ce_sae:.3f} (zero {st.ce_zero:.3f}), recovered {st.loss_recovered:.1%}")
        stats.append(st)

    out = Path(tempfile.mkdtemp())
    weights = []
    for sae in saes:
        path = out / f"sae_L{layer}_w{width}_k{sae.k}.pt"
        torch.save({"state_dict": sae.state_dict(), "scale": scale, "layer": layer, "width": width, "k": sae.k,
                    "d_model": sae.d_model, "model": harvest.model, "hook": f"gpt_neox.layers.{layer} output"}, path)
        weights.append(await File.from_local(str(path)))

    examples = None
    if n_features and label_k in ks:
        tok = AutoTokenizer.from_pretrained(harvest.model)
        feats = S.feature_examples(saes[ks.index(label_k)], acts, tokens, scale, lambda t: tok.decode([t]),
                                   n_features=n_features, n_top=n_top, n_sampled=n_sampled, seed=seed + layer)
        path = out / "examples.json"
        path.write_text(json.dumps(feats))
        examples = await File.from_local(str(path))
        print(f"  sampled {len(feats)} features for autointerp")
    return Trained(layer=layer, width=width, ks=ks, stats=stats, weights=weights, curve=curve,
                   train_s=round(train_s, 1), examples=examples, label_k=label_k if examples else 0)

@torch.no_grad()
def reconstruction_metrics(sae: TopKSAE, acts: torch.Tensor, scale: float, chunk: int = 128) -> dict:
    """Fraction of variance explained, L0, dead features, and each feature's firing rate, on held-out tokens."""
    x_all = acts[:, 1:]
    mean = x_all.reshape(-1, x_all.shape[-1]).float().mean(0) * scale
    sse = sst = l0 = 0.0
    fires = torch.zeros(sae.width, device=acts.device)
    n = 0
    for s in range(0, x_all.shape[0], chunk):
        x = x_all[s:s + chunk].reshape(-1, x_all.shape[-1]).float() * scale
        idx, vals, _ = sae.encode(x)
        x_hat = sae.decode(idx, vals)
        sse += float((x - x_hat).pow(2).sum())
        sst += float((x - mean).pow(2).sum())
        on = vals > 0
        l0 += float(on.sum())
        fires.index_add_(0, idx[on], torch.ones(int(on.sum()), device=acts.device))
        n += x.shape[0]
    density = (fires / n).cpu().numpy()
    return {"fve": 1 - sse / sst, "l0": l0 / n, "dead_frac": float((density == 0).mean()),
            "density": density, "tokens": n}

@torch.no_grad()
def spliced_losses(model, tokens: torch.Tensor, layer: int, saes: list[TopKSAE], scale: float, batch: int = 64) -> dict:
    """Next-token loss with layer ``layer``'s output left alone, replaced by each SAE's reconstruction, or zeroed.

    Position 0 is never touched, the same as in training.
    """
    block = model.gpt_neox.layers[layer]
    state: dict = {"replace": None}

    def hook(_module, _inputs, output):
        hidden = output[0] if isinstance(output, tuple) else output
        if state["replace"] is None:
            return output
        new = hidden.clone()
        new[:, 1:] = state["replace"](hidden[:, 1:])
        return (new, *output[1:]) if isinstance(output, tuple) else new

    def through(sae):
        def f(h):
            x = h.reshape(-1, h.shape[-1]).float() * scale
            return (sae(x) / scale).reshape(h.shape).to(h.dtype)
        return f

    def loss() -> float:
        total = 0.0
        for s in range(0, tokens.shape[0], batch):
            t = tokens[s:s + batch]
            total += float(model(input_ids=t, labels=t).loss) * t.shape[0]
        return total / tokens.shape[0]

    handle = block.register_forward_hook(hook)
    try:
        out = {"clean": loss()}
        state["replace"] = torch.zeros_like
        out["zero"] = loss()
        for sae in saes:
            state["replace"] = through(sae)
            out[f"k{sae.k}"] = loss()
    finally:
        handle.remove()
    return out

def loss_recovered(clean: float, spliced: float, zero: float) -> float:
    """1.0 when the reconstruction is as good as the real activations, 0.0 when it is no better than zeros."""
    return (zero - spliced) / (zero - clean)

@torch.no_grad()
def feature_examples(
    sae: TopKSAE,
    acts: torch.Tensor,
    tokens: np.ndarray,
    scale: float,
    decode_token,
    n_features: int = 40,
    n_top: int = 12,
    n_sampled: int = 8,
    n_detect: int = 8,
    before: int = 24,
    after: int = 8,
    seed: int = 0,
    chunk: int = 128,
) -> list[dict]:
    """For a random sample of live features: contexts to explain it from, a held-out test set, and a histogram.

    Claude explains a feature from its ``n_top`` strongest contexts plus ``n_sampled`` contexts drawn
    from the rest of its clear firings. Top contexts alone hide what a feature does at moderate
    strength, which is most of its firings (Paulo et al., 2024). The test set for each feature mixes ``n_detect`` contexts where it fires clearly (none of them
    among its top contexts) with ``n_detect`` random contexts where it does not fire at all. Autointerp scores an
    explanation by how well it predicts which is which.
    """
    rng = np.random.default_rng(seed)
    n_seq, ctx, _ = acts.shape
    # Pass 1: how many contexts each feature fires in.
    seq_hits = torch.zeros(sae.width, device=acts.device)
    for s in range(0, n_seq, chunk):
        x = acts[s:s + chunk, 1:].reshape(-1, acts.shape[-1]).float() * scale
        idx, vals, _ = sae.encode(x)
        per_seq = torch.zeros(min(chunk, n_seq - s), sae.width, device=acts.device)
        rows = torch.arange(per_seq.shape[0], device=acts.device).repeat_interleave((ctx - 1) * sae.k)
        per_seq.index_put_((rows, idx.flatten()), (vals.flatten() > 0).float(), accumulate=True)
        seq_hits += (per_seq > 0).float().sum(0)
    # Sample uniformly among features that fire in enough contexts to fill both example sets.
    enough = (seq_hits >= 2 * (n_top + n_sampled + n_detect)) & (seq_hits <= n_seq - 2 * n_detect)
    eligible = torch.nonzero(enough).flatten().cpu().numpy()
    feats = np.sort(rng.choice(eligible, size=min(n_features, len(eligible)), replace=False))
    # Pass 2: every activation of the sampled features, at every position.
    col = torch.full((sae.width,), -1, dtype=torch.long, device=acts.device)
    col[torch.as_tensor(feats, device=acts.device)] = torch.arange(len(feats), device=acts.device)
    f_acts = torch.zeros(n_seq, ctx, len(feats), dtype=torch.float16, device=acts.device)
    for s in range(0, n_seq, chunk):
        block = acts[s:s + chunk, 1:]
        x = block.reshape(-1, acts.shape[-1]).float() * scale
        idx, vals, _ = sae.encode(x)
        c = col[idx]
        keep = (c >= 0) & (vals > 0)
        rows = torch.arange(x.shape[0], device=acts.device).unsqueeze(1).expand_as(idx)[keep]
        flat = torch.zeros(x.shape[0], len(feats), dtype=torch.float16, device=acts.device)
        flat[rows, c[keep]] = vals[keep].half()
        f_acts[s:s + block.shape[0], 1:] = flat.reshape(block.shape[0], ctx - 1, len(feats))

    def window(seq: int, center: int, values: np.ndarray) -> dict:
        # Every excerpt has the same length, shifted inward at the context's edges, so that
        # length never tells the detection test which excerpts are positives.
        size = before + after + 1
        lo = min(max(1, center - before), ctx - size)
        hi = lo + size
        return {"tokens": [decode_token(int(t)) for t in tokens[seq, lo:hi]],
                "acts": [round(float(v), 3) for v in values[seq, lo:hi]]}

    out = []
    for j, f in enumerate(feats):
        a = f_acts[:, :, j].float()
        seq_max, seq_arg = a.max(dim=1)
        order = torch.argsort(seq_max, descending=True).cpu().numpy()
        seq_max, seq_arg, a_np = seq_max.cpu().numpy(), seq_arg.cpu().numpy(), a.cpu().numpy()
        active = order[seq_max[order] > 0]
        top = active[:n_top]
        rest = active[n_top:]
        # Held-out positives: contexts outside the top where the feature still fires clearly, at
        # least a quarter of its peak. Weaker firings are near the SAE's noise floor.
        clear = rest[seq_max[rest] >= 0.25 * seq_max[top[0]]]
        enough_clear = len(clear) >= n_sampled + n_detect
        pool = clear if enough_clear else rest[:4 * (n_sampled + n_detect)]
        picked = rng.choice(pool, size=min(n_sampled + n_detect, len(pool)), replace=False)
        sampled, test_on = np.sort(picked[:n_sampled]), picked[n_sampled:]
        test_on_set = [{**window(int(sq), int(seq_arg[sq]), a_np), "active": True} for sq in test_on]
        quiet = np.nonzero(seq_max == 0)[0]
        test_off = rng.choice(quiet, size=min(n_detect, len(quiet)), replace=False) if len(quiet) else []
        test_off_set = [{**window(int(sq), int(rng.integers(before + 1, ctx - after)), a_np), "active": False}
                        for sq in test_off]
        nz = a_np[a_np > 0]
        counts, edges = np.histogram(nz, bins=24, range=(0, float(nz.max())))
        out.append({
            "feature": int(f),
            "density": float(len(nz) / (n_seq * (ctx - 1))),
            "context_rate": float(len(active) / n_seq),
            "max_act": float(nz.max()),
            "hist": {"counts": counts.tolist(), "edges": [round(float(e), 3) for e in edges]},
            "top": [window(int(sq), int(seq_arg[sq]), a_np) for sq in top],
            "sampled_clear": bool(enough_clear),  # False: sampled from its next-strongest contexts instead
            "sampled": sorted((window(int(sq), int(seq_arg[sq]), a_np) for sq in sampled), key=lambda w: -max(w["acts"])),
            "test": test_on_set + test_off_set,
        })
    return out

EXPLAIN_SYSTEM = """\
You are an interpretability researcher. You study one feature of a sparse autoencoder trained
on a small language model's residual stream. A feature is a direction that is active on some
tokens and silent on the rest. You will see the text excerpts where the feature is most active,
then a random sample of other excerpts where it fires, with every token where it fires at a tenth of its peak or more marked as <<token>>, and the
peak activation of each excerpt on a 0-10 scale. Find what the marked tokens have in common, in the token itself and in its
context. Be specific: name the words, the syntax, or the topic. If it fires on a token only
in some contexts, say which. Your explanation will be used to predict, for unseen excerpts,
whether the feature fires anywhere in them."""

EXPLAIN_SCHEMA = {
    "type": "object",
    "properties": {
        "notes": {"type": "string", "description": "One or two sentences of reasoning about the pattern."},
        "explanation": {"type": "string", "description": "The explanation, at most 20 words."},
    },
    "required": ["notes", "explanation"],
    "additionalProperties": False,
}

DETECT_SYSTEM = """\
You test an explanation of one feature of a sparse autoencoder trained on a small language
model. You will get the explanation and numbered text excerpts. Some excerpts contain a token
where the feature fires and the rest do not. Using only the explanation, list the numbers of
the excerpts in which the feature fires."""

DETECT_SCHEMA = {
    "type": "object",
    "properties": {"active": {"type": "array", "items": {"type": "integer"}}},
    "required": ["active"],
    "additionalProperties": False,
}

def detection_score(answer_json: str, test: list[dict]) -> float:
    """Balanced accuracy of Claude's predictions against where the feature really fired."""
    predicted = {int(i) - 1 for i in json.loads(answer_json)["active"]}
    on = [i for i, ex in enumerate(test) if ex["active"]]
    off = [i for i, ex in enumerate(test) if not ex["active"]]
    hit = sum(i in predicted for i in on) / max(len(on), 1)
    reject = sum(i not in predicted for i in off) / max(len(off), 1)
    return (hit + reject) / 2

class Unanswered(Exception):
    """Claude gave no complete JSON answer: a refusal, or output cut off at max_tokens."""

async def ask_claude(system: str, prompt: str, schema: dict, max_tokens: int = 4096) -> str:
    import anthropic

    client = anthropic.AsyncAnthropic(max_retries=8)
    message = await client.messages.create(
        model=CLAUDE,
        max_tokens=max_tokens,
        system=system,
        output_config={"effort": "low", "format": {"type": "json_schema", "schema": schema}},
        messages=[{"role": "user", "content": prompt}],
    )
    # Any thinking counts toward max_tokens, so check that the answer finished before reading it.
    if message.stop_reason != "end_turn":
        raise Unanswered(f"stop_reason={message.stop_reason}")
    text = next((b.text for b in message.content if b.type == "text"), None)
    if text is None:
        raise Unanswered("no text block")
    return text

@flyte.trace
async def explain_feature(layer: int, feature: int, prompt: str) -> str:
    """Claude's explanation of one feature. Traced, so a retried task replays it instead of paying again."""
    return await ask_claude(autointerp.EXPLAIN_SYSTEM, prompt, autointerp.EXPLAIN_SCHEMA)

@flyte.trace
async def detect_feature(layer: int, feature: int, prompt: str) -> str:
    """Claude's guess of which held-out excerpts the explanation predicts. Traced, like the explanation."""
    return await ask_claude(autointerp.DETECT_SYSTEM, prompt, autointerp.DETECT_SCHEMA)

@interp_env.task(retries=1, cache=flyte.Cache("auto", salt=autointerp.VERSION))
async def label_features(layer: int, width: int, k: int, examples: File, concurrency: int = 8) -> str:
    """Explain each sampled feature, then score every explanation on the feature's held-out contexts.

    Each feature is also scored with another feature's explanation. That control shows what the
    test gives an explanation that is fluent but wrong. A feature whose calls fail is recorded
    as unlabeled, so one bad answer never fails the whole task.
    """
    with open(await examples.download()) as f:
        features = json.load(f)
    gate = asyncio.Semaphore(concurrency)
    unlabeled: list[dict] = []

    async def explain(feat: dict) -> dict | None:
        try:
            async with gate:
                answer = json.loads(await explain_feature(layer, feat["feature"], autointerp.explain_prompt(feat)))
            return {**feat, **answer}
        except Exception as exc:  # Unanswered, malformed JSON, or an API error that outlived the retries
            unlabeled.append({"feature": feat["feature"], "error": f"explain: {exc}"})
            return None

    labeled = [f for f in await asyncio.gather(*[explain(f) for f in features]) if f is not None]
    shuffled = [labeled[(i + 1) % len(labeled)]["explanation"] for i in range(len(labeled))]

    async def score(feat: dict, explanation: str) -> float | None:
        test = autointerp.detect_set(feat, seed=feat["feature"])
        try:
            async with gate:
                answer = await detect_feature(layer, feat["feature"], autointerp.detect_prompt(explanation, test))
            return autointerp.detection_score(answer, test)
        except Exception as exc:
            unlabeled.append({"feature": feat["feature"], "error": f"detect: {exc}"})
            return None

    real = await asyncio.gather(*[score(f, f["explanation"]) for f in labeled])
    control = await asyncio.gather(*[score(f, e) for f, e in zip(labeled, shuffled)])
    scored = []
    for feat, s, c in zip(labeled, real, control):
        if s is None or c is None:
            continue
        feat.update(score=s, control=c)
        scored.append(feat)
        print(f"  #{feat['feature']}: {s:.2f} (control {c:.2f})  {feat['explanation']}")
    if scored:
        print(f"layer {layer}: mean detection {sum(f['score'] for f in scored) / len(scored):.2f}, "
              f"mismatched control {sum(f['control'] for f in scored) / len(scored):.2f}")
    unlabeled = list({u["feature"]: u for u in unlabeled}.values())  # one entry per feature
    for u in unlabeled:
        print(f"  #{u['feature']} unlabeled: {u['error']}")
    return json.dumps({"layer": layer, "width": width, "k": k, "features": scored, "unlabeled": unlabeled})

@driver_env.task(report=True, produces_artifacts=True)
async def find_features(
    layers: list[int] = [1, 3, 5],
    widths: list[int] = [4096, 16384],
    ks: list[int] = [16, 32, 64],
    n_tokens: int = 8_000_000,
    steps: int = 8000,
    label_k: int = 32,
    label_width: int = 16384,
    features_per_sae: int = 40,
) -> File:
    """Harvest activations, train an SAE grid in parallel, label features with Claude, and publish the best SAE."""
    if label_width not in widths:
        label_width = max(widths)
    started = time.time()
    await flyte.report.replace.aio(reports.page_start(layers, widths, ks, n_tokens), do_flush=True)

    harvest = await harvest_activations(layers=layers, n_tokens=n_tokens)
    trained: list[Trained] = await asyncio.gather(*[
        train_sae(harvest, layer, width, ks, steps=steps, label_k=label_k,
                  n_features=features_per_sae if width == label_width else 0)
        for layer in layers for width in widths
    ])
    await flyte.report.replace.aio(reports.overview(trained, [], harvest, time.time() - started), do_flush=True)

    labeled = [json.loads(s) for s in await asyncio.gather(*[
        label_features(t.layer, t.width, t.label_k, t.examples) for t in trained if t.examples is not None
    ])]
    await flyte.report.replace.aio(reports.overview(trained, labeled, harvest, time.time() - started))
    flyte.report.get_tab("Feature browser").replace(reports.browser(labeled))
    await flyte.report.flush.aio()

    # The SAE whose features Claude explained best, by mean detection score.
    mean = lambda lab: sum(f["score"] for f in lab["features"]) / len(lab["features"])
    best_labels = max((lab for lab in labeled if lab["features"]), key=mean)
    best = next(t for t in trained if (t.layer, t.width) == (best_labels["layer"], best_labels["width"]))
    stats = next(s for s in best.stats if s.k == best_labels["k"])
    return await publish(best, stats, best_labels, flyte.ctx().action.run_name)

def browser(labeled: list[dict]) -> str:
    """Every labeled feature, most convincing first: highest detection score, then the widest margin over the control."""
    feats = [(f, lab) for lab in labeled for f in lab["features"]]
    feats.sort(key=lambda p: (-p[0]["score"], -(p[0]["score"] - p[0]["control"]), -p[0]["context_rate"]))
    good = sum(f["score"] >= SCORE_GOOD for f, _ in feats)
    return viz.page(
        "Feature browser",
        f"{len(feats)} features sampled at random from the live features of each SAE, sorted by how well Claude's "
        f"explanation predicts held-out contexts. {good} score at least {SCORE_GOOD}. Highlight strength is the "
        "feature's activation; hover a token for its value.",
        BROWSER_STYLE,
        "".join(feature_html(f, lab["layer"], lab["width"], lab["k"]) for f, lab in feats),
    )

async def publish(best: Trained, stats: SAEStats, labeled: dict, run_name: str) -> File:
    card = await artifacts.Card.create_from.aio(content=reports.card(best, stats, labeled, run_name), format="html",
                                                card_type="model")
    scores = [f["score"] for f in labeled["features"]]
    metadata = artifacts.Metadata.create_model_metadata(
        name="pythia-70m-topk-sae",
        description=(f"TopK SAE on Pythia-70m-deduped layer {stats.layer}: {stats.width} features, k={stats.k}, "
                     f"{stats.fve:.0%} variance explained, {stats.loss_recovered:.0%} loss recovered"),
        framework="PyTorch",
        model_type="TopK sparse autoencoder",
        architecture=f"d_model 512 -> {stats.width} -> 512, TopK k={stats.k}",
        task="interpretability",
        modality=("text",),
        serial_format="pt",
        card=card,
        attrs={"layer": str(stats.layer), "width": str(stats.width), "k": str(stats.k), "fve": f"{stats.fve:.4f}",
               "loss_recovered": f"{stats.loss_recovered:.4f}", "dead_frac": f"{stats.dead_frac:.4f}",
               "detection": f"{sum(scores) / len(scores):.3f}"},
    )
    return artifacts.new(best.weights[best.ks.index(stats.k)], metadata)

python find_features.py --widths 4096 16384 32768