Brunobkr/llama.cpp_AlgMor24_github

1from __future__ import annotations2 3from typing import Callable, TYPE_CHECKING4 5if TYPE_CHECKING:6    from torch import Tensor7 8from .base import ModelBase, TextModel, gguf, logger9 10 11@ModelBase.register("WavTokenizerDec")12class WavTokenizerDecModel(TextModel):13    model_arch = gguf.MODEL_ARCH.WAVTOKENIZER_DEC14 15    @classmethod16    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:17        name, gen = item18 19        if \20                name.endswith("codebook.cluster_size") or \21                name.endswith("codebook.embed_avg") or \22                name.endswith("codebook.inited"):23            logger.debug(f"Skipping {name!r}")24            return None25 26        return super().filter_tensors(item)27 28    def set_vocab(self):29        self._set_vocab_none()30 31    def set_gguf_parameters(self):32        super().set_gguf_parameters()33        self.gguf_writer.add_vocab_size         (self.hparams["vocab_size"])34        self.gguf_writer.add_features_length    (self.hparams["n_embd_features"])35        self.gguf_writer.add_feed_forward_length(self.hparams["n_ff"])36        self.gguf_writer.add_group_norm_eps     (self.hparams["group_norm_epsilon"])37        self.gguf_writer.add_group_norm_groups  (self.hparams["group_norm_groups"])38 39        self.gguf_writer.add_posnet_embedding_length(self.hparams["posnet"]["n_embd"])40        self.gguf_writer.add_posnet_block_count     (self.hparams["posnet"]["n_layer"])41 42        self.gguf_writer.add_convnext_embedding_length(self.hparams["convnext"]["n_embd"])43        self.gguf_writer.add_convnext_block_count     (self.hparams["convnext"]["n_layer"])44 45        self.gguf_writer.add_causal_attention(False)46