Felipe97/llama-cpp-compiled

1from __future__ import annotations2 3from typing import Iterable, TYPE_CHECKING4 5if TYPE_CHECKING:6    from torch import Tensor7 8from .base import ModelBase, TextModel, gguf9 10 11@ModelBase.register("DreamModel")12@ModelBase.example("Dream-org/Dream-v0-Instruct-7B")13class DreamModel(TextModel):14    model_arch = gguf.MODEL_ARCH.DREAM15 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 AutoTokenizer21        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_size26 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 tokens39                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, tokpre52 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 diffusion64        self.gguf_writer.add_causal_attention(False)65 66        # Add Dream-specific parameters67        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 model73        yield from super().modify_tensors(data_torch, name, bid)74