"""Epoch-specific virtual windows; chat quota counts actual cropped targets.""" import json from pathlib import Path import numpy as np from vimeml.training.data import SentenceWindowDataset, file_sha from vimeml.training.data_v2 import PrefixCropWindowDataset ROOT = Path(__file__).resolve().parents[3] class JoinedPrefixWindows: """Virtual concatenation of compatible frozen stores; no token copies.""" epoch_indexed = True def __init__(self, bases, primary, source_maps): self.bases, self.source_maps = bases, source_maps self.primary = primary self.ends = np.cumsum([len(base) for base in bases], dtype=np.int64) self.starts = np.r_[0, self.ends[:-1]] selected = bases[primary] for key in ("token_dir", "context_length", "split", "seed", "probability", "min_remaining_tokens"): setattr(self, key, getattr(selected, key)) self._store = None self.epoch = 0 def __len__(self): return int(self.ends[-1]) def set_epoch(self, epoch): self.epoch = epoch for base in self.bases: base.set_epoch(epoch) def _open(self): self.bases[self.primary]._open() self._store = self.bases[self.primary]._store def __getitem__(self, key): epoch, index = key if isinstance(key, tuple) else (self.epoch, key) index = int(index) if not 0 <= index < len(self): raise IndexError("Joined window outside frozen stores") source = int(np.searchsorted(self.ends, index, side="right")) sample = self.bases[source][(epoch, int(index - self.starts[source]))] return {**sample, "source_id": self.source_maps[source][sample["source_id"]], "sample_index": index} def window_lengths(self, indices): return self._lengths(indices, cropped=True) def uncropped_window_lengths(self, indices): return self._lengths(indices, cropped=False) def _lengths(self, indices, *, cropped): indices = np.asarray(indices, dtype=np.int64) if np.any(indices < 0) or np.any(indices >= len(self)): raise IndexError("Joined window outside frozen stores") lengths = np.empty(len(indices), dtype=np.int64) for base, start, end in zip(self.bases, self.starts, self.ends): positions = np.flatnonzero((indices >= start) & (indices < end)) if len(positions): local = indices[positions] - start lengths[positions] = (base.window_lengths(local) if cropped else SentenceWindowDataset.window_lengths(base, local)) return lengths def close(self): for base in self.bases: base.close() self._store = None def __getstate__(self): return {**self.__dict__, "_store": None} class MixtureWindowDataset: epoch_indexed = True def __init__(self, base, directory): self.base = base self.directory = Path(directory) self.manifest = json.loads((self.directory / "manifest.json").read_text(encoding="utf-8")) if self.manifest.get("status") != "complete" or self.manifest.get("format") != "vimeml_token_mixture_v1": raise ValueError("Incomplete or unsupported mixture.") request = self.manifest["request"] if (request["token_manifest_sha256"] != file_sha(base.token_dir / "manifest.json") or request["window_manifest_sha256"] != file_sha(ROOT / request["index_dir"].replace("\\", "/") / "manifest.json")): raise ValueError("Mixture belongs to different token data/indexes.") if (request["seed"], request["crop_probability"], request["crop_minimum"]) != ( base.seed, base.probability, base.min_remaining_tokens): raise ValueError("Mixture crop settings differ from training.") self.context_length = base.context_length self.split = "train" self.epoch = 0 self._mapping = self._mapping_epoch = self._store = None def _map(self, epoch): if not 0 <= epoch < len(self.manifest["epochs"]): raise ValueError("Epoch is outside the frozen mixture plan.") if self._mapping_epoch != epoch: self._close_map() record = self.manifest["epochs"][epoch] path = self.directory / record["file"] if path.stat().st_size != record["bytes"]: raise ValueError("Mixture mapping size differs from manifest.") self._mapping = np.load(path, mmap_mode="r", allow_pickle=False) if self._mapping.dtype != np.dtype("uint32") or self._mapping.shape != (record["windows"],): self._close_map() raise ValueError("Invalid mixture mapping dtype/shape.") self._mapping_epoch = epoch return self._mapping def __len__(self): return self.manifest["epochs"][self.epoch]["windows"] def set_epoch(self, epoch): if not 0 <= epoch < len(self.manifest["epochs"]): raise ValueError("Epoch is outside the frozen mixture plan.") self.epoch = epoch self.base.set_epoch(epoch) def _open(self): self.base._open() self._store = self.base._store def __getitem__(self, key): epoch, index = key if isinstance(key, tuple) else (self.epoch, key) physical = int(self._map(epoch)[index]) return self.base[(epoch, physical)] def window_lengths(self, indices): self.base.set_epoch(self.epoch) return self.base.window_lengths(self._map(self.epoch)[np.asarray(indices, dtype=np.int64)]) def _close_map(self): if self._mapping is not None: self._mapping._mmap.close() self._mapping = self._mapping_epoch = None def close(self): self._close_map() self.base.close() self._store = None def __getstate__(self): state = self.__dict__.copy() state["_mapping"] = state["_mapping_epoch"] = state["_store"] = None return state def mixture_dataset(config, root): settings = config["training"] directory = root / config["data_mixture"]["directory"] manifest = json.loads((directory / "manifest.json").read_text(encoding="utf-8")) records = manifest["request"].get("joined_stores") bases = [] try: if records: maps = [] for record in records: if (file_sha(root / record["token_dir"] / "manifest.json") != record["token_manifest_sha256"] or file_sha(root / record["index_dir"] / "manifest.json") != record["window_manifest_sha256"]): raise ValueError("Joined store identity differs from frozen mixture") bases.append(PrefixCropWindowDataset(root / record["token_dir"], root / record["index_dir"], seed=settings["seed"], probability=settings.get("prefix_crop_probability", 0.0), min_remaining_tokens=settings.get("prefix_crop_min_remaining_tokens", 8))) maps.append({int(k): v for k, v in record["source_map"].items()}) base = JoinedPrefixWindows(bases, manifest["request"]["primary_store"], maps) if str(base.token_dir.resolve()) != str((root / config["token_dir"]).resolve()): raise ValueError("Joined primary store must match configuration") else: base = PrefixCropWindowDataset(root / config["token_dir"], root / config["index_dir"], seed=settings["seed"], probability=settings.get("prefix_crop_probability", 0.0), min_remaining_tokens=settings.get("prefix_crop_min_remaining_tokens", 8)) bases.append(base) return MixtureWindowDataset(base, directory) except BaseException: for base in bases: base.close() raise