File size: 4,821 Bytes
e89551c | 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 | #!/usr/bin/env python3
"""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():
# The original W4 Hub bundle has code/ but keeps evidence in reports/.
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()
|