koboldcpp/conversion/mistral3.py
Takamichi Ohara f9599d0cc8
vocab : fix Nemotron-3.5-Embed tokenization (#30974)
* vocab : fix Nemotron-3.5-Embed tokenization

* Update src/llama-vocab.h

Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@huggingface.co>

* rename

* Update src/llama-vocab.cpp

Co-authored-by: Daniel Bevenius <daniel.bevenius@gmail.com>

---------

Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@huggingface.co>
Co-authored-by: Daniel Bevenius <daniel.bevenius@gmail.com>
2026-10-23 07:31:30 +02:00

105 lines
4.1 KiB
Python

from __future__ import annotations
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from torch import Tensor
from .base import ModelBase, TextModel, get_model_architecture, gguf
from .deepseek import DeepseekV2Model
from .llama import LlamaModel
@ModelBase.register(
"Mistral3ForConditionalGeneration",
"Mistral3Model",
"Ministral3ForCausalLM",
"Ministral3Model",
)
@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):
class Ministral3Model(LlamaModel):
model_arch = gguf.MODEL_ARCH.MISTRAL3
def set_gguf_parameters(self):
super().set_gguf_parameters()
rope_params = self.rope_parameters
if self.hparams.get("model_type") == "ministral3":
assert rope_params, "ministral3 must have 'rope_parameters' config"
assert rope_params["rope_type"] == "yarn", "ministral3 rope_type must be 'yarn'"
self.gguf_writer.add_rope_scaling_yarn_log_mul(rope_params["mscale_all_dim"])
self.gguf_writer.add_attn_temperature_scale(rope_params["llama_4_scaling_beta"])
class Mistral3EmbedNemotronModel(Ministral3Model):
model_arch = gguf.MODEL_ARCH.MISTRAL3
def set_vocab(self):
super().set_vocab()
self.gguf_writer.add_add_bos_token(False)
self.gguf_writer.add_add_eos_token(False)
def get_vocab_base_pre(self, tokenizer) -> str:
if self.origin_hf_arch == "Mistral3Model":
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()
class Mistral4Model(DeepseekV2Model):
model_arch = gguf.MODEL_ARCH.MISTRAL4
skip_mtp = False # model contains no MTP layers, so no need to skip
merge_expert = False # experts are already stacked as 3D
def modify_tensors(self, data_torch, name, bid):
if name.endswith(".down_proj") or name.endswith(".gate_up_proj"):
name = name + ".weight"
yield from super().modify_tensors(data_torch, name, bid)
model_arch = gguf.MODEL_ARCH.MISTRAL3 # unused
impl: 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"):
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"):
self.impl = Mistral3Model.Mistral3EmbedNemotronModel(dir_model, *args, **kwargs)
else:
self.impl = Mistral3Model.Ministral3Model(dir_model, *args, **kwargs)
def set_vocab(self):
self.impl.set_vocab()
def set_gguf_parameters(self):
self.impl.set_gguf_parameters()
def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None):
yield from self.impl.modify_tensors(data_torch, name, bid)
def prepare_tensors(self):
self.impl.prepare_tensors()
def write_vocab(self):
self.impl.write_vocab()
def write(self):
self.impl.write()
@ModelBase.register("VoxtralRealtimeForConditionalGeneration")
@ModelBase.example("mistralai/Voxtral-Mini-4B-Realtime-2602")
class VoxtralRealtimeModel(LlamaModel):
model_arch = gguf.MODEL_ARCH.MISTRAL3
def set_gguf_parameters(self):
super().set_gguf_parameters()
self.gguf_writer.add_realtime_delay_tokens(self.hparams["default_num_delay_tokens"])