char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Xuan-Son Nguyenconvert: add @ModelBase.example (#27208)fa88ae936

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
26.9 KiB641 linesraw
1from __future__ import annotations
2
3import json
4import os
5
6from pathlib import Path
7from typing import Any, Callable, Iterable, TYPE_CHECKING
8
9import torch
10
11if TYPE_CHECKING:
12    from torch import Tensor
13
14from .base import ModelBase, SentencePieceTokenTypes, TextModel, gguf, logger
15
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.BERT
21
22    def __init__(self, *args, **kwargs):
23        super().__init__(*args, **kwargs)
24        self.vocab_size = None
25
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 AutoConfig
29                cls_out_labels = None
30        self.cls_out_labels = cls_out_labels
31
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 embeddings
45        # though currently we are passing all zeros to the token_type embeddings
46        # "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 vocab
50        def phantom(tok, toktype):
51            if toktype == gguf.TokenType.CONTROL:
52                return tok
53            if tok.startswith("##"):
54                return tok[2:]
55            return "\u2581" + tok
56        assert len(tokens) == len(toktypes)
57        tokens = list(map(phantom, tokens, toktypes))
58
59        # add vocab to gguf
60        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 tokens
66        special_vocab = gguf.SpecialVocab(self.dir_model, n_vocab=len(tokens))
67        special_vocab.add_to_gguf(self.gguf_writer)
68
69    @classmethod
70    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
71        name, gen = item
72
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 layer
83        if name in ("embeddings.position_ids", "pooler.dense.weight", "pooler.dense.bias"):
84            return None
85
86        if name.startswith("cls.predictions"):
87            return None
88
89        if name.startswith("cls.seq_relationship"):
90            return None
91
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 matrix
107        if (pad_token_id := self.hparams.get("pad_token_id")) is not None:
108            self._position_offset = 1 + pad_token_id
109            if "max_position_embeddings" in self.hparams:
110                self.hparams["max_position_embeddings"] -= self._position_offset
111        else:
112            self._position_offset = None
113
114    def _xlmroberta_set_vocab(self) -> None:
115        # to avoid TypeError: Descriptors cannot be created directly
116        # exception when importing sentencepiece_model_pb2
117        os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"
118        from sentencepiece import SentencePieceProcessor
119        from sentencepiece import sentencepiece_model_pb2 as model
120
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 b64decode
133            from transformers import AutoTokenizer
134            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  # UNIGRAM
152
153            add_prefix = sentencepiece_model.normalizer_spec.add_dummy_prefix
154            remove_whitespaces = sentencepiece_model.normalizer_spec.remove_extra_whitespaces
155            precompiled_charsmap = sentencepiece_model.normalizer_spec.precompiled_charsmap
156
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_size
164        toktypes: list[int] = [SentencePieceTokenTypes.UNUSED] * vocab_size
165
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.NORMAL
173                if tokenizer.IsUnknown(token_id):
174                    toktype = SentencePieceTokenTypes.UNKNOWN
175                elif tokenizer.IsControl(token_id):
176                    toktype = SentencePieceTokenTypes.CONTROL
177                elif tokenizer.IsUnused(token_id):
178                    toktype = SentencePieceTokenTypes.UNUSED
179                elif tokenizer.IsByte(token_id):
180                    toktype = SentencePieceTokenTypes.BYTE
181
182                tokens[token_id] = text
183                scores[token_id] = score
184                toktypes[token_id] = toktype
185        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.NORMAL
197                    if token_id == unk_token_id:
198                        toktype = SentencePieceTokenTypes.UNKNOWN
199                    elif token_id in tokenizer.all_special_ids:  # ty: ignore[unresolved-attribute]
200                        toktype = SentencePieceTokenTypes.CONTROL
201                    elif token_id in added_vocab.values():
202                        toktype = SentencePieceTokenTypes.USER_DEFINED
203                    # No reliable way to detect this, but jina doesn't have any
204                    # elif tokenizer.IsByte(token_id):
205                    #     toktype = SentencePieceTokenTypes.BYTE
206
207                    tokens[token_id] = text
208                    scores[token_id] = score
209                    toktypes[token_id] = toktype
210
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.model
224                tokens[250001] = b'<mask>'
225                scores[250001] = 0.0
226                toktypes[250001] = SentencePieceTokenTypes.CONTROL
227
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.BERT
247
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    @classmethod
254    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
255        name, gen = item
256
257        if name.startswith("distilbert."):
258            name = name[11:]
259
260        # These layers act as MLM head, so we don't need them
261        if name.startswith("vocab_"):
262            return None
263
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.BERT
271
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 matrix
276        if (pad_token_id := self.hparams.get("pad_token_id")) is not None:
277            self._position_offset = 1 + pad_token_id
278            if "max_position_embeddings" in self.hparams:
279                self.hparams["max_position_embeddings"] -= self._position_offset
280        else:
281            self._position_offset = None
282
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 embeddings
290            # though currently we are passing all zeros to the token_type embeddings
291            # "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    @classmethod
298    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
299        name, gen = item
300
301        # if name starts with "roberta.", remove the prefix
302        # e.g. https://huggingface.co/BAAI/bge-reranker-v2-m3/tree/main
303        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 tensor
310        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.BERT
321
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_BERT
329
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 version
347        assert self.hparams["causal"] is False
348        # no bias tensors unless MoE
349        assert self.hparams["qkv_proj_bias"] == self.is_moe
350        assert self.hparams["mlp_fc1_bias"]  == self.is_moe
351        assert self.hparams["mlp_fc2_bias"]  == self.is_moe
352
353        # norm at end of layer
354        assert self.hparams["prenorm"] is False
355        # standard RoPE
356        assert self.hparams["rotary_emb_fraction"] == 1.0
357        assert self.hparams["rotary_emb_interleaved"] is False
358        assert self.hparams["rotary_emb_scale_base"] is None
359
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    @classmethod
366    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
367        name, gen = item
368
369        # If the tensor is an experts bias tensor, skip it.
370        if "mlp.experts.bias" in name:
371            return None
372
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 True
401        if toktyp == "WordPiece":
402            return False
403        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_BERT
410
411    def set_gguf_parameters(self):
412        super().set_gguf_parameters()
413
414        # NeoBERT uses 2/3 of the intermediate size as feed forward length
415        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 NeoBERT
417        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 NeoBERT
420        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-use
424
425    @classmethod
426    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
427        name, gen = item
428
429        if name.startswith("decoder."):
430            return None
431
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.EUROBERT
442
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    @classmethod
458    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
459        name, gen = item
460
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.BERT
471    _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_names
481            self.model_arch = gguf.MODEL_ARCH.JINA_BERT_V3
482
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    @classmethod
504    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
505        name, gen = item
506
507        # if name starts with "roberta.", remove the prefix
508        # e.g. https://huggingface.co/BAAI/bge-reranker-v2-m3/tree/main
509        if name.startswith("roberta."):
510            name = name[8:]
511
512        # jina-embeddings-v3
513        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 tensor
522        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                return
529
530            num_loras = data_torch.size(0)
531            assert num_loras == len(self._lora_names)
532
533            # Split out each LoRA in their own GGUF
534            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 shape
538                if new_name == "token_embd.weight.lora_b":
539                    data = data.T
540                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            return
545
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-v3
552        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_V2
574
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 = None
584            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_BERT
602
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). The
617        # original ModernBERT uses GELU (-> GeGLU); some derivatives such as IBM
618        # Granite Embedding 97m R2 use SiLU (-> SwiGLU). Persist this so the
619        # 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    @classmethod
624    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
625        name, gen = item
626
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)