| diff --git 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",
|
|
|
|
|
|
|
|
|
| @@ -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)
|
|
|
|
|
|
|
|
|
| @@ -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,
|
|
|
|
|
|
|
|
|
| @@ -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);
|
| }
|
|
|
