| |
| """Assemble the native-feature boundary-to-all conditional-probe dataset. |
| |
| Self-Forcing and HY-static already use structured temporal/spatial sampling and |
| are linked read-only. Causal-Forcing is normalized from a fresh recorder run |
| whose features are stored as [temporal, spatial, channel]. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import os |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
|
|
|
|
| ROLE_MAP = {7: "early", 14: "middle", 22: "late", 29: "final"} |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--self_dir", type=Path, required=True) |
| parser.add_argument("--causal_root", type=Path, required=True) |
| parser.add_argument("--hy_dir", type=Path, required=True) |
| parser.add_argument("--output_root", type=Path, required=True) |
| parser.add_argument("--num_prompts", type=int, default=10) |
| parser.add_argument("--chunks", type=int, default=4) |
| parser.add_argument("--steps", type=int, default=4) |
| return parser.parse_args() |
|
|
|
|
| def ensure_link(path: Path, target: Path) -> None: |
| target = target.resolve() |
| if path.is_symlink(): |
| if path.resolve() != target: |
| raise ValueError(f"Existing symlink {path} points to {path.resolve()}, not {target}") |
| return |
| if path.exists(): |
| raise FileExistsError(path) |
| path.symlink_to(target, target_is_directory=True) |
|
|
|
|
| def atomic_save(path: Path, value: dict) -> None: |
| temporary = path.with_suffix(path.suffix + ".tmp") |
| torch.save(value, temporary) |
| os.replace(temporary, path) |
|
|
|
|
| def coords_from_indices(indices: torch.Tensor) -> np.ndarray: |
| flat = indices.detach().cpu().numpy().astype(np.int64).reshape(-1) |
| plane = 30 * 52 |
| temporal = flat // plane |
| remainder = flat % plane |
| coords = np.stack([temporal, remainder // 52, remainder % 52], axis=1) |
| slots = sorted(int(value) for value in np.unique(temporal)) |
| if slots != [0, 1, 2]: |
| raise ValueError(f"Expected three temporal slots, got {slots}") |
| reference = coords[coords[:, 0] == slots[-1], 1:] |
| for slot in slots: |
| if not np.array_equal(coords[coords[:, 0] == slot, 1:], reference): |
| raise ValueError(f"Causal slot {slot} does not share the structured spatial grid") |
| return coords |
|
|
|
|
| def normalize_causal(path: Path, chunks: int, steps: int) -> dict: |
| state = torch.load(path, map_location="cpu", weights_only=False) |
| features = {} |
| for layer, role in ROLE_MAP.items(): |
| chunk_rows = [] |
| for chunk in range(chunks): |
| step_rows = [] |
| for step in range(steps): |
| value = state["features"][f"{layer}:{chunk}:{step}"] |
| step_rows.append(value.reshape(-1, value.shape[-1]).to(torch.float16)) |
| chunk_rows.append(torch.stack(step_rows, dim=0)) |
| features[role] = torch.stack(chunk_rows, dim=0).contiguous() |
| index_key = f"{next(iter(ROLE_MAP))}:0:0" |
| coords = coords_from_indices(state["feature_indices"][index_key]) |
| token_counts = {int(value.shape[2]) for value in features.values()} |
| if token_counts != {len(coords)}: |
| raise ValueError(f"Feature/coordinate mismatch: tokens={token_counts}, coords={len(coords)}") |
| return { |
| "prompt_id": int(state["prompt_id"]), |
| "prompt": state["prompt"], |
| "seed": int(state["seed"]), |
| "model_family": "causal_forcing", |
| "model_variant": "dmd4_boundary_structured", |
| "features": features, |
| "timesteps": np.asarray([1000.0, 937.5, 833.3333, 625.0], dtype=np.float32), |
| "coords": coords, |
| "grid_shape": np.asarray([3, 30, 52], dtype=np.int64), |
| } |
|
|
|
|
| def main() -> None: |
| args = parse_args() |
| output = args.output_root.resolve() |
| output.mkdir(parents=True, exist_ok=True) |
| ensure_link(output / "self_forcing", args.self_dir) |
| ensure_link(output / "hy_worldplay", args.hy_dir) |
| causal_output = output / "causal_forcing" |
| causal_output.mkdir(exist_ok=True) |
| inventory = [] |
| for prompt_id in range(args.num_prompts): |
| source = ( |
| args.causal_root |
| / "runs" |
| / f"prompt_{prompt_id:04d}" |
| / "feature_snapshots.pt" |
| ) |
| if not source.exists(): |
| raise FileNotFoundError(source) |
| item = normalize_causal(source, args.chunks, args.steps) |
| if item["prompt_id"] != prompt_id: |
| raise ValueError(f"Prompt mismatch in {source}: {item['prompt_id']}") |
| destination = causal_output / f"prompt_{prompt_id:04d}.pt" |
| atomic_save(destination, item) |
| inventory.append({ |
| "prompt_id": prompt_id, |
| "source": str(source), |
| "destination": str(destination), |
| "tokens": int(item["features"]["early"].shape[2]), |
| }) |
| manifest = { |
| "dataset_version": 1, |
| "chunk_pairing": "boundary_to_all", |
| "prompt_ids": list(range(args.num_prompts)), |
| "chunks": args.chunks, |
| "steps": args.steps, |
| "self_forcing": str(args.self_dir.resolve()), |
| "hy_worldplay": str(args.hy_dir.resolve()), |
| "causal_source": str(args.causal_root.resolve()), |
| "causal_inventory": inventory, |
| } |
| (output / "manifest.json").write_text( |
| json.dumps(manifest, indent=2, ensure_ascii=False) + "\n", |
| encoding="utf-8", |
| ) |
| print(f"[complete] {output} prompts={args.num_prompts}", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
