from transformers import BertConfig, BertForMaskedLM
model = build_model(args.tokenizer_dir, args.max_position_embeddings).to(device)

def build_model(tokenizer_dir: str, max_position_embeddings: int = 512):
    tok_path_json = os.path.join(tokenizer_dir, "tokenizer.json")
    tok = Tokenizer.from_file(tok_path_json)
    vocab_size = tok.get_vocab_size()
    cfg = BertConfig(
        vocab_size=vocab_size,
        hidden_size=512,
        num_hidden_layers=6,
        num_attention_heads=8,
        intermediate_size=2048,
        max_position_embeddings=max_position_embeddings,
        hidden_act="gelu",
        layer_norm_eps=1e-12,
        type_vocab_size=2,
        pad_token_id=0,
    )
    model = BertForMaskedLM(cfg)
    return model