| """Load auxiliary heads from an smerkd/arrow-heads snapshot, without a backbone.""" |
| from pathlib import Path |
| import json |
| import numpy as np |
| from safetensors.torch import load_file |
| from arrow_train.heads import TokenHeads |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
|
|
| def load_token_heads(variant='E1', device='cpu'): |
| if variant not in ('E0b', 'E1'): |
| raise ValueError('Token-head variant must be E0b or E1') |
| folder = ROOT / 'weights' / variant |
| cfg = json.loads((folder / 'head_config.json').read_text()) |
| model = TokenHeads(tuple(cfg['names']), cfg['hidden'], cfg['descriptor_dim']) |
| model.load_state_dict(load_file(str(folder / 'heads.safetensors')), strict=True) |
| return model.eval().to(device) |
|
|
| def load_gaussian_head(device='cpu'): |
| from arrow_train.gaussians import GaussianHead, GaussianConfig |
| from arrow_train.language import LanguageDecoder |
| folder = ROOT / 'weights/G1d' |
| cfg = json.loads((folder / 'head_config.json').read_text()) |
| model = GaussianHead(cfg=GaussianConfig(**cfg)) |
| with np.load(ROOT / 'assets/label_embeddings.npz') as z: |
| model.lang = LanguageDecoder(cfg['lang_dim'], z['targets']) |
| model.load_state_dict(load_file(str(folder / 'gaussian_head.safetensors')), strict=True) |
| return model.eval().to(device) |
|
|
