reader.py

1"""Bounded-index reader: validate every shard/count before yielding any examples."""
2import hashlib
3import json
4from pathlib import Path
5import re
6import struct
7from .tokenizer import State, VERSION, windows
8from .limits import MAX_SHARDS, MAX_MANIFEST_BYTES
9from .convert import canonical
10
11SPLITS = ('train', 'validation', 'test')
12
13
14def _integer(value, minimum=0):
15    if type(value) is not int or value < minimum:
16        raise ValueError('invalid manifest count')
17    return value
18
19
20def _read_shard(root, shard):
21    contents = {}
22    for name, expected in shard['files'].items():
23        path = root/name
24        if path.stat().st_size != expected['bytes']:
25            raise ValueError('shard size mismatch')
26        with path.open('rb') as stream:
27            blob = stream.read(expected['bytes'] + 1)
28        if len(blob) != expected['bytes'] or hashlib.sha256(blob).hexdigest() != expected['sha256']:
29            raise ValueError('shard integrity failure')
30        contents[name.split('.', 1)[1]] = blob
31    tokens = contents['tokens.u8']
32    raw_offsets = contents['offsets.u64le']
33    if len(tokens) != shard['tokens'] or len(raw_offsets) != 8*(shard['games']+1):
34        raise ValueError('shard token/game count mismatch')
35    offsets = struct.unpack('<' + 'Q'*(len(raw_offsets)//8), raw_offsets)
36    metadata = [json.loads(line) for line in contents['metadata.jsonl'].splitlines()]
37    if (len(metadata) != shard['games'] or offsets[0] != 0 or offsets[-1] != len(tokens)
38            or any(a >= b for a,b in zip(offsets, offsets[1:]))):
39        raise ValueError('invalid game boundaries or metadata count')
40    return tokens, offsets, metadata
41
42
43def _manifest(root):
44    with (root/'manifest.json').open('rb') as stream:
45        raw = stream.read(MAX_MANIFEST_BYTES+1)
46    if len(raw) > MAX_MANIFEST_BYTES:
47        raise ValueError('manifest exceeds v1 byte limit')
48    manifest = json.loads(raw)
49    fingerprint = manifest.get('fingerprint')
50    if (not isinstance(fingerprint, str) or not re.fullmatch('[0-9a-f]{64}', fingerprint)
51            or hashlib.sha256(canonical(manifest['provenance'])).hexdigest() != fingerprint):
52        raise ValueError('provenance fingerprint mismatch or invalid format')
53    if (manifest['format'] != 'chess-dataset-v1' or manifest['provenance']['contract'] != VERSION
54            or manifest['provenance']['config']['count_only']):
55        raise ValueError('unsupported format, token contract or count-only output')
56    limit = manifest['provenance']['config'].get('max_shards', MAX_SHARDS)
57    if not 1 <= _integer(limit, 1) <= MAX_SHARDS or len(manifest['shards']) > limit:
58        raise ValueError('manifest shard limit exceeded')
59    totals, seen = {}, set()
60    for shard in manifest['shards']:
61        split = shard['split']
62        if split not in SPLITS:
63            raise ValueError('unknown shard split')
64        n, t = _integer(shard['games'], 1), _integer(shard['tokens'], 2)
65        stems, suffixes = set(), set()
66        if len(shard['files']) != 3:
67            raise ValueError('expected three shard files')
68        for name, expected in shard['files'].items():
69            match = re.fullmatch(r'(train|validation|test)-[0-9]{6}\.(tokens\.u8|offsets\.u64le|metadata\.jsonl)', name)
70            if not match or match[1] != split or name in seen:
71                raise ValueError('duplicate or invalid shard file')
72            seen.add(name); stems.add(name.split('.', 1)[0]); suffixes.add(match[2])
73            _integer(expected['bytes'], 1)
74            if not re.fullmatch('[0-9a-f]{64}', expected['sha256']):
75                raise ValueError('invalid shard hash')
76        if len(stems) != 1 or suffixes != {'tokens.u8','offsets.u64le','metadata.jsonl'}:
77            raise ValueError('inconsistent shard file set')
78        total = totals.setdefault(split, {'games':0, 'tokens':0})
79        total['games'] += n; total['tokens'] += t
80    for split, total in manifest['splits'].items():
81        if split not in SPLITS or set(total) != {'games','tokens'}:
82            raise ValueError('invalid split counts')
83        _integer(total['games']); _integer(total['tokens'])
84    for count in manifest['counts'].values():
85        _integer(count)
86    if (manifest['splits'] != totals
87            or _integer(manifest['accepted_tokens']) != sum(s['tokens'] for s in totals.values())
88            or manifest['counts'].get('accepted',0) != sum(s['games'] for s in totals.values())):
89        raise ValueError('aggregate manifest count mismatch')
90    # Full preflight validates counts against actual data in ALL splits. It retains
91    # only one shard at a time; selected shards are read again when yielding.
92    for shard in manifest['shards']:
93        _read_shard(root, shard)
94    return manifest
95
96
97def games(directory, split='train'):
98    root = Path(directory)
99    if split not in SPLITS:
100        raise ValueError('unknown split')
101    try:
102        manifest = _manifest(root)
103    except (KeyError, TypeError, AttributeError, json.JSONDecodeError, UnicodeError, struct.error) as error:
104        raise ValueError('invalid manifest/shard schema') from error
105    for shard in manifest['shards']:
106        if shard['split'] != split:
107            continue
108        tokens, offsets, metadata = _read_shard(root, shard)
109        for start, end, meta in zip(offsets, offsets[1:], metadata):
110            sequence = list(tokens[start:end])
111            state = State()
112            for token in sequence:
113                state.feed(token)
114            state.finish()
115            yield sequence, meta
116
117
118def training_windows(directory, split='train', length=1024, stride=None):
119    """Metadata stays outside model inputs; each example belongs to exactly one game."""
120    for sequence, metadata in games(directory, split):
121        yield from windows(sequence, length, stride)
122