"""Load the Slayer149 ARC-E fine-tune v2 with its native PyTorch architecture."""
import json
from pathlib import Path
import torch
from safetensors.torch import load_model as load_safetensors_model
from tokenizers import Tokenizer
from balanced_model import BalancedLM

def load_model(directory=".", device="cpu"):
    directory = Path(directory)
    config = json.loads((directory / "config.json").read_text())
    model = BalancedLM(config)
    load_safetensors_model(model, str(directory / "model.safetensors"), strict=True)
    model.eval().to(device)
    tokenizer = Tokenizer.from_file(str(directory / "tokenizer.json"))
    return model, tokenizer