#!/usr/bin/env python3 """ make_belief_cache_v2.py ═══════════════════════════════════════════════════════════════════════════════ Cache pre-VLM features for ablation matrix M0–M14 (CoT-Pool plan, Phase 0). Modes ───── --cache_mode mean_pool (legacy, sanity-equivalent to v1) output: beliefs [N, D] fp16 --cache_mode dual_pool (M1: image vs text mean, separately) output: beliefs_img [N, D] fp16 beliefs_text [N, D] fp16 --cache_mode per_frame (M3-M5: time-axis preserved, spatial pooled) output: beliefs_frame [N, F, D] fp16 (F = MAX_FRAMES = 8) valid_frames [N, F] bool beliefs_text [N, D] fp16 (auxiliary text pool) --cache_mode spatial4x4 (M6-M11: time + 4×4 spatial per frame) output: beliefs_grid [N, F, 16, D] fp16 (16 = 4×4 spatial pooled) valid_frames [N, F] bool beliefs_text [N, D] fp16 All modes additionally save: tta_means [N] fp32, tta_vars [N] fp32, schema_version=2, cache_mode, hidden_dim, n_frames. Why fp16? • Belief vectors come from a bf16/fp16 forward; fp32 storage is wasteful. • Halves disk + IO; trainer can promote to fp32 at use-time if needed. Storage budget (217k samples, D=2048, F=8) mean_pool ≈ 1.7 GB dual_pool ≈ 3.4 GB per_frame ≈ 13.5 GB spatial4x4 ≈ 113 GB (use mmap; do NOT load fully into RAM) Index invariant (same as v1) cache[i] corresponds to manifest sample i in data/policy_labels/{split}.json["samples"][i]. Usage ───── cd PROJECT_ROOT python -m training.Policy.make_belief_cache_v2 \\ --sft_checkpoint checkpoints/SFT/sft_v2/best \\ --cache_mode spatial4x4 \\ --label_dir data/policy_labels \\ --out_dir data/belief_cache_v2 \\ --batch_size 4 """ from __future__ import annotations import argparse import json import logging from pathlib import Path from typing import Dict, List, Optional, Tuple import torch import torch.nn.functional as F from torch.amp import autocast from torch.utils.data import DataLoader from tqdm import tqdm import sys sys.path.insert(0, str(Path(__file__).resolve().parents[2])) from training.Policy.policy_model import PolicyModel from training.Policy.policy_dataset import PolicyDataset, policy_collate_fn, MAX_FRAMES logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") logger = logging.getLogger("Policy.make_cache_v2") SCHEMA_VERSION = 2 # ───────────────────────────────────────────────────────────────────────────── # Helpers — per-image token slicing # ───────────────────────────────────────────────────────────────────────────── def _get_spatial_merge_size(model: PolicyModel) -> int: """Read spatial_merge_size from VLM vision config. Qwen2.5-VL = 2.""" base = model.sft.get_base_model() cfg = getattr(base, "config", None) vc = getattr(cfg, "vision_config", None) if cfg is not None else None sms = getattr(vc, "spatial_merge_size", None) if vc is not None else None if sms is None: logger.warning("Could not read vision_config.spatial_merge_size; " "defaulting to 2 (Qwen2.5-VL).") sms = 2 return int(sms) def _per_image_token_counts(image_grid_thw: torch.Tensor, spatial_merge_size: int) -> List[int]: """ For each image i in this batch, how many LLM-visible visual tokens it emits. count_i = t_i * h_i * w_i // (spatial_merge_size**2) """ counts: List[int] = [] sms2 = spatial_merge_size * spatial_merge_size for row in image_grid_thw.tolist(): t, h, w = row[0], row[1], row[2] c = (t * h * w) // sms2 counts.append(int(c)) return counts def _spatial_pool_image(tokens: torch.Tensor, h_post: int, w_post: int, out_hw: int = 4) -> torch.Tensor: """ tokens : [n_tok, D] flattened post-merger spatial sequence for ONE image h_post : post-merger height = h // spatial_merge_size w_post : post-merger width = w // spatial_merge_size out_hw : target spatial side (4 → 4×4 = 16 outputs) Returns : [out_hw*out_hw, D] """ n_tok, D = tokens.shape assert n_tok == h_post * w_post, \ f"token count {n_tok} != h_post*w_post={h_post * w_post}" # → [1, D, h_post, w_post] grid = tokens.transpose(0, 1).reshape(1, D, h_post, w_post) pooled = F.adaptive_avg_pool2d(grid.float(), (out_hw, out_hw)) # promote to fp32 for AAP # → [out_hw*out_hw, D] pooled = pooled.reshape(D, out_hw * out_hw).transpose(0, 1) return pooled.to(tokens.dtype) # ───────────────────────────────────────────────────────────────────────────── # Per-sample feature extraction # ───────────────────────────────────────────────────────────────────────────── def _split_sample_visual_tokens( hidden_states_b: torch.Tensor, # [L, D] one sample's tokens input_ids_b: torch.Tensor, # [L] attention_mask_b: torch.Tensor, # [L] image_grid_thw_b: torch.Tensor, # [n_img_in_sample, 3] image_token_id: int, spatial_merge_size: int, ) -> Tuple[List[torch.Tensor], List[Tuple[int, int]]]: """ Split a single sample's hidden states into per-image chunks. Returns ------- chunks : list of length n_img, each [count_i, D] (image-token hiddens) shapes : list of (h_post, w_post) per image """ # 1. Find positions of image_token_id within VALID region. valid = attention_mask_b > 0 is_img = (input_ids_b == image_token_id) & valid img_positions = torch.nonzero(is_img, as_tuple=False).squeeze(-1) n_img_tokens = int(img_positions.numel()) counts = _per_image_token_counts(image_grid_thw_b, spatial_merge_size) expected_total = sum(counts) if n_img_tokens != expected_total: raise RuntimeError( f"Visual-token count mismatch: input_ids has {n_img_tokens} " f"image-token positions, but image_grid_thw expects {expected_total}. " f"image_grid_thw rows: {image_grid_thw_b.tolist()}" ) # 2. Slice hidden_states at those positions (already contiguous per Qwen layout). img_hidden = hidden_states_b[img_positions] # [n_img_tokens, D] # 3. Partition into per-image chunks; remember (h_post, w_post). chunks: List[torch.Tensor] = [] shapes: List[Tuple[int, int]] = [] cursor = 0 for i, c in enumerate(counts): chunks.append(img_hidden[cursor:cursor + c]) t = int(image_grid_thw_b[i, 0].item()) h = int(image_grid_thw_b[i, 1].item()) w = int(image_grid_thw_b[i, 2].item()) # Qwen2.5-VL still images: t==1, post-merger spatial = (h//sms, w//sms). # If t > 1 (rare for our pipeline of single frames), we collapse t into # the "n_tok" sequence and re-derive spatial as h_post*w_post*t per image. # For our use case t=1 always — assert and proceed. if t != 1: raise RuntimeError( f"Unexpected image_grid_thw t={t} (>1). This pipeline assumes " f"per-frame image inputs, not video tensors." ) h_post = h // spatial_merge_size w_post = w // spatial_merge_size shapes.append((h_post, w_post)) cursor += c return chunks, shapes def _extract_features_for_batch( model: PolicyModel, inputs: Dict[str, torch.Tensor], cache_mode: str, spatial_merge_size: int, image_token_id: int, n_frames: int, ) -> Dict[str, torch.Tensor]: """ Run one VLM forward and return (CPU, fp16 where appropriate) tensors for the requested cache_mode. All outputs have leading dim B. Returns dict with keys depending on cache_mode (see file header). """ # Move tensors to device moved: Dict[str, torch.Tensor] = {} for k, v in inputs.items(): if not isinstance(v, torch.Tensor): moved[k] = v continue if k == "pixel_values": moved[k] = v.to(model.device, dtype=model.sft.dtype, non_blocking=True) else: moved[k] = v.to(model.device, non_blocking=True) base = model.sft.get_base_model() core = getattr(base, "model", None) # Run base text+vision encoder; get last hidden state with autocast(device_type="cuda", dtype=model._amp_dtype, enabled=True): if core is not None: out = core( input_ids = moved["input_ids"], attention_mask = moved.get("attention_mask"), pixel_values = moved.get("pixel_values"), image_grid_thw = moved.get("image_grid_thw"), use_cache = False, return_dict = True, ) hs = out.last_hidden_state if hasattr(out, "last_hidden_state") else out[0] else: out = base( input_ids = moved["input_ids"], attention_mask = moved.get("attention_mask"), pixel_values = moved.get("pixel_values"), image_grid_thw = moved.get("image_grid_thw"), use_cache = False, return_dict = True, output_hidden_states = True, ) hs = out.hidden_states[-1] # TTA for downstream compatibility — uses the canonical pooled belief. belief_canon = model.sft.belief_aggregator( hs, moved.get("attention_mask"), moved.get("input_ids"), ) # belief_aggregator may produce 2D for dual_pool — but we use the # ORIGINAL training strategy here (whatever the SFT ckpt has). The # tta_head was trained against THAT strategy, so feed it the canonical. tta_mean, tta_logvar = model.sft.tta_head(belief_canon) tta_var = torch.exp(tta_logvar.float().clamp(-20.0, 20.0)) tta_mean = tta_mean.float() B = hs.shape[0] D = hs.shape[-1] attn = moved.get("attention_mask") ids = moved.get("input_ids") igt = moved.get("image_grid_thw") # [total_images_in_batch, 3] out_dict: Dict[str, torch.Tensor] = { "tta_means": tta_mean.detach().cpu(), "tta_vars": tta_var.detach().cpu(), } # ── mean_pool (legacy) ──────────────────────────────────────────────────── if cache_mode == "mean_pool": if attn is not None: m = attn.unsqueeze(-1).to(hs.dtype) beliefs = (hs * m).sum(dim=1) / m.sum(dim=1).clamp(min=1e-6) else: beliefs = hs.mean(dim=1) out_dict["beliefs"] = beliefs.detach().to(torch.float16).cpu() return out_dict # ── dual_pool (image-mean, text-mean) ───────────────────────────────────── if cache_mode == "dual_pool": is_img = (ids == image_token_id) if attn is not None: valid = attn > 0 is_img = is_img & valid is_text = (~is_img) & valid else: is_text = ~is_img def _mm(mask_b: torch.Tensor) -> torch.Tensor: m = mask_b.unsqueeze(-1).to(hs.dtype) s = (hs * m).sum(dim=1) denom = m.sum(dim=1).clamp(min=1e-6) return s / denom b_img = _mm(is_img) b_txt = _mm(is_text) out_dict["beliefs_img"] = b_img.detach().to(torch.float16).cpu() out_dict["beliefs_text"] = b_txt.detach().to(torch.float16).cpu() return out_dict # ── per_frame / spatial4x4 — both need per-image splitting ──────────────── if cache_mode in ("per_frame", "spatial4x4"): if igt is None: raise RuntimeError( f"cache_mode={cache_mode} requires image_grid_thw, but the " f"processor did not emit it (no images in batch?)." ) # We need to know which (sample, frame) slot each row of image_grid_thw # belongs to. The processor concatenates images in batch order; per # sample the count equals number of frames passed in. Recover via the # number of distinct image-token RUNS in that sample's input_ids. # Simpler & more robust: per sample count = number of PIL images we # passed. But here we no longer have access to that; recover from # contiguous groups in input_ids. # # For Qwen2.5-VL each image's tokens form a contiguous run prefixed # and suffixed by special <|vision_start|>/<|vision_end|> tokens. We # only need image_token_id runs to count images per sample. igt_cursor = 0 beliefs_frame: Optional[torch.Tensor] = None beliefs_grid: Optional[torch.Tensor] = None if cache_mode == "per_frame": beliefs_frame = torch.zeros(B, n_frames, D, dtype=torch.float16) else: # spatial4x4 beliefs_grid = torch.zeros(B, n_frames, 16, D, dtype=torch.float16) valid_frames = torch.zeros(B, n_frames, dtype=torch.bool) beliefs_text = torch.zeros(B, D, dtype=torch.float16) for b in range(B): ids_b = ids[b] attn_b = attn[b] if attn is not None else torch.ones_like(ids_b) hs_b = hs[b] # Count contiguous runs of image_token_id (= number of images in this sample) valid = attn_b > 0 is_img_b = (ids_b == image_token_id) & valid # diff to find run boundaries x = is_img_b.to(torch.int8) diff = torch.cat([x.new_zeros(1), x[1:] - x[:-1]]) n_runs = int((diff == 1).sum().item()) if n_runs == 0: # No images for this sample — leave zeros, valid_frames stays False # Still compute text mean. m_text = valid.unsqueeze(-1).to(hs_b.dtype) t_mean = (hs_b * m_text).sum(dim=0) / m_text.sum(dim=0).clamp(min=1e-6) beliefs_text[b] = t_mean.detach().to(torch.float16).cpu() continue # Slice this sample's image_grid_thw rows igt_b = igt[igt_cursor:igt_cursor + n_runs] igt_cursor += n_runs chunks, shapes = _split_sample_visual_tokens( hs_b, ids_b, attn_b, igt_b, image_token_id, spatial_merge_size, ) n_imgs_use = min(len(chunks), n_frames) for f in range(n_imgs_use): tok_f = chunks[f] h_post, w_post = shapes[f] if cache_mode == "per_frame": pooled = tok_f.float().mean(dim=0).to(torch.float16) beliefs_frame[b, f] = pooled.detach().cpu() else: # spatial4x4 grid = _spatial_pool_image(tok_f, h_post, w_post, out_hw=4) beliefs_grid[b, f] = grid.detach().to(torch.float16).cpu() valid_frames[b, f] = True # text mean (non-image valid tokens) is_text_b = (~is_img_b) & valid m_text = is_text_b.unsqueeze(-1).to(hs_b.dtype) denom = m_text.sum(dim=0).clamp(min=1e-6) t_mean = (hs_b * m_text).sum(dim=0) / denom beliefs_text[b] = t_mean.detach().to(torch.float16).cpu() if cache_mode == "per_frame": out_dict["beliefs_frame"] = beliefs_frame else: out_dict["beliefs_grid"] = beliefs_grid out_dict["valid_frames"] = valid_frames out_dict["beliefs_text"] = beliefs_text return out_dict raise ValueError(f"Unknown cache_mode: {cache_mode}") # ───────────────────────────────────────────────────────────────────────────── # Cache builder # ───────────────────────────────────────────────────────────────────────────── def _flush_chunk(accumulators: Dict[str, List[torch.Tensor]], chunk_dir: Path, chunk_idx: int) -> int: """Concat the in-memory batches and atomically save one chunk file. Returns number of samples in the chunk.""" if not accumulators: return 0 part = {k: torch.cat(v, dim=0) for k, v in accumulators.items()} n = next(iter(part.values())).shape[0] tmp = chunk_dir / f"chunk_{chunk_idx:05d}.pt.tmp" fin = chunk_dir / f"chunk_{chunk_idx:05d}.pt" torch.save(part, tmp) tmp.rename(fin) return int(n) def _scan_chunks(chunk_dir: Path) -> Tuple[int, int]: """Return (n_chunks, n_samples_total) present on disk (sorted).""" if not chunk_dir.exists(): return 0, 0 files = sorted(chunk_dir.glob("chunk_*.pt")) # Drop stray .tmp for t in chunk_dir.glob("*.tmp"): t.unlink(missing_ok=True) n_samples = 0 for f in files: try: d = torch.load(f, map_location="cpu", weights_only=True) n_samples += int(next(iter(d.values())).shape[0]) except Exception as e: logger.warning(f" [resume] chunk {f.name} unreadable ({e}); dropping") f.unlink(missing_ok=True) return len(list(chunk_dir.glob("chunk_*.pt"))), n_samples def _merge_chunks(chunk_dir: Path) -> Dict[str, torch.Tensor]: """Load all chunks in order and concatenate into a single cache dict.""" files = sorted(chunk_dir.glob("chunk_*.pt")) if not files: return {} acc: Dict[str, List[torch.Tensor]] = {} for f in files: d = torch.load(f, map_location="cpu", weights_only=True) for k, v in d.items(): acc.setdefault(k, []).append(v) return {k: torch.cat(lst, dim=0) for k, lst in acc.items()} @torch.no_grad() def build_cache( model: PolicyModel, loader: DataLoader, split_name: str, cache_mode: str, spatial_merge_size: int, image_token_id: int, n_frames: int, chunk_dir: Optional[Path] = None, chunk_size: int = 200, expected_n: Optional[int] = None, ) -> Dict[str, torch.Tensor]: """ If chunk_dir is provided, save a chunk every `chunk_size` batches and resume by scanning existing chunks. `expected_n` is the total sample count (used to sanity-check resume alignment). """ model.eval() batch_size = loader.batch_size or 1 # ── Resume detection ──────────────────────────────────────────────────── start_batch = 0 chunk_idx = 0 if chunk_dir is not None: chunk_dir.mkdir(parents=True, exist_ok=True) n_chunks, n_done = _scan_chunks(chunk_dir) if n_chunks > 0: # Each chunk (except possibly the last from a previous partial run) # contains `chunk_size * batch_size` samples. We skip exactly that # many batches so the DataLoader resumes at the next untouched one. start_batch = n_chunks * chunk_size chunk_idx = n_chunks logger.info( f" [resume] found {n_chunks} chunk(s) with {n_done} samples; " f"skipping first {start_batch} batches" ) if expected_n is not None and n_done >= expected_n: logger.info(f" [resume] chunks already cover all {expected_n} " f"samples; merging") return _merge_chunks(chunk_dir) accumulators: Dict[str, List[torch.Tensor]] = {} batches_since_flush = 0 processed_batches = 0 pbar = tqdm(loader, desc=f"cache[{cache_mode}]{split_name}", ncols=80, leave=True) for bi, batch in enumerate(pbar): if bi < start_batch: # Still need to let DataLoader workers produce the item (cheap — CPU # image load only — and keeps ordering deterministic). continue inputs = model._build_inputs(batch["images"], batch["metadata"]) feats = _extract_features_for_batch( model, inputs, cache_mode, spatial_merge_size, image_token_id, n_frames, ) for k, v in feats.items(): accumulators.setdefault(k, []).append(v) batches_since_flush += 1 processed_batches += 1 if chunk_dir is not None and batches_since_flush >= chunk_size: n_flush = _flush_chunk(accumulators, chunk_dir, chunk_idx) pbar.set_postfix_str(f"chunk={chunk_idx} +{n_flush}") accumulators = {} batches_since_flush = 0 chunk_idx += 1 # Final partial chunk if chunk_dir is not None and accumulators: n_flush = _flush_chunk(accumulators, chunk_dir, chunk_idx) logger.info(f" [chunk] final partial flushed (+{n_flush})") accumulators = {} chunk_idx += 1 # ── Assemble final cache ──────────────────────────────────────────────── if chunk_dir is not None: cache = _merge_chunks(chunk_dir) else: cache = {k: torch.cat(lst, dim=0) for k, lst in accumulators.items()} # NaN/Inf sanity for k, t in cache.items(): if t.dtype.is_floating_point: n_nan = int(torch.isnan(t).sum().item()) n_inf = int(torch.isinf(t).sum().item()) if n_nan or n_inf: logger.warning( f" {split_name}/{k}: {n_nan} NaN, {n_inf} Inf " f"(out of {t.numel()} elems)" ) n = next(iter(cache.values())).shape[0] nbytes = sum(t.element_size() * t.numel() for t in cache.values()) logger.info( f" {split_name}: cached {n} samples " f"keys={list(cache.keys())} size={nbytes / 1e9:.2f} GB" ) return cache # ───────────────────────────────────────────────────────────────────────────── # Main # ───────────────────────────────────────────────────────────────────────────── def main(): ap = argparse.ArgumentParser("make_belief_cache_v2") ap.add_argument("--sft_checkpoint", default="checkpoints/SFT/sft_v2/best") ap.add_argument("--label_dir", default="data/policy_labels") ap.add_argument("--out_dir", default="data/belief_cache_v2") ap.add_argument("--cache_mode", required=True, choices=["mean_pool", "dual_pool", "per_frame", "spatial4x4"]) ap.add_argument("--batch_size", type=int, default=4, help="Smaller for spatial4x4 (more GPU memory for hidden states)") ap.add_argument("--num_workers", type=int, default=2) ap.add_argument("--splits", nargs="+", default=["train", "val"]) ap.add_argument("--split", default=None, help="Shortcut for a single split; overrides --splits when set") ap.add_argument("--manifest", default=None, help="Explicit manifest path; overrides label_dir/{split}.json") ap.add_argument("--out", default=None, help="Explicit output .pt path; overrides out_dir/cache_mode/{split}.pt") ap.add_argument("--n_frames", type=int, default=MAX_FRAMES, help="Number of frames per clip (8, 16, 24, ...)") ap.add_argument("--sampling", default="original", choices=["original", "uniform", "last_biased", "last_2s"], help="Frame-index resampling scheme (cf. plan Stage K)") ap.add_argument("--source_filter", default="all", choices=["all", "nexar", "multisrc", "dada", "dad"], help="Restrict samples to a data source (Stage K multi-source variants)") ap.add_argument("--debug", action="store_true", help="Smoke-test on 16 samples per split") ap.add_argument("--debug_samples", type=int, default=16) ap.add_argument("--overwrite", action="store_true") ap.add_argument("--chunk_size", type=int, default=200, help="Flush a chunk to disk every N batches (resume-safe). " "0 disables chunked save.") ap.add_argument("--keep_chunks", action="store_true", help="Keep {out}.chunks/ dir after successful merge " "(default: delete on success).") args = ap.parse_args() if args.split is not None: args.splits = [args.split] odir = Path(args.out_dir) / args.cache_mode odir.mkdir(parents=True, exist_ok=True) # Monkey-patch module-level MAX_FRAMES so _extract_features_for_batch sees it # (per_frame / spatial4x4 preallocate buffers based on this). import training.Policy.policy_dataset as pds pds.MAX_FRAMES = args.n_frames logger.info("Loading SFTModel (frozen) for feature extraction...") model = PolicyModel(args.sft_checkpoint, use_bf16=True) sms = _get_spatial_merge_size(model) img_tok_id = model.sft.belief_aggregator.image_token_id if img_tok_id is None: img_tok_id = 151655 logger.info(f" spatial_merge_size = {sms}") logger.info(f" image_token_id = {img_tok_id}") logger.info(f" hidden_dim = {model.hidden_dim}") logger.info(f" cache_mode = {args.cache_mode}") logger.info(f" n_frames = {args.n_frames}") logger.info(f" sampling = {args.sampling}") logger.info(f" source_filter = {args.source_filter}") for split in args.splits: if args.manifest is not None: label_path = Path(args.manifest) else: label_path = Path(args.label_dir) / f"{split}.json" if not label_path.exists(): logger.warning(f" {label_path} not found — skipping {split}") continue if args.out is not None: out_path = Path(args.out) out_path.parent.mkdir(parents=True, exist_ok=True) else: out_path = odir / f"{split}.pt" if out_path.exists() and not args.overwrite: logger.info(f" Cache exists: {out_path} — skip (use --overwrite to rebuild)") continue ds = PolicyDataset( manifests = [label_path], split = split, debug = args.debug, debug_samples = args.debug_samples, n_frames = args.n_frames, sampling = args.sampling, source_filter = args.source_filter, ) if len(ds) == 0: logger.warning(f" {split}: dataset empty after filtering — skipping") continue loader = DataLoader( ds, batch_size = args.batch_size, shuffle = False, num_workers = args.num_workers, collate_fn = policy_collate_fn, pin_memory = True, ) chunk_dir = None if args.chunk_size > 0: chunk_dir = out_path.parent / (out_path.stem + ".chunks") cache = build_cache( model, loader, split, args.cache_mode, sms, img_tok_id, args.n_frames, chunk_dir=chunk_dir, chunk_size=args.chunk_size, expected_n=len(ds), ) # Preserve sample IDs / labels in meta for downstream alignment ids = [s.get("video_id") for s in ds.samples] labels = [int(s.get("action_label", -1)) for s in ds.samples] meta = { "schema_version": SCHEMA_VERSION, "cache_mode": args.cache_mode, "hidden_dim": model.hidden_dim, "n_frames": args.n_frames, "sampling": args.sampling, "source_filter": args.source_filter, "n_samples": int(next(iter(cache.values())).shape[0]), "spatial_merge_size": sms, "image_token_id": int(img_tok_id), "sft_checkpoint": str(args.sft_checkpoint), "label_path": str(label_path), "ids": ids, "action_labels": labels, } cache_to_save = {k: v for k, v in cache.items() if k != "__meta__"} cache_to_save["meta"] = meta tmp_path = out_path.with_suffix(out_path.suffix + ".tmp") torch.save(cache_to_save, tmp_path) tmp_path.rename(out_path) logger.info(f" Saved → {out_path}") with open(out_path.with_suffix(".meta.json"), "w") as f: meta_slim = {k: v for k, v in meta.items() if k not in ("ids", "action_labels")} meta_slim["n_ids"] = len(ids) json.dump(meta_slim, f, indent=2) if chunk_dir is not None and chunk_dir.exists() and not args.keep_chunks: import shutil shutil.rmtree(chunk_dir) logger.info(f" Removed chunk dir {chunk_dir}") logger.info("\nbelief_cache_v2 complete.") if __name__ == "__main__": main()