File size: 656 Bytes
7d11e8f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 | """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
|