Blame for load_model.py

"""Load this repository's custom PyTorch causal language model."""

import json
from pathlib import Path

import torch

from model import CausalLM
from moe import TransformerConfig


def load_model(directory, device="cpu"):
    """Load verified tensor weights and return a model in evaluation mode."""
    directory = Path(directory)
    config = TransformerConfig(**json.loads((directory / "config.json").read_text()))
    model = CausalLM(config)
    weights = torch.load(
        directory / "model_state.pt", map_location="cpu", weights_only=True
    )
    model.load_state_dict(weights, strict=True)
    return model.eval().to(device)