"""Load the local safetensors release without Transformers remote code."""
import json
from pathlib import Path
from types import SimpleNamespace
import torch
from safetensors.torch import load_model as load_safetensors
from tokenizers import Tokenizer
from modeling_gollem import GPT

def load_model(directory='.',device='cpu'):
    directory=Path(directory)
    cfg=SimpleNamespace(**json.loads((directory/'config.json').read_text()))
    model=GPT(cfg.vocab,cfg.n_layer,cfg.n_embd,cfg.n_head,cfg.block,cfg)
    load_safetensors(model,str(directory/'model.safetensors'),strict=True)
    return model.eval().to(device),Tokenizer.from_file(str(directory/'tokenizer.json'))