File size: 609 Bytes
fdec8d5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 | import json
import torch
from safetensors.torch import load_file
from model import ModernBertConfig, ModernBertForMLM
def load_bertc(model_dir="."):
with open(f"{model_dir}/config.json") as f:
cfg = ModernBertConfig(**json.load(f))
model = ModernBertForMLM(cfg)
state = load_file(f"{model_dir}/model.safetensors")
model.load_state_dict(state, strict=True)
model.eval()
return model
if __name__ == "__main__":
model = load_bertc(".")
ids = torch.tensor([[1, 2, 3]], dtype=torch.long)
with torch.no_grad():
out = model(ids)
print(out["logits"].shape)
|