上传 global_step_224 at 2025-11-18 18:51:35

added_tokens.json

@@ -1,26 +0,0 @@
-{
-  "</tool_call>": 151658,
-  "<tool_call>": 151657,
-  "<|beginoftext|>": 151665,
-  "<|box_end|>": 151649,
-  "<|box_start|>": 151648,
-  "<|endoftext|>": 151643,
-  "<|file_sep|>": 151664,
-  "<|fim_middle|>": 151660,
-  "<|fim_pad|>": 151662,
-  "<|fim_prefix|>": 151659,
-  "<|fim_suffix|>": 151661,
-  "<|im_end|>": 151645,
-  "<|im_start|>": 151644,
-  "<|image_pad|>": 151655,
-  "<|mask|>": 151666,
-  "<|object_ref_end|>": 151647,
-  "<|object_ref_start|>": 151646,
-  "<|quad_end|>": 151651,
-  "<|quad_start|>": 151650,
-  "<|repo_name|>": 151663,
-  "<|video_pad|>": 151656,
-  "<|vision_end|>": 151653,
-  "<|vision_pad|>": 151654,
-  "<|vision_start|>": 151652
-}

chat_template.jinja

@@ -1,54 +0,0 @@
-{%- if tools %}
-    {{- '<|im_start|>system\n' }}
-    {%- if messages[0]['role'] == 'system' %}
-        {{- messages[0]['content'] }}
-    {%- else %}
-        {{- 'You are a helpful assistant.' }}
-    {%- endif %}
-    {{- "\n\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
-    {%- for tool in tools %}
-        {{- "\n" }}
-        {{- tool | tojson }}
-    {%- endfor %}
-    {{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n" }}
-{%- else %}
-    {%- if messages[0]['role'] == 'system' %}
-        {{- '<|im_start|>system\n' + messages[0]['content'] + '<|im_end|>\n' }}
-    {%- else %}
-        {{- '<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n' }}
-    {%- endif %}
-{%- endif %}
-{%- for message in messages %}
-    {%- if (message.role == "user") or (message.role == "system" and not loop.first) or (message.role == "assistant" and not message.tool_calls) %}
-        {{- '<|im_start|>' + message.role + '\n' + message.content + '<|im_end|>' + '\n' }}
-    {%- elif message.role == "assistant" %}
-        {{- '<|im_start|>' + message.role }}
-        {%- if message.content %}
-            {{- '\n' + message.content }}
-        {%- endif %}
-        {%- for tool_call in message.tool_calls %}
-            {%- if tool_call.function is defined %}
-                {%- set tool_call = tool_call.function %}
-            {%- endif %}
-            {{- '\n<tool_call>\n{"name": "' }}
-            {{- tool_call.name }}
-            {{- '", "arguments": ' }}
-            {{- tool_call.arguments | tojson }}
-            {{- '}\n</tool_call>' }}
-        {%- endfor %}
-        {{- '<|im_end|>\n' }}
-    {%- elif message.role == "tool" %}
-        {%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != "tool") %}
-            {{- '<|im_start|>user' }}
-        {%- endif %}
-        {{- '\n<tool_response>\n' }}
-        {{- message.content }}
-        {{- '\n</tool_response>' }}
-        {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
-            {{- '<|im_end|>\n' }}
-        {%- endif %}
-    {%- endif %}
-{%- endfor %}
-{%- if add_generation_prompt %}
-    {{- '<|im_start|>assistant\n' }}
-{%- endif %}

config.json

@@ -1,35 +0,0 @@
-{
-  "architectures": [
-    "DreamModel"
-  ],
-  "attention_dropout": 0.0,
-  "auto_map": {
-    "AutoConfig": "configuration_dream.DreamConfig",
-    "AutoModel": "modeling_dream.DreamModel"
-  },
-  "bos_token_id": 151643,
-  "dtype": "float32",
-  "eos_token_id": 151643,
-  "hidden_act": "silu",
-  "hidden_size": 3584,
-  "initializer_range": 0.02,
-  "intermediate_size": 18944,
-  "mask_token_id": 151666,
-  "max_position_embeddings": 131072,
-  "max_window_layers": 28,
-  "model_type": "Dream",
-  "num_attention_heads": 28,
-  "num_hidden_layers": 28,
-  "num_key_value_heads": 4,
-  "pad_token_id": 151643,
-  "rms_norm_eps": 1e-06,
-  "rope_scaling": null,
-  "rope_theta": 1000000.0,
-  "sliding_window": null,
-  "tie_word_embeddings": false,
-  "transformers_version": "4.57.1",
-  "use_cache": true,
-  "use_mrope": false,
-  "use_sliding_window": false,
-  "vocab_size": 152064
-}

configuration_dream.py

@@ -1,86 +0,0 @@
-# coding=utf-8
-# Copyright 2024 The Dream team, HKUNLP Group and the HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Dream model configuration"""
-
-from transformers.configuration_utils import PretrainedConfig
-from transformers.modeling_rope_utils import rope_config_validation
-from transformers.utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-class DreamConfig(PretrainedConfig):
-    model_type = "Dream"
-    keys_to_ignore_at_inference = ["past_key_values"]
-
-    def __init__(
-        self,
-        vocab_size=151936,
-        hidden_size=4096,
-        intermediate_size=22016,
-        num_hidden_layers=32,
-        num_attention_heads=32,
-        num_key_value_heads=32,
-        hidden_act="silu",
-        max_position_embeddings=32768,
-        initializer_range=0.02,
-        rms_norm_eps=1e-6,
-        use_cache=False,  # cache not used in diffusion
-        tie_word_embeddings=False,
-        rope_theta=10000.0,
-        rope_scaling=None,
-        use_sliding_window=False,
-        sliding_window=4096,
-        max_window_layers=28,
-        attention_dropout=0.0,
-        mask_token_id=151666,
-        pad_token_id=151643,
-        **kwargs,
-    ):
-        self.vocab_size = vocab_size
-        self.max_position_embeddings = max_position_embeddings
-        self.hidden_size = hidden_size
-        self.intermediate_size = intermediate_size
-        self.num_hidden_layers = num_hidden_layers
-        self.num_attention_heads = num_attention_heads
-        self.use_sliding_window = use_sliding_window
-        self.sliding_window = sliding_window if use_sliding_window else None
-        self.max_window_layers = max_window_layers
-
-        # for backward compatibility
-        if num_key_value_heads is None:
-            num_key_value_heads = num_attention_heads
-
-        self.num_key_value_heads = num_key_value_heads
-        self.hidden_act = hidden_act
-        self.initializer_range = initializer_range
-        self.rms_norm_eps = rms_norm_eps
-        self.use_cache = use_cache
-        self.rope_theta = rope_theta
-        self.rope_scaling = rope_scaling
-        self.attention_dropout = attention_dropout
-        # Validate the correctness of rotary position embeddings parameters
-        # BC: if there is a 'type' field, move it to 'rope_type'.
-        if self.rope_scaling is not None and "type" in self.rope_scaling:
-            self.rope_scaling["rope_type"] = self.rope_scaling["type"]
-        rope_config_validation(self)
-        
-        super().__init__(
-            tie_word_embeddings=tie_word_embeddings,
-            **kwargs,
-        )
-        self.mask_token_id = mask_token_id
-        self.pad_token_id = pad_token_id

generation_config.json

@@ -1,11 +0,0 @@
-{
-  "_from_model_config": true,
-  "bos_token_id": 151643,
-  "eos_token_id": 151643,
-  "pad_token_id": 151643,
-  "temperature": null,
-  "top_k": null,
-  "top_p": null,
-  "transformers_version": "4.57.1",
-  "typical_p": null
-}

generation_utils.py

@@ -1,464 +0,0 @@
-# coding=utf-8
-# Copyright 2024 The Dream team, HKUNLP Group and the HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import warnings
-import copy
-from dataclasses import dataclass
-from typing import Any, Dict, Optional, Tuple, Union
-
-import torch
-import torch.distributions as dists
-from torch.nn import functional as F
-from transformers import __version__
-from transformers.generation.configuration_utils import (
-    GenerationConfig
-)
-from transformers.utils import (
-    ModelOutput,
-    is_torchdynamo_compiling,
-    logging,
-)
-
-logger = logging.get_logger(__name__)
-
-
-def top_p_logits(logits, top_p=None):
-    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
-    cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
-    sorted_indices_to_remove = cumulative_probs > top_p
-    # Shift the indices to the right to keep the first token above the threshold
-    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
-    sorted_indices_to_remove[..., 0] = 0
-
-    mask = torch.zeros_like(logits, dtype=torch.bool, device=logits.device)
-    mask = mask.scatter_(-1, sorted_indices, sorted_indices_to_remove)
-    logits = logits.masked_fill(mask, torch.finfo(logits.dtype).min)
-    return logits
-
-def top_k_logits(logits, top_k=None):
-    top_k = min(top_k, logits.size(-1))  # Safety check
-    # Remove all tokens with a probability less than the last token of the top-k
-    indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
-    logits = logits.masked_fill(indices_to_remove, torch.finfo(logits.dtype).min)
-    return logits
-
-
-def sample_tokens(logits, temperature=0.0, top_p=None, top_k=None, margin_confidence=False, neg_entropy=False):
-
-    if temperature > 0:
-        logits = logits / temperature
-    if top_p is not None and top_p < 1:
-        logits = top_p_logits(logits, top_p)
-    if top_k is not None:
-        logits = top_k_logits(logits, top_k)
-    probs = torch.softmax(logits, dim=-1)
-
-    if temperature > 0:
-        try:
-            x0 = dists.Categorical(probs=probs).sample()
-            confidence = torch.gather(probs, -1, x0.unsqueeze(-1)).squeeze(-1)
-        except:
-            confidence, x0 = probs.max(dim=-1)
-    else:
-        confidence, x0 = probs.max(dim=-1)
-    
-    if margin_confidence:
-        sorted_probs, _ = torch.sort(probs, dim=-1, descending=True)
-        # Extract top1 and top2 probabilities
-        top1_probs = sorted_probs[:, 0] 
-        top2_probs = sorted_probs[:, 1] 
-        # Calculate confidence as top1 - top2
-        confidence = top1_probs - top2_probs 
-    
-    if neg_entropy:
-        epsilon = 1e-10
-        log_probs = torch.log(probs + epsilon)
-        confidence = torch.sum(probs * log_probs, dim=-1)
-    
-    return confidence, x0
-
-
-@dataclass
-class DreamModelOutput(ModelOutput):
-    sequences: torch.LongTensor = None
-    history: Optional[Tuple[torch.FloatTensor]] = None
-
-
-class DreamGenerationConfig(GenerationConfig):
-    def __init__(self, **kwargs):
-        self.temperature: float = kwargs.pop("temperature", 0.0)
-        self.top_p: Optional[float] = kwargs.pop("top_p", None)
-        self.top_k: Optional[int] = kwargs.pop("top_k", None)
-        self.max_length = kwargs.pop("max_length", 20)
-        self.max_new_tokens = kwargs.pop("max_new_tokens", None)
-        # diffusion specific params
-        self.eps: float = kwargs.pop("eps", 1e-3)
-        self.steps: int = kwargs.pop("steps", 512)
-        self.alg: str = kwargs.pop("alg", 'origin')
-        self.alg_temp: Optional[float] = kwargs.pop("alg_temp", None)
-
-        # Parameters that define the output variables of `generate`
-        self.num_return_sequences: int = kwargs.pop("num_return_sequences", 1)
-        self.return_dict_in_generate: bool = kwargs.pop("return_dict_in_generate", False)
-        self.output_history: bool = kwargs.pop("output_history", False)
-
-        # Special tokens that can be used at generation time
-        self.mask_token_id = kwargs.pop("mask_token_id", None)
-        self.pad_token_id = kwargs.pop("pad_token_id", None)
-        self.bos_token_id = kwargs.pop("bos_token_id", None)
-        self.eos_token_id = kwargs.pop("eos_token_id", None)
-
-        # Wild card
-        self.generation_kwargs = kwargs.pop("generation_kwargs", {})
-
-        # The remaining attributes do not parametrize `.generate()`, but are informative and/or used by the hub
-        # interface.
-        self._from_model_config = kwargs.pop("_from_model_config", False)
-        self._commit_hash = kwargs.pop("_commit_hash", None)
-        self.transformers_version = kwargs.pop("transformers_version", __version__)
-
-        # Additional attributes without default values
-        if not self._from_model_config:
-            # we don't want to copy values from the model config if we're initializing a `GenerationConfig` from a
-            # model's default configuration file
-            for key, value in kwargs.items():
-                try:
-                    setattr(self, key, value)
-                except AttributeError as err:
-                    logger.error(f"Can't set {key} with value {value} for {self}")
-                    raise err
-
-        # Validate the values of the attributes
-        self.validate(is_init=True)
-
-    def validate(self, is_init=False):
-        pass
-
-class DreamGenerationMixin:
-    @staticmethod
-    def _expand_inputs_for_generation(
-        expand_size: int = 1,
-        input_ids: Optional[torch.LongTensor] = None,
-        attention_mask: Optional[torch.LongTensor] = None
-    ) -> Tuple[torch.LongTensor, Dict[str, Any]]:
-        """Expands tensors from [batch_size, ...] to [batch_size * expand_size, ...]"""
-        # Do not call torch.repeat_interleave if expand_size is 1 because it clones
-        # the input tensor and thus requires more memory although no change is applied
-        if expand_size == 1:
-            return input_ids, attention_mask
-        if input_ids is not None:
-            input_ids = input_ids.repeat_interleave(expand_size, dim=0)
-        if attention_mask is not None:
-            attention_mask = attention_mask.repeat_interleave(expand_size, dim=0)
-        return input_ids, attention_mask
-
-    def _validate_generated_length(self, generation_config, input_ids_length, has_default_max_length):
-        """Performs validation related to the resulting generated length"""
-
-        # Can't throw warnings/exceptions during compilation
-        if is_torchdynamo_compiling():
-            return
-
-        # 1. Max length warnings related to poor parameterization
-        if has_default_max_length and generation_config.max_new_tokens is None and generation_config.max_length == 20:
-            # 20 is the default max_length of the generation config
-            warnings.warn(
-                f"Using the model-agnostic default `max_length` (={generation_config.max_length}) to control the "
-                "generation length. We recommend setting `max_new_tokens` to control the maximum length of the "
-                "generation.",
-                UserWarning,
-            )
-        if input_ids_length >= generation_config.max_length:
-            input_ids_string = "input_ids"
-            raise ValueError(
-                f"Input length of {input_ids_string} is {input_ids_length}, but `max_length` is set to"
-                f" {generation_config.max_length}. This can lead to unexpected behavior. You should consider"
-                " increasing `max_length` or, better yet, setting `max_new_tokens`."
-            )
-
-    def _prepare_generated_length(
-        self,
-        generation_config,
-        has_default_max_length,
-        input_ids_length,
-    ):
-        """Prepared max and min length in generation configs to avoid clashes between similar attributes"""
-
-        if generation_config.max_new_tokens is not None:
-            if not has_default_max_length and generation_config.max_length is not None:
-                logger.warning(
-                    f"Both `max_new_tokens` (={generation_config.max_new_tokens}) and `max_length`(="
-                    f"{generation_config.max_length}) seem to have been set. `max_new_tokens` will take precedence. "
-                    "Please refer to the documentation for more information. "
-                    "(https://huggingface.co/docs/transformers/main/en/main_classes/text_generation)"
-                )
-            generation_config.max_length = generation_config.max_new_tokens + input_ids_length
-
-        elif has_default_max_length:
-            if generation_config.max_length == DreamGenerationConfig().max_length:
-                generation_config.max_length = generation_config.max_length + input_ids_length
-                max_position_embeddings = getattr(self.config, "max_position_embeddings", None)
-                if max_position_embeddings is not None:
-                    generation_config.max_length = min(generation_config.max_length, max_position_embeddings)
-
-        return generation_config
-
-    def _prepare_generation_config(
-        self, generation_config: Optional[DreamGenerationConfig], **kwargs: Dict
-    ) -> DreamGenerationConfig:
-        """
-        Prepares the base generation config, then applies any generation configuration options from kwargs. This
-        function handles retrocompatibility with respect to configuration files.
-        """
-        # priority: `generation_config` argument > `model.generation_config` (the default generation config)
-        using_model_generation_config = False
-        if generation_config is None:
-            generation_config = DreamGenerationConfig.from_model_config(self.config)
-            using_model_generation_config = True
-
-        # `torch.compile` can't compile `copy.deepcopy`, arguments in `kwargs` that are part of `generation_config`
-        # will mutate the object with `.update`. As such, passing these arguments through `kwargs` is disabled -- an
-        # exception will be raised in `_validate_model_kwargs`
-        if not is_torchdynamo_compiling():
-            generation_config = copy.deepcopy(generation_config)
-            _kwargs = generation_config.update(**kwargs)
-            # If `generation_config` is provided, let's fallback ALL special tokens to the default values for the model
-            if not using_model_generation_config:
-                if generation_config.bos_token_id is None:
-                    generation_config.bos_token_id = self.generation_config.bos_token_id
-                if generation_config.eos_token_id is None:
-                    generation_config.eos_token_id = self.generation_config.eos_token_id
-                if generation_config.pad_token_id is None:
-                    generation_config.pad_token_id = self.generation_config.pad_token_id
-                if generation_config.mask_token_id is None:
-                    generation_config.mask_token_id = self.generation_config.mask_token_id
-
-        return generation_config
-
-    def _prepare_special_tokens(
-        self,
-        generation_config: DreamGenerationConfig,
-        device: Optional[Union[torch.device, str]] = None,
-    ):
-        """
-        Prepares the special tokens for generation, overwriting the generation config with their processed versions
-        converted to tensor.
-
-        Note that `generation_config` is changed in place and stops being serializable after this method is called.
-        That is no problem if called within `generate` (`generation_config` is a local copy that doesn't leave the
-        function). However, if called outside `generate`, consider creating a copy of `generation_config` first.
-        """
-
-        # Convert special tokens to tensors
-        def _tensor_or_none(token, device=None):
-            if token is None:
-                return token
-
-            device = device if device is not None else self.device
-            if isinstance(token, torch.Tensor):
-                return token.to(device)
-            return torch.tensor(token, device=device, dtype=torch.long)
-
-        bos_token_tensor = _tensor_or_none(generation_config.bos_token_id, device=device)
-        eos_token_tensor = _tensor_or_none(generation_config.eos_token_id, device=device)
-        pad_token_tensor = _tensor_or_none(generation_config.pad_token_id, device=device)
-        mask_token_tensor = _tensor_or_none(generation_config.mask_token_id, device=device)
-
-        # We can have more than one eos token. Always treat it as a 1D tensor (when it exists).
-        if eos_token_tensor is not None and eos_token_tensor.ndim == 0:
-            eos_token_tensor = eos_token_tensor.unsqueeze(0)
-
-        # Set pad token if unset (and there are conditions to do so)
-        if pad_token_tensor is None and eos_token_tensor is not None:
-            pad_token_tensor = eos_token_tensor[0]
-            logger.warning(f"Setting `pad_token_id` to `eos_token_id`:{pad_token_tensor} for open-end generation.")
-
-        # Update generation config with the updated special tokens tensors
-        # NOTE: this must be written into a different attribute name than the one holding the original special tokens
-        # (in their non-tensor form), in order to enable end-to-end compilation. See
-        # https://pytorch.org/docs/stable/torch.compiler_cudagraph_trees.html#limitations
-        generation_config._bos_token_tensor = bos_token_tensor
-        generation_config._eos_token_tensor = eos_token_tensor
-        generation_config._pad_token_tensor = pad_token_tensor
-        generation_config._mask_token_tensor = mask_token_tensor
-
-    @torch.no_grad()
-    def diffusion_generate(
-        self,
-        inputs: Optional[torch.Tensor] = None,
-        generation_config: Optional[DreamGenerationConfig] = None,
-        **kwargs,
-    ) -> Union[DreamModelOutput, torch.LongTensor]:
-        # 1. Handle `generation_config` and kwargs that might update it, and validate the `.generate()` call
-        generation_config = self._prepare_generation_config(generation_config, **kwargs)
-        generation_tokens_hook_func = kwargs.pop("generation_tokens_hook_func", lambda step, x, logits: x)
-        generation_logits_hook_func = kwargs.pop("generation_logits_hook_func", lambda step, x, logits: logits)
-
-        # 2. Define model inputs
-        assert inputs is not None
-        input_ids = inputs
-        device = input_ids.device
-        attention_mask = kwargs.pop("attention_mask", None)
-        self._prepare_special_tokens(generation_config, device=device)
-
-        # 3. Prepare `max_length`.
-        input_ids_length = input_ids.shape[-1]
-        has_default_max_length = kwargs.get("max_length") is None and generation_config.max_length is not None
-        generation_config = self._prepare_generated_length(
-            generation_config=generation_config,
-            has_default_max_length=has_default_max_length,
-            input_ids_length=input_ids_length,
-        )
-
-        self._validate_generated_length(generation_config, input_ids_length, has_default_max_length)
-        
-        # 4. Check input_ids
-        if not is_torchdynamo_compiling() and self.device.type != input_ids.device.type:
-            warnings.warn(
-                "You are calling .generate() with the `input_ids` being on a device type different"
-                f" than your model's device. `input_ids` is on {input_ids.device.type}, whereas the model"
-                f" is on {self.device.type}. You may experience unexpected behaviors or slower generation."
-                " Please make sure that you have put `input_ids` to the"
-                f" correct device by calling for example input_ids = input_ids.to('{self.device.type}') before"
-                " running `.generate()`.",
-                UserWarning,
-            )
-        if (
-            hasattr(generation_config, "pad_token_id") and
-            torch.any(input_ids == generation_config.pad_token_id) and 
-            attention_mask is None
-        ):
-            warnings.warn(
-                "Padding was detected but no attention mask is passed here. For correct "
-                "generation results, please set `attention_mask` when batch-padding inputs.",
-                UserWarning,
-            )
-
-        input_ids, attention_mask = self._expand_inputs_for_generation(
-            expand_size=generation_config.num_return_sequences,
-            input_ids=input_ids,
-            attention_mask=attention_mask 
-        )
-
-        result = self._sample(
-            input_ids,
-            attention_mask=attention_mask,
-            generation_config=generation_config,
-            generation_tokens_hook_func=generation_tokens_hook_func,
-            generation_logits_hook_func=generation_logits_hook_func
-        )
-        return result
-
-    def _sample(
-        self,
-        input_ids: torch.LongTensor,
-        attention_mask: Optional[torch.LongTensor],
-        generation_config: DreamGenerationConfig,
-        generation_tokens_hook_func,
-        generation_logits_hook_func
-    ) -> Union[DreamModelOutput, torch.LongTensor]:
-        # init values
-        output_history = generation_config.output_history
-        return_dict_in_generate = generation_config.return_dict_in_generate
-        max_length = generation_config.max_length
-        mask_token_id = generation_config.mask_token_id
-        steps = generation_config.steps
-        eps = generation_config.eps
-        alg = generation_config.alg
-        alg_temp = generation_config.alg_temp
-        temperature = generation_config.temperature
-        top_p = generation_config.top_p
-        top_k = generation_config.top_k
-
-        histories = [] if (return_dict_in_generate and output_history) else None
-
-        # pad input_ids to max_length
-        x = F.pad(input_ids, (0, max_length - input_ids.shape[1]), value=mask_token_id)
-
-        if attention_mask is not None and torch.any(attention_mask == 0.0):
-            # we do not mask the [MASK] tokens so value = 1.0
-            attention_mask = F.pad(attention_mask, (0, max_length - attention_mask.shape[1]), value=1.0)
-            tok_idx = attention_mask.long().cumsum(-1) - 1
-            tok_idx.masked_fill_(attention_mask == 0, 1)
-            # attention_mask is of shape [B, N]
-            # broadcast to [B, 1, N, N]
-            attention_mask = torch.logical_and(
-                attention_mask.unsqueeze(1).unsqueeze(-2),
-                attention_mask.unsqueeze(1).unsqueeze(-1),
-            )
-        else:
-            tok_idx = None
-            attention_mask = "full"
-
-        timesteps = torch.linspace(1, eps, steps + 1, device=x.device)
-
-        # this allows user-defined token control of the intermediate steps
-        x = generation_tokens_hook_func(None, x, None)
-        for i in range(steps):
-            mask_index = (x == mask_token_id)
-            logits = self(x, attention_mask, tok_idx).logits
-            logits = torch.cat([logits[:,:1], logits[:, :-1]], dim=1)
-
-            # this allows user-defined logits control of the intermediate steps
-            logits = generation_logits_hook_func(i, x, logits)
-
-            mask_logits = logits[mask_index]
-            t = timesteps[i]
-            s = timesteps[i + 1]
-        
-            if alg == 'origin':
-                p_transfer = 1 - s / t if i < steps - 1 else 1
-                x0 = torch.zeros_like(x[mask_index], device=self.device, dtype=torch.long) + mask_token_id
-                transfer_index_t_s = torch.rand(*x0.shape, device=self.device) < p_transfer
-                _, x0[transfer_index_t_s]= sample_tokens(mask_logits[transfer_index_t_s], temperature=temperature, top_p=top_p, top_k=top_k)
-                x[mask_index] = x0.clone()
-            else:
-                if alg == 'maskgit_plus':
-                    confidence, x0 = sample_tokens(mask_logits, temperature=temperature, top_p=top_p, top_k=top_k)
-                elif alg == 'topk_margin':
-                    confidence, x0 = sample_tokens(mask_logits, temperature=temperature, top_p=top_p, top_k=top_k, margin_confidence=True)
-                elif alg == 'entropy':
-                    confidence, x0 = sample_tokens(mask_logits, temperature, top_p=top_p, top_k=top_k, neg_entropy=True)
-                else:
-                    raise RuntimeError(f"Unknown alg: {alg}")
-                num_mask_token = mask_index.sum() / mask_index.shape[0]
-                number_transfer_tokens = int(num_mask_token * (1 - s / t)) if i < steps - 1 else int(num_mask_token)
-                full_confidence = torch.full_like(x, -torch.inf, device=self.device, dtype=logits.dtype)
-                full_confidence[mask_index] = confidence
-                if number_transfer_tokens > 0:
-                    if alg_temp is None or alg_temp == 0:
-                        _, transfer_index = torch.topk(full_confidence, number_transfer_tokens)
-                    else:
-                        full_confidence = full_confidence / alg_temp
-                        full_confidence = F.softmax(full_confidence, dim=-1)
-                        transfer_index = torch.multinomial(full_confidence, num_samples=number_transfer_tokens)
-                    x_ = torch.zeros_like(x, device=self.device, dtype=torch.long) + mask_token_id
-                    x_[mask_index] = x0.clone()
-                    row_indices = torch.arange(x.size(0), device=self.device).unsqueeze(1).expand_as(transfer_index)
-                    x[row_indices,transfer_index] = x_[row_indices,transfer_index]
-
-            # this allows user-defined token control of the intermediate steps
-            x = generation_tokens_hook_func(i, x, logits)
-
-            if histories is not None:
-                histories.append(x.clone())
-        
-        if return_dict_in_generate:
-            return DreamModelOutput(
-                sequences=x,
-                history=histories,
-            )
-        else:
-            return x

model-00001-of-00007.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:c3449e42c526a70c2d2c4ebe35cda78feb712ad0e6533b7e72e3cc9602595b02
-size 4976687216

model-00002-of-00007.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:d2833d533deb3b2622af4e6d765332273627caa7b3497bd92407835b17313216
-size 4778622352

model-00003-of-00007.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:ea00fe3e1dfade1ddfc346fb323b3896c55a43acd6d19b412c23fe1c7894335b
-size 4932743960

model-00004-of-00007.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:b95d2ef0932290d91ef9ff09beebda252903c0b7f104dca00ca97b8ef75473f8
-size 4932743992

model-00005-of-00007.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:d2eae009ce5581bb8327285d6fb73ae3090fdbc1d2b9e08964ddc2457b81c3a1
-size 4998852296

model-00006-of-00007.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:8a196a869bd9874e52fe10f65812cb4ed21908d5bd3b9df61e3e0a3d8d7444e7
-size 3662865184

model-00007-of-00007.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:82090eff08e75305f3af8be82b2dee4f07048c159f9e637eebd8299309a0474e
-size 2179989632

model.safetensors.index.json

@@ -1,347 +0,0 @@
-{
-  "metadata": {
-    "total_parameters": 1269269438,
-    "total_size": 30462466048
-  },
-  "weight_map": {
-    "lm_head.weight": "model-00007-of-00007.safetensors",
-    "model.embed_tokens.weight": "model-00001-of-00007.safetensors",
-    "model.layers.0.input_layernorm.weight": "model-00001-of-00007.safetensors",
-    "model.layers.0.mlp.down_proj.weight": "model-00001-of-00007.safetensors",
-    "model.layers.0.mlp.gate_proj.weight": "model-00001-of-00007.safetensors",
-    "model.layers.0.mlp.up_proj.weight": "model-00001-of-00007.safetensors",
-    "model.layers.0.post_attention_layernorm.weight": "model-00001-of-00007.safetensors",
-    "model.layers.0.self_attn.k_proj.bias": "model-00001-of-00007.safetensors",
-    "model.layers.0.self_attn.k_proj.weight": "model-00001-of-00007.safetensors",
-    "model.layers.0.self_attn.o_proj.weight": "model-00001-of-00007.safetensors",
-    "model.layers.0.self_attn.q_proj.bias": "model-00001-of-00007.safetensors",
-    "model.layers.0.self_attn.q_proj.weight": "model-00001-of-00007.safetensors",
-    "model.layers.0.self_attn.v_proj.bias": "model-00001-of-00007.safetensors",
-    "model.layers.0.self_attn.v_proj.weight": "model-00001-of-00007.safetensors",
-    "model.layers.1.input_layernorm.weight": "model-00001-of-00007.safetensors",
-    "model.layers.1.mlp.down_proj.weight": "model-00001-of-00007.safetensors",
-    "model.layers.1.mlp.gate_proj.weight": "model-00001-of-00007.safetensors",
-    "model.layers.1.mlp.up_proj.weight": "model-00001-of-00007.safetensors",
-    "model.layers.1.post_attention_layernorm.weight": "model-00001-of-00007.safetensors",
-    "model.layers.1.self_attn.k_proj.bias": "model-00001-of-00007.safetensors",
-    "model.layers.1.self_attn.k_proj.weight": "model-00001-of-00007.safetensors",
-    "model.layers.1.self_attn.o_proj.weight": "model-00001-of-00007.safetensors",
-    "model.layers.1.self_attn.q_proj.bias": "model-00001-of-00007.safetensors",
-    "model.layers.1.self_attn.q_proj.weight": "model-00001-of-00007.safetensors",
-    "model.layers.1.self_attn.v_proj.bias": "model-00001-of-00007.safetensors",
-    "model.layers.1.self_attn.v_proj.weight": "model-00001-of-00007.safetensors",
-    "model.layers.10.input_layernorm.weight": "model-00003-of-00007.safetensors",
-    "model.layers.10.mlp.down_proj.weight": "model-00003-of-00007.safetensors",
-    "model.layers.10.mlp.gate_proj.weight": "model-00003-of-00007.safetensors",
-    "model.layers.10.mlp.up_proj.weight": "model-00003-of-00007.safetensors",
-    "model.layers.10.post_attention_layernorm.weight": "model-00003-of-00007.safetensors",
-    "model.layers.10.self_attn.k_proj.bias": "model-00003-of-00007.safetensors",
-    "model.layers.10.self_attn.k_proj.weight": "model-00003-of-00007.safetensors",
-    "model.layers.10.self_attn.o_proj.weight": "model-00003-of-00007.safetensors",
-    "model.layers.10.self_attn.q_proj.bias": "model-00003-of-00007.safetensors",
-    "model.layers.10.self_attn.q_proj.weight": "model-00003-of-00007.safetensors",
-    "model.layers.10.self_attn.v_proj.bias": "model-00003-of-00007.safetensors",
-    "model.layers.10.self_attn.v_proj.weight": "model-00003-of-00007.safetensors",
-    "model.layers.11.input_layernorm.weight": "model-00003-of-00007.safetensors",
-    "model.layers.11.mlp.down_proj.weight": "model-00003-of-00007.safetensors",
-    "model.layers.11.mlp.gate_proj.weight": "model-00003-of-00007.safetensors",
-    "model.layers.11.mlp.up_proj.weight": "model-00003-of-00007.safetensors",
-    "model.layers.11.post_attention_layernorm.weight": "model-00003-of-00007.safetensors",
-    "model.layers.11.self_attn.k_proj.bias": "model-00003-of-00007.safetensors",
-    "model.layers.11.self_attn.k_proj.weight": "model-00003-of-00007.safetensors",
-    "model.layers.11.self_attn.o_proj.weight": "model-00003-of-00007.safetensors",
-    "model.layers.11.self_attn.q_proj.bias": "model-00003-of-00007.safetensors",
-    "model.layers.11.self_attn.q_proj.weight": "model-00003-of-00007.safetensors",
-    "model.layers.11.self_attn.v_proj.bias": "model-00003-of-00007.safetensors",
-    "model.layers.11.self_attn.v_proj.weight": "model-00003-of-00007.safetensors",
-    "model.layers.12.input_layernorm.weight": "model-00003-of-00007.safetensors",
-    "model.layers.12.mlp.down_proj.weight": "model-00003-of-00007.safetensors",
-    "model.layers.12.mlp.gate_proj.weight": "model-00003-of-00007.safetensors",
-    "model.layers.12.mlp.up_proj.weight": "model-00003-of-00007.safetensors",
-    "model.layers.12.post_attention_layernorm.weight": "model-00003-of-00007.safetensors",
-    "model.layers.12.self_attn.k_proj.bias": "model-00003-of-00007.safetensors",
-    "model.layers.12.self_attn.k_proj.weight": "model-00003-of-00007.safetensors",
-    "model.layers.12.self_attn.o_proj.weight": "model-00003-of-00007.safetensors",
-    "model.layers.12.self_attn.q_proj.bias": "model-00003-of-00007.safetensors",
-    "model.layers.12.self_attn.q_proj.weight": "model-00003-of-00007.safetensors",
-    "model.layers.12.self_attn.v_proj.bias": "model-00003-of-00007.safetensors",
-    "model.layers.12.self_attn.v_proj.weight": "model-00003-of-00007.safetensors",
-    "model.layers.13.input_layernorm.weight": "model-00004-of-00007.safetensors",
-    "model.layers.13.mlp.down_proj.weight": "model-00004-of-00007.safetensors",
-    "model.layers.13.mlp.gate_proj.weight": "model-00003-of-00007.safetensors",
-    "model.layers.13.mlp.up_proj.weight": "model-00004-of-00007.safetensors",
-    "model.layers.13.post_attention_layernorm.weight": "model-00004-of-00007.safetensors",
-    "model.layers.13.self_attn.k_proj.bias": "model-00003-of-00007.safetensors",
-    "model.layers.13.self_attn.k_proj.weight": "model-00003-of-00007.safetensors",
-    "model.layers.13.self_attn.o_proj.weight": "model-00003-of-00007.safetensors",
-    "model.layers.13.self_attn.q_proj.bias": "model-00003-of-00007.safetensors",
-    "model.layers.13.self_attn.q_proj.weight": "model-00003-of-00007.safetensors",
-    "model.layers.13.self_attn.v_proj.bias": "model-00003-of-00007.safetensors",
-    "model.layers.13.self_attn.v_proj.weight": "model-00003-of-00007.safetensors",
-    "model.layers.14.input_layernorm.weight": "model-00004-of-00007.safetensors",
-    "model.layers.14.mlp.down_proj.weight": "model-00004-of-00007.safetensors",
-    "model.layers.14.mlp.gate_proj.weight": "model-00004-of-00007.safetensors",
-    "model.layers.14.mlp.up_proj.weight": "model-00004-of-00007.safetensors",
-    "model.layers.14.post_attention_layernorm.weight": "model-00004-of-00007.safetensors",
-    "model.layers.14.self_attn.k_proj.bias": "model-00004-of-00007.safetensors",
-    "model.layers.14.self_attn.k_proj.weight": "model-00004-of-00007.safetensors",
-    "model.layers.14.self_attn.o_proj.weight": "model-00004-of-00007.safetensors",
-    "model.layers.14.self_attn.q_proj.bias": "model-00004-of-00007.safetensors",
-    "model.layers.14.self_attn.q_proj.weight": "model-00004-of-00007.safetensors",
-    "model.layers.14.self_attn.v_proj.bias": "model-00004-of-00007.safetensors",
-    "model.layers.14.self_attn.v_proj.weight": "model-00004-of-00007.safetensors",
-    "model.layers.15.input_layernorm.weight": "model-00004-of-00007.safetensors",
-    "model.layers.15.mlp.down_proj.weight": "model-00004-of-00007.safetensors",
-    "model.layers.15.mlp.gate_proj.weight": "model-00004-of-00007.safetensors",
-    "model.layers.15.mlp.up_proj.weight": "model-00004-of-00007.safetensors",
-    "model.layers.15.post_attention_layernorm.weight": "model-00004-of-00007.safetensors",
-    "model.layers.15.self_attn.k_proj.bias": "model-00004-of-00007.safetensors",
-    "model.layers.15.self_attn.k_proj.weight": "model-00004-of-00007.safetensors",
-    "model.layers.15.self_attn.o_proj.weight": "model-00004-of-00007.safetensors",
-    "model.layers.15.self_attn.q_proj.bias": "model-00004-of-00007.safetensors",
-    "model.layers.15.self_attn.q_proj.weight": "model-00004-of-00007.safetensors",
-    "model.layers.15.self_attn.v_proj.bias": "model-00004-of-00007.safetensors",
-    "model.layers.15.self_attn.v_proj.weight": "model-00004-of-00007.safetensors",
-    "model.layers.16.input_layernorm.weight": "model-00004-of-00007.safetensors",
-    "model.layers.16.mlp.down_proj.weight": "model-00004-of-00007.safetensors",
-    "model.layers.16.mlp.gate_proj.weight": "model-00004-of-00007.safetensors",
-    "model.layers.16.mlp.up_proj.weight": "model-00004-of-00007.safetensors",
-    "model.layers.16.post_attention_layernorm.weight": "model-00004-of-00007.safetensors",
-    "model.layers.16.self_attn.k_proj.bias": "model-00004-of-00007.safetensors",
-    "model.layers.16.self_attn.k_proj.weight": "model-00004-of-00007.safetensors",
-    "model.layers.16.self_attn.o_proj.weight": "model-00004-of-00007.safetensors",
-    "model.layers.16.self_attn.q_proj.bias": "model-00004-of-00007.safetensors",
-    "model.layers.16.self_attn.q_proj.weight": "model-00004-of-00007.safetensors",
-    "model.layers.16.self_attn.v_proj.bias": "model-00004-of-00007.safetensors",
-    "model.layers.16.self_attn.v_proj.weight": "model-00004-of-00007.safetensors",
-    "model.layers.17.input_layernorm.weight": "model-00004-of-00007.safetensors",
-    "model.layers.17.mlp.down_proj.weight": "model-00004-of-00007.safetensors",
-    "model.layers.17.mlp.gate_proj.weight": "model-00004-of-00007.safetensors",
-    "model.layers.17.mlp.up_proj.weight": "model-00004-of-00007.safetensors",
-    "model.layers.17.post_attention_layernorm.weight": "model-00004-of-00007.safetensors",
-    "model.layers.17.self_attn.k_proj.bias": "model-00004-of-00007.safetensors",
-    "model.layers.17.self_attn.k_proj.weight": "model-00004-of-00007.safetensors",
-    "model.layers.17.self_attn.o_proj.weight": "model-00004-of-00007.safetensors",
-    "model.layers.17.self_attn.q_proj.bias": "model-00004-of-00007.safetensors",
-    "model.layers.17.self_attn.q_proj.weight": "model-00004-of-00007.safetensors",
-    "model.layers.17.self_attn.v_proj.bias": "model-00004-of-00007.safetensors",
-    "model.layers.17.self_attn.v_proj.weight": "model-00004-of-00007.safetensors",
-    "model.layers.18.input_layernorm.weight": "model-00005-of-00007.safetensors",
-    "model.layers.18.mlp.down_proj.weight": "model-00005-of-00007.safetensors",
-    "model.layers.18.mlp.gate_proj.weight": "model-00004-of-00007.safetensors",
-    "model.layers.18.mlp.up_proj.weight": "model-00004-of-00007.safetensors",
-    "model.layers.18.post_attention_layernorm.weight": "model-00005-of-00007.safetensors",
-    "model.layers.18.self_attn.k_proj.bias": "model-00004-of-00007.safetensors",
-    "model.layers.18.self_attn.k_proj.weight": "model-00004-of-00007.safetensors",
-    "model.layers.18.self_attn.o_proj.weight": "model-00004-of-00007.safetensors",
-    "model.layers.18.self_attn.q_proj.bias": "model-00004-of-00007.safetensors",
-    "model.layers.18.self_attn.q_proj.weight": "model-00004-of-00007.safetensors",
-    "model.layers.18.self_attn.v_proj.bias": "model-00004-of-00007.safetensors",
-    "model.layers.18.self_attn.v_proj.weight": "model-00004-of-00007.safetensors",
-    "model.layers.19.input_layernorm.weight": "model-00005-of-00007.safetensors",
-    "model.layers.19.mlp.down_proj.weight": "model-00005-of-00007.safetensors",
-    "model.layers.19.mlp.gate_proj.weight": "model-00005-of-00007.safetensors",
-    "model.layers.19.mlp.up_proj.weight": "model-00005-of-00007.safetensors",
-    "model.layers.19.post_attention_layernorm.weight": "model-00005-of-00007.safetensors",
-    "model.layers.19.self_attn.k_proj.bias": "model-00005-of-00007.safetensors",
-    "model.layers.19.self_attn.k_proj.weight": "model-00005-of-00007.safetensors",
-    "model.layers.19.self_attn.o_proj.weight": "model-00005-of-00007.safetensors",
-    "model.layers.19.self_attn.q_proj.bias": "model-00005-of-00007.safetensors",
-    "model.layers.19.self_attn.q_proj.weight": "model-00005-of-00007.safetensors",
-    "model.layers.19.self_attn.v_proj.bias": "model-00005-of-00007.safetensors",
-    "model.layers.19.self_attn.v_proj.weight": "model-00005-of-00007.safetensors",
-    "model.layers.2.input_layernorm.weight": "model-00001-of-00007.safetensors",
-    "model.layers.2.mlp.down_proj.weight": "model-00001-of-00007.safetensors",
-    "model.layers.2.mlp.gate_proj.weight": "model-00001-of-00007.safetensors",
-    "model.layers.2.mlp.up_proj.weight": "model-00001-of-00007.safetensors",
-    "model.layers.2.post_attention_layernorm.weight": "model-00001-of-00007.safetensors",
-    "model.layers.2.self_attn.k_proj.bias": "model-00001-of-00007.safetensors",
-    "model.layers.2.self_attn.k_proj.weight": "model-00001-of-00007.safetensors",
-    "model.layers.2.self_attn.o_proj.weight": "model-00001-of-00007.safetensors",
-    "model.layers.2.self_attn.q_proj.bias": "model-00001-of-00007.safetensors",
-    "model.layers.2.self_attn.q_proj.weight": "model-00001-of-00007.safetensors",
-    "model.layers.2.self_attn.v_proj.bias": "model-00001-of-00007.safetensors",
-    "model.layers.2.self_attn.v_proj.weight": "model-00001-of-00007.safetensors",
-    "model.layers.20.input_layernorm.weight": "model-00005-of-00007.safetensors",
-    "model.layers.20.mlp.down_proj.weight": "model-00005-of-00007.safetensors",
-    "model.layers.20.mlp.gate_proj.weight": "model-00005-of-00007.safetensors",
-    "model.layers.20.mlp.up_proj.weight": "model-00005-of-00007.safetensors",
-    "model.layers.20.post_attention_layernorm.weight": "model-00005-of-00007.safetensors",
-    "model.layers.20.self_attn.k_proj.bias": "model-00005-of-00007.safetensors",
-    "model.layers.20.self_attn.k_proj.weight": "model-00005-of-00007.safetensors",
-    "model.layers.20.self_attn.o_proj.weight": "model-00005-of-00007.safetensors",
-    "model.layers.20.self_attn.q_proj.bias": "model-00005-of-00007.safetensors",
-    "model.layers.20.self_attn.q_proj.weight": "model-00005-of-00007.safetensors",
-    "model.layers.20.self_attn.v_proj.bias": "model-00005-of-00007.safetensors",
-    "model.layers.20.self_attn.v_proj.weight": "model-00005-of-00007.safetensors",
-    "model.layers.21.input_layernorm.weight": "model-00005-of-00007.safetensors",
-    "model.layers.21.mlp.down_proj.weight": "model-00005-of-00007.safetensors",
-    "model.layers.21.mlp.gate_proj.weight": "model-00005-of-00007.safetensors",
-    "model.layers.21.mlp.up_proj.weight": "model-00005-of-00007.safetensors",
-    "model.layers.21.post_attention_layernorm.weight": "model-00005-of-00007.safetensors",
-    "model.layers.21.self_attn.k_proj.bias": "model-00005-of-00007.safetensors",
-    "model.layers.21.self_attn.k_proj.weight": "model-00005-of-00007.safetensors",
-    "model.layers.21.self_attn.o_proj.weight": "model-00005-of-00007.safetensors",
-    "model.layers.21.self_attn.q_proj.bias": "model-00005-of-00007.safetensors",
-    "model.layers.21.self_attn.q_proj.weight": "model-00005-of-00007.safetensors",
-    "model.layers.21.self_attn.v_proj.bias": "model-00005-of-00007.safetensors",
-    "model.layers.21.self_attn.v_proj.weight": "model-00005-of-00007.safetensors",
-    "model.layers.22.input_layernorm.weight": "model-00005-of-00007.safetensors",
-    "model.layers.22.mlp.down_proj.weight": "model-00005-of-00007.safetensors",
-    "model.layers.22.mlp.gate_proj.weight": "model-00005-of-00007.safetensors",
-    "model.layers.22.mlp.up_proj.weight": "model-00005-of-00007.safetensors",
-    "model.layers.22.post_attention_layernorm.weight": "model-00005-of-00007.safetensors",
-    "model.layers.22.self_attn.k_proj.bias": "model-00005-of-00007.safetensors",
-    "model.layers.22.self_attn.k_proj.weight": "model-00005-of-00007.safetensors",
-    "model.layers.22.self_attn.o_proj.weight": "model-00005-of-00007.safetensors",
-    "model.layers.22.self_attn.q_proj.bias": "model-00005-of-00007.safetensors",
-    "model.layers.22.self_attn.q_proj.weight": "model-00005-of-00007.safetensors",
-    "model.layers.22.self_attn.v_proj.bias": "model-00005-of-00007.safetensors",
-    "model.layers.22.self_attn.v_proj.weight": "model-00005-of-00007.safetensors",
-    "model.layers.23.input_layernorm.weight": "model-00005-of-00007.safetensors",
-    "model.layers.23.mlp.down_proj.weight": "model-00005-of-00007.safetensors",
-    "model.layers.23.mlp.gate_proj.weight": "model-00005-of-00007.safetensors",
-    "model.layers.23.mlp.up_proj.weight": "model-00005-of-00007.safetensors",
-    "model.layers.23.post_attention_layernorm.weight": "model-00005-of-00007.safetensors",
-    "model.layers.23.self_attn.k_proj.bias": "model-00005-of-00007.safetensors",
-    "model.layers.23.self_attn.k_proj.weight": "model-00005-of-00007.safetensors",
-    "model.layers.23.self_attn.o_proj.weight": "model-00005-of-00007.safetensors",
-    "model.layers.23.self_attn.q_proj.bias": "model-00005-of-00007.safetensors",
-    "model.layers.23.self_attn.q_proj.weight": "model-00005-of-00007.safetensors",
-    "model.layers.23.self_attn.v_proj.bias": "model-00005-of-00007.safetensors",
-    "model.layers.23.self_attn.v_proj.weight": "model-00005-of-00007.safetensors",
-    "model.layers.24.input_layernorm.weight": "model-00006-of-00007.safetensors",
-    "model.layers.24.mlp.down_proj.weight": "model-00006-of-00007.safetensors",
-    "model.layers.24.mlp.gate_proj.weight": "model-00006-of-00007.safetensors",
-    "model.layers.24.mlp.up_proj.weight": "model-00006-of-00007.safetensors",
-    "model.layers.24.post_attention_layernorm.weight": "model-00006-of-00007.safetensors",
-    "model.layers.24.self_attn.k_proj.bias": "model-00005-of-00007.safetensors",
-    "model.layers.24.self_attn.k_proj.weight": "model-00005-of-00007.safetensors",
-    "model.layers.24.self_attn.o_proj.weight": "model-00006-of-00007.safetensors",
-    "model.layers.24.self_attn.q_proj.bias": "model-00005-of-00007.safetensors",
-    "model.layers.24.self_attn.q_proj.weight": "model-00005-of-00007.safetensors",
-    "model.layers.24.self_attn.v_proj.bias": "model-00005-of-00007.safetensors",
-    "model.layers.24.self_attn.v_proj.weight": "model-00005-of-00007.safetensors",
-    "model.layers.25.input_layernorm.weight": "model-00006-of-00007.safetensors",
-    "model.layers.25.mlp.down_proj.weight": "model-00006-of-00007.safetensors",
-    "model.layers.25.mlp.gate_proj.weight": "model-00006-of-00007.safetensors",
-    "model.layers.25.mlp.up_proj.weight": "model-00006-of-00007.safetensors",
-    "model.layers.25.post_attention_layernorm.weight": "model-00006-of-00007.safetensors",
-    "model.layers.25.self_attn.k_proj.bias": "model-00006-of-00007.safetensors",
-    "model.layers.25.self_attn.k_proj.weight": "model-00006-of-00007.safetensors",
-    "model.layers.25.self_attn.o_proj.weight": "model-00006-of-00007.safetensors",
-    "model.layers.25.self_attn.q_proj.bias": "model-00006-of-00007.safetensors",
-    "model.layers.25.self_attn.q_proj.weight": "model-00006-of-00007.safetensors",
-    "model.layers.25.self_attn.v_proj.bias": "model-00006-of-00007.safetensors",
-    "model.layers.25.self_attn.v_proj.weight": "model-00006-of-00007.safetensors",
-    "model.layers.26.input_layernorm.weight": "model-00006-of-00007.safetensors",
-    "model.layers.26.mlp.down_proj.weight": "model-00006-of-00007.safetensors",
-    "model.layers.26.mlp.gate_proj.weight": "model-00006-of-00007.safetensors",
-    "model.layers.26.mlp.up_proj.weight": "model-00006-of-00007.safetensors",
-    "model.layers.26.post_attention_layernorm.weight": "model-00006-of-00007.safetensors",
-    "model.layers.26.self_attn.k_proj.bias": "model-00006-of-00007.safetensors",
-    "model.layers.26.self_attn.k_proj.weight": "model-00006-of-00007.safetensors",
-    "model.layers.26.self_attn.o_proj.weight": "model-00006-of-00007.safetensors",
-    "model.layers.26.self_attn.q_proj.bias": "model-00006-of-00007.safetensors",
-    "model.layers.26.self_attn.q_proj.weight": "model-00006-of-00007.safetensors",
-    "model.layers.26.self_attn.v_proj.bias": "model-00006-of-00007.safetensors",
-    "model.layers.26.self_attn.v_proj.weight": "model-00006-of-00007.safetensors",
-    "model.layers.27.input_layernorm.weight": "model-00006-of-00007.safetensors",
-    "model.layers.27.mlp.down_proj.weight": "model-00006-of-00007.safetensors",
-    "model.layers.27.mlp.gate_proj.weight": "model-00006-of-00007.safetensors",
-    "model.layers.27.mlp.up_proj.weight": "model-00006-of-00007.safetensors",
-    "model.layers.27.post_attention_layernorm.weight": "model-00006-of-00007.safetensors",
-    "model.layers.27.self_attn.k_proj.bias": "model-00006-of-00007.safetensors",
-    "model.layers.27.self_attn.k_proj.weight": "model-00006-of-00007.safetensors",
-    "model.layers.27.self_attn.o_proj.weight": "model-00006-of-00007.safetensors",
-    "model.layers.27.self_attn.q_proj.bias": "model-00006-of-00007.safetensors",
-    "model.layers.27.self_attn.q_proj.weight": "model-00006-of-00007.safetensors",
-    "model.layers.27.self_attn.v_proj.bias": "model-00006-of-00007.safetensors",
-    "model.layers.27.self_attn.v_proj.weight": "model-00006-of-00007.safetensors",
-    "model.layers.3.input_layernorm.weight": "model-00002-of-00007.safetensors",
-    "model.layers.3.mlp.down_proj.weight": "model-00002-of-00007.safetensors",
-    "model.layers.3.mlp.gate_proj.weight": "model-00002-of-00007.safetensors",
-    "model.layers.3.mlp.up_proj.weight": "model-00002-of-00007.safetensors",
-    "model.layers.3.post_attention_layernorm.weight": "model-00002-of-00007.safetensors",
-    "model.layers.3.self_attn.k_proj.bias": "model-00002-of-00007.safetensors",
-    "model.layers.3.self_attn.k_proj.weight": "model-00002-of-00007.safetensors",
-    "model.layers.3.self_attn.o_proj.weight": "model-00002-of-00007.safetensors",
-    "model.layers.3.self_attn.q_proj.bias": "model-00002-of-00007.safetensors",
-    "model.layers.3.self_attn.q_proj.weight": "model-00002-of-00007.safetensors",
-    "model.layers.3.self_attn.v_proj.bias": "model-00002-of-00007.safetensors",
-    "model.layers.3.self_attn.v_proj.weight": "model-00002-of-00007.safetensors",
-    "model.layers.4.input_layernorm.weight": "model-00002-of-00007.safetensors",
-    "model.layers.4.mlp.down_proj.weight": "model-00002-of-00007.safetensors",
-    "model.layers.4.mlp.gate_proj.weight": "model-00002-of-00007.safetensors",
-    "model.layers.4.mlp.up_proj.weight": "model-00002-of-00007.safetensors",
-    "model.layers.4.post_attention_layernorm.weight": "model-00002-of-00007.safetensors",
-    "model.layers.4.self_attn.k_proj.bias": "model-00002-of-00007.safetensors",
-    "model.layers.4.self_attn.k_proj.weight": "model-00002-of-00007.safetensors",
-    "model.layers.4.self_attn.o_proj.weight": "model-00002-of-00007.safetensors",
-    "model.layers.4.self_attn.q_proj.bias": "model-00002-of-00007.safetensors",
-    "model.layers.4.self_attn.q_proj.weight": "model-00002-of-00007.safetensors",
-    "model.layers.4.self_attn.v_proj.bias": "model-00002-of-00007.safetensors",
-    "model.layers.4.self_attn.v_proj.weight": "model-00002-of-00007.safetensors",
-    "model.layers.5.input_layernorm.weight": "model-00002-of-00007.safetensors",
-    "model.layers.5.mlp.down_proj.weight": "model-00002-of-00007.safetensors",
-    "model.layers.5.mlp.gate_proj.weight": "model-00002-of-00007.safetensors",
-    "model.layers.5.mlp.up_proj.weight": "model-00002-of-00007.safetensors",
-    "model.layers.5.post_attention_layernorm.weight": "model-00002-of-00007.safetensors",
-    "model.layers.5.self_attn.k_proj.bias": "model-00002-of-00007.safetensors",
-    "model.layers.5.self_attn.k_proj.weight": "model-00002-of-00007.safetensors",
-    "model.layers.5.self_attn.o_proj.weight": "model-00002-of-00007.safetensors",
-    "model.layers.5.self_attn.q_proj.bias": "model-00002-of-00007.safetensors",
-    "model.layers.5.self_attn.q_proj.weight": "model-00002-of-00007.safetensors",
-    "model.layers.5.self_attn.v_proj.bias": "model-00002-of-00007.safetensors",
-    "model.layers.5.self_attn.v_proj.weight": "model-00002-of-00007.safetensors",
-    "model.layers.6.input_layernorm.weight": "model-00002-of-00007.safetensors",
-    "model.layers.6.mlp.down_proj.weight": "model-00002-of-00007.safetensors",
-    "model.layers.6.mlp.gate_proj.weight": "model-00002-of-00007.safetensors",
-    "model.layers.6.mlp.up_proj.weight": "model-00002-of-00007.safetensors",
-    "model.layers.6.post_attention_layernorm.weight": "model-00002-of-00007.safetensors",
-    "model.layers.6.self_attn.k_proj.bias": "model-00002-of-00007.safetensors",
-    "model.layers.6.self_attn.k_proj.weight": "model-00002-of-00007.safetensors",
-    "model.layers.6.self_attn.o_proj.weight": "model-00002-of-00007.safetensors",
-    "model.layers.6.self_attn.q_proj.bias": "model-00002-of-00007.safetensors",
-    "model.layers.6.self_attn.q_proj.weight": "model-00002-of-00007.safetensors",
-    "model.layers.6.self_attn.v_proj.bias": "model-00002-of-00007.safetensors",
-    "model.layers.6.self_attn.v_proj.weight": "model-00002-of-00007.safetensors",
-    "model.layers.7.input_layernorm.weight": "model-00002-of-00007.safetensors",
-    "model.layers.7.mlp.down_proj.weight": "model-00002-of-00007.safetensors",
-    "model.layers.7.mlp.gate_proj.weight": "model-00002-of-00007.safetensors",
-    "model.layers.7.mlp.up_proj.weight": "model-00002-of-00007.safetensors",
-    "model.layers.7.post_attention_layernorm.weight": "model-00002-of-00007.safetensors",
-    "model.layers.7.self_attn.k_proj.bias": "model-00002-of-00007.safetensors",
-    "model.layers.7.self_attn.k_proj.weight": "model-00002-of-00007.safetensors",
-    "model.layers.7.self_attn.o_proj.weight": "model-00002-of-00007.safetensors",
-    "model.layers.7.self_attn.q_proj.bias": "model-00002-of-00007.safetensors",
-    "model.layers.7.self_attn.q_proj.weight": "model-00002-of-00007.safetensors",
-    "model.layers.7.self_attn.v_proj.bias": "model-00002-of-00007.safetensors",
-    "model.layers.7.self_attn.v_proj.weight": "model-00002-of-00007.safetensors",
-    "model.layers.8.input_layernorm.weight": "model-00003-of-00007.safetensors",
-    "model.layers.8.mlp.down_proj.weight": "model-00003-of-00007.safetensors",
-    "model.layers.8.mlp.gate_proj.weight": "model-00003-of-00007.safetensors",
-    "model.layers.8.mlp.up_proj.weight": "model-00003-of-00007.safetensors",
-    "model.layers.8.post_attention_layernorm.weight": "model-00003-of-00007.safetensors",
-    "model.layers.8.self_attn.k_proj.bias": "model-00002-of-00007.safetensors",
-    "model.layers.8.self_attn.k_proj.weight": "model-00002-of-00007.safetensors",
-    "model.layers.8.self_attn.o_proj.weight": "model-00002-of-00007.safetensors",
-    "model.layers.8.self_attn.q_proj.bias": "model-00002-of-00007.safetensors",
-    "model.layers.8.self_attn.q_proj.weight": "model-00002-of-00007.safetensors",
-    "model.layers.8.self_attn.v_proj.bias": "model-00002-of-00007.safetensors",
-    "model.layers.8.self_attn.v_proj.weight": "model-00002-of-00007.safetensors",
-    "model.layers.9.input_layernorm.weight": "model-00003-of-00007.safetensors",
-    "model.layers.9.mlp.down_proj.weight": "model-00003-of-00007.safetensors",
-    "model.layers.9.mlp.gate_proj.weight": "model-00003-of-00007.safetensors",
-    "model.layers.9.mlp.up_proj.weight": "model-00003-of-00007.safetensors",
-    "model.layers.9.post_attention_layernorm.weight": "model-00003-of-00007.safetensors",
-    "model.layers.9.self_attn.k_proj.bias": "model-00003-of-00007.safetensors",
-    "model.layers.9.self_attn.k_proj.weight": "model-00003-of-00007.safetensors",
-    "model.layers.9.self_attn.o_proj.weight": "model-00003-of-00007.safetensors",
-    "model.layers.9.self_attn.q_proj.bias": "model-00003-of-00007.safetensors",
-    "model.layers.9.self_attn.q_proj.weight": "model-00003-of-00007.safetensors",
-    "model.layers.9.self_attn.v_proj.bias": "model-00003-of-00007.safetensors",
-    "model.layers.9.self_attn.v_proj.weight": "model-00003-of-00007.safetensors",
-    "model.norm.weight": "model-00006-of-00007.safetensors"
-  }
-}

modeling_dream.py

@@ -1,824 +0,0 @@
-# coding=utf-8
-# Copyright 2024 The Dream team, HKUNLP Group and the HuggingFace Inc. team. All rights reserved.
-#
-# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
-# and OPT and Qwen implementations in this library. It has been modified from its
-# original forms to accommodate minor architectural differences compared
-# to GPT-NeoX and OPT and Qwen used by the Meta AI and Qwen team that trained the model.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""PyTorch Dream model."""
-
-import math
-from typing import List, Optional, Tuple, Union
-import os
-import torch
-import torch.utils.checkpoint
-from torch import nn
-
-from transformers.activations import ACT2FN
-from transformers.cache_utils import Cache, DynamicCache
-from transformers.modeling_outputs import (
-    BaseModelOutput,
-    MaskedLMOutput,
-)
-from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
-from transformers.modeling_utils import PreTrainedModel
-from transformers.utils import (
-    add_start_docstrings,
-    add_start_docstrings_to_model_forward,
-    is_flash_attn_2_available,
-    is_flash_attn_greater_or_equal_2_10,
-    logging,
-)
-from transformers import PretrainedConfig
-from .configuration_dream import DreamConfig
-from .generation_utils import DreamGenerationMixin, DreamGenerationConfig
-
-if is_flash_attn_2_available():
-    from transformers.modeling_flash_attention_utils import _flash_attention_forward
-
-
-logger = logging.get_logger(__name__)
-
-
-_CHECKPOINT_FOR_DOC = "Dream-7B"
-_CONFIG_FOR_DOC = "DreamConfig"
-
-
-# Copied from transformers.models.llama.modeling_llama.LlamaRMSNorm with Llama->Dream
-class DreamRMSNorm(nn.Module):
-    def __init__(self, hidden_size, eps=1e-6):
-        """
-        DreamRMSNorm is equivalent to T5LayerNorm
-        """
-        super().__init__()
-        self.weight = nn.Parameter(torch.ones(hidden_size))
-        self.variance_epsilon = eps
-
-    def forward(self, hidden_states):
-        input_dtype = hidden_states.dtype
-        hidden_states = hidden_states.to(torch.float32)
-        variance = hidden_states.pow(2).mean(-1, keepdim=True)
-        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
-        return self.weight * hidden_states.to(input_dtype)
-
-    def extra_repr(self):
-        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
-
-
-# Copied from transformers.models.llama.modeling_llama.LlamaRotaryEmbedding with Llama->Dream
-class DreamRotaryEmbedding(nn.Module):
-    def __init__(
-        self,
-        dim=None,
-        max_position_embeddings=2048,
-        base=10000,
-        device=None,
-        scaling_factor=1.0,
-        rope_type="default",
-        config: Optional[DreamConfig] = None,
-    ):
-        super().__init__()
-        # TODO (joao): remove the `if` below, only used for BC
-        self.rope_kwargs = {}
-        if config is None:
-            logger.warning_once(
-                "`DreamRotaryEmbedding` can now be fully parameterized by passing the model config through the "
-                "`config` argument. All other arguments will be removed in v4.46"
-            )
-            self.rope_kwargs = {
-                "rope_type": rope_type,
-                "factor": scaling_factor,
-                "dim": dim,
-                "base": base,
-                "max_position_embeddings": max_position_embeddings,
-            }
-            self.rope_type = rope_type
-            self.max_seq_len_cached = max_position_embeddings
-            self.original_max_seq_len = max_position_embeddings
-        else:
-            # BC: "rope_type" was originally "type"
-            if config.rope_scaling is not None:
-                self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
-            else:
-                self.rope_type = "default"
-            self.max_seq_len_cached = config.max_position_embeddings
-            self.original_max_seq_len = config.max_position_embeddings
-
-        self.config = config
-        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
-
-        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device, **self.rope_kwargs)
-        self.register_buffer("inv_freq", inv_freq, persistent=False)
-        self.original_inv_freq = self.inv_freq
-
-    def reset_parameters(self):
-        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, self.inv_freq.device, **self.rope_kwargs)
-        self.register_buffer("inv_freq", inv_freq, persistent=False)
-        self.original_inv_freq = self.inv_freq
-    
-
-    def _dynamic_frequency_update(self, position_ids, device):
-        """
-        dynamic RoPE layers should recompute `inv_freq` in the following situations:
-        1 - growing beyond the cached sequence length (allow scaling)
-        2 - the current sequence length is in the original scale (avoid losing precision with small sequences)
-        """
-        seq_len = torch.max(position_ids) + 1
-        if seq_len > self.max_seq_len_cached:  # growth
-            inv_freq, self.attention_scaling = self.rope_init_fn(
-                self.config, device, seq_len=seq_len, **self.rope_kwargs
-            )
-            self.register_buffer("inv_freq", inv_freq, persistent=False)  # TODO joao: may break with compilation
-            self.max_seq_len_cached = seq_len
-
-        if seq_len < self.original_max_seq_len and self.max_seq_len_cached > self.original_max_seq_len:  # reset
-            self.register_buffer("inv_freq", self.original_inv_freq, persistent=False)
-            self.max_seq_len_cached = self.original_max_seq_len
-
-    @torch.no_grad()
-    def forward(self, x, position_ids):
-        if "dynamic" in self.rope_type:
-            self._dynamic_frequency_update(position_ids, device=x.device)
-
-        # Core RoPE block
-        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
-        position_ids_expanded = position_ids[:, None, :].float()
-        # Force float32 (see https://github.com/huggingface/transformers/pull/29285)
-        device_type = x.device.type
-        device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu"
-        with torch.autocast(device_type=device_type, enabled=False):
-            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
-            emb = torch.cat((freqs, freqs), dim=-1)
-            cos = emb.cos()
-            sin = emb.sin()
-
-        # Advanced RoPE types (e.g. yarn) apply a post-processing scaling factor, equivalent to scaling attention
-        cos = cos * self.attention_scaling
-        sin = sin * self.attention_scaling
-
-        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
-
-
-# Copied from transformers.models.llama.modeling_llama.rotate_half
-def rotate_half(x):
-    """Rotates half the hidden dims of the input."""
-    x1 = x[..., : x.shape[-1] // 2]
-    x2 = x[..., x.shape[-1] // 2 :]
-    return torch.cat((-x2, x1), dim=-1)
-
-
-# Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb
-def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
-    """Applies Rotary Position Embedding to the query and key tensors.
-
-    Args:
-        q (`torch.Tensor`): The query tensor.
-        k (`torch.Tensor`): The key tensor.
-        cos (`torch.Tensor`): The cosine part of the rotary embedding.
-        sin (`torch.Tensor`): The sine part of the rotary embedding.
-        position_ids (`torch.Tensor`, *optional*):
-            Deprecated and unused.
-        unsqueeze_dim (`int`, *optional*, defaults to 1):
-            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
-            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
-            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
-            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
-            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
-            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
-    Returns:
-        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
-    """
-    cos = cos.unsqueeze(unsqueeze_dim)
-    sin = sin.unsqueeze(unsqueeze_dim)
-    q_embed = (q * cos) + (rotate_half(q) * sin)
-    k_embed = (k * cos) + (rotate_half(k) * sin)
-    return q_embed, k_embed
-
-
-# Copied from transformers.models.mistral.modeling_mistral.MistralMLP with Mistral->Dream
-class DreamMLP(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.hidden_size = config.hidden_size
-        self.intermediate_size = config.intermediate_size
-        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
-        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
-        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
-        self.act_fn = ACT2FN[config.hidden_act]
-
-    def forward(self, hidden_state):
-        return self.down_proj(self.act_fn(self.gate_proj(hidden_state)) * self.up_proj(hidden_state))
-
-
-# Copied from transformers.models.llama.modeling_llama.repeat_kv
-def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
-    """
-    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
-    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
-    """
-    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
-    if n_rep == 1:
-        return hidden_states
-    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
-    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
-
-
-class DreamAttention(nn.Module):
-    """
-    Multi-headed attention from 'Attention Is All You Need' paper. Modified to use sliding window attention: Longformer
-    and "Generating Long Sequences with Sparse Transformers".
-    """
-
-    def __init__(self, config: DreamConfig, layer_idx: Optional[int] = None):
-        super().__init__()
-        self.config = config
-        self.layer_idx = layer_idx
-        if layer_idx is None:
-            logger.warning_once(
-                f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will "
-                "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "
-                "when creating this class."
-            )
-
-        self.hidden_size = config.hidden_size
-        self.num_heads = config.num_attention_heads
-        self.head_dim = self.hidden_size // self.num_heads
-        self.num_key_value_heads = config.num_key_value_heads
-        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
-        self.max_position_embeddings = config.max_position_embeddings
-        self.rope_theta = config.rope_theta
-        self.is_causal = False
-        self.attention_dropout = config.attention_dropout
-
-        if (self.head_dim * self.num_heads) != self.hidden_size:
-            raise ValueError(
-                f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
-                f" and `num_heads`: {self.num_heads})."
-            )
-        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=True)
-        self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True)
-        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True)
-        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
-
-        self.rotary_emb = DreamRotaryEmbedding(config=self.config)
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        past_key_value: Optional[Cache] = None,
-        output_attentions: bool = False,
-        use_cache: bool = False,
-        cache_position: Optional[torch.LongTensor] = None,
-        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,  # will become mandatory in v4.46
-    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
-        bsz, q_len, _ = hidden_states.size()
-
-        query_states = self.q_proj(hidden_states)
-        key_states = self.k_proj(hidden_states)
-        value_states = self.v_proj(hidden_states)
-
-        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
-        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
-        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
-
-        if position_embeddings is None:
-            logger.warning_once(
-                "The attention layers in this model are transitioning from computing the RoPE embeddings internally "
-                "through `position_ids` (2D tensor with the indexes of the tokens), to using externally computed "
-                "`position_embeddings` (Tuple of tensors, containing cos and sin). In v4.46 `position_ids` will be "
-                "removed and `position_embeddings` will be mandatory."
-            )
-            cos, sin = self.rotary_emb(value_states, position_ids)
-        else:
-            cos, sin = position_embeddings
-        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
-
-        if past_key_value is not None:
-            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}  # Specific to RoPE models
-            key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
-
-        # repeat k/v heads if n_kv_heads < n_heads
-        key_states = repeat_kv(key_states, self.num_key_value_groups)
-        value_states = repeat_kv(value_states, self.num_key_value_groups)
-
-        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
-        if attention_mask is not None:  # no matter the length, we just slice it
-            causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
-            attn_weights = attn_weights + causal_mask
-
-        # upcast attention to fp32
-        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
-        attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)
-        attn_output = torch.matmul(attn_weights, value_states)
-
-        if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
-            raise ValueError(
-                f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
-                f" {attn_output.size()}"
-            )
-
-        attn_output = attn_output.transpose(1, 2).contiguous()
-        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
-
-        attn_output = self.o_proj(attn_output)
-
-        if not output_attentions:
-            attn_weights = None
-
-        return attn_output, attn_weights, past_key_value
-    
-
-class DreamSdpaAttention(DreamAttention):
-    """
-    Dream attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from
-    `DreamAttention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to
-    SDPA API.
-    """
-
-    # Adapted from DreamAttention.forward
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        past_key_value: Optional[Cache] = None,
-        output_attentions: bool = False,
-        use_cache: bool = False,
-        cache_position: Optional[torch.LongTensor] = None,
-        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,  # will become mandatory in v4.46
-    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
-        if output_attentions:
-            # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented.
-            logger.warning_once(
-                "DreamModel is using DreamSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, "
-                'but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
-            )
-            return super().forward(
-                hidden_states=hidden_states,
-                attention_mask=attention_mask,
-                position_ids=position_ids,
-                past_key_value=past_key_value,
-                output_attentions=output_attentions,
-                use_cache=use_cache,
-            )
-
-        bsz, q_len, _ = hidden_states.size()
-
-        query_states = self.q_proj(hidden_states)
-        key_states = self.k_proj(hidden_states)
-        value_states = self.v_proj(hidden_states)
-
-        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
-        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
-        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
-
-        if position_embeddings is None:
-            logger.warning_once(
-                "The attention layers in this model are transitioning from computing the RoPE embeddings internally "
-                "through `position_ids` (2D tensor with the indexes of the tokens), to using externally computed "
-                "`position_embeddings` (Tuple of tensors, containing cos and sin). In v4.46 `position_ids` will be "
-                "removed and `position_embeddings` will be mandatory."
-            )
-            cos, sin = self.rotary_emb(value_states, position_ids)
-        else:
-            cos, sin = position_embeddings
-        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
-
-        if past_key_value is not None:
-            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}  # Specific to RoPE models
-            key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
-
-        key_states = repeat_kv(key_states, self.num_key_value_groups)
-        value_states = repeat_kv(value_states, self.num_key_value_groups)
-
-        # causal_mask = attention_mask
-        # if attention_mask is not None:  # no matter the length, we just slice it
-        #     causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
-
-        # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask,
-        # Reference: https://github.com/pytorch/pytorch/issues/112577.
-        if query_states.device.type == "cuda" and attention_mask is not None:
-            query_states = query_states.contiguous()
-            key_states = key_states.contiguous()
-            value_states = value_states.contiguous()
-
-        # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an inline conditional assignment
-        # in SDPA to support both torch.compile's dynamic shapes and full graph options. An inline conditional prevents dynamic shapes from compiling.
-        # The q_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a causal mask in case q_len == 1.
-        # is_causal = True if causal_mask is None and q_len > 1 else False
-
-        attn_output = torch.nn.functional.scaled_dot_product_attention(
-            query_states,
-            key_states,
-            value_states,
-            attn_mask=attention_mask if isinstance(attention_mask, torch.Tensor) else None,
-            dropout_p=self.attention_dropout if self.training else 0.0,
-            is_causal=False, # hard coded
-        )
-
-        attn_output = attn_output.transpose(1, 2).contiguous()
-        attn_output = attn_output.view(bsz, q_len, self.hidden_size)
-
-        attn_output = self.o_proj(attn_output)
-
-        return attn_output, None, past_key_value
-
-
-class DreamDecoderLayer(nn.Module):
-    def __init__(self, config: DreamConfig, layer_idx: int):
-        super().__init__()
-        self.hidden_size = config.hidden_size
-
-        if config.sliding_window and config._attn_implementation != "flash_attention_2":
-            logger.warning_once(
-                f"Sliding Window Attention is enabled but not implemented for `{config._attn_implementation}`; "
-                "unexpected results may be encountered."
-            )
-        
-        # self.self_attn = Dream_ATTENTION_CLASSES[config._attn_implementation](config, layer_idx)
-        self.self_attn = DreamSdpaAttention(config, layer_idx)
-
-        self.mlp = DreamMLP(config)
-        self.input_layernorm = DreamRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.post_attention_layernorm = DreamRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        past_key_value: Optional[Tuple[torch.Tensor]] = None,
-        output_attentions: Optional[bool] = False,
-        use_cache: Optional[bool] = False,
-        cache_position: Optional[torch.LongTensor] = None,
-        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,  # will become mandatory in v4.46
-        **kwargs,
-    ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
-        """
-        Args:
-            hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
-            attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
-                `(batch, sequence_length)` where padding elements are indicated by 0.
-            output_attentions (`bool`, *optional*):
-                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
-                returned tensors for more detail.
-            use_cache (`bool`, *optional*):
-                If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
-                (see `past_key_values`).
-            past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
-            cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
-                Indices depicting the position of the input sequence tokens in the sequence.
-            position_embeddings (`Tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*):
-                Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`,
-                with `head_dim` being the embedding dimension of each attention head.
-            kwargs (`dict`, *optional*):
-                Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code
-                into the model
-        """
-
-        residual = hidden_states
-
-        hidden_states = self.input_layernorm(hidden_states)
-
-        # Self Attention
-        hidden_states, self_attn_weights, present_key_value = self.self_attn(
-            hidden_states=hidden_states,
-            attention_mask=attention_mask,
-            position_ids=position_ids,
-            past_key_value=past_key_value,
-            output_attentions=output_attentions,
-            use_cache=use_cache,
-            cache_position=cache_position,
-            position_embeddings=position_embeddings,
-        )
-        hidden_states = residual + hidden_states
-
-        # Fully Connected
-        residual = hidden_states
-        hidden_states = self.post_attention_layernorm(hidden_states)
-        hidden_states = self.mlp(hidden_states)
-        hidden_states = residual + hidden_states
-
-        outputs = (hidden_states,)
-
-        if output_attentions:
-            outputs += (self_attn_weights,)
-
-        if use_cache:
-            outputs += (present_key_value,)
-
-        return outputs
-
-class DreamPreTrainedModel(PreTrainedModel):
-    config_class = DreamConfig
-    base_model_prefix = "model"
-    supports_gradient_checkpointing = True
-    _no_split_modules = ["DreamDecoderLayer"]
-    _skip_keys_device_placement = "past_key_values"
-    _supports_flash_attn_2 = True
-    _supports_sdpa = True
-    _supports_cache_class = True
-    _supports_quantized_cache = True
-    _supports_static_cache = True
-
-    def _init_weights(self, module):
-        std = self.config.initializer_range
-        if isinstance(module, nn.Linear):
-            module.weight.data.normal_(mean=0.0, std=std)
-            if module.bias is not None:
-                module.bias.data.zero_()
-        elif isinstance(module, nn.Embedding):
-            module.weight.data.normal_(mean=0.0, std=std)
-            if module.padding_idx is not None:
-                module.weight.data[module.padding_idx].zero_()
-
-    @classmethod
-    def from_pretrained(
-        cls,
-        pretrained_model_name_or_path: Optional[Union[str, os.PathLike]],
-        *model_args,
-        config: Optional[Union[PretrainedConfig, str, os.PathLike]] = None,
-        cache_dir: Optional[Union[str, os.PathLike]] = None,
-        ignore_mismatched_sizes: bool = False,
-        force_download: bool = False,
-        local_files_only: bool = False,
-        token: Optional[Union[str, bool]] = None,
-        revision: str = "main",
-        use_safetensors: Optional[bool] = None,
-        weights_only: bool = True,
-        **kwargs,
-    ):
-        _model = super().from_pretrained(
-            pretrained_model_name_or_path,
-            *model_args,
-            config=config,
-            cache_dir=cache_dir,
-            ignore_mismatched_sizes=ignore_mismatched_sizes,
-            force_download=force_download,
-            local_files_only=local_files_only,
-            token=token,
-            revision=revision,
-            use_safetensors=use_safetensors,
-            weights_only=weights_only,
-            **kwargs,
-        )
-        # NOTE(Lin): we need to override the generation config
-        # because the generation config loaded in `from_pretrained` 
-        # does not include all the attributes of DreamGenerationConfig
-        resume_download = kwargs.get("resume_download", None)
-        proxies = kwargs.get("proxies", None)
-        subfolder = kwargs.get("subfolder", "")
-        from_auto_class = kwargs.get("_from_auto", False)
-        from_pipeline = kwargs.get("_from_pipeline", None)
-        _model.generation_config = DreamGenerationConfig.from_pretrained(
-            pretrained_model_name_or_path,
-            cache_dir=cache_dir,
-            force_download=force_download,
-            resume_download=resume_download,
-            proxies=proxies,
-            local_files_only=local_files_only,
-            token=token,
-            revision=revision,
-            subfolder=subfolder,
-            _from_auto=from_auto_class,
-            _from_pipeline=from_pipeline,
-        )
-        return _model
-
-class DreamBaseModel(DreamPreTrainedModel):
-    """
-    Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`DreamDecoderLayer`]
-
-    Args:
-        config: DreamConfig
-    """
-
-    def __init__(self, config: DreamConfig):
-        super().__init__(config)
-        self.padding_idx = config.pad_token_id
-        self.vocab_size = config.vocab_size
-
-        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
-        self.layers = nn.ModuleList(
-            [DreamDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
-        )
-        self._attn_implementation = config._attn_implementation
-        self.norm = DreamRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.rotary_emb = DreamRotaryEmbedding(config=config)
-
-        self.gradient_checkpointing = False
-        # Initialize weights and apply final processing
-        self.post_init()
-
-    def get_input_embeddings(self):
-        return self.embed_tokens
-
-    def set_input_embeddings(self, value):
-        self.embed_tokens = value
-
-    def forward(
-        self,
-        input_ids: torch.LongTensor = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        past_key_values: Optional[List[torch.FloatTensor]] = None,
-        inputs_embeds: Optional[torch.FloatTensor] = None,
-        use_cache: Optional[bool] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-        cache_position: Optional[torch.LongTensor] = None,
-    ) -> Union[Tuple, BaseModelOutput]:
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
-            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
-        )
-        use_cache = use_cache if use_cache is not None else self.config.use_cache
-
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
-        if (input_ids is None) ^ (inputs_embeds is not None):
-            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
-
-        if self.gradient_checkpointing and self.training:
-            if use_cache:
-                logger.warning_once(
-                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
-                )
-                use_cache = False
-
-        if inputs_embeds is None:
-            inputs_embeds = self.embed_tokens(input_ids)
-        
-        if use_cache and past_key_values is None:
-            past_key_values = DynamicCache()
-
-        if cache_position is None:
-            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
-            cache_position = torch.arange(
-                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
-            )
-
-        if position_ids is None:
-            position_ids = cache_position.unsqueeze(0)
-
-        hidden_states = inputs_embeds
-
-        # create position embeddings to be shared across the decoder layers
-        position_embeddings = self.rotary_emb(hidden_states, position_ids)
-
-        # decoder layers
-        all_hidden_states = () if output_hidden_states else None
-        all_self_attns = () if output_attentions else None
-
-        for decoder_layer in self.layers:
-            if output_hidden_states:
-                all_hidden_states += (hidden_states,)
-
-            if self.gradient_checkpointing and self.training:
-                layer_outputs = self._gradient_checkpointing_func(
-                    decoder_layer.__call__,
-                    hidden_states,
-                    attention_mask,
-                    position_ids,
-                    past_key_values,
-                    output_attentions,
-                    use_cache,
-                    cache_position,
-                    position_embeddings,
-                )
-            else:
-                layer_outputs = decoder_layer(
-                    hidden_states,
-                    attention_mask=attention_mask,
-                    position_ids=position_ids,
-                    past_key_value=past_key_values,
-                    output_attentions=output_attentions,
-                    use_cache=use_cache,
-                    cache_position=cache_position,
-                    position_embeddings=position_embeddings,
-                )
-
-            hidden_states = layer_outputs[0]
-
-            if output_attentions:
-                all_self_attns += (layer_outputs[1],)
-
-        hidden_states = self.norm(hidden_states)
-
-        # add hidden states from the last decoder layer
-        if output_hidden_states:
-            all_hidden_states += (hidden_states,)
-
-        if not return_dict:
-            return tuple(v for v in [hidden_states, all_hidden_states, all_self_attns] if v is not None)
-        return BaseModelOutput(
-            last_hidden_state=hidden_states,
-            hidden_states=all_hidden_states,
-            attentions=all_self_attns,
-        )
-
-
-class DreamModel(DreamGenerationMixin, DreamPreTrainedModel):
-    _tied_weights_keys = ["lm_head.weight"]
-
-    def __init__(self, config):
-        super().__init__(config)
-        self.model = DreamBaseModel(config)
-        self.vocab_size = config.vocab_size
-        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
-
-        # Initialize weights and apply final processing
-        self.post_init()
-
-    def reset_rope_parameters(self):
-        self.model.rotary_emb.reset_parameters()
-        for layer in self.model.layers:
-            layer.self_attn.rotary_emb.reset_parameters()
-
-    def get_input_embeddings(self):
-        return self.model.embed_tokens
-
-    def set_input_embeddings(self, value):
-        self.model.embed_tokens = value
-
-    def get_output_embeddings(self):
-        return self.lm_head
-
-    def set_output_embeddings(self, new_embeddings):
-        self.lm_head = new_embeddings
-
-    def set_decoder(self, decoder):
-        self.model = decoder
-
-    def get_decoder(self):
-        return self.model
-
-    def forward(
-        self,
-        input_ids: torch.LongTensor = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        past_key_values: Optional[List[torch.FloatTensor]] = None,
-        inputs_embeds: Optional[torch.FloatTensor] = None,
-        labels: Optional[torch.LongTensor] = None,
-        use_cache: Optional[bool] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-        cache_position: Optional[torch.LongTensor] = None,
-        num_logits_to_keep: int = 0,
-        **loss_kwargs,
-    ) -> Union[Tuple, MaskedLMOutput]:
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
-            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
-        )
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
-        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
-        outputs = self.model(
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            position_ids=position_ids,
-            past_key_values=past_key_values,
-            inputs_embeds=inputs_embeds,
-            use_cache=use_cache,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            return_dict=return_dict,
-            cache_position=cache_position,
-        )
-
-        hidden_states = outputs[0]
-        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
-        logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :])
-
-        loss = None
-        if labels is not None:
-            loss = self.loss_function(logits, labels, self.vocab_size, **loss_kwargs)
-
-        if not return_dict:
-            output = (logits,) + outputs[1:]
-            return (loss,) + output if loss is not None else output
-
-        return MaskedLMOutput(
-            loss=loss,
-            logits=logits,
-            hidden_states=outputs.hidden_states,
-            attentions=outputs.attentions,
-        )

