Felipe97/llama-cpp-compiled

1from __future__ import annotations2 3from typing import Any, Callable, Iterable, TYPE_CHECKING4 5import torch6 7if TYPE_CHECKING:8    from torch import Tensor9 10from .base import MmprojModel, ModelBase, TextModel, gguf11 12from .gemma import ConformerAudioModel13 14 15@ModelBase.register("Lfm2ForCausalLM", "LFM2ForCausalLM")16@ModelBase.example("LiquidAI/LFM2-1.2B", "LiquidAI/LFM2.5-350M")17class LFM2Model(TextModel):18    model_arch = gguf.MODEL_ARCH.LFM219 20    def _add_feed_forward_length(self):21        ff_dim = self.find_hparam(["block_ff_dim", "intermediate_size"])22        auto_adjust_ff_dim = self.hparams["block_auto_adjust_ff_dim"]23        ffn_dim_multiplier = self.hparams["block_ffn_dim_multiplier"]24        multiple_of = self.hparams["block_multiple_of"]25 26        if auto_adjust_ff_dim:27            ff_dim = int(2 * ff_dim / 3)28            # custom dim factor multiplier29            if ffn_dim_multiplier is not None:30                ff_dim = int(ffn_dim_multiplier * ff_dim)31            ff_dim = multiple_of * ((ff_dim + multiple_of - 1) // multiple_of)32 33        self.gguf_writer.add_feed_forward_length(ff_dim)34 35    def set_gguf_parameters(self):36        # set num_key_value_heads only for attention layers37        self.hparams["num_key_value_heads"] = [38            self.hparams["num_key_value_heads"] if layer_type != "conv" else 039            for layer_type in self.hparams["layer_types"]40        ]41 42        super().set_gguf_parameters()43        self.gguf_writer.add_vocab_size(self.hparams["vocab_size"])44        self.gguf_writer.add_shortconv_l_cache(self.hparams["conv_L_cache"])45        self.gguf_writer.add_layer_norm_rms_eps(self.hparams["norm_eps"])46        self._add_feed_forward_length()47 48    @classmethod49    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:50        name, gen = item51 52        if ConformerAudioModel.is_audio_tensor(name):53            # skip multimodal tensors54            return None55 56        name = name.replace("lfm.", "model.")      # audio57 58        return super().filter_tensors((name, gen))59 60    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:61        # conv op requires 2d tensor62        if 'conv.conv' in name:63            data_torch = data_torch.squeeze(1)64 65        yield from super().modify_tensors(data_torch, name, bid)66 67 68@ModelBase.register("Lfm2Model", "Lfm2BidirectionalModel")69@ModelBase.example("LiquidAI/LFM2.5-ColBERT-350M", "LiquidAI/LFM2.5-Embedding-350M")70class LFM2ColBertModel(LFM2Model):71    model_arch = gguf.MODEL_ARCH.LFM272    dense_tensor_name = "dense_2"73 74    def set_gguf_parameters(self):75        super().set_gguf_parameters()76        if self.hf_arch == "Lfm2BidirectionalModel":77            self.gguf_writer.add_causal_attention(False)78        self._try_set_pooling_type()79 80    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:81        if not name.startswith(self.dense_tensor_name):82            name = "model." + name83 84        yield from super().modify_tensors(data_torch, name, bid)85 86    def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:87        # optional dense tensor is stored in a separate safetensors file88        from safetensors.torch import load_file89        tensors_file = self.dir_model / "1_Dense" / "model.safetensors"90        if not tensors_file.is_file():91            return92        tensor = load_file(tensors_file)["linear.weight"]93        self.gguf_writer.add_embedding_length_out(tensor.shape[0])94        yield f"{self.dense_tensor_name}.weight", tensor.clone()95 96 97@ModelBase.register("Lfm2MoeForCausalLM")98@ModelBase.example("LiquidAI/LFM2-8B-A1B")99class LFM2MoeModel(TextModel):100    model_arch = gguf.MODEL_ARCH.LFM2MOE101 102    def set_gguf_parameters(self):103        # set num_key_value_heads only for attention layers104        self.hparams["num_key_value_heads"] = [105            self.hparams["num_key_value_heads"] if layer_type == "full_attention" else 0106            for layer_type in self.hparams["layer_types"]107        ]108 109        super().set_gguf_parameters()110 111        self.gguf_writer.add_expert_feed_forward_length(self.hparams["moe_intermediate_size"])112        self.gguf_writer.add_leading_dense_block_count(self.hparams["num_dense_layers"])113        self.gguf_writer.add_expert_gating_func(gguf.ExpertGatingFuncType.SIGMOID)114 115        self.gguf_writer.add_vocab_size(self.hparams["vocab_size"])116        self.gguf_writer.add_shortconv_l_cache(self.hparams["conv_L_cache"])117 118    # cache for experts weights for merging119    _experts_cache: dict[int, dict[str, Tensor]] = {}120 121    @classmethod122    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:123        name, gen = item124 125        if name.endswith(".expert_bias"):126            name = name.replace(".expert_bias", ".expert_bias.bias")127 128        return super().filter_tensors((name, gen))129 130    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:131        # conv op requires 2d tensor132        if 'conv.conv' in name:133            data_torch = data_torch.squeeze(1)134 135        # merge expert weights136        if 'experts' in name:137            n_experts = self.find_hparam(["num_local_experts", "num_experts"])138            assert bid is not None139 140            expert_cache = self._experts_cache.setdefault(bid, {})141            expert_cache[name] = data_torch142            expert_weights = ["w1", "w2", "w3"]143 144            # not enough expert weights to merge145            if len(expert_cache) < n_experts * len(expert_weights):146                return147 148            for w_name in expert_weights:149                datas: list[Tensor] = []150 151                for xid in range(n_experts):152                    ename = f"model.layers.{bid}.feed_forward.experts.{xid}.{w_name}.weight"153                    datas.append(expert_cache[ename])154                    del expert_cache[ename]155 156                data_torch = torch.stack(datas, dim=0)157                merged_name = f"layers.{bid}.feed_forward.experts.{w_name}.weight"158 159                yield from super().modify_tensors(data_torch, merged_name, bid)160 161            del self._experts_cache[bid]162            return163 164        yield from super().modify_tensors(data_torch, name, bid)165 166    def prepare_tensors(self):167        super().prepare_tensors()168        assert not self._experts_cache169 170 171@ModelBase.register("Lfm2VlForConditionalGeneration")172@ModelBase.example("LiquidAI/LFM2-VL-450M")173class LFM2VLModel(MmprojModel):174    def __init__(self, *args, **kwargs):175        super().__init__(*args, **kwargs)176        assert self.hparams_vision is not None177        # TODO(tarek): for dynamic resolution image_size is not specified, setting here for compatibility178        self.hparams_vision["image_size"] = 256179 180    def set_gguf_parameters(self):181        super().set_gguf_parameters()182        self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.LFM2)183        self.gguf_writer.add_vision_attention_layernorm_eps(self.find_vparam(["layer_norm_eps"]))184        self.gguf_writer.add_vision_projector_scale_factor(self.global_config.get("downsample_factor", 2))185        self.gguf_writer.add_vision_use_gelu(True)186        # python notation, e.g. for vision_feature_layer == -1, we pick last layer -> vision_feature_layers_to_drop = 0187        vision_feature_layers_to_drop = -(self.global_config.get("vision_feature_layer", -1) + 1)188        self.gguf_writer.add_vision_block_count(self.find_vparam(self.n_block_keys) - vision_feature_layers_to_drop)189 190    @classmethod191    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:192        name, gen = item193 194        name = name.replace("model.vision_tower.", "vision_tower.")195        name = name.replace("model.multi_modal_projector.", "multi_modal_projector.")196 197        return super().filter_tensors((name, gen))198 199    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:200        if "patch_embedding.weight" in name:201            data_torch = data_torch.view(data_torch.shape[0], 16, 16, 3).permute(0, 3, 1, 2)202 203        yield from super().modify_tensors(data_torch, name, bid)204 205 206@ModelBase.register("Lfm2AudioForConditionalGeneration")207@ModelBase.example("LiquidAI/LFM2.5-Audio-1.5B", "LiquidAI/LFM2-Audio-1.5B")208class LFM2AudioModel(ConformerAudioModel):209    has_vision_encoder = False210    has_audio_encoder = True211    model_name = "Lfm2AudioEncoder"212 213    def get_audio_config(self) -> dict[str, Any] | None:214        return self.global_config.get("encoder")215 216    def set_gguf_parameters(self):217        assert self.hparams_audio is not None218        self.hparams_audio["hidden_size"] = self.hparams_audio["d_model"]219        self.hparams_audio["intermediate_size"] = self.hparams_audio["d_model"]220        self.hparams_audio["num_attention_heads"] = self.hparams_audio["n_heads"]221        super().set_gguf_parameters()222        self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.LFM2A)223        self.gguf_writer.add_audio_num_mel_bins(self.hparams_audio["feat_in"])224        self.gguf_writer.add_audio_attention_layernorm_eps(1e-5)225 226    @classmethod227    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:228        name, gen = item229 230        # skip language model tensors231        if name.startswith("lfm."):232            return None233 234        # for training only235        if any(p in name for p in ["audio_loss_weight"]):236            return None237 238        # for audio output239        if any(p in name for p in ["codebook_offsets", "depth_embeddings", "depth_linear", "depthformer"]):240            return None241 242        return super().filter_tensors(item)243 244 245@ModelBase.register("Lfm25AudioTokenizer")246@ModelBase.example("LiquidAI/LFM2.5-Audio-1.5B")247class LFM25AudioTokenizer(LFM2Model):248    model_arch = gguf.MODEL_ARCH.LFM2249 250    def set_vocab(self):251        self._set_vocab_none()252 253    def set_gguf_parameters(self):254        super().set_gguf_parameters()255        self.gguf_writer.add_sliding_window(self.hparams["sliding_window"])256        self.gguf_writer.add_embedding_length_out(self.hparams["output_size"])257 258    @classmethod259    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:260        name, gen = item261 262        # skip language model tensors263        if name == "istft.window" or name.startswith("emb.emb"):264            return None265 266        if name.startswith("lin"):267            name = name.replace("lin", "dense_2_out")268 269        return super().filter_tensors((name, gen))270