Brunobkr/llama.cpp_AlgMor24_github

1from __future__ import annotations2 3import json4 5from typing import Iterable, TYPE_CHECKING6 7if TYPE_CHECKING:8    from torch import Tensor9 10from .base import MmprojModel, ModelBase, gguf, logger11 12from .llama import LlamaModel13 14 15@ModelBase.register(16    "LlavaForConditionalGeneration", # pixtral17    "Mistral3ForConditionalGeneration", # mistral small 3.118)19class LlavaVisionModel(MmprojModel):20    img_break_tok_id = -121    use_break_tok = True22 23    def __init__(self, *args, **kwargs):24        super().__init__(*args, **kwargs)25        if self.hparams.get("model_type") == "pixtral":26            # layer_norm_eps is not in config.json, it is hard-coded in modeling_pixtral.py27            self.hparams["layer_norm_eps"] = self.hparams.get("layer_norm_eps", 1e-5)28            if self.use_break_tok:29                self.img_break_tok_id = self.get_token_id("[IMG_BREAK]")30        elif self.is_mistral_format:31            # hparams is already vision config here so norm_eps is only defined in global_config.32            self.hparams["norm_eps"] = self.global_config.get("norm_eps", None)33            assert self.hparams["norm_eps"] is not None, "norm_eps not found in params.json"34            if self.use_break_tok:35                self.img_break_tok_id = self.find_vparam(["image_break_token_id"])36 37                # params.json may ship -1 placeholders (Mistral Medium 3.5)38                # resolve the real id from the bundled tokenizer in that case39                if self.img_break_tok_id < 0:40                    self.img_break_tok_id = self.get_mistral_token_id("[IMG_BREAK]")41        else:42            raise ValueError(f"Unsupported model type: {self.hparams['model_type']}")43        logger.info(f"Image break token id: {self.img_break_tok_id}")44 45    def get_token_id(self, token: str) -> int:46        tokenizer_config_file = self.dir_model / 'tokenizer_config.json'47        with open(tokenizer_config_file, "r", encoding="utf-8") as f:48            added_tokens_decoder = json.load(f).get('added_tokens_decoder') or {}49            for id_, token_data in added_tokens_decoder.items():50                if token_data.get("content") == token:51                    return int(id_)52            # fallthrough to tokenizer.json53        with open(self.dir_model / "tokenizer.json", "r", encoding="utf-8") as f:54            tokenizer_json = json.load(f)55            for token_data in tokenizer_json["added_tokens"]:56                if token_data["content"] == token:57                    return int(token_data["id"])58        raise ValueError(f"Token '{token}' not found in tokenizer config.")59 60    def get_mistral_token_id(self, token: str) -> int:61        # mistral native format ships tekken.json or a versioned spm tokenizer62        tekken_file = self.dir_model / "tekken.json"63        if tekken_file.is_file():64            with open(tekken_file, "r", encoding="utf-8") as f:65                data = json.load(f)66            for entry in data.get("special_tokens", []):67                if entry.get("token_str") == token:68                    return int(entry["rank"])69        tokenizer_json_file = self.dir_model / "tokenizer.json"70        if tokenizer_json_file.is_file():71            with open(tokenizer_json_file, "r", encoding="utf-8") as f:72                data = json.load(f)73            for entry in data.get("added_tokens", []):74                if entry.get("content") == token:75                    return int(entry["id"])76        raise ValueError(f"Token '{token}' not found in mistral tokenizer files.")77 78    def set_gguf_parameters(self):79        super().set_gguf_parameters()80        hparams = self.hparams81        if hparams.get("model_type") == "pixtral":82            self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.PIXTRAL)83            self.gguf_writer.add_vision_attention_layernorm_eps(hparams["layer_norm_eps"])84 85            # hidden_act86            if hparams["hidden_act"] == "silu":87                self.gguf_writer.add_vision_use_silu(True)88            elif hparams["hidden_act"] == "gelu":89                self.gguf_writer.add_vision_use_gelu(True)90            else:91                raise ValueError(f"Unsupported hidden_act: {hparams['hidden_act']}")92 93            # spatial_merge_size94            if "spatial_merge_size" in self.global_config:95                self.gguf_writer.add_vision_spatial_merge_size(self.global_config["spatial_merge_size"])96 97    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:98        n_head = (99            self.hparams["num_attention_heads"] if not self.is_mistral_format else self.find_vparam(["num_attention_heads"])100        )101        n_kv_head = n_head102 103        valid_prefixes = (104            "multi_modal_projector.",105            "vision_tower.",106            "vision_encoder.",107            "vision_language_adapter.",108            "patch_merger.",109            "pre_mm_projector_norm",110        )111 112        if any(name.startswith(prefix) for prefix in valid_prefixes):113            # process vision tensors114            if name.endswith(("q_proj.weight", "q_proj.bias")) and not self.is_mistral_format:115                data_torch = LlamaModel.permute(data_torch, n_head, n_head)116            if name.endswith(("k_proj.weight", "k_proj.bias")) and not self.is_mistral_format:117                data_torch = LlamaModel.permute(data_torch, n_head, n_kv_head)118            yield from super().modify_tensors(data_torch, name, bid)119            return120 121        embed_key = "embed_tokens.weight" if not self.is_mistral_format else "tok_embeddings.weight"122        if self.img_break_tok_id > 0 and embed_key in name:123            logger.info(f"Extracting [IMG_BREAK] token embedding from {name}")124            # for pixtral model, we need to extract the [IMG_BREAK] token embedding125            img_break_embd = data_torch[self.img_break_tok_id]126            name = gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.V_TOK_EMBD_IMG_BREAK]127            yield from super().modify_tensors(img_break_embd, name, bid)128 129        return # skip other tensors130