| """Task sources over the pretokenized caches and a seeded mixture loader.
|
|
|
| Every micro-batch holds a single task (so lengths are similar and padding is small) and is
|
| drawn from an RNG seeded by (seed, step, micro_index): sampling is a pure function of the step,
|
| which makes runs reproducible and resumable without storing data-loader state.
|
| """
|
| from __future__ import annotations
|
|
|
| import json
|
| import queue
|
| import threading
|
| from pathlib import Path
|
|
|
| import numpy as np
|
| import torch
|
|
|
| from ..tokenizers.multimodal import Sample, UnifiedVocab
|
| from ..tokenizers.text import TextTokenizer
|
| from .collator import collate
|
|
|
| BASE_TASKS = ("text", "t2i", "i2t", "i2i", "t2v", "v2t", "vcont")
|
| CHAT_TASKS = ("chat_text", "chat_mm")
|
| TASK_MODALITY = {"text": "text", "chat_text": "text", "t2i": "image", "i2t": "image", "i2i": "image",
|
| "t2v": "video", "v2t": "video", "vcont": "video", "chat_mm": "multimodal"}
|
|
|
| CHAT_PROMPTS = {
|
| "t2i": ["Generate an image of {c}.", "Draw {c}.", "Create a picture of {c}.", "Show me {c}.", "Can you make an image of {c}?"],
|
| "t2v": ["Make a video of {c}.", "Generate a video of {c}.", "Create a short clip of {c}.", "Animate {c}."],
|
| "i2t": ["Describe this image.", "What do you see in this picture?", "What is in this image?"],
|
| "v2t": ["Describe this video.", "What is happening in this video?", "What happens in this clip?"],
|
| "vcont": ["Continue the movement in this video.", "Continue this video.", "What happens next? Show me."],
|
| "i2v": ["Animate this image: {c}.", "Bring this picture to life: {c}."],
|
| }
|
| CHAT_KINDS = ("t2i", "t2v", "i2t", "v2t", "vcont", "i2i", "i2v")
|
|
|
|
|
| def _load_media(cache_dir: Path, name: str, split: str, vocab: UnifiedVocab, keys: tuple[str, ...]) -> tuple[dict, dict] | None:
|
| """Load <name>_<split>.npz/.json, dropping rows whose codes are out of range (corruption guard)."""
|
| npz, js = cache_dir / f"{name}_{split}.npz", cache_dir / f"{name}_{split}.json"
|
| if not npz.exists():
|
| return None
|
| try:
|
| with np.load(npz) as z:
|
| arrs = {k: z[k] for k in keys}
|
| meta = json.loads(js.read_text())
|
| except (OSError, ValueError, KeyError) as e:
|
| raise RuntimeError(f"cache {npz} is unreadable ({e}); regenerate with scripts/generate_dataset.py") from e
|
| ok = np.ones(len(arrs[keys[0]]), dtype=bool)
|
| for a in arrs.values():
|
| ok &= (a.reshape(len(a), -1) < vocab.visual_size).all(1)
|
| if not ok.all():
|
| print(f"[data] {npz.name}: dropping {int((~ok).sum())} corrupted rows")
|
| keep = np.nonzero(ok)[0]
|
| return {k: a[keep] for k, a in arrs.items()}, {k: [v[i] for i in keep] for k, v in meta.items()}
|
|
|
|
|
| class TaskSources:
|
| """Builds one training Sample for any task from the caches of a given split."""
|
|
|
| def __init__(self, cache_dir: str | Path, split: str, vocab: UnifiedVocab, tok: TextTokenizer, seq_len: int, caption_dropout: float = 0.1):
|
| cache_dir = Path(cache_dir)
|
| self.vocab, self.tok, self.seq_len = vocab, tok, seq_len
|
|
|
| self.caption_dropout = caption_dropout if split == "train" else 0.0
|
| tp = cache_dir / f"text_{split}.bin"
|
| self.text = np.memmap(tp, dtype=np.uint16, mode="r") if tp.exists() else None
|
| self.images = _load_media(cache_dir, "images", split, vocab, ("codes",))
|
| self.edits = _load_media(cache_dir, "edits", split, vocab, ("src", "tgt"))
|
| self.videos = _load_media(cache_dir, "videos", split, vocab, ("codes",))
|
| cp = cache_dir / f"chat_{split}.json"
|
| self.chat = json.loads(cp.read_text()) if cp.exists() else None
|
|
|
| cap_idx = lambda m: np.array([i for i, c in enumerate(m[1]["captions"]) if c], dtype=np.int64) if m else np.zeros(0, np.int64)
|
| self.img_cap, self.vid_cap = cap_idx(self.images), cap_idx(self.videos)
|
| need = {"t2i": len(self.img_cap), "i2t": len(self.img_cap), "t2v": len(self.vid_cap), "v2t": len(self.vid_cap),
|
| "vcont": self.videos is not None, "i2i": self.edits is not None, "i2v": len(self.vid_cap)}
|
| self.chat_kinds = [k for k in CHAT_KINDS if need[k]]
|
|
|
| def available(self) -> list[str]:
|
| have = {"text": self.text is not None and len(self.text) > self.seq_len + 1, "chat_text": bool(self.chat),
|
| "t2i": len(self.img_cap) > 0, "i2t": len(self.img_cap) > 0, "i2i": self.edits is not None,
|
| "t2v": len(self.vid_cap) > 0, "v2t": len(self.vid_cap) > 0, "vcont": self.videos is not None,
|
| "chat_mm": bool(self.chat_kinds)}
|
| return [t for t, ok in have.items() if ok]
|
|
|
| def enc(self, s: str) -> list[int]:
|
| return self.tok.encode(s)
|
|
|
| def _prompt(self, caps: list[str], rng: np.random.Generator) -> list[int]:
|
| return [] if rng.random() < self.caption_dropout else self.enc(str(rng.choice(caps)))
|
|
|
| def sample(self, task: str, rng: np.random.Generator) -> Sample:
|
| v = self.vocab
|
| if task == "text":
|
| off = int(rng.integers(0, len(self.text) - self.seq_len - 1))
|
| return Sample(np.asarray(self.text[off : off + self.seq_len + 1], dtype=np.int32), 0)
|
| if task in ("t2i", "i2t"):
|
| (arr, meta), i = self.images, int(self.img_cap[rng.integers(len(self.img_cap))])
|
| caps = meta["captions"][i]
|
| if task == "t2i":
|
| return v.build("t2i", self._prompt(caps, rng), image=arr["codes"][i])
|
| return v.build("i2t", self.enc(caps[0]), image=arr["codes"][i])
|
| if task == "i2i":
|
| (arr, meta), i = self.edits, int(rng.integers(len(self.edits[0]["src"])))
|
| return v.build("i2i", self.enc(meta["instructions"][i]), image=arr["src"][i], image2=arr["tgt"][i])
|
| if task in ("t2v", "v2t", "vcont"):
|
| pool = self.vid_cap if task != "vcont" else np.arange(len(self.videos[0]["codes"]))
|
| (arr, meta), i = self.videos, int(pool[rng.integers(len(pool))])
|
| caps = meta["captions"][i]
|
| if task == "t2v":
|
| return v.build("t2v", self._prompt(caps, rng), video=arr["codes"][i])
|
| if task == "v2t":
|
| return v.build("v2t", self.enc(caps[0]), video=arr["codes"][i])
|
| return v.build("vcont", video=arr["codes"][i])
|
| if task == "chat_text":
|
| r = self.chat[int(rng.integers(len(self.chat)))]
|
| s = v.chat(self.enc(r["user"]), self.enc(r["assistant"]))
|
| return Sample(s.tokens[: self.seq_len + 1], s.target_start)
|
| if task == "chat_mm":
|
| return self.chat_mm(str(rng.choice(self.chat_kinds)), rng)
|
| raise KeyError(task)
|
|
|
| def chat_mm(self, kind: str, rng: np.random.Generator) -> Sample:
|
| """Instruction-formatted multimodal conversation built from the media caches."""
|
| v, enc = self.vocab, self.enc
|
| pick = lambda xs: str(rng.choice(xs))
|
| if kind in ("t2i", "i2t"):
|
| (arr, meta), i = self.images, int(self.img_cap[rng.integers(len(self.img_cap))])
|
| cap, img = meta["captions"][i][0], arr["codes"][i]
|
| if kind == "t2i":
|
| return v.chat(enc(pick(CHAT_PROMPTS["t2i"]).format(c=pick(meta["captions"][i]))), enc(f"Here is an image of {cap}.") + v.image_tokens(img))
|
| return v.chat(v.image_tokens(img) + enc(" " + pick(CHAT_PROMPTS["i2t"])), enc(f"The image shows {cap}."))
|
| if kind == "i2i":
|
| (arr, meta), i = self.edits, int(rng.integers(len(self.edits[0]["src"])))
|
| instr = meta["instructions"][i]
|
| return v.chat(v.image_tokens(arr["src"][i]) + enc(" " + instr[0].upper() + instr[1:] + "."), enc("Here is the edited image.") + v.image_tokens(arr["tgt"][i]))
|
| pool = self.vid_cap if kind != "vcont" else np.arange(len(self.videos[0]["codes"]))
|
| (arr, meta), i = self.videos, int(pool[rng.integers(len(pool))])
|
| cap, vid = (meta["captions"][i] or [""])[0], arr["codes"][i]
|
| if kind == "t2v":
|
| return v.chat(enc(pick(CHAT_PROMPTS["t2v"]).format(c=pick(meta["captions"][i]))), enc(f"Here is a video of {cap}.") + v.video_tokens(vid))
|
| if kind == "v2t":
|
| return v.chat(v.video_tokens(vid) + enc(" " + pick(CHAT_PROMPTS["v2t"])), enc(f"The video shows {cap}."))
|
| if kind == "vcont":
|
| k = len(vid) // 2
|
| return v.chat(v.video_tokens(vid[:k]) + enc(" " + pick(CHAT_PROMPTS["vcont"])), enc("Here is the continuation.") + v.video_tokens(vid[k:], first_index=k))
|
|
|
| return v.chat(v.image_tokens(vid[0]) + enc(" " + pick(CHAT_PROMPTS["i2v"]).format(c=cap)), enc("Here is the animation.") + v.video_tokens(vid[1:], first_index=1))
|
|
|
|
|
| def split_task(name: str) -> tuple[str, str]:
|
| """'user/t2i' -> ('user', 't2i'); 't2i' -> ('', 't2i') (source prefix selects the cache)."""
|
| src, _, task = name.rpartition("/")
|
| return src, task
|
|
|
|
|
| def available_tasks(sources: dict[str, TaskSources]) -> list[str]:
|
| return [f"{k}/{t}" if k else t for k, src in sources.items() for t in src.available()]
|
|
|
|
|
| def compile_bucket_length(length: int, context: int) -> int:
|
| """At most four fixed sequence sizes at context 1024; no content truncation."""
|
| if not 0 < length <= context:
|
| raise ValueError("sample length is outside the model context")
|
| return min(context, max(128, 1 << (length - 1).bit_length()))
|
|
|
|
|
| class MixtureLoader:
|
| """Yields (task, input_ids, labels) micro-batches; one task per micro-batch, ~`micro_tokens` tokens each.
|
|
|
| `sources` maps a source name to its TaskSources ('' = built-in cache, 'user' = imported data);
|
| mixture keys address them as 't2i' or 'user/t2i', so every source/task share is an explicit weight."""
|
|
|
| def __init__(self, sources: dict[str, TaskSources] | TaskSources, weights: dict[str, float], micro_tokens: int, seed: int, prefetch: int = 6,
|
| compile_buckets: bool = False):
|
| sources = sources if isinstance(sources, dict) else {"": sources}
|
| missing = [t for t, w in weights.items() if w > 0 and t not in available_tasks(sources)]
|
| if missing:
|
| raise ValueError(f"mixture uses tasks with no cache: {missing}; available: {available_tasks(sources)}")
|
| self.sources, self.seed, self.micro_tokens = sources, seed, micro_tokens
|
| self.vocab = next(iter(sources.values())).vocab
|
| self.tasks = [t for t, w in weights.items() if w > 0]
|
| p = np.array([weights[t] for t in self.tasks], dtype=np.float64)
|
| self.p = p / p.sum()
|
| self.prefetch = prefetch
|
| self.compile_buckets = compile_buckets
|
|
|
| def batch(self, step: int, micro: int) -> tuple[str, torch.Tensor, torch.Tensor]:
|
| rng = np.random.default_rng([self.seed, step, micro])
|
| task = self.tasks[int(rng.choice(len(self.tasks), p=self.p))]
|
| samples: list[Sample] = []
|
| longest = 0
|
| src, base = split_task(task)
|
| while True:
|
| s = self.sources[src].sample(base, rng)
|
| proposed = max(longest, len(s.tokens) - 1 if self.compile_buckets else len(s.tokens))
|
| padded = compile_bucket_length(proposed, self.sources[src].seq_len) if self.compile_buckets else proposed
|
| if self.compile_buckets and padded > self.micro_tokens:
|
| raise ValueError("compiled batch token budget is smaller than its padded context")
|
| if samples and (len(samples) + 1) * padded > self.micro_tokens:
|
| break
|
| longest = proposed
|
| samples.append(s)
|
| if self.compile_buckets:
|
| length = compile_bucket_length(longest, self.sources[src].seq_len)
|
|
|
| inp, lab = collate(samples, self.vocab.pad, target_length=length,
|
| batch_size=self.micro_tokens // length)
|
| else:
|
| inp, lab = collate(samples, self.vocab.pad)
|
| return task, inp, lab
|
|
|
| def iterate(self, start_step: int, grad_accum: int):
|
| """Background-thread prefetching iterator of (step, micro, task, inp, lab)."""
|
| q: queue.Queue = queue.Queue(maxsize=self.prefetch)
|
| stop = threading.Event()
|
|
|
| def work():
|
| step = start_step
|
| while not stop.is_set():
|
| for m in range(grad_accum):
|
| item = (step, m, *self.batch(step, m))
|
| while not stop.is_set():
|
| try:
|
| q.put(item, timeout=0.5)
|
| break
|
| except queue.Full:
|
| continue
|
| step += 1
|
|
|
| th = threading.Thread(target=work, daemon=True)
|
| th.start()
|
| try:
|
| while True:
|
| yield q.get()
|
| finally:
|
| stop.set()
|
|
|
|
|
| def fixed_eval_batches(sources: dict[str, TaskSources] | TaskSources, tasks: list[str], micro_tokens: int, n_batches: int, seed: int = 999) -> dict[str, list]:
|
| """Deterministic held-out batches per task (same every evaluation)."""
|
| sources = sources if isinstance(sources, dict) else {"": sources}
|
| out = {}
|
| for ti, task in enumerate(tasks):
|
| if task not in available_tasks(sources):
|
| continue
|
| ml = MixtureLoader(sources, {task: 1.0}, micro_tokens, seed)
|
| out[task] = [ml.batch(10_000 + ti, i)[1:] for i in range(n_batches)]
|
| return out
|
|
|
