| |
| """CPU helpers: assemble released runtime files or recreate WT2 test tokens.""" |
| from __future__ import annotations |
| import argparse |
| import hashlib |
| import json |
| from pathlib import Path |
| import shutil |
| import sys |
|
|
| sys.dont_write_bytecode = True |
| from run_recall_native_prefill_wt2_test_v1 import PINS, PINS_SHA256, need, put, sha |
|
|
|
|
| def fixed_pins(): |
| need(sha(PINS) == PINS_SHA256, 'Frozen test pins changed') |
| return json.loads(PINS.read_text()) |
|
|
|
|
| def prepare_runtime(args): |
| fixed = fixed_pins()['models'][args.model_kind] |
| need(not args.out_dir.exists(), 'Fresh runtime output required') |
| source = args.download_root.resolve() |
| planned = {} |
| for relative, expected in fixed['frozen_files'].items(): |
| |
| code_candidate = source / 'code' / relative |
| path = code_candidate if code_candidate.is_file() else source / relative |
| need(path.is_file() and sha(path) == expected, 'Published runtime/evidence differs: ' + relative) |
| planned[relative] = path |
| adapter = source / fixed['adapter_path'] |
| if not adapter.is_file() and args.model_kind == 'w4': |
| adapter = source / 'adapter/adapter_fp16.pt' |
| need(adapter.is_file() and sha(adapter) == fixed['adapter']['sha256'], 'Published adapter differs') |
| planned[fixed['adapter_path']] = adapter |
| args.out_dir.mkdir(parents=True) |
| for relative, path in planned.items(): |
| target = args.out_dir / relative |
| target.parent.mkdir(parents=True, exist_ok=True) |
| shutil.copyfile(path, target) |
| receipt = {'complete': True, 'model_kind': args.model_kind, 'pins_sha256': sha(PINS), |
| 'source': str(source), 'runtime_root': str(args.out_dir.resolve()), |
| 'files': {relative: {'bytes': path.stat().st_size, 'sha256': sha(path)} |
| for relative, path in planned.items()}, |
| 'weight_files_copied': False, 'cuda_initialized': False, 'torch_imported': 'torch' in sys.modules} |
| put(args.out_dir / 'test_runtime_receipt.json', receipt) |
| print(json.dumps(receipt, indent=2)) |
|
|
|
|
| def recreate_tokens(args): |
| pins = fixed_pins() |
| release = args.release_root.resolve() |
| fixed = pins['models'][args.model_kind] |
| for relative in ('mamba2_recall/__init__.py', 'mamba2_recall/runtime.py', 'mamba2_recall/calibration.py'): |
| need(sha(release / relative) == fixed['frozen_files'][relative], 'Frozen tokenizer/data source changed') |
| need('mamba2_recall' not in sys.modules, 'A different runtime was imported') |
| sys.path.insert(0, str(release)) |
| from mamba2_recall import runtime, calibration |
| import torch |
| need(not torch.cuda.is_initialized(), 'Token reconstruction must be CPU only') |
| tokenizer = runtime.SentencePieceTokenizer(args.tokenizer_source_dir) |
| ids, metadata = calibration.load_wikitext_tokens(tokenizer, 'test') |
| need(metadata == pins['test_dataset'], 'Pinned official test tokenization differs') |
| raw = ids.numpy().astype('<i8', copy=False).tobytes() |
| need(hashlib.sha256(raw).hexdigest() == metadata['token_stream_sha256_int64le'], 'Raw token identity differs') |
| args.out_dir.mkdir(parents=True, exist_ok=True) |
| path = args.out_dir / 'tokens.int64le' |
| if path.exists(): |
| need(sha(path) == metadata['token_stream_sha256_int64le'], 'Existing raw token file differs') |
| else: |
| path.write_bytes(raw) |
| need(not torch.cuda.is_initialized(), 'CPU helper initialized CUDA') |
| receipt = {'complete': True, 'model_kind': args.model_kind, 'dataset': metadata, |
| 'file': str(path.resolve()), 'bytes': len(raw), 'sha256': sha(path), |
| 'cuda_initialized': False, 'weight_checkpoint_loaded': False} |
| print(json.dumps(receipt, indent=2)) |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| sub = parser.add_subparsers(dest='command', required=True) |
| runtime = sub.add_parser('runtime', help='Small frozen files from the original published bundle; no weights copied') |
| runtime.add_argument('--model-kind', choices=('fp16', 'w4'), required=True) |
| runtime.add_argument('--download-root', type=Path, required=True) |
| runtime.add_argument('--out-dir', type=Path, required=True) |
| tokens = sub.add_parser('tokens', help='Recreate unpublished raw tokens for the CPU evidence audit') |
| tokens.add_argument('--model-kind', choices=('fp16', 'w4'), required=True) |
| tokens.add_argument('--release-root', type=Path, required=True) |
| tokens.add_argument('--tokenizer-source-dir', type=Path, required=True) |
| tokens.add_argument('--out-dir', type=Path, required=True) |
| args = parser.parse_args() |
| (prepare_runtime if args.command == 'runtime' else recreate_tokens)(args) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
