File size: 6,370 Bytes
81cffe7 | 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 | """Prepare the V4.3 three-store mixture: old web + expansion + chat.
The old web store and its frozen keep mask are reused unchanged from the V4.1/V4.2
mixture. The expansion store's keep mask excludes windows whose sentence matches a
frozen benchmark provenance text (when the evaluation inputs are supplied).
"""
import argparse
import hashlib
import json
from pathlib import Path
import numpy as np
import sys
ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT / 'src'))
OLD_MIX = ROOT / 'artifacts/training-data/corpus-v41-mixed-c128'
NEW = ROOT / 'artifacts/training-data/corpus-v43-mixed-c128'
EPOCHS = 30
def sha(path):
digest = hashlib.sha256()
with Path(path).open('rb') as stream:
for block in iter(lambda: stream.read(1 << 20), b''):
digest.update(block)
return digest.hexdigest()
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('--benchmark-inputs', type=Path, default=None,
help='Frozen evaluation track JSONs; provenance texts are excluded from the expansion pool')
parser.add_argument('--verify-only', action='store_true')
args = parser.parse_args()
from vimeml.training.data import write_json
from vimeml.training.v43_data import prepare_maps_v43, window_sentences, sentence_keys
if (NEW / 'manifest.json').exists():
result = prepare_maps_v43(NEW, root=ROOT)
print(json.dumps({'status': 'verified', 'epochs': len(result['epochs']),
'dataset_sha256': sha(NEW / 'manifest.json')}))
return
if args.verify_only:
raise ValueError('No V4.3 mixture to verify')
old_request = json.loads((OLD_MIX / 'request.json').read_text('utf-8'))
old_store, chat_store = old_request['joined_stores'][0], old_request['joined_stores'][1]
exp_token = ROOT / 'artifacts/token-data/corpus-v43-exp8k'
exp_index = ROOT / 'artifacts/training-data/corpus-v43-exp-c128'
if not (exp_token / 'manifest.json').is_file() or not (exp_index / 'manifest.json').is_file():
raise ValueError('Expansion token store is not built yet')
NEW.mkdir(parents=True, exist_ok=True)
exp_keep = np.ones(json.loads((exp_index / 'manifest.json').read_text('utf-8'))['splits']['train']['windows'],
dtype=bool)
excluded = 0
if args.benchmark_inputs and args.benchmark_inputs.is_dir():
from vimeml.training.data_v2 import PrefixCropWindowDataset
import unicodedata
keys = sentence_keys(_benchmark_texts(args.benchmark_inputs))
excluded_lines = set()
dedup_manifest = json.loads((ROOT / 'outputs/corpus-expansion-v43-dedup/manifest.json').read_text('utf-8'))
with (ROOT / 'outputs/corpus-expansion-v43-dedup/train.txt').open(encoding='utf-8') as stream:
for line_number, line in enumerate(stream):
text = line.rstrip('\r\n')
if text and (' '.join(unicodedata.normalize('NFKC', text).split()).encode() and
hashlib.blake2b(' '.join(unicodedata.normalize('NFKC', text).split()).encode(),
digest_size=16).digest() in keys):
excluded_lines.add(line_number)
base = PrefixCropWindowDataset(exp_token, exp_index, seed=old_request['seed'],
probability=0.0, min_remaining_tokens=8)
try:
keep = np.ones(len(base), dtype=bool)
for start in range(0, len(base), 262144):
ids = np.arange(start, min(len(base), start + 262144), dtype=np.int64)
sentences = window_sentences(base, ids)
keep[start:start + len(ids)] = ~np.isin(sentences, np.fromiter(excluded_lines, dtype=np.int64))
excluded = int((~keep).sum())
exp_keep = keep
finally:
base.close()
np.save(NEW / 'expansion.keep.npy', exp_keep, allow_pickle=False)
request = {
'joined_stores': [
{**old_store, 'keep': 'artifacts/training-data/corpus-v41-mixed-c128/old.keep.npy',
'keep_sha256': old_request['old_keep_sha256']},
{'token_dir': 'artifacts/token-data/corpus-v43-exp8k',
'index_dir': 'artifacts/training-data/corpus-v43-exp-c128',
'token_manifest_sha256': sha(exp_token / 'manifest.json'),
'window_manifest_sha256': sha(exp_index / 'manifest.json'),
'source_map': {'0': 4},
'keep': 'artifacts/training-data/corpus-v43-mixed-c128/expansion.keep.npy',
'keep_sha256': sha(NEW / 'expansion.keep.npy')},
{**chat_store},
],
'primary_store': 0,
'token_manifest_sha256': old_store['token_manifest_sha256'],
'window_manifest_sha256': old_store['window_manifest_sha256'],
'index_dir': old_store['index_dir'],
'epochs': EPOCHS, 'seed': old_request['seed'],
'crop_probability': old_request['crop_probability'],
'crop_minimum': old_request['crop_minimum'], 'chat_fraction': old_request['chat_fraction'],
'old_keep': 'artifacts/training-data/corpus-v41-mixed-c128/old.keep.npy',
'old_keep_sha256': old_request['old_keep_sha256'],
'derived_from': 'V4.1/V4.2 data pools plus the V4.3 deduplicated expansion store',
'expansion_benchmark_excluded_windows': excluded,
'test_or_blind_read': False, 'numpy_version': np.__version__,
}
(NEW / 'request.json').write_text(json.dumps(request, ensure_ascii=False, indent=1) + '\n', encoding='utf-8')
result = prepare_maps_v43(NEW, root=ROOT)
if len(result['epochs']) != EPOCHS:
raise ValueError('prepare_maps_v43 did not produce the full plan')
print(json.dumps({'status': 'prepared', 'epochs': EPOCHS,
'benchmark_excluded_windows': excluded,
'dataset_sha256': sha(NEW / 'manifest.json')}))
def _benchmark_texts(directory):
rows = []
for path in sorted(Path(directory).glob('*.json')):
data = json.loads(path.read_text('utf-8'))
rows.extend(row.get('provenance', {}).get('original_message')
for row in data if row.get('provenance', {}).get('original_message'))
return [text for text in rows if text]
if __name__ == '__main__':
main()
|