File size: 6,414 Bytes
99297d7 | 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 | """V4.3 mixture support: two web stores with keep masks plus the chat bank."""
import json
from pathlib import Path
import numpy as np
from vimeml.training.data import file_sha, write_json
from vimeml.training.data_v2 import PrefixCropWindowDataset
from vimeml.training.data_mixture import JoinedPrefixWindows
from vimeml.posttraining.streams import ConversionWindows
ROOT = Path(__file__).resolve().parents[3]
def window_sentences(dataset, ids):
position = np.searchsorted(dataset.end, ids, side='right')
previous = np.zeros(len(ids), dtype=np.int64)
present = position > 0
previous[present] = dataset.end[position[present] - 1] - dataset.sentences[position[present] - 1] - 1
sentences = ids - previous
candidates = np.flatnonzero(position < len(dataset.end))
extra = candidates[ids[candidates] >= dataset.first[position[candidates]]]
sentences[extra] = dataset.sentences[position[extra]]
return sentences
def sentence_keys(paths):
"""Canonical text keys of every sentence in the listed TXT files."""
import hashlib
import unicodedata
keys = set()
for path in paths:
with Path(path).open(encoding='utf-8') as stream:
for line in stream:
text = line.rstrip('\r\n')
if text:
keys.add(hashlib.blake2b(' '.join(unicodedata.normalize('NFKC', text).split()).encode(),
digest_size=16).digest())
return keys
def prepare_maps_v43(directory, root=ROOT):
"""Rebuild epoch maps for the three-store V4.3 mixture, verifying SHAs."""
directory, root = Path(directory), Path(root)
request = json.loads((directory / 'request.json').read_text('utf-8'))
expected = (json.loads((directory / 'manifest.json').read_text('utf-8'))
if (directory / 'manifest.json').exists() else None)
if expected and all((directory / r['file']).exists() and file_sha(directory / r['file']) == r['sha256']
for r in expected['epochs']):
return expected
bases = []
for record in request['joined_stores']:
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 bank identity changed')
bases.append(PrefixCropWindowDataset(root / record['token_dir'], root / record['index_dir'],
seed=request['seed'], probability=request['crop_probability'],
min_remaining_tokens=request['crop_minimum']))
web_parts = []
for record, base in zip(request['joined_stores'][:-1], bases[:-1]):
keep_path = root / record['keep']
if file_sha(keep_path) != record['keep_sha256']:
raise ValueError(f"Store keep mask changed: {record['token_dir']}")
keep = np.load(keep_path, mmap_mode='r', allow_pickle=False)
chunks = []
for start in range(0, len(base), 262144):
ids = np.arange(start, min(len(base), start + 262144), dtype=np.int64)
chunks.append(ids[keep[window_sentences(base, ids)]])
web_parts.append(np.concatenate(chunks))
web = np.concatenate(web_parts)
joined = JoinedPrefixWindows(bases, request['primary_store'],
[{int(k): v for k, v in r['source_map'].items()} for r in request['joined_stores']])
chat_base = bases[-1]
try:
chat_start = sum(len(base) for base in bases[:-1])
chat = np.arange(chat_start, chat_start + len(chat_base), dtype=np.int64)
records = []
for epoch in range(request['epochs']):
joined.set_epoch(epoch)
web_tokens = 0
for start in range(0, len(web), 262144):
web_tokens += int(joined.window_lengths(web[start:start + 262144]).sum())
chat_lengths = joined.window_lengths(chat)
target = round(web_tokens * request['chat_fraction'] / (1 - request['chat_fraction']))
rng = np.random.default_rng(request['seed'] + epoch)
chosen, used = [], 0
while used < target:
order = rng.permutation(len(chat))
cumulative = np.cumsum(chat_lengths[order], dtype=np.int64)
take = int(np.searchsorted(cumulative, target - used, side='right'))
if take:
chosen.append(chat[order[:take]])
used += int(cumulative[take - 1])
if take < len(order):
break
chosen = np.concatenate(chosen)
mapping = np.concatenate((web, chosen)).astype('<u4')
rng.shuffle(mapping)
path = directory / f'epoch-{epoch}.npy'
temp = path.with_suffix('.tmp.npy')
np.save(temp, mapping, allow_pickle=False)
temp.replace(path)
record = {'epoch': epoch, 'file': path.name, 'bytes': path.stat().st_size,
'sha256': file_sha(path), 'windows': len(mapping),
'web_windows': len(web), 'chat_windows': len(chosen),
'web_tokens': web_tokens, 'chat_tokens': used,
'effective_tokens': web_tokens + used,
'chat_fraction': used / (web_tokens + used),
'chat_effective_passes': used / int(chat_lengths.sum())}
if expected and record != expected['epochs'][epoch]:
raise ValueError(f'Rebuilt map differs from frozen epoch {epoch}')
records.append(record)
print(json.dumps({'map_epoch': epoch + 1, 'windows': len(mapping),
'chat_fraction': record['chat_fraction']}), flush=True)
result = {'format': 'vimeml_token_mixture_v1', 'status': 'complete', 'request': request,
'epochs': records,
'policy': 'V4.3: retained old LM windows and deduplicated expansion windows once per epoch; '
'chat-only repeated shuffled cycles at 5% actual cropped web targets; maps SHA256-pinned'}
if expected and result != expected:
raise ValueError('Rebuilt mixture manifest differs')
if expected is None:
write_json(directory / 'manifest.json', result)
return result
finally:
joined.close()
|