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