File size: 3,236 Bytes
64b9802 | 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 | """Load an explicit released world-model variant; never substitute adapters."""
import hashlib
import json
from pathlib import Path, PurePosixPath
def safe_path(root, relative):
rel = PurePosixPath(relative)
if rel.is_absolute() or not rel.parts or '..' in rel.parts or '\\' in relative:
raise ValueError(f'Invalid relative path: {relative!r}')
root = Path(root).resolve()
path = (root / relative).resolve()
if path != root and root not in path.parents:
raise ValueError('Path escapes bundle root')
return path
def file_sha256(path):
value = hashlib.sha256()
with Path(path).open('rb') as stream:
for block in iter(lambda: stream.read(8 * 1024 * 1024), b''):
value.update(block)
return value.hexdigest()
def verify_bundle(root):
root = Path(root)
manifest = json.loads((root / 'manifest.json').read_text())
if manifest.get('schema_version') != 1 or not manifest.get('files'):
raise ValueError('Unsupported or empty manifest')
seen = set()
for item in manifest['files']:
name = item['path']
if name in seen:
raise ValueError(f'Duplicate file: {name}')
seen.add(name)
path = safe_path(root, name)
if not path.is_file() or path.stat().st_size != item['bytes']:
raise ValueError(f'Missing file or wrong size: {name}')
if file_sha256(path) != item['sha256']:
raise ValueError(f'Checksum mismatch: {name}')
return manifest
def load_world_model(root, dataset, variant='main', device='cpu', verify=True):
"""Load the predictor only. Codec/actions/T5 preprocessing are separate.
FP32 and SDPA match the audited loaders. Dataset/variant selection is
explicit; missing text adapters are an error, never silently zero-filled.
"""
import torch
from safetensors.torch import load_file
from transformers import AutoConfig, AutoModelForCausalLM
from .text_conditioning import TextWorldModel
root = Path(root)
if verify:
verify_bundle(root)
settings = json.loads(safe_path(root, f'{dataset}/futureworlds_config.json').read_text())
if settings['dataset'] != dataset or variant not in settings['variants']:
raise ValueError('Dataset or variant does not match bundle configuration')
entry = settings['variants'][variant]
folder = safe_path(root, entry['path'])
cfg = AutoConfig.from_pretrained(folder, local_files_only=True, trust_remote_code=False)
core = AutoModelForCausalLM.from_config(cfg, torch_dtype=torch.float32, attn_implementation='sdpa', trust_remote_code=False)
core.load_state_dict(load_file(str(folder / 'model.safetensors')), strict=True)
if settings['text_conditioning']['enabled']:
tc = settings['text_conditioning']
model = TextWorldModel(core, text_width=tc['width'], layers=tc['layers'], inner=tc['inner'])
adapters = folder / 'text_adapters.safetensors'
if not adapters.is_file():
raise FileNotFoundError(f'Required trained text adapters missing: {adapters}')
model.adapters.load_state_dict(load_file(str(adapters)), strict=True)
else:
model = core
return model.to(device).eval()
|