experiment/training/build_gen_feature_cache.py

"""Build a gen-only SAE feature cache (drop-in for Train_Probe_AllLayers
output, but pooled over the assistant's generated caption tokens only).

The existing ``training/multilayer_sae_ckpt/feature_cache/`` was written by
``training/Train_Probe_AllLayers.py``'s ``topk_sae_max_pool``, which
max-pools over the *entire* sequence (image patches + prompt + caption +
padding). That mismatches the probes trained by ``train_probe_gen.py``,
which pool only over the assistant's generated tokens via ``gen_mask``.

This script reuses the exact pipeline of ``train_probe_gen.py``:
generate caption → teacher-force [prompt+caption] → capture hidden
states → ``FrozenSAEEncoder`` → ``masked_max_pool`` (or
``masked_mean_pool``) over ``gen_mask``. The result is saved in the same
``feature_cache/`` directory layout (``labels.pt`` + ``layer_LL.pt`` +
``meta.json``) so ``Select_Features_F1_AllLayers.py`` and
``mechanistic_interp/select_features_activation.py`` consume it
unchanged.

Single-GPU or torchrun multi-GPU. Each rank writes its features into a
shared tensor list, and rank 0 gathers + saves.

Usage:
    python -m experiment.training.build_gen_feature_cache \\
        --config experiment/adv_config.json \\
        --cache_out training/multilayer_sae_ckpt/feature_cache_gen

    torchrun --nproc_per_node=8 -m experiment.training.build_gen_feature_cache \\
        --config experiment/adv_config.json \\
        --cache_out training/multilayer_sae_ckpt/feature_cache_gen
"""
from __future__ import annotations

import argparse
import json
import os
import sys

import torch
import torch.distributed as dist
from torch.utils.data import DataLoader, DistributedSampler
from transformers import AutoModelForPreTraining, AutoProcessor
from tqdm import tqdm

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../.."))

from sae.Training_Utils import str_to_torch_dtype
from experiment.config.train_config import TrainConfig
from experiment.config.relation_config import get_relation_config
from experiment.data.datasets import FinetuneDataset, finetune_dataset_extra_kwargs
from experiment.training.finetune_adv import (
    FrozenSAEEncoder,
    HiddenStateCapture,
    count_lm_layers,
)
from experiment.training.train_probe_gen import _generate_captions
from experiment.training.gen_features import (
    build_position_masks,
    left_pad_collate,
    masked_max_pool,
    masked_mean_pool,
)


def parse_local_args():
    p = argparse.ArgumentParser(add_help=False)
    p.add_argument("--config", required=True)
    p.add_argument("--cache_out", required=True,
                   help="Output directory (will be created).")
    p.add_argument("--relation", default=None)
    p.add_argument("--pool", choices=("max", "mean"), default="max")
    p.add_argument("--split", default="train",
                   help="HF split fed into FinetuneDataset.")
    p.add_argument("--max_samples", type=int, default=0,
                   help="0 = all (overrides config.max_train_samples).")
    extra, remaining = p.parse_known_args()
    sys.argv = [sys.argv[0]] + remaining
    return extra


