File size: 7,587 Bytes
d911efa | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 | #!/usr/bin/env python3
"""Lazy container-backed loader for the multi-caption S1->S10 curriculum."""
from __future__ import annotations
import hashlib
import json
import os
from pathlib import Path
import random
import numpy as np
from timed_prompt import repair_prompt
TAR_MEMBER_MAP = "/e/scratch/reformo/schuhmann1_moss/out/m2_600m_ladder/tar_member_map/manifest.json"
FAMILIES = ("general_script", "bude_v1", "bude_v1_1", "timbre", "voice_tagging", "gemma", "transcript_only")
DESC_DTYPE = np.dtype([
("base", "<u4"), ("reference", "<u4"), ("annotation", "<u4"),
("family", "u1"), ("surface", "u1"), ("mode", "u1"), ("reserved", "u1"),
])
MISSING = np.iinfo(np.uint32).max
class JsonlOffsets:
def __init__(self, records: str | Path, offsets: str | Path):
self.path = Path(records)
self.offsets = np.load(offsets, mmap_mode="r", allow_pickle=False)
self.fd = os.open(self.path, os.O_RDONLY)
assert self.offsets.dtype == np.int64 and int(self.offsets[0]) == 0
assert int(self.offsets[-1]) == self.path.stat().st_size
def row(self, index: int) -> dict:
start, end = int(self.offsets[index]), int(self.offsets[index + 1])
return json.loads(os.pread(self.fd, end - start, start))
class PhysicalBase(JsonlOffsets):
def __init__(self, manifest: str | Path):
path = Path(manifest)
spec = json.loads(path.read_text())
root = path.parent
super().__init__(root / spec["records"], root / spec["offsets"])
def shuffled(tags: list[str], uid: str, stage: str, presentation: int, family: str) -> list[str]:
result = list(tags)
seed = int.from_bytes(hashlib.blake2b(
f"20260923|{stage}|{uid}|{family}|{presentation}".encode(), digest_size=8).digest(), "little")
random.Random(seed).shuffle(result)
return result
def caption_prompt(caption: str, transcript: str) -> str:
caption = " ".join(str(caption).replace('"', "'").split())
assert caption
return f'CAPTION: {caption}\nTRANSCRIPT: "{transcript}"'
class CaptionCurriculumRecords:
def __init__(self, manifest: str | Path):
self.manifest_path = Path(manifest)
self.spec = json.loads(self.manifest_path.read_text())
assert self.spec["status"] == "complete"
assert self.spec["format"] == "m2-caption-curriculum-v1"
root = self.manifest_path.parent
self.phase = self.spec["stage"]
self.descriptors = np.load(root / self.spec["descriptors"], mmap_mode="r", allow_pickle=False)
self.order = np.load(root / self.spec["order"], mmap_mode="r", allow_pickle=False)
assert self.descriptors.dtype == DESC_DTYPE
assert len(self.descriptors) == len(self.order) == self.spec["presentations"]
self.base = PhysicalBase(self.spec["base_manifest"])
assets = json.loads(Path(self.spec["assets_manifest"]).read_text())
assert assets["status"] == "complete"
asset_root = Path(self.spec["assets_manifest"]).parent
self.annotations = JsonlOffsets(asset_root / assets["annotations"]["records"],
asset_root / assets["annotations"]["offsets"])
self.references = JsonlOffsets(asset_root / assets["references"]["records"],
asset_root / assets["references"]["offsets"])
import sys
sys.path.insert(0, "/e/scratch/reformo/schuhmann1_moss/code/m2_20k_score")
from dataset import ShardReader
from tar_member_map import TarMemberMap
class MappedShardReader(ShardReader):
def __init__(self):
super().__init__(capacity=16)
self.packed_members = TarMemberMap(TAR_MEMBER_MAP)
def member(self, path, name):
offset, size = self.packed_members.locate(path, name)
return self.bytes_at(path, offset, size)
self.reader = MappedShardReader()
def __len__(self):
return len(self.order)
def example(self, index: int, packer):
descriptor = self.descriptors[int(self.order[index])]
base = self.base.row(int(descriptor["base"]))
uid = base["uid"]
family = FAMILIES[int(descriptor["family"])]
mode = "reference" if int(descriptor["mode"]) else "instruction"
meta = dict(self.reader.record(base["record_locator"]))
assert meta["uid"] == uid
meta["frames"] = int(meta.get("moss_frames", base["target_locator"]["full_frames"]))
target = self.reader.codes(base["target_locator"])
assert len(target) == meta["frames"]
transcript = str(meta["text"])
annotation = None
if int(descriptor["annotation"]) != MISSING:
annotation = self.annotations.row(int(descriptor["annotation"]))
assert annotation["uid"] == uid
prompt_info = {"timed": False, "eligible": False, "duration_tags": 0,
"format_id": "m2-multiformat-caption-v1-20260923"}
form = "B"
if family == "general_script":
meta["prompt"], prompt_info = repair_prompt(
meta, caption=None, form="A", mode=mode,
presentation=int(index), phase=self.phase)
form = "A"
elif family == "bude_v1":
assert base.get("bude_caption")
meta["prompt"] = caption_prompt(base["bude_caption"], transcript)
elif family == "bude_v1_1":
meta["prompt"] = caption_prompt(annotation["bude"], transcript)
elif family in ("timbre", "voice_tagging"):
prefix = "timbre" if family == "timbre" else "voice"
tags = shuffled(annotation[prefix + "_tags"], uid, self.phase, int(index), family)
prose = annotation[prefix + "_prose"]
surface = int(descriptor["surface"])
if surface == 0:
assert tags
body = ", ".join(tags)
elif surface == 1:
body = prose
else:
assert tags and prose
body = "Tags: " + ", ".join(tags) + ". Caption: " + prose
meta["prompt"] = caption_prompt(body, transcript)
elif family == "gemma":
meta["prompt"] = caption_prompt(annotation["gemma_caption"], transcript)
elif family == "transcript_only":
meta["prompt"] = f'TRANSCRIPT: "{transcript}"'
form = "A"
else:
raise AssertionError(family)
reference = None
if mode == "reference":
assert int(descriptor["reference"]) != MISSING
reference_row = self.references.row(int(descriptor["reference"]))
assert reference_row["uid"] == uid and reference_row["reference_uid"] != uid
reference = self.reader.codes(reference_row["locator"])
assert 0 < len(reference) <= 37
result = packer.pack_mode(meta, target, mode, reference)
result["accounting"] = {
"uid": uid, "mode": mode, "form": form, "family": family,
"frames": len(target), "target_audio_hours": float(meta["dur_s"]) / 3600,
"reference_frames": len(reference) if reference is not None else 0,
"prompt_timed": bool(prompt_info["timed"]),
"prompt_eligible": bool(prompt_info["eligible"]),
"duration_tags": int(prompt_info["duration_tags"]),
"prompt_format_id": (prompt_info["format_id"] if family == "general_script"
else "m2-multiformat-caption-v1-20260923"),
}
return result
|