| """Load an exact completed-5B release, with hash and endpoint verification. |
| |
| This is a ResearchLM checkpoint, not a Transformers AutoModel checkpoint. |
| Review the bundled research_code and install its pinned CUDA dependencies first. |
| """ |
| from __future__ import annotations |
|
|
| import hashlib |
| import json |
| from pathlib import Path, PurePosixPath |
| import sys |
|
|
|
|
| def sha256(path): |
| h = hashlib.sha256() |
| with open(path, 'rb') as stream: |
| for part in iter(lambda: stream.read(8 << 20), b''): |
| h.update(part) |
| return h.hexdigest() |
|
|
|
|
| def verify_release(folder): |
| folder = Path(folder).resolve() |
| manifest = json.loads((folder / 'release_manifest.json').read_text()) |
| for name, expected in manifest['files'].items(): |
| relative = PurePosixPath(name) |
| |
| |
| if relative.is_absolute() or '..' in relative.parts: |
| raise ValueError(f'Invalid release path: {name}') |
| path = folder / relative |
| if not path.is_file(): |
| raise ValueError(f'Invalid release file: {name}') |
| if path.stat().st_size != expected['bytes'] or sha256(path) != expected['sha256']: |
| raise ValueError(f'Release file differs: {name}') |
| provenance = json.loads((folder / 'provenance.json').read_text()) |
| if provenance['step'] != 9537 or provenance['tokens'] != 5000134656: |
| raise ValueError('A completed 5B endpoint is required') |
| if provenance['corpus_sha256'] != '1786fc7cf798784edf2e0ee9400456ca93add1cbc5688cd8e83865cb230bc475': |
| raise ValueError('Unexpected training corpus') |
| return provenance |
|
|
|
|
| def load_release(folder, *, device='cuda', dtype=None): |
| import torch |
| from safetensors.torch import load_file |
|
|
| folder = Path(folder).resolve() |
| provenance = verify_release(folder) |
| code = folder / 'research_code' |
| existing = sys.modules.get('pretraining_5b.model') |
| if existing is not None and Path(existing.__file__).resolve() != (code / 'pretraining_5b/model.py').resolve(): |
| raise RuntimeError('Use a fresh process: a different pretraining_5b.model is already imported') |
| sys.path.insert(0, str(code)) |
| from pretraining_5b.model import ModelConfig, ResearchLM |
|
|
| config = json.loads((folder / 'research_config.json').read_text()) |
| weights = load_file(str(folder / 'model.safetensors'), device='cpu') |
| aliases = json.loads((folder / 'tied_weights.json').read_text()) |
| for alias, source in aliases.items(): |
| weights[alias] = weights[source] |
| model = ResearchLM(ModelConfig(**config)) |
| model.load_state_dict(weights, strict=True) |
| if model.parameter_count() != provenance['parameter_count']: |
| raise ValueError('Model parameter count differs from the admitted release') |
| model = model.to(device=device, dtype=dtype or torch.bfloat16).eval() |
| return model, provenance |
|
|
