load_laya.py

"""Self-contained native Laya inference loader; no original base checkpoint needed."""
from pathlib import Path
import importlib.util
import json

def load_laya(directory=None, device="cpu"):
    import torch
    from transformers import AutoConfig, AutoModel, AutoTokenizer
    from safetensors.torch import load_file
    directory = Path(directory or Path(__file__).parent).resolve()
    spec = importlib.util.spec_from_file_location("_exported_laya_native", directory/"native_common.py")
    common = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(common)
    config = AutoConfig.from_pretrained(directory/"encoder", local_files_only=True)
    config.reference_compile = False
    encoder = AutoModel.from_config(config, attn_implementation="sdpa")
    native = json.loads((directory/"rl_agent_config.json").read_text())
    model = common.DecisionModel(encoder, native.get("head_layers", 2), len(native.get("act_costs", {}))+1)
    model.load_state_dict(load_file(str(directory/"model.safetensors")), strict=True)
    model.to(device).eval()
    tokenizer = AutoTokenizer.from_pretrained(directory/"tokenizer", local_files_only=True)
    return model, tokenizer, common