diff --git a/conversion/__init__.py b/conversion/__init__.py index 3976e03..4c7d5cc 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -213,6 +213,8 @@ TEXT_MODEL_MAP: dict[str, str] = { "Ministral3Model": "mistral3", "Mistral3ForConditionalGeneration": "mistral3", "Mistral3Model": "mistral3", + "Mistral3ForSequenceClassification": "mistral3", + "Ministral3ForSequenceClassification": "mistral3", "MistralLarge4ForConditionalGeneration": "mistral4", "VoxtralRealtimeForConditionalGeneration": "mistral3", "MistralForCausalLM": "llama", diff --git a/conversion/mistral3.py b/conversion/mistral3.py index 6c771b2..e77bef4 100644 --- a/conversion/mistral3.py +++ b/conversion/mistral3.py @@ -16,6 +16,8 @@ from .llama import LlamaModel "Mistral3Model", "Ministral3ForCausalLM", "Ministral3Model", + "Mistral3ForSequenceClassification", # LOCAL PATCH + "Ministral3ForSequenceClassification", # LOCAL PATCH ) @ModelBase.example("mistralai/Mistral-Small-3.1-24B-Instruct-2503", "hf-tiny-v2/tiny-random-Ministral3ForCausalLM", "nvidia/Nemotron-3-Embed-1B-BF16", "nvidia/Nemotron-3.5-Embed-1B-BF16") class Mistral3Model(TextModel): @@ -40,13 +42,28 @@ class Mistral3Model(TextModel): self.gguf_writer.add_add_eos_token(False) def get_vocab_base_pre(self, tokenizer) -> str: - if self.origin_hf_arch == "Mistral3Model": + if self.origin_hf_arch in ("Mistral3Model", "Mistral3ForSequenceClassification", "Ministral3ForSequenceClassification"): return "nemotron-3.5-embed" return super().get_vocab_base_pre(tokenizer) def set_gguf_parameters(self): super().set_gguf_parameters() self._try_set_pooling_type() + if self.is_rerank: # LOCAL PATCH: Nemotron rerank -> native RANK pooling (mean + cls.output) + self.gguf_writer.add_pooling_type(gguf.PoolingType.RANK) + self.gguf_writer.add_classifier_output_labels(["score"]) + self.gguf_writer.add_chat_template([{"name": "rerank", "template": "query: {query}\n\npassage: {document}"}]) + + @property + def is_rerank(self) -> bool: + return "sequenceclassification" in (self.origin_hf_arch or "").lower() + + def modify_tensors(self, data_torch, name, bid): # LOCAL PATCH: rerank score head -> cls.output + if name.endswith("score.weight"): + if self.is_rerank: + yield (gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.CLS_OUT] + ".weight", data_torch) + return + yield from super().modify_tensors(data_torch, name, bid) class Mistral4Model(DeepseekV2Model): model_arch = gguf.MODEL_ARCH.MISTRAL4 @@ -64,14 +81,14 @@ class Mistral3Model(TextModel): def __init__(self, dir_model, *args, **kwargs): hparams = ModelBase.load_hparams(dir_model, is_mistral_format=False) hf_arch = get_model_architecture(hparams, self.model_type) - if hf_arch in ("Ministral3Model", "Mistral3Model"): + if hf_arch in ("Ministral3Model", "Mistral3Model", "Mistral3ForSequenceClassification", "Ministral3ForSequenceClassification"): text_config = hparams.get("text_config", hparams) if text_config.get("is_causal") is not False: raise NotImplementedError(f"Expected a Nemotron Embed model for {hf_arch}, but is_causal is not False") super().__init__(dir_model, *args, **kwargs) if self.hparams.get("model_type") == "mistral4": self.impl = Mistral3Model.Mistral4Model(dir_model, *args, **kwargs) - elif self.hf_arch in ("Ministral3Model", "Mistral3Model"): + elif self.hf_arch in ("Ministral3Model", "Mistral3Model", "Mistral3ForSequenceClassification", "Ministral3ForSequenceClassification"): self.impl = Mistral3Model.Mistral3EmbedNemotronModel(dir_model, *args, **kwargs) else: self.impl = Mistral3Model.Ministral3Model(dir_model, *args, **kwargs) diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 1d67dc8..238b26a 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -5820,6 +5820,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { ], MODEL_ARCH.MISTRAL3: [ MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.CLS_OUT, MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, diff --git a/src/models/mistral3.cpp b/src/models/mistral3.cpp index af058d5..bbdad09 100644 --- a/src/models/mistral3.cpp +++ b/src/models/mistral3.cpp @@ -20,6 +20,11 @@ void llama_model_mistral3::load_arch_hparams(llama_model_loader & ml) { } } + // rank pooling (Nemotron rerank): mean pooling before cls_out + if (hparams.pooling_type_cls == LLAMA_POOLING_TYPE_UNSPECIFIED) { + hparams.pooling_type_cls = LLAMA_POOLING_TYPE_MEAN; + } + switch (hparams.n_layer()) { case 26: type = LLM_TYPE_3B; break; case 34: type = LLM_TYPE_8B; break; @@ -38,6 +43,9 @@ void llama_model_mistral3::load_arch_tensors(llama_model_loader & ml) { output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED); // if output is NULL, init from the input tok embed + // optional classifier head (rerank) + cls_out = create_tensor(tn(LLM_TENSOR_CLS_OUT, "weight"), {n_embd, hparams.n_cls_out}, TENSOR_NOT_REQUIRED); + if (output == NULL) { output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED); }