File size: 1,443 Bytes
f8de5bd | 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 | """Load a merged Drex DLM checkpoint.
The backbone shards live next to ``head.pt``. There is no adapter to merge.
"""
from pathlib import Path
import torch
from .model import DecisionModel, load_tokenizer
def read_meta(path):
path = Path(path)
return torch.load(path / "head.pt", map_location="cpu", weights_only=False)
def load(path, device=None, dtype=None):
"""Return ``(tokenizer, model)`` in eval mode.
``device`` defaults to CUDA, then Apple MPS, then CPU. ``dtype`` defaults
to bfloat16 on CUDA and MPS, and float32 on CPU.
"""
path = Path(path)
meta = read_meta(path)
if device is None:
if torch.cuda.is_available():
device = "cuda"
elif getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
device = "mps"
else:
device = "cpu"
if dtype is None:
dtype = torch.bfloat16 if device in ("cuda", "mps") else torch.float32
tok = load_tokenizer(str(path))
model = DecisionModel(
str(path),
tok,
device,
lora=None,
head_dim=int(meta.get("head_dim", 256)),
option_isolation=bool(meta.get("option_isolation", False)),
dtype=dtype,
state_bidir=bool(meta.get("state_bidir", True)),
)
model.head.load_state_dict(meta["head"])
model.head.temperature = float(meta.get("temperature", 1.0))
model.eval()
return tok, model
|