special_tokens_map.json

@@ -1,34 +0,0 @@
-{
-  "additional_special_tokens": [
-    "<|beginoftext|>",
-    "<|mask|>"
-  ],
-  "bos_token": {
-    "content": "<|beginoftext|>",
-    "lstrip": false,
-    "normalized": false,
-    "rstrip": false,
-    "single_word": false
-  },
-  "eos_token": {
-    "content": "<|endoftext|>",
-    "lstrip": false,
-    "normalized": false,
-    "rstrip": false,
-    "single_word": false
-  },
-  "mask_token": {
-    "content": "<|mask|>",
-    "lstrip": false,
-    "normalized": false,
-    "rstrip": false,
-    "single_word": false
-  },
-  "pad_token": {
-    "content": "<|endoftext|>",
-    "lstrip": false,
-    "normalized": false,
-    "rstrip": false,
-    "single_word": false
-  }
-}

tokenization_dream.py

@@ -1,340 +0,0 @@
-# coding=utf-8
-# Copyright 2024 The Dream team, HKUNLP Group and The HuggingFace Inc. team. All rights reserved.
-#
-# This code is based on Qwen's implementations in this library.
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Tokenization classes for Dream."""
-
-import json
-import os
-import unicodedata
-from functools import lru_cache
-from typing import Optional, Tuple
-
-import regex as re
-
-from transformers.tokenization_utils import AddedToken, PreTrainedTokenizer
-from transformers.utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-VOCAB_FILES_NAMES = {
-    "vocab_file": "vocab.json",
-    "merges_file": "merges.txt",
-}
-
-
-MAX_MODEL_INPUT_SIZES = {"dream/dream-tokenizer": 32768}
-
-PRETOKENIZE_REGEX = r"""(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"""
-
-
-@lru_cache()
-# Copied from transformers.models.gpt2.tokenization_gpt2.bytes_to_unicode
-def bytes_to_unicode():
-    """
-    Returns list of utf-8 byte and a mapping to unicode strings. We specifically avoids mapping to whitespace/control
-    characters the bpe code barfs on.
-
-    The reversible bpe codes work on unicode strings. This means you need a large # of unicode characters in your vocab
-    if you want to avoid UNKs. When you're at something like a 10B token dataset you end up needing around 5K for
-    decent coverage. This is a significant percentage of your normal, say, 32K bpe vocab. To avoid that, we want lookup
-    tables between utf-8 bytes and unicode strings.
-    """
-    bs = (
-        list(range(ord("!"), ord("~") + 1)) + list(range(ord("¡"), ord("¬") + 1)) + list(range(ord("®"), ord("ÿ") + 1))
-    )
-    cs = bs[:]
-    n = 0
-    for b in range(2**8):
-        if b not in bs:
-            bs.append(b)
-            cs.append(2**8 + n)
-            n += 1
-    cs = [chr(n) for n in cs]
-    return dict(zip(bs, cs))
-
-
-# Copied from transformers.models.gpt2.tokenization_gpt2.get_pairs
-def get_pairs(word):
-    """
-    Return set of symbol pairs in a word.
-
-    Word is represented as tuple of symbols (symbols being variable-length strings).
-    """
-    pairs = set()
-    prev_char = word[0]
-    for char in word[1:]:
-        pairs.add((prev_char, char))
-        prev_char = char
-    return pairs
-
-
-class DreamTokenizer(PreTrainedTokenizer):
-    """
-    Construct a Dream tokenizer. Based on byte-level Byte-Pair-Encoding.
-
-    Same with GPT2Tokenizer, this tokenizer has been trained to treat spaces like parts of the tokens so a word will
-    be encoded differently whether it is at the beginning of the sentence (without space) or not:
-
-    ```python
-    >>> from transformers import AutoTokenizer
-
-    >>> tokenizer = AutoTokenizer.from_pretrained("Dream-org/Dream-v0-Base-7B", trust_remote_code=True)
-    >>> tokenizer("Hello world")["input_ids"]
-    [9707, 1879]
-
-    >>> tokenizer(" Hello world")["input_ids"]
-    [21927, 1879]
-    ```
-    This is expected.
-
-    You should not use GPT2Tokenizer instead, because of the different pretokenization rules.
-
-    This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to
-    this superclass for more information regarding those methods.
-
-    Args:
-        vocab_file (`str`):
-            Path to the vocabulary file.
-        merges_file (`str`):
-            Path to the merges file.
-        errors (`str`, *optional*, defaults to `"replace"`):
-            Paradigm to follow when decoding bytes to UTF-8. See
-            [bytes.decode](https://docs.python.org/3/library/stdtypes.html#bytes.decode) for more information.
-        unk_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
-            The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
-            token instead.
-        bos_token (`str`, *optional*):
-            The beginning of sequence token. Not applicable for this tokenizer.
-        eos_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
-            The end of sequence token.
-        pad_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
-            The token used for padding, for example when batching sequences of different lengths.
-        clean_up_tokenization_spaces (`bool`, *optional*, defaults to `False`):
-            Whether or not the model should cleanup the spaces that were added when splitting the input text during the
-            tokenization process. Not applicable to this tokenizer, since tokenization does not add spaces.
-        split_special_tokens (`bool`, *optional*, defaults to `False`):
-            Whether or not the special tokens should be split during the tokenization process. The default behavior is
-            to not split special tokens. This means that if `<|endoftext|>` is the `eos_token`, then `tokenizer.tokenize("<|endoftext|>") =
-            ['<|endoftext|>`]. Otherwise, if `split_special_tokens=True`, then `tokenizer.tokenize("<|endoftext|>")` will be give `['<',
-            '|', 'endo', 'ft', 'ext', '|', '>']`. This argument is only supported for `slow` tokenizers for the moment.
-    """
-
-    vocab_files_names = VOCAB_FILES_NAMES
-    model_input_names = ["input_ids", "attention_mask"]
-
-    def __init__(
-        self,
-        vocab_file,
-        merges_file,
-        errors="replace",
-        unk_token="<|endoftext|>",
-        bos_token=None,
-        eos_token="<|endoftext|>",
-        pad_token="<|endoftext|>",
-        clean_up_tokenization_spaces=False,
-        split_special_tokens=False,
-        **kwargs,
-    ):
-        # Dream vocab does not contain control tokens; added tokens need to be special
-        bos_token = (
-            AddedToken(bos_token, lstrip=False, rstrip=False, special=True, normalized=False)
-            if isinstance(bos_token, str)
-            else bos_token
-        )
-        eos_token = (
-            AddedToken(eos_token, lstrip=False, rstrip=False, special=True, normalized=False)
-            if isinstance(eos_token, str)
-            else eos_token
-        )
-        unk_token = (
-            AddedToken(unk_token, lstrip=False, rstrip=False, special=True, normalized=False)
-            if isinstance(unk_token, str)
-            else unk_token
-        )
-        pad_token = (
-            AddedToken(pad_token, lstrip=False, rstrip=False, special=True, normalized=False)
-            if isinstance(pad_token, str)
-            else pad_token
-        )
-
-        with open(vocab_file, encoding="utf-8") as vocab_handle:
-            self.encoder = json.load(vocab_handle)
-        self.decoder = {v: k for k, v in self.encoder.items()}
-        self.errors = errors  # how to handle errors in decoding
-        self.byte_encoder = bytes_to_unicode()
-        self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}
-        bpe_merges = []
-        with open(merges_file, encoding="utf-8") as merges_handle:
-            for i, line in enumerate(merges_handle):
-                line = line.strip()
-                if (i == 0 and line.startswith("#version:")) or not line:
-                    continue
-                bpe_merges.append(tuple(line.split()))
-        self.bpe_ranks = dict(zip(bpe_merges, range(len(bpe_merges))))
-        # NOTE: the cache can grow without bound and will get really large for long running processes
-        # (esp. for texts of language that do not use space between word, e.g. Chinese); technically
-        # not a memory leak but appears as one.
-        # GPT2Tokenizer has the same problem, so let's be consistent.
-        self.cache = {}
-
-        self.pat = re.compile(PRETOKENIZE_REGEX)
-
-        if kwargs.get("add_prefix_space", False):
-            logger.warning_once(
-                f"{self.__class__.__name} does not support `add_prefix_space`, setting it to True has no effect."
-            )
-
-        super().__init__(
-            errors=errors,
-            bos_token=bos_token,
-            eos_token=eos_token,
-            pad_token=pad_token,
-            unk_token=unk_token,
-            clean_up_tokenization_spaces=clean_up_tokenization_spaces,
-            split_special_tokens=split_special_tokens,
-            **kwargs,
-        )
-
-    @property
-    def vocab_size(self) -> int:
-        return len(self.encoder)
-
-    # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer.get_vocab
-    def get_vocab(self):
-        return dict(self.encoder, **self.added_tokens_encoder)
-
-    # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer.bpe
-    def bpe(self, token):
-        if token in self.cache:
-            return self.cache[token]
-        word = tuple(token)
-        pairs = get_pairs(word)
-
-        if not pairs:
-            return token
-
-        while True:
-            bigram = min(pairs, key=lambda pair: self.bpe_ranks.get(pair, float("inf")))
-            if bigram not in self.bpe_ranks:
-                break
-            first, second = bigram
-            new_word = []
-            i = 0
-            while i < len(word):
-                try:
-                    j = word.index(first, i)
-                except ValueError:
-                    new_word.extend(word[i:])
-                    break
-                else:
-                    new_word.extend(word[i:j])
-                    i = j
-
-                if word[i] == first and i < len(word) - 1 and word[i + 1] == second:
-                    new_word.append(first + second)
-                    i += 2
-                else:
-                    new_word.append(word[i])
-                    i += 1
-            new_word = tuple(new_word)
-            word = new_word
-            if len(word) == 1:
-                break
-            else:
-                pairs = get_pairs(word)
-        word = " ".join(word)
-        self.cache[token] = word
-        return word
-
-    # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer._tokenize
-    def _tokenize(self, text):
-        """Tokenize a string."""
-        bpe_tokens = []
-        for token in re.findall(self.pat, text):
-            token = "".join(
-                self.byte_encoder[b] for b in token.encode("utf-8")
-            )  # Maps all our bytes to unicode strings, avoiding control tokens of the BPE (spaces in our case)
-            bpe_tokens.extend(bpe_token for bpe_token in self.bpe(token).split(" "))
-        return bpe_tokens
-
-    # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer._convert_token_to_id
-    def _convert_token_to_id(self, token):
-        """Converts a token (str) in an id using the vocab."""
-        return self.encoder.get(token, self.encoder.get(self.unk_token))
-
-    # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer._convert_id_to_token
-    def _convert_id_to_token(self, index):
-        """Converts an index (integer) in a token (str) using the vocab."""
-        return self.decoder.get(index)
-
-    # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer.convert_tokens_to_string
-    def convert_tokens_to_string(self, tokens):
-        """Converts a sequence of tokens (string) in a single string."""
-        text = "".join(tokens)
-        text = bytearray([self.byte_decoder[c] for c in text]).decode("utf-8", errors=self.errors)
-        return text
-
-    def decode(
-        self,
-        token_ids,
-        skip_special_tokens: bool = False,
-        clean_up_tokenization_spaces: Optional[bool] = False,
-        spaces_between_special_tokens: bool = False,
-        **kwargs,
-    ) -> str:
-        # `spaces_between_special_tokens` defaults to True for _decode in slow tokenizers
-        # and cannot be configured elsewhere, but it should default to False for DreamTokenizer
-        return super().decode(
-            token_ids,
-            skip_special_tokens=skip_special_tokens,
-            clean_up_tokenization_spaces=clean_up_tokenization_spaces,
-            spaces_between_special_tokens=spaces_between_special_tokens,
-            **kwargs,
-        )
-
-    # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer.save_vocabulary
-    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
-        if not os.path.isdir(save_directory):
-            logger.error(f"Vocabulary path ({save_directory}) should be a directory")
-            return
-        vocab_file = os.path.join(
-            save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
-        )
-        merge_file = os.path.join(
-            save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["merges_file"]
-        )
-
-        with open(vocab_file, "w", encoding="utf-8") as f:
-            f.write(json.dumps(self.encoder, indent=2, sort_keys=True, ensure_ascii=False) + "\n")
-
-        index = 0
-        with open(merge_file, "w", encoding="utf-8") as writer:
-            writer.write("#version: 0.2\n")
-            for bpe_tokens, token_index in sorted(self.bpe_ranks.items(), key=lambda kv: kv[1]):
-                if index != token_index:
-                    logger.warning(
-                        f"Saving vocabulary to {merge_file}: BPE merge indices are not consecutive."
-                        " Please check that the tokenizer is not corrupted!"
-                    )
-                    index = token_index
-                writer.write(" ".join(bpe_tokens) + "\n")
-                index += 1
-
-        return vocab_file, merge_file
-
-    def prepare_for_tokenization(self, text, **kwargs):
-        text = unicodedata.normalize("NFC", text)
-        return (text, kwargs)

tokenizer_config.json

@@ -1,219 +0,0 @@
-{
-  "add_bos_token": false,
-  "add_prefix_space": false,
-  "added_tokens_decoder": {
-    "151643": {
-      "content": "<|endoftext|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "151644": {
-      "content": "<|im_start|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "151645": {
-      "content": "<|im_end|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "151646": {
-      "content": "<|object_ref_start|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "151647": {
-      "content": "<|object_ref_end|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "151648": {
-      "content": "<|box_start|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "151649": {
-      "content": "<|box_end|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "151650": {
-      "content": "<|quad_start|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "151651": {
-      "content": "<|quad_end|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "151652": {
-      "content": "<|vision_start|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "151653": {
-      "content": "<|vision_end|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "151654": {
-      "content": "<|vision_pad|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "151655": {
-      "content": "<|image_pad|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "151656": {
-      "content": "<|video_pad|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "151657": {
-      "content": "<tool_call>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": false
-    },
-    "151658": {
-      "content": "</tool_call>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": false
-    },
-    "151659": {
-      "content": "<|fim_prefix|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": false
-    },
-    "151660": {
-      "content": "<|fim_middle|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": false
-    },
-    "151661": {
-      "content": "<|fim_suffix|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": false
-    },
-    "151662": {
-      "content": "<|fim_pad|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": false
-    },
-    "151663": {
-      "content": "<|repo_name|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": false
-    },
-    "151664": {
-      "content": "<|file_sep|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": false
-    },
-    "151665": {
-      "content": "<|beginoftext|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "151666": {
-      "content": "<|mask|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    }
-  },
-  "additional_special_tokens": [
-    "<|beginoftext|>",
-    "<|mask|>"
-  ],
-  "auto_map": {
-    "AutoTokenizer": [
-      "tokenization_dream.DreamTokenizer",
-      null
-    ]
-  },
-  "bos_token": "<|beginoftext|>",
-  "clean_up_tokenization_spaces": false,
-  "eos_token": "<|endoftext|>",
-  "errors": "replace",
-  "extra_special_tokens": {},
-  "mask_token": "<|mask|>",
-  "model_max_length": 131072,
-  "pad_token": "<|endoftext|>",
-  "split_special_tokens": false,
-  "tokenizer_class": "DreamTokenizer",
-  "unk_token": null
-}

training_state.pt

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:c2e57be8b2bb7e72d34621d36abbb2af709e6f116b91e82672b48b8607f92c1c
-size 1495