def main():
    extra = parse_local_args()
    config = TrainConfig.load(extra.config)
    if extra.relation:
        config.relation = extra.relation
    config.resolve_from_relation()

    relation_config = get_relation_config(config.relation)
    adv_cfg = config.adv
    model_dtype = str_to_torch_dtype(config.dtype)
    assert adv_cfg.sae_checkpoint, "adv.sae_checkpoint must be set in config"

    use_ddp = "LOCAL_RANK" in os.environ
    if use_ddp:
        local_rank = int(os.environ["LOCAL_RANK"])
        dist.init_process_group(backend="nccl")
        torch.cuda.set_device(local_rank)
        device = torch.device(f"cuda:{local_rank}")
        is_main = local_rank == 0
        world = dist.get_world_size()
    else:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        local_rank, is_main, world = 0, True, 1

    pool_fn = masked_max_pool if extra.pool == "max" else masked_mean_pool

    if is_main:
        print(f"[build_gen_cache] relation={config.relation}  "
              f"sae={adv_cfg.sae_checkpoint}  pool={extra.pool}  "
              f"split={extra.split}  out={extra.cache_out}")
        os.makedirs(extra.cache_out, exist_ok=True)

    model = AutoModelForPreTraining.from_pretrained(
        config.model_name, torch_dtype=model_dtype, device_map={"": device},
    )
    processor = AutoProcessor.from_pretrained(config.model_name)
    pad_token_id = processor.tokenizer.pad_token_id or processor.tokenizer.eos_token_id
    for p in model.parameters():
        p.requires_grad_(False)
    model.eval()

    n_layers = count_lm_layers(model)
    layer_indices = adv_cfg.probe_layers if adv_cfg.probe_layers else list(range(n_layers))
    if is_main:
        print(f"[build_gen_cache] layers: {layer_indices}")

    sae = FrozenSAEEncoder.from_checkpoint(adv_cfg.sae_checkpoint, device)
    d_sae = sae.encoder.weight.shape[0]

    if extra.max_samples > 0:
        max_samples = extra.max_samples
    else:
        max_samples = config.max_train_samples
    dataset = FinetuneDataset(
        processor=processor,
        prompt_config=config.prompts,
        dataset_id=config.dataset_id,
        scene_col=relation_config.scene_key,
        object_col=relation_config.object_key,
        csv_path=config.csv_path,
        image_dir=config.image_dir,
        max_samples=max_samples,
        split=extra.split,
        upsample_categories=None,
        **finetune_dataset_extra_kwargs(config),
    )
    if use_ddp:
        sampler = DistributedSampler(dataset, shuffle=False, drop_last=False)
        loader = DataLoader(
            dataset, batch_size=config.batch_size, sampler=sampler,
            num_workers=config.num_workers, pin_memory=True, drop_last=False,
        )
    else:
        loader = DataLoader(
            dataset, batch_size=config.batch_size, shuffle=False,
            num_workers=config.num_workers, pin_memory=True, drop_last=False,
        )

    capture = HiddenStateCapture(model, layer_indices)
    n_image_patches = None

    layer_feats: dict[int, list[torch.Tensor]] = {l: [] for l in layer_indices}
    has_object_list: list[torch.Tensor] = []
    is_scene_list: list[torch.Tensor] = []

    for batch in tqdm(loader, desc=f"rank{local_rank}", disable=not is_main):
        has_object = batch.pop("has_object").to(device)
        is_scene = batch.pop("is_scene").to(device)
        batch.pop("labels", None)
        batch = {k: v.to(device) for k, v in batch.items()}

        # 1) Free-generate captions per row.
        _, full_seqs, prompt_lens, gen_lens = _generate_captions(
            model,
            pixel_values=batch["pixel_values"],
            input_ids=batch["input_ids"],
            attention_mask=batch["attention_mask"],
            max_new_tokens=adv_cfg.max_new_tokens_train,
            do_sample=adv_cfg.gen_do_sample,
            temperature=adv_cfg.gen_temperature,
            pad_token_id=pad_token_id,
        )

        # 2) Teacher-force [prompt + caption], capture hidden states.
        tf_ids, tf_attn = left_pad_collate(full_seqs, pad_id=pad_token_id)
        tf_ids = tf_ids.to(device)
        tf_attn = tf_attn.to(device)
        with torch.no_grad():
            with capture:
                model(
                    pixel_values=batch["pixel_values"],
                    input_ids=tf_ids,
                    attention_mask=tf_attn,
                    use_cache=False,
                )

        # 3) Build gen_mask. Infer n_image_patches on first batch.
        S = next(iter(capture.hidden_states.values())).shape[1]
        if n_image_patches is None:
            n_image_patches = S - tf_ids.shape[1] + 1
            if is_main:
                print(f"[build_gen_cache] inferred n_image_patches={n_image_patches} "
                      f"(S={S}, L_max={tf_ids.shape[1]})")
        prompt_real_lens_t = torch.tensor(prompt_lens, device=device)
        gen_lens_t = torch.tensor(gen_lens, device=device)
        _, gen_mask = build_position_masks(
            prompt_real_lens_t, gen_lens_t, n_image_patches, S,
        )
        keep = gen_lens_t > 0
        if not keep.any():
            continue

        # 4) SAE-encode every position, gen-pool, store.
        for l, h in capture.hidden_states.items():
            latents = sae(h)                       # (B, S, d_sae) float32 (gated)
            feat = pool_fn(latents, gen_mask)      # (B, d_sae)
            # masked_max_pool returns -inf where mask is empty; we already
            # dropped those rows via `keep`. Cast to float16 to mirror the
            # original cache (compact + fast to load).
            layer_feats[l].append(feat[keep].detach().to(torch.float16).cpu())

        has_object_list.append(has_object[keep].detach().cpu().float())
        is_scene_list.append(is_scene[keep].detach().cpu().float())

    local_layer_feats = {l: torch.cat(layer_feats[l]) for l in layer_indices}
    local_has_object = torch.cat(has_object_list) if has_object_list else torch.zeros(0)
    local_is_scene = torch.cat(is_scene_list) if is_scene_list else torch.zeros(0)

    if use_ddp:
        # Gather across ranks (concatenate per-rank tensors).
        def _gather_concat(t: torch.Tensor) -> torch.Tensor:
            parts = [None] * world
            dist.all_gather_object(parts, t)
            return torch.cat(parts) if is_main else None

        all_has_object = _gather_concat(local_has_object)
        all_is_scene = _gather_concat(local_is_scene)
        all_layer_feats = {}
        for l in layer_indices:
            all_layer_feats[l] = _gather_concat(local_layer_feats[l])
    else:
        all_has_object = local_has_object
        all_is_scene = local_is_scene
        all_layer_feats = local_layer_feats

    if not is_main:
        if use_ddp:
            dist.barrier()
            dist.destroy_process_group()
        return

    # ── Save in feature_cache format (labels=has_object, mirrors original) ──
    out = extra.cache_out
    os.makedirs(out, exist_ok=True)
    torch.save(all_has_object, os.path.join(out, "labels.pt"))
    torch.save(all_is_scene, os.path.join(out, "is_scene.pt"))
    for l in layer_indices:
        torch.save(all_layer_feats[l], os.path.join(out, f"layer_{l:02d}.pt"))
    meta = {
        "d_sae": int(d_sae),
        "n_layers": int(n_layers),
        "probe_label_mode": "object_only",
        "dataset_id": config.dataset_id,
        "pool_over": "gen",
        "pool": extra.pool,
        "sae_ckpt": adv_cfg.sae_checkpoint,
        "split": extra.split,
        "relation": config.relation,
        "n_samples": int(all_has_object.numel()),
        "n_pos": int(all_has_object.sum()),
        "n_neg": int((1 - all_has_object).sum()),
        "n_image_patches": int(n_image_patches),
        "source": "build_gen_feature_cache.py (gen_mask, FrozenSAEEncoder)",
    }
    with open(os.path.join(out, "meta.json"), "w") as f:
        json.dump(meta, f, indent=2)
    print(f"[build_gen_cache] wrote cache → {out}")
    print(f"  labels: N={meta['n_samples']}  pos={meta['n_pos']}  neg={meta['n_neg']}")
    print(f"  per-layer shape: ({meta['n_samples']}, {d_sae}) float16")

    if use_ddp:
        dist.barrier()
        dist.destroy_process_group()


if __name__ == "__main__":
    main()