Upload inference/generation_utils.py with huggingface_hub

inference/generation_utils.py

@@ -0,0 +1,313 @@
+# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
+# Copyright 2025 NVIDIA CORPORATION & AFFILIATES
+#
+# 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.
+#
+# SPDX-License-Identifier: Apache-2.0
+# Modified from Dream repos: https://github.com/HKUNLP/Dream
+
+
+
+from dataclasses import dataclass
+from collections.abc import Iterable
+from typing import Any, Dict, Optional, Tuple, Union
+
+import torch
+try:
+    import torch_npu
+except ImportError as e:
+    pass
+import torch.distributions as dists
+from torch.nn import functional as F
+
+from transformers.cache_utils import Cache, DynamicCache
+
+
+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
+
+
+class BlockDynamicCache(DynamicCache):
+    """
+    When `skip_cache_update` is True, this class does NOT update the cached key and value states.
+    Instead, it concatenates the current states with the original cached states along the sequence dimension
+    and returns the result. 
+
+    Example:
+
+        ```python
+        >>> past_key_values = BlockDynamicCache()
+        >>> past_key_values.skip_cache_update = True
+        >>> outputs.past_key_values
+        ```
+    """
+    def __init__(self, _distributed_cache_data: Optional[Iterable] = None) -> None:
+        """
+        Initialize a BlockDynamicCache instance.
+
+        skip_cache_update is False by default.
+        """
+        super().__init__(_distributed_cache_data)
+        self.skip_cache_update = False
+        
+    def update(
+        self,
+        key_states: torch.Tensor,
+        value_states: torch.Tensor,
+        layer_idx: int,
+        cache_kwargs: Optional[dict[str, Any]] = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Updates the cache with the new `key_states` and `value_states` for the layer `layer_idx`.
+
+        Behavior depends on the `skip_cache_update` flag:
+        - If `skip_cache_update` is True:
+            * Does NOT update the stored cache.
+            * Concatenates the current `key_states` and `value_states` 
+              with the original cached states along the sequence dimension.
+            * Returns the concatenated result.
+        - If `skip_cache_update` is False:
+            * Uses the parent class update logic to update the cache.
+
+        Parameters:
+            key_states (`torch.Tensor`):
+                The new key states to cache.
+            value_states (`torch.Tensor`):
+                The new value states to cache.
+            layer_idx (`int`):
+                The index of the layer to cache the states for.
+            cache_kwargs (`dict[str, Any]`, `optional`):
+                Additional arguments for the cache subclass. No additional arguments are used in `DynamicCache`.
+
+        Returns:
+            Tuple[torch.Tensor, torch.Tensor]:
+                The updated key and value states after concatenation or update.
+                When `skip_cache_update=True`, returns the concatenated tensor without modifying cache.
+                When `skip_cache_update=False`, returns the result from the parent class.
+        """
+        if self.skip_cache_update:
+            key_cache = torch.cat([self.key_cache[layer_idx], key_states], dim=-2)
+            value_cache = torch.cat([self.value_cache[layer_idx], value_states], dim=-2)
+            return key_cache, value_cache
+        return super().update(key_states, value_states, layer_idx, cache_kwargs)
+
+    
+
+@torch.no_grad()
+def diffusion_generate(
+    model,
+    inputs: Optional[torch.Tensor] = None,
+    top_p: Optional[int] = None,
+    top_k: Optional[int] = None,
+    threshold: Optional[float] = 0.9,
+    num_small_blocks: Optional[int] = 1,
+    **kwargs,
+):
+    block_length=kwargs.pop("block_length", 32)
+    attention_mask = kwargs.pop("attention_mask", None)
+    alg = kwargs.get("alg", 'origin')
+    temperature = kwargs.get("temperature", 0.0)
+    mask_token_id = kwargs.get("mask_token_id", None)
+    eos_token_id = kwargs.get("eos_token_id", None)
+
+    if mask_token_id is None:
+        raise ValueError("mask_token_id must be provided")
+
+    if eos_token_id is None:
+        raise ValueError("eos_token_id must be provided")
+
+    if inputs is None:
+        raise ValueError("inputs must be provided")
+
+    if attention_mask is None:
+        raise ValueError("attention_mask must be provided")
+
+
+    input_ids = inputs
+
+    if type(kwargs.get('max_new_tokens', None)) is int:
+        max_length = kwargs.get('max_new_tokens') + input_ids.shape[-1]
+    elif kwargs.get('max_length', None) is None:
+        raise ValueError("Pass max_new_tokens or max_length")
+
+    prompt_length = input_ids.shape[1]
+    if (max_length - prompt_length) % block_length != 0:
+        raise ValueError(
+            f"The token length ({max_length - prompt_length}) "
+            f"cannot be evenly divided by the block length ({block_length})."
+        )
+
+    num_blocks = (max_length - prompt_length) // block_length
+    device = model.device
+    position_ids = torch.arange(max_length, device=device).unsqueeze(0)
+    # pad input_ids to max_length
+    x = F.pad(input_ids, (0, max_length - prompt_length), value=mask_token_id)
+
+    # Initialize cache for the prompt
+    past_key_values = BlockDynamicCache()
+
+    causal_mask = torch.tril(torch.ones(max_length, max_length, device=device, dtype=torch.bool))[None, None, :, :]
+    
+    padding_mask = F.pad(attention_mask, (0, max_length - attention_mask.shape[1]), value=1.0)
+    position_ids = padding_mask.long().cumsum(-1) - 1
+    position_ids.masked_fill_(padding_mask == 0, 1)
+    # [B, N] --> [B, 1, N, N]
+    padding_mask = torch.logical_and(
+        padding_mask.unsqueeze(1).unsqueeze(-2),
+        padding_mask.unsqueeze(1).unsqueeze(-1),
+    )
+    attention_mask = padding_mask & causal_mask
+
+    
+    # Prefill stage
+    if prompt_length > 0:
+        cur_x = x[:, :prompt_length]
+        cur_attn_mask = attention_mask[:, :, :prompt_length, :prompt_length]
+        cur_position_ids = position_ids[:, :prompt_length]
+        output = model(cur_x,
+                attention_mask=cur_attn_mask,
+                position_ids=cur_position_ids,
+                past_key_values=past_key_values,
+                use_cache=True
+                        )
+        past_key_values = output.past_key_values
+    
+        logits = output.logits[:, -1:]
+        confidence, x0 = sample_tokens(logits, temperature=temperature, top_p=top_p, top_k=top_k)
+        x[:, prompt_length:prompt_length + 1] = x0
+        
+    # Process each block
+    for num_block in range(num_blocks):
+        block_start = prompt_length + num_block * block_length
+        block_end = prompt_length + (num_block + 1) * block_length
+        cur_x = x[:, block_start:block_end]
+        cur_attn_mask = attention_mask[:, :, block_start:block_end, :block_end]
+        cur_padding_mask = padding_mask[:, :, block_start:block_end, :block_end]
+        cur_position_ids = position_ids[:, block_start:block_end]    
+        # Use cache for generation
+        small_block_length = block_length // num_small_blocks
+
+        if block_length % num_small_blocks != 0:
+            raise ValueError(
+                f"block_length ({block_length}) must be divisible by num_small_blocks ({num_small_blocks})."
+            )
+
+        # Just concatenates current key value states, do not update key value cache
+        past_key_values.skip_cache_update = True
+        for small_block_idx in range(num_small_blocks):
+            small_block_start = small_block_idx * small_block_length
+            small_block_end = small_block_start + small_block_length
+
+            while True:
+                sub_mask_index = (cur_x[:, small_block_start:small_block_end] == mask_token_id)
+                if sub_mask_index.sum() == 0:
+                    break
+                
+                output = model(cur_x,
+                        attention_mask=cur_padding_mask,
+                        position_ids=cur_position_ids,
+                        past_key_values=past_key_values,
+                        use_cache=True)
+                logits = output.logits
+                logits = torch.cat([logits[:, :1], logits[:, :-1]], dim=1)
+                logits = logits[:, small_block_start:small_block_end]
+
+                confidence, x0 = sample_tokens(logits, temperature=temperature, top_p=top_p, top_k=top_k,
+                                            neg_entropy=(alg == 'entropy'), margin_confidence=(alg == 'topk_margin'))
+                confidence = torch.where(sub_mask_index, confidence, -torch.inf)
+                transfer_index = (F.one_hot(torch.max(confidence, dim=1)[1], num_classes=small_block_length) == 1)
+                if alg == 'confidence_threshold':
+                    transfer_index |= (confidence > threshold)
+                cur_x[:, small_block_start:small_block_end][transfer_index] = x0[transfer_index]
+
+            if eos_token_id and (x[:, prompt_length:] == eos_token_id).any(dim=1).all():
+                return x
+            
+        # Store kv cache
+        past_key_values.skip_cache_update = False
+        output = model(cur_x, 
+                        attention_mask=cur_attn_mask,
+                position_ids=cur_position_ids,
+                past_key_values=past_key_values,
+                use_cache=True,
+                        )
+        past_key_values = output.past_key_values
+        if num_block < num_blocks - 1:
+            logits = output.logits[:, -1:]
+            confidence, x0 = sample_tokens(logits, temperature=temperature, top_p=top_p, top_k=top_k)
+            x[:, block_end:block_end + 1] = x0
+
+    return x
+
+
+
+
+
+
+
+
+
+