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)