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