File size: 1,034 Bytes
37e177f | 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 | from pathlib import Path
import torch
from huggingface_hub import snapshot_download
import laya
from tinycenn_lm.laya_lab.core import LayaLabConfig
from tinycenn_lm.laya_lab.factory import make_replacement
def load_model(repo_id, device=None, token=None):
root = Path(snapshot_download(
repo_id,
repo_type="model",
token=token,
allow_patterns=["adapter.pt", "report.json", "model_meta.json"],
))
payload = torch.load(root / "adapter.pt", map_location="cpu", weights_only=False)
cfg = LayaLabConfig(**payload["lab_config"])
agent = laya.load(cfg.model_id, device=device, token=token)
for layer_id, spec in payload["adapters"].items():
idx = int(layer_id)
layer = agent.model.encoder.layers[idx]
replacement = make_replacement(layer.attn, cfg)
replacement.load_state_dict(spec["state_dict"], strict=True)
replacement.to(agent.device).eval().requires_grad_(False)
layer.attn = replacement
agent.model.eval()
return agent
|