Change to using token classification model.

bioscan-ml · Dec 17, 2024 · 7b62c5e · 7b62c5e
1 parent 60b64c6
commit 7b62c5e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/bioscanclip/model/dna_encoder.py b/bioscanclip/model/dna_encoder.py
@@ -22,7 +22,7 @@ def load_pre_trained_bioscan_bert_trained_with_5m(bioscan_bert_checkpoint, k=4):
     ckpt = torch.load(bioscan_bert_checkpoint, map_location=device)
     model_ckpt = remove_extra_pre_fix(ckpt["model"])
     bert_config = BertConfig(**ckpt["bert_config"])
-    model = BertForMaskedLM(bert_config)
+    model = BertForTokenClassification(bert_config)
     model.load_state_dict(model_ckpt)
     return model.to(device)