| """Prompt trajectories and fixed FFFF targets for Predictor-v4 Stage 2.""" |
|
|
| from __future__ import annotations |
|
|
| import json |
| from pathlib import Path |
| from typing import Any |
|
|
| import torch |
| from safetensors import safe_open |
| from torch.utils.data import Dataset |
|
|
|
|
| class PredictorV4TrajectoryDataset(Dataset): |
| """One deterministic seven-chunk trajectory per training prompt.""" |
|
|
| def __init__( |
| self, |
| cases_path: str | Path, |
| *, |
| data_root: str | Path, |
| max_cases: int | None = None, |
| ) -> None: |
| self.cases_path = Path(cases_path).resolve() |
| self.data_root = Path(data_root).resolve() |
| with self.cases_path.open("r", encoding="utf-8") as handle: |
| cases = [json.loads(line) for line in handle if line.strip()] |
| if max_cases is not None: |
| cases = cases[: int(max_cases)] |
| if not cases: |
| raise ValueError(f"No trajectory cases in {self.cases_path}") |
| for position, case in enumerate(cases): |
| if int(case["case_id"]) != position: |
| raise ValueError("Stage-2 case_id values must be dense and ordered") |
| if int(case.get("seed", -1)) != 0: |
| raise ValueError(f"case {position} does not use latent seed 0") |
| if not str(case.get("prompt", "")).strip(): |
| raise ValueError(f"case {position} has an empty prompt") |
| self.cases = cases |
|
|
| def __len__(self) -> int: |
| return len(self.cases) |
|
|
| def __getitem__(self, index: int) -> dict[str, Any]: |
| case = self.cases[index] |
| return { |
| "case_id": int(case["case_id"]), |
| "prompt": str(case["prompt"]), |
| "seed": int(case["seed"]), |
| } |
|
|
| def offline_step_path(self, case_id: int, chunk_id: int) -> Path: |
| path = ( |
| self.data_root |
| / "steps" |
| / f"case_{int(case_id):06d}" |
| / f"chunk_{int(chunk_id):02d}.safetensors" |
| ) |
| if not path.is_file(): |
| raise FileNotFoundError(path) |
| return path |
|
|
|
|
| def trajectory_collate(items: list[dict[str, Any]]) -> dict[str, Any]: |
| if len(items) != 1: |
| raise ValueError( |
| "Stage-2 rollout currently requires one trajectory per rank" |
| ) |
| return items[0] |
|
|
|
|
| def load_offline_ffff_target( |
| path: str | Path, |
| *, |
| step_id: int, |
| device: torch.device, |
| ) -> dict[str, torch.Tensor]: |
| """Load one immutable Full-trajectory hidden/flow target.""" |
|
|
| path = Path(path) |
| names = { |
| "hidden": f"step_{int(step_id)}_final_hidden", |
| "flow": f"step_{int(step_id)}_flow", |
| "timestep": f"step_{int(step_id)}_timestep", |
| } |
| with safe_open(str(path), framework="pt", device="cpu") as handle: |
| missing = set(names.values()).difference(handle.keys()) |
| if missing: |
| raise KeyError(f"{path} lacks fixed FFFF tensors {sorted(missing)}") |
| result = {key: handle.get_tensor(name) for key, name in names.items()} |
| if result["hidden"].dtype != torch.bfloat16: |
| raise TypeError(f"{path}: target hidden must be BF16") |
| if result["flow"].dtype != torch.bfloat16: |
| raise TypeError(f"{path}: target flow must be BF16") |
| if result["timestep"].dtype != torch.int64: |
| raise TypeError(f"{path}: legacy offline timestep must be INT64") |
| return { |
| key: value.to(device=device, non_blocking=True) |
| for key, value in result.items() |
| } |
|
|
|
|
| __all__ = [ |
| "PredictorV4TrajectoryDataset", |
| "load_offline_ffff_target", |
| "trajectory_collate", |
| ] |
|
|
