| """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 | |