bdbj/Dream-Coder-v0-Instruct-7B-SM
194
1# coding=utf-82# Copyright 2024 The Dream team, HKUNLP Group and the HuggingFace Inc. team. All rights reserved.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8# http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15 16import warnings17import copy18from dataclasses import dataclass19from typing import Any, Dict, Optional, Tuple, Union20 21import torch22import torch.distributions as dists23from torch.nn import functional as F24from transformers import __version__25from transformers.generation.configuration_utils import (26 GenerationConfig27)28from transformers.utils import (29 ModelOutput,30 is_torchdynamo_compiling,31 logging,32)33from .softmasking_utils import SMArgs, get_mixing_factors_for_softmasking34 35logger = logging.get_logger(__name__)36 37 38def top_p_logits(logits, top_p=None):39 sorted_logits, sorted_indices = torch.sort(logits, descending=True)40 cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)41 sorted_indices_to_remove = cumulative_probs > top_p42 # Shift the indices to the right to keep the first token above the threshold43 sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()44 sorted_indices_to_remove[..., 0] = 045 46 mask = torch.zeros_like(logits, dtype=torch.bool, device=logits.device)47 mask = mask.scatter_(-1, sorted_indices, sorted_indices_to_remove)48 logits = logits.masked_fill(mask, torch.finfo(logits.dtype).min)49 return logits50 51def top_k_logits(logits, top_k=None):52 top_k = min(top_k, logits.size(-1)) # Safety check53 # Remove all tokens with a probability less than the last token of the top-k54 indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]55 logits = logits.masked_fill(indices_to_remove, torch.finfo(logits.dtype).min)56 return logits57 58 59def sample_tokens(logits, temperature=0.0, top_p=None, top_k=None, margin_confidence=False, neg_entropy=False):60 original_dtype = logits.dtype61 logits = logits.to(torch.float32)62 if temperature > 0:63 logits = logits / temperature64 if top_p is not None and top_p < 1:65 logits = top_p_logits(logits, top_p)66 if top_k is not None:67 logits = top_k_logits(logits, top_k)68 probs = torch.softmax(logits, dim=-1)69 70 if temperature > 0:71 x0 = dists.Categorical(probs=probs).sample()72 confidence = torch.gather(probs, -1, x0.unsqueeze(-1)).squeeze(-1)73 else:74 confidence, x0 = probs.max(dim=-1)75 76 if margin_confidence:77 sorted_probs, _ = torch.sort(probs, dim=-1, descending=True)78 # Extract top1 and top2 probabilities79 top1_probs = sorted_probs[:, 0]80 top2_probs = sorted_probs[:, 1]81 # Calculate confidence as top1 - top282 confidence = top1_probs - top2_probs83 84 if neg_entropy:85 epsilon = 1e-1086 log_probs = torch.log(probs + epsilon)87 confidence = torch.sum(probs * log_probs, dim=-1)88 89 return confidence.to(original_dtype), x090 91 92@dataclass93class DreamModelOutput(ModelOutput):94 sequences: torch.LongTensor = None95 history: Optional[Tuple[torch.FloatTensor]] = None96 97 98class DreamGenerationConfig(GenerationConfig):99 def __init__(self, **kwargs):100 self.temperature: float = kwargs.pop("temperature", 0.0)101 self.top_p: Optional[float] = kwargs.pop("top_p", None)102 self.top_k: Optional[int] = kwargs.pop("top_k", None)103 self.max_length = kwargs.pop("max_length", 20)104 self.max_new_tokens = kwargs.pop("max_new_tokens", None)105 # diffusion specific params106 self.eps: float = kwargs.pop("eps", 1e-3)107 self.steps: int = kwargs.pop("steps", 512)108 self.alg: str = kwargs.pop("alg", 'origin')109 self.alg_temp: Optional[float] = kwargs.pop("alg_temp", None)110 self.eos_penalty: Optional[float] = kwargs.pop("eos_penalty", 0.0)111 112 # Parameters that define the output variables of `generate`113 self.num_return_sequences: int = kwargs.pop("num_return_sequences", 1)114 self.return_dict_in_generate: bool = kwargs.pop("return_dict_in_generate", False)115 self.output_history: bool = kwargs.pop("output_history", False)116 117 # Special tokens that can be used at generation time118 self.mask_token_id = kwargs.pop("mask_token_id", None)119 self.pad_token_id = kwargs.pop("pad_token_id", None)120 self.bos_token_id = kwargs.pop("bos_token_id", None)121 self.eos_token_id = kwargs.pop("eos_token_id", None)122 123 # Wild card124 self.generation_kwargs = kwargs.pop("generation_kwargs", {})125 126 # The remaining attributes do not parametrize `.generate()`, but are informative and/or used by the hub127 # interface.128 self._from_model_config = kwargs.pop("_from_model_config", False)129 self._commit_hash = kwargs.pop("_commit_hash", None)130 self.transformers_version = kwargs.pop("transformers_version", __version__)131 132 # Additional attributes without default values133 if not self._from_model_config:134 # we don't want to copy values from the model config if we're initializing a `GenerationConfig` from a135 # model's default configuration file136 for key, value in kwargs.items():137 try:138 setattr(self, key, value)139 except AttributeError as err:140 logger.error(f"Can't set {key} with value {value} for {self}")141 raise err142 143 # Validate the values of the attributes144 self.validate(is_init=True)145 146 def validate(self, is_init=False):147 pass148 149class DreamGenerationMixin:150 @staticmethod151 def _expand_inputs_for_generation(152 expand_size: int = 1,153 input_ids: Optional[torch.LongTensor] = None,154 attention_mask: Optional[torch.LongTensor] = None155 ) -> Tuple[torch.LongTensor, Dict[str, Any]]:156 """Expands tensors from [batch_size, ...] to [batch_size * expand_size, ...]"""157 # Do not call torch.repeat_interleave if expand_size is 1 because it clones158 # the input tensor and thus requires more memory although no change is applied159 if expand_size == 1:160 return input_ids, attention_mask161 if input_ids is not None:162 input_ids = input_ids.repeat_interleave(expand_size, dim=0)163 if attention_mask is not None:164 attention_mask = attention_mask.repeat_interleave(expand_size, dim=0)165 return input_ids, attention_mask166 167 def _validate_generated_length(self, generation_config, input_ids_length, has_default_max_length):168 """Performs validation related to the resulting generated length"""169 170 # Can't throw warnings/exceptions during compilation171 if is_torchdynamo_compiling():172 return173 174 # 1. Max length warnings related to poor parameterization175 if has_default_max_length and generation_config.max_new_tokens is None and generation_config.max_length == 20:176 # 20 is the default max_length of the generation config177 warnings.warn(178 f"Using the model-agnostic default `max_length` (={generation_config.max_length}) to control the "179 "generation length. We recommend setting `max_new_tokens` to control the maximum length of the "180 "generation.",181 UserWarning,182 )183 if input_ids_length >= generation_config.max_length:184 input_ids_string = "input_ids"185 raise ValueError(186 f"Input length of {input_ids_string} is {input_ids_length}, but `max_length` is set to"187 f" {generation_config.max_length}. This can lead to unexpected behavior. You should consider"188 " increasing `max_length` or, better yet, setting `max_new_tokens`."189 )190 191 def _prepare_generated_length(192 self,193 generation_config,194 has_default_max_length,195 input_ids_length,196 ):197 """Prepared max and min length in generation configs to avoid clashes between similar attributes"""198 199 if generation_config.max_new_tokens is not None:200 if not has_default_max_length and generation_config.max_length is not None:201 logger.warning(202 f"Both `max_new_tokens` (={generation_config.max_new_tokens}) and `max_length`(="203 f"{generation_config.max_length}) seem to have been set. `max_new_tokens` will take precedence. "204 "Please refer to the documentation for more information. "205 "(https://huggingface.co/docs/transformers/main/en/main_classes/text_generation)"206 )207 generation_config.max_length = generation_config.max_new_tokens + input_ids_length208 209 elif has_default_max_length:210 if generation_config.max_length == DreamGenerationConfig().max_length:211 generation_config.max_length = generation_config.max_length + input_ids_length212 max_position_embeddings = getattr(self.config, "max_position_embeddings", None)213 if max_position_embeddings is not None:214 generation_config.max_length = min(generation_config.max_length, max_position_embeddings)215 216 return generation_config217 218 def _prepare_generation_config(219 self, generation_config: Optional[DreamGenerationConfig], **kwargs: Dict220 ) -> DreamGenerationConfig:221 """222 Prepares the base generation config, then applies any generation configuration options from kwargs. This223 function handles retrocompatibility with respect to configuration files.224 """225 # priority: `generation_config` argument > `model.generation_config` (the default generation config)226 using_model_generation_config = False227 if generation_config is None:228 generation_config = DreamGenerationConfig.from_model_config(self.config)229 using_model_generation_config = True230 231 # `torch.compile` can't compile `copy.deepcopy`, arguments in `kwargs` that are part of `generation_config`232 # will mutate the object with `.update`. As such, passing these arguments through `kwargs` is disabled -- an233 # exception will be raised in `_validate_model_kwargs`234 if not is_torchdynamo_compiling():235 generation_config = copy.deepcopy(generation_config)236 _kwargs = generation_config.update(**kwargs)237 # If `generation_config` is provided, let's fallback ALL special tokens to the default values for the model238 if not using_model_generation_config:239 if generation_config.bos_token_id is None:240 generation_config.bos_token_id = self.generation_config.bos_token_id241 if generation_config.eos_token_id is None:242 generation_config.eos_token_id = self.generation_config.eos_token_id243 if generation_config.pad_token_id is None:244 generation_config.pad_token_id = self.generation_config.pad_token_id245 if generation_config.mask_token_id is None:246 generation_config.mask_token_id = self.generation_config.mask_token_id247 248 return generation_config249 250 def _prepare_special_tokens(251 self,252 generation_config: DreamGenerationConfig,253 device: Optional[Union[torch.device, str]] = None,254 ):255 """256 Prepares the special tokens for generation, overwriting the generation config with their processed versions257 converted to tensor.258 259 Note that `generation_config` is changed in place and stops being serializable after this method is called.260 That is no problem if called within `generate` (`generation_config` is a local copy that doesn't leave the261 function). However, if called outside `generate`, consider creating a copy of `generation_config` first.262 """263 264 # Convert special tokens to tensors265 def _tensor_or_none(token, device=None):266 if token is None:267 return token268 269 device = device if device is not None else self.device270 if isinstance(token, torch.Tensor):271 return token.to(device)272 return torch.tensor(token, device=device, dtype=torch.long)273 274 bos_token_tensor = _tensor_or_none(generation_config.bos_token_id, device=device)275 eos_token_tensor = _tensor_or_none(generation_config.eos_token_id, device=device)276 pad_token_tensor = _tensor_or_none(generation_config.pad_token_id, device=device)277 mask_token_tensor = _tensor_or_none(generation_config.mask_token_id, device=device)278 279 # We can have more than one eos token. Always treat it as a 1D tensor (when it exists).280 if eos_token_tensor is not None and eos_token_tensor.ndim == 0:281 eos_token_tensor = eos_token_tensor.unsqueeze(0)282 283 # Set pad token if unset (and there are conditions to do so)284 if pad_token_tensor is None and eos_token_tensor is not None:285 pad_token_tensor = eos_token_tensor[0]286 logger.warning(f"Setting `pad_token_id` to `eos_token_id`:{pad_token_tensor} for open-end generation.")287 288 # Update generation config with the updated special tokens tensors289 # NOTE: this must be written into a different attribute name than the one holding the original special tokens290 # (in their non-tensor form), in order to enable end-to-end compilation. See291 # https://pytorch.org/docs/stable/torch.compiler_cudagraph_trees.html#limitations292 generation_config._bos_token_tensor = bos_token_tensor293 generation_config._eos_token_tensor = eos_token_tensor294 generation_config._pad_token_tensor = pad_token_tensor295 generation_config._mask_token_tensor = mask_token_tensor296 297 @torch.no_grad()298 def diffusion_generate(299 self,300 inputs: Optional[torch.Tensor] = None,301 generation_config: Optional[DreamGenerationConfig] = None,302 **kwargs,303 ) -> Union[DreamModelOutput, torch.LongTensor]:304 # 1. Handle `generation_config` and kwargs that might update it, and validate the `.generate()` call305 generation_config = self._prepare_generation_config(generation_config, **kwargs)306 generation_tokens_hook_func = kwargs.pop("generation_tokens_hook_func", lambda step, x, logits: x)307 generation_logits_hook_func = kwargs.pop("generation_logits_hook_func", lambda step, x, logits: logits)308 309 # 2. Define model inputs310 assert inputs is not None311 input_ids = inputs312 device = input_ids.device313 attention_mask = kwargs.pop("attention_mask", None)314 self._prepare_special_tokens(generation_config, device=device)315 316 # 3. Prepare `max_length`.317 input_ids_length = input_ids.shape[-1]318 has_default_max_length = kwargs.get("max_length") is None and generation_config.max_length is not None319 generation_config = self._prepare_generated_length(320 generation_config=generation_config,321 has_default_max_length=has_default_max_length,322 input_ids_length=input_ids_length,323 )324 325 self._validate_generated_length(generation_config, input_ids_length, has_default_max_length)326 327 # 4. Check input_ids328 if not is_torchdynamo_compiling() and self.device.type != input_ids.device.type:329 warnings.warn(330 "You are calling .generate() with the `input_ids` being on a device type different"331 f" than your model's device. `input_ids` is on {input_ids.device.type}, whereas the model"332 f" is on {self.device.type}. You may experience unexpected behaviors or slower generation."333 " Please make sure that you have put `input_ids` to the"334 f" correct device by calling for example input_ids = input_ids.to('{self.device.type}') before"335 " running `.generate()`.",336 UserWarning,337 )338 if (339 hasattr(generation_config, "pad_token_id") and340 torch.any(input_ids == generation_config.pad_token_id) and341 attention_mask is None342 ):343 warnings.warn(344 "Padding was detected but no attention mask is passed here. For correct "345 "generation results, please set `attention_mask` when batch-padding inputs.",346 UserWarning,347 )348 349 input_ids, attention_mask = self._expand_inputs_for_generation(350 expand_size=generation_config.num_return_sequences,351 input_ids=input_ids,352 attention_mask=attention_mask353 )354 355 ###### LOAD IN Softmasking PARAMETERS ######356 sm_args = SMArgs(357 sm_alg=kwargs.pop("transparency_alg", "none"),358 sm_schedule=kwargs.pop("transparency_scheduling", "none"),359 scale=kwargs.pop("transparency_scale", 0.0),360 steepness=kwargs.pop("transparency_steepness", 0.0),361 offset=kwargs.pop("transparency_centre", 0.0),362 mixinputs_k=kwargs.pop("mixinputs_k", 1),363 mixinputs_temp=kwargs.pop("mixture_temp", 1.0),364 )365 #############################################366 367 print(f"SM parameters used in generation: {sm_args}")368 369 result = self._sample(370 input_ids,371 attention_mask=attention_mask,372 generation_config=generation_config,373 generation_tokens_hook_func=generation_tokens_hook_func,374 generation_logits_hook_func=generation_logits_hook_func,375 sm_args=sm_args376 )377 return result378 379 def _sample(380 self,381 input_ids: torch.LongTensor,382 attention_mask: Optional[torch.LongTensor],383 generation_config: DreamGenerationConfig,384 generation_tokens_hook_func,385 generation_logits_hook_func,386 sm_args: SMArgs,387 ) -> Union[DreamModelOutput, torch.LongTensor]:388 # init values389 output_history = generation_config.output_history390 return_dict_in_generate = generation_config.return_dict_in_generate391 max_length = generation_config.max_length392 mask_token_id = generation_config.mask_token_id393 pad_token_id = generation_config.pad_token_id394 steps = generation_config.steps395 eps = generation_config.eps396 alg = generation_config.alg397 alg_temp = generation_config.alg_temp398 temperature = generation_config.temperature399 top_p = generation_config.top_p400 top_k = generation_config.top_k401 eos_penalty = generation_config.eos_penalty402 403 histories = [] if (return_dict_in_generate and output_history) else None404 405 # pad input_ids to max_length406 x = F.pad(input_ids, (0, max_length - input_ids.shape[1]), value=mask_token_id)407 408 if attention_mask is not None and torch.any(attention_mask == 0.0):409 # we do not mask the [MASK] tokens so value = 1.0410 attention_mask = F.pad(attention_mask, (0, max_length - attention_mask.shape[1]), value=1.0)411 tok_idx = attention_mask.long().cumsum(-1) - 1412 tok_idx.masked_fill_(attention_mask == 0, 1)413 # attention_mask is of shape [B, N]414 # broadcast to [B, 1, N, N]415 attention_mask = torch.logical_and(416 attention_mask.unsqueeze(1).unsqueeze(-2),417 attention_mask.unsqueeze(1).unsqueeze(-1),418 )419 else:420 tok_idx = None421 attention_mask = "full"422 423 timesteps = torch.linspace(1, eps, steps + 1, device=x.device)424 425 # this allows user-defined token control of the intermediate steps426 x = generation_tokens_hook_func(None, x, None)427 428 # Initialize necessary SM variables429 inputs_embeds = None430 embed_weights = self.get_input_embeddings().weight # (V,D)431 max_gen_length = (x == mask_token_id).sum().item()432 433 for i in range(steps):434 mask_index = (x == mask_token_id)435 if inputs_embeds is None:436 logits = self(x, attention_mask, tok_idx).logits437 else:438 logits = self(inputs_embeds=inputs_embeds, attention_mask=attention_mask, position_ids=tok_idx).logits439 logits = torch.cat([logits[:,:1], logits[:, :-1]], dim=1)440 441 # this allows user-defined logits control of the intermediate steps442 logits = generation_logits_hook_func(i, x, logits)443 444 mask_logits = logits[mask_index]445 t = timesteps[i]446 s = timesteps[i + 1]447 448 mask_logits[:,pad_token_id] += eos_penalty * torch.log(1-t+eps)449 if alg == 'origin':450 p_transfer = 1 - s / t if i < steps - 1 else 1451 x0 = torch.zeros_like(x[mask_index], device=self.device, dtype=torch.long) + mask_token_id452 transfer_index_t_s = torch.rand(*x0.shape, device=self.device) < p_transfer453 _, x0[transfer_index_t_s]= sample_tokens(mask_logits[transfer_index_t_s], temperature=temperature, top_p=top_p, top_k=top_k)454 x[mask_index] = x0.clone()455 else:456 if alg == 'maskgit_plus':457 confidence, x0 = sample_tokens(mask_logits, temperature=temperature, top_p=top_p, top_k=top_k)458 elif alg == 'topk_margin':459 confidence, x0 = sample_tokens(mask_logits, temperature=temperature, top_p=top_p, top_k=top_k, margin_confidence=True)460 elif alg == 'entropy':461 confidence, x0 = sample_tokens(mask_logits, temperature, top_p=top_p, top_k=top_k, neg_entropy=True)462 else:463 raise RuntimeError(f"Unknown alg: {alg}")464 num_mask_token = mask_index.sum() / mask_index.shape[0]465 number_transfer_tokens = int(num_mask_token * (1 - s / t)) if i < steps - 1 else int(num_mask_token)466 full_confidence = torch.full_like(x, -torch.inf, device=self.device, dtype=logits.dtype)467 full_confidence[mask_index] = confidence468 if number_transfer_tokens > 0:469 if alg_temp is None or alg_temp == 0:470 _, transfer_index = torch.topk(full_confidence, number_transfer_tokens)471 else:472 full_confidence = full_confidence / alg_temp473 full_confidence = F.softmax(full_confidence, dim=-1)474 transfer_index = torch.multinomial(full_confidence, num_samples=number_transfer_tokens)475 x_ = torch.zeros_like(x, device=self.device, dtype=torch.long) + mask_token_id476 x_[mask_index] = x0.clone()477 row_indices = torch.arange(x.size(0), device=self.device).unsqueeze(1).expand_as(transfer_index)478 x[row_indices,transfer_index] = x_[row_indices,transfer_index]479 480 # this allows user-defined token control of the intermediate steps481 x = generation_tokens_hook_func(i, x, logits)482 483 # DO SOFTMASKING MIXING484 if sm_args.sm_alg != "none":485 p_sm = get_mixing_factors_for_softmasking(486 x, 487 logits, 488 mask_token_id, 489 max_gen_length, 490 sm_args491 )492 inputs_embeds = torch.matmul(p_sm, embed_weights) # (B,T,D)493 494 if histories is not None:495 histories.append(x.clone())496 497 if return_dict_in_generate:498 return DreamModelOutput(499 sequences=x,500 history=histories,501 )502 else:503 return x