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 →
2.8 KiB73 linesraw
1from __future__ import annotations
2
3from typing import Iterable, TYPE_CHECKING
4
5if TYPE_CHECKING:
6    from torch import Tensor
7
8from .base import ModelBase, TextModel, gguf
9
10
11@ModelBase.register("DreamModel")
12@ModelBase.example("Dream-org/Dream-v0-Instruct-7B")
13class DreamModel(TextModel):
14    model_arch = gguf.MODEL_ARCH.DREAM
15
16    def get_vocab_base(self) -> tuple[list[str], list[int], str]:
17        tokens: list[str] = []
18        toktypes: list[int] = []
19
20        from transformers import AutoTokenizer
21        tokenizer = AutoTokenizer.from_pretrained(self.dir_model, trust_remote_code=True)
22
23        vocab_dict = tokenizer.get_vocab()  # ty: ignore[unresolved-attribute]
24        vocab_size = self.hparams.get("vocab_size", len(vocab_dict))
25        assert max(vocab_dict.values()) < vocab_size
26
27        tokpre = self.get_vocab_base_pre(tokenizer)
28
29        reverse_vocab = {id_: encoded_tok for encoded_tok, id_ in vocab_dict.items()}
30        added_vocab = tokenizer.get_added_vocab()  # ty: ignore[unresolved-attribute]
31
32        for i in range(vocab_size):
33            if i not in reverse_vocab:
34                tokens.append(f"[PAD{i}]")
35                toktypes.append(gguf.TokenType.UNUSED)
36            elif reverse_vocab[i] in added_vocab:
37                tokens.append(reverse_vocab[i])
38                # Check if it's a special token - treat special tokens as CONTROL tokens
39                if hasattr(tokenizer, 'added_tokens_decoder') and i in tokenizer.added_tokens_decoder:
40                    if tokenizer.added_tokens_decoder[i].special:
41                        toktypes.append(gguf.TokenType.CONTROL)
42                    else:
43                        toktypes.append(gguf.TokenType.USER_DEFINED)
44                else:
45                    # Fallback: treat all added vocab as control tokens for special tokens like <|im_start|>
46                    toktypes.append(gguf.TokenType.CONTROL)
47            else:
48                tokens.append(reverse_vocab[i])
49                toktypes.append(gguf.TokenType.NORMAL)
50
51        return tokens, toktypes, tokpre
52
53    def set_vocab(self):
54        try:
55            self._set_vocab_sentencepiece()
56        except FileNotFoundError:
57            self._set_vocab_gpt2()
58
59    def set_gguf_parameters(self):
60        super().set_gguf_parameters()
61        self._try_set_pooling_type()
62
63        # Dream models use non-causal attention for diffusion
64        self.gguf_writer.add_causal_attention(False)
65
66        # Add Dream-specific parameters
67        mask_token_id = self.hparams.get("mask_token_id")
68        if mask_token_id is not None:
69            self.gguf_writer.add_mask_token_id(mask_token_id)
70
71    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
72        # Dream model tensors should be mapped directly since it's the base model
73        yield from super().modify_tensors(data_torch, name, bid)