Felipe97/llama-cpp-compiled

1from __future__ import annotations2 3import json4import os5 6from pathlib import Path7from typing import Any, Callable, Iterable, TYPE_CHECKING8 9import torch10 11if TYPE_CHECKING:12    from torch import Tensor13 14from .base import ModelBase, SentencePieceTokenTypes, TextModel, gguf, logger15 16 17@ModelBase.register("BertModel", "BertForMaskedLM", "CamembertModel", "BertForSequenceClassification")18@ModelBase.example("BAAI/bge-small-en-v1.5", "dangvantuan/sentence-camembert-base")19class BertModel(TextModel):20    model_arch = gguf.MODEL_ARCH.BERT21 22    def __init__(self, *args, **kwargs):23        super().__init__(*args, **kwargs)24        self.vocab_size = None25 26        if cls_out_labels := self.hparams.get("id2label"):27            if len(cls_out_labels) == 2 and cls_out_labels[0] == "LABEL_0":28                # Remove dummy labels added by AutoConfig29                cls_out_labels = None30        self.cls_out_labels = cls_out_labels31 32    def set_gguf_parameters(self):33        super().set_gguf_parameters()34        self.gguf_writer.add_causal_attention(False)35        self._try_set_pooling_type()36 37        if self.cls_out_labels:38            self.gguf_writer.add_classifier_output_labels([v for k, v in sorted(self.cls_out_labels.items())])39 40    def set_vocab(self):41        tokens, toktypes, tokpre = self.get_vocab_base()42        self.vocab_size = len(tokens)43 44        # we need this to validate the size of the token_type embeddings45        # though currently we are passing all zeros to the token_type embeddings46        # "Sequence A" or "Sequence B"47        self.gguf_writer.add_token_type_count(self.hparams.get("type_vocab_size", 1))48 49        # convert to phantom space vocab50        def phantom(tok, toktype):51            if toktype == gguf.TokenType.CONTROL:52                return tok53            if tok.startswith("##"):54                return tok[2:]55            return "\u2581" + tok56        assert len(tokens) == len(toktypes)57        tokens = list(map(phantom, tokens, toktypes))58 59        # add vocab to gguf60        self.gguf_writer.add_tokenizer_model("bert")61        self.gguf_writer.add_tokenizer_pre(tokpre)62        self.gguf_writer.add_token_list(tokens)63        self.gguf_writer.add_token_types(toktypes)64 65        # handle special tokens66        special_vocab = gguf.SpecialVocab(self.dir_model, n_vocab=len(tokens))67        special_vocab.add_to_gguf(self.gguf_writer)68 69    @classmethod70    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:71        name, gen = item72 73        if name.startswith("bert."):74            name = name[5:]75 76        if name.endswith(".gamma"):77            name = name[:-6] + ".weight"78 79        if name.endswith(".beta"):80            name = name[:-5] + ".bias"81 82        # we are only using BERT for embeddings so we don't need the pooling layer83        if name in ("embeddings.position_ids", "pooler.dense.weight", "pooler.dense.bias"):84            return None85 86        if name.startswith("cls.predictions"):87            return None88 89        if name.startswith("cls.seq_relationship"):90            return None91 92        return super().filter_tensors((name, gen))93 94    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:95        if self.cls_out_labels:96            # For BertForSequenceClassification (direct projection layer)97            if name == "classifier.weight":98                name = "classifier.out_proj.weight"99 100            if name == "classifier.bias":101                name = "classifier.out_proj.bias"102 103        yield from super().modify_tensors(data_torch, name, bid)104 105    def _xlmroberta_tokenizer_init(self) -> None:106        # we need the pad_token_id to know how to chop down position_embd matrix107        if (pad_token_id := self.hparams.get("pad_token_id")) is not None:108            self._position_offset = 1 + pad_token_id109            if "max_position_embeddings" in self.hparams:110                self.hparams["max_position_embeddings"] -= self._position_offset111        else:112            self._position_offset = None113 114    def _xlmroberta_set_vocab(self) -> None:115        # to avoid TypeError: Descriptors cannot be created directly116        # exception when importing sentencepiece_model_pb2117        os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"118        from sentencepiece import SentencePieceProcessor119        from sentencepiece import sentencepiece_model_pb2 as model120 121        tokenizer_path = self.dir_model / 'sentencepiece.bpe.model'122 123        tokenizer_json = {}124        tokenizer_config_json = {}125        if not tokenizer_path.is_file():126            tokenizer_path = self.dir_model / 'tokenizer.json'127            tokenizer_config_path = self.dir_model / 'tokenizer_config.json'128 129            if not tokenizer_path.is_file():130                raise FileNotFoundError(f"File not found: {tokenizer_path}")131 132            from base64 import b64decode133            from transformers import AutoTokenizer134            tokenizer = AutoTokenizer.from_pretrained(self.dir_model)135 136            with open(tokenizer_path, "r", encoding="utf-8") as fp:137                tokenizer_json = json.load(fp)138 139            if tokenizer_config_path.is_file():140                with open(tokenizer_config_path, "r", encoding="utf-8") as fp:141                    tokenizer_config_json = json.load(fp)142 143            add_prefix = tokenizer.add_prefix_space  # ty: ignore[unresolved-attribute]144            remove_whitespaces = tokenizer.clean_up_tokenization_spaces  # ty: ignore[unresolved-attribute]145            precompiled_charsmap = b64decode(tokenizer_json["normalizer"]["precompiled_charsmap"])146 147            vocab_size = max(self.hparams.get("vocab_size", 0), tokenizer.vocab_size)  # ty: ignore[unresolved-attribute]148        else:149            sentencepiece_model = model.ModelProto()  # pyright: ignore[reportAttributeAccessIssue] # ty: ignore[unresolved-attribute]150            sentencepiece_model.ParseFromString(open(tokenizer_path, "rb").read())151            assert sentencepiece_model.trainer_spec.model_type == 1  # UNIGRAM152 153            add_prefix = sentencepiece_model.normalizer_spec.add_dummy_prefix154            remove_whitespaces = sentencepiece_model.normalizer_spec.remove_extra_whitespaces155            precompiled_charsmap = sentencepiece_model.normalizer_spec.precompiled_charsmap156 157            tokenizer = SentencePieceProcessor()158            tokenizer.LoadFromFile(str(tokenizer_path))159 160            vocab_size = max(self.hparams.get("vocab_size", 0), tokenizer.vocab_size())161 162        tokens: list[bytes] = [f"[PAD{i}]".encode("utf-8") for i in range(vocab_size)]163        scores: list[float] = [-10000.0] * vocab_size164        toktypes: list[int] = [SentencePieceTokenTypes.UNUSED] * vocab_size165 166        if isinstance(tokenizer, SentencePieceProcessor):167            for token_id in range(tokenizer.vocab_size()):168                piece = tokenizer.IdToPiece(token_id)169                text = piece.encode("utf-8")170                score = tokenizer.GetScore(token_id)171 172                toktype = SentencePieceTokenTypes.NORMAL173                if tokenizer.IsUnknown(token_id):174                    toktype = SentencePieceTokenTypes.UNKNOWN175                elif tokenizer.IsControl(token_id):176                    toktype = SentencePieceTokenTypes.CONTROL177                elif tokenizer.IsUnused(token_id):178                    toktype = SentencePieceTokenTypes.UNUSED179                elif tokenizer.IsByte(token_id):180                    toktype = SentencePieceTokenTypes.BYTE181 182                tokens[token_id] = text183                scores[token_id] = score184                toktypes[token_id] = toktype185        else:186            added_vocab = tokenizer.get_added_vocab()  # ty: ignore[unresolved-attribute]187            unk_token = tokenizer_config_json.get("unk_token")188            unk_token_id = added_vocab.get(unk_token, tokenizer_json["model"].get("unk_id", 3))  # ty: ignore[no-matching-overload]189 190            for token_id in range(tokenizer.vocab_size):  # ty: ignore[unresolved-attribute]191                piece = tokenizer._convert_id_to_token(token_id)  # ty: ignore[unresolved-attribute]192                if (piece := tokenizer._convert_id_to_token(token_id)) is not None:  # ty: ignore[unresolved-attribute]193                    text = piece.encode("utf-8")194                    score = tokenizer_json["model"]["vocab"][token_id][1]195 196                    toktype = SentencePieceTokenTypes.NORMAL197                    if token_id == unk_token_id:198                        toktype = SentencePieceTokenTypes.UNKNOWN199                    elif token_id in tokenizer.all_special_ids:  # ty: ignore[unresolved-attribute]200                        toktype = SentencePieceTokenTypes.CONTROL201                    elif token_id in added_vocab.values():202                        toktype = SentencePieceTokenTypes.USER_DEFINED203                    # No reliable way to detect this, but jina doesn't have any204                    # elif tokenizer.IsByte(token_id):205                    #     toktype = SentencePieceTokenTypes.BYTE206 207                    tokens[token_id] = text208                    scores[token_id] = score209                    toktypes[token_id] = toktype210 211        if isinstance(tokenizer, SentencePieceProcessor):212            # realign tokens (see HF tokenizer code)213            tokens = [b'<s>', b'<pad>', b'</s>', b'<unk>'] + tokens[3:-1]214            scores = [0.0, 0.0, 0.0, 0.0] + scores[3:-1]215            toktypes = [216                SentencePieceTokenTypes.CONTROL,217                SentencePieceTokenTypes.CONTROL,218                SentencePieceTokenTypes.CONTROL,219                SentencePieceTokenTypes.UNKNOWN,220            ] + toktypes[3:-1]221 222            if self.model_arch == gguf.MODEL_ARCH.NOMIC_BERT_MOE:223                # Add mask token missing from sentencepiece.bpe.model224                tokens[250001] = b'<mask>'225                scores[250001] = 0.0226                toktypes[250001] = SentencePieceTokenTypes.CONTROL227 228        self.gguf_writer.add_tokenizer_model("t5")229        self.gguf_writer.add_tokenizer_pre("default")230        self.gguf_writer.add_token_list(tokens)231        self.gguf_writer.add_token_scores(scores)232        self.gguf_writer.add_token_types(toktypes)233        self.gguf_writer.add_add_space_prefix(add_prefix)234        self.gguf_writer.add_token_type_count(self.hparams.get("type_vocab_size", 1))235        self.gguf_writer.add_remove_extra_whitespaces(remove_whitespaces)236        if precompiled_charsmap:237            self.gguf_writer.add_precompiled_charsmap(precompiled_charsmap)238 239        special_vocab = gguf.SpecialVocab(self.dir_model, n_vocab=len(tokens))240        special_vocab.add_to_gguf(self.gguf_writer)241 242 243@ModelBase.register("DistilBertModel", "DistilBertForMaskedLM", "DistilBertForSequenceClassification")244@ModelBase.example("distilbert/distilbert-base-uncased")245class DistilBertModel(BertModel):246    model_arch = gguf.MODEL_ARCH.BERT247 248    def set_gguf_parameters(self):249        self.gguf_writer.add_layer_norm_eps(1e-12)250        logger.info("gguf: layer norm epsilon = 1e-12")251        super().set_gguf_parameters()252 253    @classmethod254    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:255        name, gen = item256 257        if name.startswith("distilbert."):258            name = name[11:]259 260        # These layers act as MLM head, so we don't need them261        if name.startswith("vocab_"):262            return None263 264        return super().filter_tensors((name, gen))265 266 267@ModelBase.register("RobertaModel", "RobertaForSequenceClassification")268@ModelBase.example("sentence-transformers/stsb-roberta-base")269class RobertaModel(BertModel):270    model_arch = gguf.MODEL_ARCH.BERT271 272    def __init__(self, *args, **kwargs):273        super().__init__(*args, **kwargs)274 275        # we need the pad_token_id to know how to chop down position_embd matrix276        if (pad_token_id := self.hparams.get("pad_token_id")) is not None:277            self._position_offset = 1 + pad_token_id278            if "max_position_embeddings" in self.hparams:279                self.hparams["max_position_embeddings"] -= self._position_offset280        else:281            self._position_offset = None282 283    def set_vocab(self):284        """Support BPE tokenizers for roberta models"""285        bpe_tok_path = self.dir_model / "tokenizer.json"286        if bpe_tok_path.exists():287            self._set_vocab_gpt2()288 289            # we need this to validate the size of the token_type embeddings290            # though currently we are passing all zeros to the token_type embeddings291            # "Sequence A" or "Sequence B"292            self.gguf_writer.add_token_type_count(self.hparams.get("type_vocab_size", 1))293 294        else:295            return super().set_vocab()296 297    @classmethod298    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:299        name, gen = item300 301        # if name starts with "roberta.", remove the prefix302        # e.g. https://huggingface.co/BAAI/bge-reranker-v2-m3/tree/main303        if name.startswith("roberta."):304            name = name[8:]305 306        return super().filter_tensors((name, gen))307 308    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:309        # position embeddings start at pad_token_id + 1, so just chop down the weight tensor310        if name == "embeddings.position_embeddings.weight":311            if self._position_offset is not None:312                data_torch = data_torch[self._position_offset:,:]313 314        yield from super().modify_tensors(data_torch, name, bid)315 316 317@ModelBase.register("NomicBertModel")318@ModelBase.example("nomic-ai/nomic-embed-text-v1.5")319class NomicBertModel(BertModel):320    model_arch = gguf.MODEL_ARCH.BERT321 322    def __init__(self, dir_model: Path, ftype: gguf.LlamaFileType, fname_out: Path, **kwargs: Any):323        hparams = kwargs.pop("hparams", None)324        if hparams is None:325            hparams = ModelBase.load_hparams(dir_model, False)326 327        self.is_moe = bool(hparams.get("moe_every_n_layers"))328        self.model_arch = gguf.MODEL_ARCH.NOMIC_BERT_MOE if self.is_moe else gguf.MODEL_ARCH.NOMIC_BERT329 330        super().__init__(dir_model, ftype, fname_out, hparams=hparams, **kwargs)331 332        self._tokenizer_is_xlmroberta = self._is_tokenizer_xlmroberta()333        if self._tokenizer_is_xlmroberta:334            self._xlmroberta_tokenizer_init()335 336        npos, mtp = self.hparams["n_positions"], self.hparams.get("max_trained_positions", 2048)337        if npos == 8192 and mtp == 2048:338            self.hparams["n_positions"] = 2048  # nomic-embed-text v1 and v1.5 are trained for 2048 tokens.339        elif npos == 2048 and mtp == 2048:340            self.hparams["n_positions"] = 512   # nomic-embed-text-v2-moe is trained for 512 tokens.341        else:342            raise ValueError(f"unrecognized parameters: n_positions={npos}, max_trained_positions={mtp}")343 344        assert self.hparams["activation_function"] == "gelu" if self.is_moe else "swiglu"345 346        # this doesn't do anything in the HF version347        assert self.hparams["causal"] is False348        # no bias tensors unless MoE349        assert self.hparams["qkv_proj_bias"] == self.is_moe350        assert self.hparams["mlp_fc1_bias"]  == self.is_moe351        assert self.hparams["mlp_fc2_bias"]  == self.is_moe352 353        # norm at end of layer354        assert self.hparams["prenorm"] is False355        # standard RoPE356        assert self.hparams["rotary_emb_fraction"] == 1.0357        assert self.hparams["rotary_emb_interleaved"] is False358        assert self.hparams["rotary_emb_scale_base"] is None359 360    def set_vocab(self) -> None:361        if self._tokenizer_is_xlmroberta:362            return self._xlmroberta_set_vocab()363        return super().set_vocab()364 365    @classmethod366    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:367        name, gen = item368 369        # If the tensor is an experts bias tensor, skip it.370        if "mlp.experts.bias" in name:371            return None372 373        return super().filter_tensors(item)374 375    def modify_tensors(self, data_torch: torch.Tensor, name: str, bid: int | None) -> Iterable[tuple[str, torch.Tensor]]:376        if "mlp.experts.mlp.w1" in name:377            n_experts = self.find_hparam(["num_local_experts", "num_experts"])378            data_torch = data_torch.view(n_experts, self.hparams["n_inner"], self.hparams["n_embd"])379            name += ".weight"380 381        if "mlp.experts.mlp.w2" in name:382            n_experts = self.find_hparam(["num_local_experts", "num_experts"])383            data_torch = data_torch.view(n_experts, self.hparams["n_inner"], self.hparams["n_embd"])384            data_torch = data_torch.transpose(1, 2)385            name += ".weight"386 387        yield from super().modify_tensors(data_torch, name, bid)388 389    def set_gguf_parameters(self):390        super().set_gguf_parameters()391        if self.is_moe:392            self.gguf_writer.add_moe_every_n_layers(self.hparams["moe_every_n_layers"])393            self.gguf_writer.add_expert_used_count(self.hparams["moe_top_k"])394 395    def _is_tokenizer_xlmroberta(self) -> bool:396        with open(self.dir_model / "tokenizer.json") as f:397            tokenizer_json = json.load(f)398        toktyp = tokenizer_json["model"]["type"]399        if toktyp == "Unigram":400            return True401        if toktyp == "WordPiece":402            return False403        raise ValueError(f"unknown tokenizer: {toktyp}")404 405 406@ModelBase.register("NeoBERT", "NeoBERTLMHead", "NeoBERTForSequenceClassification")407@ModelBase.example("chandar-lab/NeoBERT")408class NeoBert(BertModel):409    model_arch = gguf.MODEL_ARCH.NEO_BERT410 411    def set_gguf_parameters(self):412        super().set_gguf_parameters()413 414        # NeoBERT uses 2/3 of the intermediate size as feed forward length415        self.gguf_writer.add_feed_forward_length(int(2 * self.hparams["intermediate_size"] / 3))416        self.gguf_writer.add_rope_freq_base(10000.0)  # default value for NeoBERT417        self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.NONE)418 419        f_rms_eps = self.hparams.get("norm_eps", 1e-6)  # default value for NeoBERT420        self.gguf_writer.add_layer_norm_rms_eps(f_rms_eps)421        logger.info(f"gguf: rms norm epsilon = {f_rms_eps}")422 423        self.gguf_writer.add_pooling_type(gguf.PoolingType.CLS) # https://huggingface.co/chandar-lab/NeoBERT#how-to-use424 425    @classmethod426    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:427        name, gen = item428 429        if name.startswith("decoder."):430            return None431 432        if name.startswith("model."):433            name = name[6:]434 435        return super().filter_tensors((name, gen))436 437 438@ModelBase.register("EuroBertModel", "JinaEmbeddingsV5Model")439@ModelBase.example("hf-tiny-v2/tiny-random-EuroBertModel", "jinaai/jina-embeddings-v5-text-nano")440class EuroBertModel(TextModel):441    model_arch = gguf.MODEL_ARCH.EUROBERT442 443    def set_vocab(self):444        self.gguf_writer.add_add_bos_token(False)445        self._set_vocab_gpt2()446 447    def set_gguf_parameters(self):448        super().set_gguf_parameters()449 450        # EuroBert is bidirectional (encoder)451        self.gguf_writer.add_causal_attention(False)452 453        self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.NONE)454 455        self._try_set_pooling_type()456 457    @classmethod458    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:459        name, gen = item460 461        if name.startswith("model."):462            name = name[6:]463 464        return super().filter_tensors((name, gen))465 466 467@ModelBase.register("XLMRobertaModel", "XLMRobertaForSequenceClassification")468@ModelBase.example("BAAI/bge-m3")469class XLMRobertaModel(BertModel):470    model_arch = gguf.MODEL_ARCH.BERT471    _lora_files = {}472    _lora_names = []473 474    def __init__(self, dir_model: Path, ftype: gguf.LlamaFileType, fname_out: Path, **kwargs: Any):475        hparams = kwargs.pop("hparams", None)476        if hparams is None:477            hparams = ModelBase.load_hparams(dir_model, False)478 479        if lora_names := hparams.get("lora_adaptations"):480            self._lora_names = lora_names481            self.model_arch = gguf.MODEL_ARCH.JINA_BERT_V3482 483        super().__init__(dir_model, ftype, fname_out, hparams=hparams, **kwargs)484        self._xlmroberta_tokenizer_init()485 486    def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:487        if self._lora_names:488            for name in self._lora_names:489                fname = self.add_prefix_to_filename(self.fname_out, f"lora-{name}-")490                self._lora_files[name] = gguf.GGUFWriter(fname, arch=gguf.MODEL_ARCH_NAMES[self.model_arch], endianess=self.endianess, use_temp_file=self.use_temp_file, dry_run=self.dry_run)491 492        return super().generate_extra_tensors()493 494    def set_type(self):495        for lora_writer in self._lora_files.values():496            lora_writer.add_type(gguf.GGUFType.ADAPTER)497            lora_writer.add_string(gguf.Keys.Adapter.TYPE, "lora")498        super().set_type()499 500    def set_vocab(self):501        self._xlmroberta_set_vocab()502 503    @classmethod504    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:505        name, gen = item506 507        # if name starts with "roberta.", remove the prefix508        # e.g. https://huggingface.co/BAAI/bge-reranker-v2-m3/tree/main509        if name.startswith("roberta."):510            name = name[8:]511 512        # jina-embeddings-v3513        if ".parametrizations." in name:514            name = name.replace(".parametrizations.", ".")515            if name.endswith(".original"):516                name = name[:-9]517 518        return super().filter_tensors((name, gen))519 520    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:521        # position embeddings start at pad_token_id + 1, so just chop down the weight tensor522        if name == "embeddings.position_embeddings.weight":523            if self._position_offset is not None:524                data_torch = data_torch[self._position_offset:,:]525 526        if name.endswith(".0.lora_A") or name.endswith(".0.lora_B"):527            if name.startswith("pooler.dense"):528                return529 530            num_loras = data_torch.size(0)531            assert num_loras == len(self._lora_names)532 533            # Split out each LoRA in their own GGUF534            for i, lora_writer in enumerate(self._lora_files.values()):535                new_name = self.map_tensor_name(name[:-9]) + name[-7:].lower()536                data = data_torch[i, :, :]537                # Transpose/flip token_embd/types into correct shape538                if new_name == "token_embd.weight.lora_b":539                    data = data.T540                elif new_name.startswith("token_types.weight."):541                    new_name = new_name[:-1] + ("a" if new_name[-1:] == "b" else "b")542                lora_writer.add_tensor(new_name, data.float().numpy(), raw_dtype=gguf.GGMLQuantizationType.F32)543 544            return545 546        yield from super().modify_tensors(data_torch, name, bid)547 548    def set_gguf_parameters(self):549        super().set_gguf_parameters()550 551        # jina-embeddings-v3552        lora_alpha = self.hparams.get("lora_alpha")553        if lora_prompt_prefixes := self.hparams.get("task_instructions"):554            assert self._lora_files and all(lora_name in lora_prompt_prefixes for lora_name in self._lora_files.keys())555        for lora_name, lora_writer in self._lora_files.items():556            lora_writer.add_float32(gguf.Keys.Adapter.LORA_ALPHA, lora_alpha if lora_alpha is not None else 1.0)557            lora_writer.add_string(gguf.Keys.Adapter.LORA_TASK_NAME, lora_name)558            if lora_prompt_prefixes:559                lora_writer.add_string(gguf.Keys.Adapter.LORA_PROMPT_PREFIX, lora_prompt_prefixes[lora_name])560 561    def write(self):562        super().write()563        for lora_writer in self._lora_files.values():564            lora_writer.write_header_to_file()565            lora_writer.write_kv_data_to_file()566            lora_writer.write_tensors_to_file(progress=True)567            lora_writer.close()568 569 570@ModelBase.register("JinaBertModel", "JinaBertForMaskedLM")571@ModelBase.example("jinaai/jina-embeddings-v2-base-en")572class JinaBertV2Model(BertModel):573    model_arch = gguf.MODEL_ARCH.JINA_BERT_V2574 575    def set_vocab(self):576        tokenizer_class = 'BertTokenizer'577        with open(self.dir_model / "tokenizer_config.json", "r", encoding="utf-8") as f:578            tokenizer_class = json.load(f)['tokenizer_class']579 580        if tokenizer_class == 'BertTokenizer':581            super().set_vocab()582        elif tokenizer_class == 'RobertaTokenizer':583            pre_tokenizer_type = None584            tokenizer_json_path = self.dir_model / "tokenizer.json"585            if tokenizer_json_path.is_file():586                with open(tokenizer_json_path, "r", encoding="utf-8") as f:587                    pre_tokenizer_type = json.load(f).get("pre_tokenizer", {}).get("type")588 589            if pre_tokenizer_type == "Whitespace":590                self._set_vocab_whitespace()591            else:592                self._set_vocab_gpt2()593            self.gguf_writer.add_token_type_count(2)594        else:595            raise NotImplementedError(f'Tokenizer {tokenizer_class} is not supported for JinaBertModel')596 597 598@ModelBase.register("ModernBertModel", "ModernBertForMaskedLM", "ModernBertForSequenceClassification")599@ModelBase.example("answerdotai/ModernBERT-base")600class ModernBertModel(BertModel):601    model_arch = gguf.MODEL_ARCH.MODERN_BERT602 603    def set_vocab(self):604        self.gguf_writer.add_add_bos_token(True)605        self.gguf_writer.add_add_eos_token(True)606        self.gguf_writer.add_add_sep_token(True)607        self._set_vocab_gpt2()608 609    def set_gguf_parameters(self):610        super().set_gguf_parameters()611        self.gguf_writer.add_sliding_window(self.hparams["local_attention"])612        if (sliding_window_pattern := self.hparams.get("global_attn_every_n_layers")) is not None:613            self.gguf_writer.add_sliding_window_pattern(sliding_window_pattern)614        self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.NONE)615        self.gguf_writer.add_vocab_size(self.hparams["vocab_size"])616        # FFN activation: ModernBert uses a GLU pair (ffn_up output is 2*n_ff). The617        # original ModernBERT uses GELU (-> GeGLU); some derivatives such as IBM618        # Granite Embedding 97m R2 use SiLU (-> SwiGLU). Persist this so the619        # llama.cpp graph can pick the matching activation.620        if hidden_act := self.hparams.get("hidden_activation"):621            self.gguf_writer.add_hidden_act(hidden_act)622 623    @classmethod624    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:625        name, gen = item626 627        if name.startswith("model."):628            name = name[6:]629 630        return super().filter_tensors((name, gen))631 632    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:633        if self.cls_out_labels:634            # For BertForSequenceClassification (direct projection layer)635            if name == "classifier.weight":636                name = "classifier.out_proj.weight"637 638            if name == "classifier.bias":639                name = "classifier.out_proj.bias"640 641        yield from super().modify_tensors(data_torch, name, bid)642