File size: 5,791 Bytes
50bfecf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 | 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);
}
|