Brunobkr/llama.cpp_AlgMor24_github

1from __future__ import annotations2 3import json4import math5import re6 7from typing import Callable, Iterable, TYPE_CHECKING8 9import torch10 11if TYPE_CHECKING:12    from torch import Tensor13 14from .base import MmprojModel, ModelBase, TextModel, gguf15 16 17@ModelBase.register("Ernie4_5_ForCausalLM", "Ernie4_5ForCausalLM")18class Ernie4_5Model(TextModel):19    model_arch = gguf.MODEL_ARCH.ERNIE4_520 21    def set_vocab(self):22        self._set_vocab_sentencepiece()23 24        tokenizer_config_file = self.dir_model / 'tokenizer_config.json'25        if tokenizer_config_file.is_file():26            with open(tokenizer_config_file, "r", encoding="utf-8") as f:27                tokenizer_config_json = json.load(f)28                if "add_prefix_space" in tokenizer_config_json:29                    self.gguf_writer.add_add_space_prefix(tokenizer_config_json["add_prefix_space"])30 31    def set_gguf_parameters(self):32        super().set_gguf_parameters()33 34    @classmethod35    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:36        name, gen = item37 38        if "ernie." in name:39            name = name.replace("ernie.", "model.")40 41        return super().filter_tensors((name, gen))42 43    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:44        num_heads = self.hparams["num_attention_heads"]45        num_kv_heads = self.hparams["num_key_value_heads"]46        if (head_dim := self.hparams.get("head_dim")) is None:47            head_dim = self.hparams["hidden_size"] // num_heads48 49        # split the qkv weights50        # qkv_proj shape: [(num_heads + 2 * num_kv_heads) * head_dim, hidden_size]51        if "qkv_proj" in name:52            name_q = name.replace("qkv_proj.weight", "q_proj.weight")53            name_k = name.replace("qkv_proj.weight", "k_proj.weight")54            name_v = name.replace("qkv_proj.weight", "v_proj.weight")55            total_q_dim = num_heads * head_dim56            total_k_dim = num_kv_heads * head_dim57            total_v_dim = num_kv_heads * head_dim58            q_proj_weight, k_proj_weight, v_proj_weight = data_torch.split([total_q_dim, total_k_dim, total_v_dim], dim=0)59            yield from super().modify_tensors(q_proj_weight, name_q, bid)60            yield from super().modify_tensors(k_proj_weight, name_k, bid)61            yield from super().modify_tensors(v_proj_weight, name_v, bid)62        # split the up_gate_proj into gate and up63        # up_gate_proj shape: [2 * intermediate_size, hidden_size]64        elif "up_gate_proj" in name:65            name_up = name.replace("up_gate_proj.weight", "up_proj.weight")66            name_gate = name.replace("up_gate_proj.weight", "gate_proj.weight")67            dim_half = data_torch.shape[0] // 268            gate_proj_weight, up_proj_weight = data_torch.split(dim_half, dim=0)69            yield from super().modify_tensors(gate_proj_weight, name_gate, bid)70            yield from super().modify_tensors(up_proj_weight, name_up, bid)71        else:72            yield from super().modify_tensors(data_torch, name, bid)73 74 75@ModelBase.register("Ernie4_5_MoeForCausalLM")76class Ernie4_5MoeModel(Ernie4_5Model):77    model_arch = gguf.MODEL_ARCH.ERNIE4_5_MOE78    _experts: list[dict[str, Tensor]] | None = None79 80    def __init__(self, *args, **kwargs):81        super().__init__(*args, **kwargs)82        self._experts = [{} for _ in range(self.block_count)]83 84    def set_gguf_parameters(self):85        super().set_gguf_parameters()86        self.gguf_writer.add_expert_count(self.hparams["moe_num_experts"])87        self.gguf_writer.add_expert_used_count(self.hparams["moe_k"])88        self.gguf_writer.add_interleave_moe_layer_step(self.hparams["moe_layer_interval"])89        self.gguf_writer.add_leading_dense_block_count(self.hparams["moe_layer_start_index"])90        if (moe_intermediate_size := self.hparams.get("moe_intermediate_size")) is not None:91            self.gguf_writer.add_expert_feed_forward_length(moe_intermediate_size)92        if (shared_expert_count := self.hparams.get('moe_num_shared_experts')) is not None:93            self.gguf_writer.add_expert_shared_count(shared_expert_count)94            if shared_expert_count > 0 and (shared_expert_intermediate_size := self.hparams.get('intermediate_size')) is not None and (num_key_value_heads := self.hparams.get('num_key_value_heads')) is not None:95                self.gguf_writer.add_expert_shared_feed_forward_length(shared_expert_intermediate_size // num_key_value_heads)96 97    @classmethod98    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:99        name, gen = item100 101        # skip Multi-Token Prediction (MTP) layers (again, same as DeepseekV2)102        match = re.match(r"model.mtp_block.(\d+)", name)103        if match:104            return None105 106        # skip all other MTP tensors for now107        match = re.match(r"model.mtp_emb_norm.(\d+)", name)108        if match:109            return None110 111        match = re.match(r"model.mtp_hidden_norm.(\d+)", name)112        if match:113            return None114 115        match = re.match(r"model.mtp_linear_proj.(\d+)", name)116        if match:117            return None118 119        return super().filter_tensors(item)120 121    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:122        # process the experts separately123        if name.find("mlp.experts") != -1:124            n_experts = self.hparams["moe_num_experts"]125            assert bid is not None126 127            if self._experts is None:128                self._experts = [{} for _ in range(self.block_count)]129 130            self._experts[bid][name] = data_torch131 132            if len(self._experts[bid]) >= n_experts * 3:133                # merge the experts into a single 3d tensor134                for w_name in ["gate_proj", "up_proj", "down_proj"]:135                    datas: list[Tensor] = []136 137                    for xid in range(n_experts):138                        ename_to_retrieve = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight"139                        datas.append(self._experts[bid][ename_to_retrieve])140                        del self._experts[bid][ename_to_retrieve]141 142                    data_torch = torch.stack(datas, dim=0)143                    merged_name = f"model.layers.{bid}.mlp.experts.{w_name}.weight"144                    yield from super().modify_tensors(data_torch, merged_name, bid)145        else:146            yield from ModelBase.modify_tensors(self, data_torch, name, bid)147 148    def prepare_tensors(self):149        super().prepare_tensors()150 151        if self._experts is not None:152            # flatten `list[dict[str, Tensor]]` into `list[str]`153            experts = [k for d in self._experts for k in d.keys()]154            if len(experts) > 0:155                raise ValueError(f"Unprocessed experts: {experts}")156 157 158@ModelBase.register("PaddleOCRVLForConditionalGeneration")159class PaddleOCRModel(Ernie4_5Model):160    model_arch = gguf.MODEL_ARCH.PADDLEOCR161 162 163@ModelBase.register("PaddleOCRVisionModel")164class PaddleOCRVisionModel(MmprojModel):165    # PaddleOCR-VL uses a modified version of Siglip166    min_pixels: int = 0167    max_pixels: int = 0168 169    def __init__(self, *args, **kwargs):170        super().__init__(*args, **kwargs)171        assert self.hparams_vision is not None172        self.min_pixels = self.preprocessor_config["min_pixels"]173        self.max_pixels = self.preprocessor_config["max_pixels"]174        self.hparams_vision["image_size"] = int(math.sqrt(self.max_pixels))175 176    def set_gguf_parameters(self):177        super().set_gguf_parameters()178        assert self.hparams_vision is not None179        hparams = self.hparams_vision180        self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.PADDLEOCR)181        self.gguf_writer.add_vision_max_pixels(self.max_pixels)182        self.gguf_writer.add_vision_min_pixels(self.min_pixels)183        self.gguf_writer.add_vision_use_gelu(True)184        self.gguf_writer.add_vision_attention_layernorm_eps(hparams.get("rms_norm_eps", 1e-6))185 186    @classmethod187    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:188        name, gen = item189 190        if "vision_model" not in name and "mlp_AR" not in name:191            return None192        name = name.replace("visual.", "model.")193        if "packing_position_embedding" in name:194            # unused195            return None196        if "vision_model.head" in name:197            # we don't yet support image embeddings for this model198            return None199 200        return super().filter_tensors((name, gen))201