File size: 1,815 Bytes
5b7b27a | 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 | #!/usr/bin/env python3
"""Prepare one provenance-bound numeric split without initializing CUDA."""
import argparse
import json
import os
from pathlib import Path
import sys
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
from mamba2_recall import resurface_data as data
from mamba2_recall.runtime import SentencePieceTokenizer
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('--split', choices=('train', 'dev', 'confirm'), required=True)
parser.add_argument('--source-dir', type=Path, required=True)
parser.add_argument('--data-root', type=Path, default=ROOT/'training_data'/'numeric_v1')
args = parser.parse_args()
if os.environ.get('CUDA_VISIBLE_DEVICES') != '':
raise RuntimeError('Set CUDA_VISIBLE_DEVICES= for CPU data preparation')
protocol = ROOT/'docs'/'PROTOCOL.md'
tokenizer = SentencePieceTokenizer(args.source_dir)
manifest = data.prepare_split(args.data_root, args.split, tokenizer,
protocol_path=protocol, protocol_sha256=data.sha_file(protocol),
tokenizer_sha256=data.TOKENIZER_SHA)
path = args.data_root/args.split/'manifest.json'
loaded, _, _ = (data.load_training(args.data_root, data.sha_file(path), tokenizer)
if args.split == 'train' else
data.load_evaluation(args.data_root, args.split, data.sha_file(path), tokenizer))
if loaded != manifest:
raise RuntimeError('Prepared split failed exact readback')
print(json.dumps({'split': args.split, 'rows': manifest['row_count'],
'manifest_sha256': data.sha_file(path),
'protocol_sha256': data.sha_file(protocol),
'cuda_initialized': False}), flush=True)
if __name__ == '__main__':
main()
|