File size: 5,302 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 | #!/usr/bin/env python3
"""M2-600M cascade adapter with deterministic timed-prompt repair.
Audio/code targets and manifests remain immutable. Every presentation passes
through ``timed_prompt.repair_prompt``; repair facts are returned in accounting
so the trainer can prove the observed timed/untimed mixture.
"""
import json
import os
import numpy as np
from timed_prompt import repair_prompt
MODES = ('instruction', 'reference')
TAR_MEMBER_MAP = '/e/scratch/reformo/schuhmann1_moss/out/m2_600m_ladder/tar_member_map/manifest.json'
class CascadeRecords:
def __init__(self, manifest):
from pathlib import Path
self.manifest_path = Path(manifest)
spec = json.loads(self.manifest_path.read_text())
assert spec['status'] == 'complete' and spec['presentations'] > 0
root = self.manifest_path.parent
self.spec = spec
self.phase = str(spec.get('phase') or spec.get('ladder_stage') or self.manifest_path.parent.name)
self.offsets = np.load(root / spec['offsets'], mmap_mode='r', allow_pickle=False)
assert len(self.offsets) == spec['presentations'] + 1
assert self.offsets.dtype == np.int64 and int(self.offsets[0]) == 0
self.order = np.load(root / spec['order'], mmap_mode='r', allow_pickle=False)
assert len(self.order) == spec['presentations']
assert self.order.dtype in (np.dtype('uint32'), np.dtype('int64'))
# Full uniqueness is established once by the manifest verifier. Never
# materialize 48M Python ints/sets independently on all 128 DDP ranks.
assert int(self.order.min()) == 0 and int(self.order.max()) == spec['presentations'] - 1
self.path = root / spec['records']
self.fd = os.open(self.path, os.O_RDONLY)
assert int(self.offsets[-1]) == self.path.stat().st_size
if 'records_bytes' in spec:
assert self.path.stat().st_size == spec['records_bytes']
assert (root / spec['offsets']).stat().st_size == spec['offsets_bytes']
assert (root / spec['order']).stat().st_size == spec['order_bytes']
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):
# ShardReader uses the same bounded capacity for Parquet
# reference tables as for TAR member dictionaries. TAR
# dictionaries are overridden below, but Parquet tables still
# need a non-zero LRU (capacity=0 caused an immediate eviction
# followed by move_to_end(KeyError) in S1 smoke 1892952).
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 row(self, index):
j = int(self.order[index])
start, end = int(self.offsets[j]), int(self.offsets[j + 1])
return json.loads(os.pread(self.fd, end - start, start))
def example(self, index, packer):
row = self.row(index)
mode = row['mode']
form = row.get('form', 'A')
assert mode in MODES, row['uid']
assert form in ('A', 'B'), row['uid']
meta = self.reader.record(row['record_locator'])
assert meta['uid'] == row['uid']
if 'moss_frames' in meta:
meta['frames'] = int(meta['moss_frames'])
else:
# Sidecar element JSONs carry no moss_frames key (pre-upload state);
# frame count comes from the manifest locator (validated in codes()).
meta['frames'] = int(row['target_locator']['full_frames'])
codes = self.reader.codes(row['target_locator'])
assert len(codes) == meta['frames']
meta = dict(meta)
meta['prompt'], prompt_info = repair_prompt(
meta, caption=row.get('bude_caption'), form=form, mode=mode,
presentation=int(index), phase=self.phase)
reference = None
if mode == 'reference':
assert row.get('reference_uid') and row['reference_uid'] != row['uid']
assert row.get('reference_locator')
reference = self.reader.codes(row['reference_locator'])
assert 0 < len(reference) <= 37
result = packer.pack_mode(meta, codes, mode, reference)
result['accounting'] = {'uid': row['uid'], 'mode': mode, 'form': form,
'frames': len(codes),
'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']}
return result
|