load_model.py

import json
from pathlib import Path
import torch
from safetensors.torch import load_file
from tokenizers import Tokenizer
from balanced_model import BalancedLM

def load_model(directory,device='cpu'):
 path=Path(directory);m=BalancedLM(json.loads((path/'config.json').read_text()))
 state=load_file(str(path/'model.safetensors'));state['lm_head.weight']=state['model.embed_tokens.weight']
 m.load_state_dict(state,strict=True);return m.eval().to(device),Tokenizer.from_file(str(path/'tokenizer.json'))