| 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')) | |