fix: JinaBertForMaskedLM registration

ggml-org · Apr 23, 2024 · e946cb0 · e946cb0
1 parent 71ff763
commit e946cb0
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/convert-hf-to-gguf.py b/convert-hf-to-gguf.py
@@ -2715,7 +2715,7 @@ def write_tensors(self):
             self.gguf_writer.add_tensor(new_name, data)
 
 
-@Model.register("JinaBertModel")
+@Model.register("JinaBertModel", "JinaBertForMaskedLM")
 class JinaBertModel(BertModel):
     model_arch = gguf.MODEL_ARCH.JINA_BERT
 
@@ -2742,7 +2742,6 @@ def set_vocab(self, *args, **kwargs):
         self.gguf_writer.add_add_eos_token(True)
 
 
-JinaBertForMaskedML = JinaBertModel
 
 ###### CONVERSION LOGIC ######