Aluode/PerceptionLabPortable
0
1# coding=utf-82# Copyright 2021 The Google AI Flax Team Authors, and The HuggingFace Inc. team.3# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.4#5# Licensed under the Apache License, Version 2.0 (the "License");6# you may not use this file except in compliance with the License.7# You may obtain a copy of the License at8#9# http://www.apache.org/licenses/LICENSE-2.010#11# Unless required by applicable law or agreed to in writing, software12# distributed under the License is distributed on an "AS IS" BASIS,13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.14# See the License for the specific language governing permissions and15# limitations under the License.16 17 18import copy19import inspect20import warnings21from functools import partial22from typing import Any, Optional, Union23 24import flax25import jax26import jax.numpy as jnp27import numpy as np28from jax import lax29 30from ..models.auto import (31 FLAX_MODEL_FOR_CAUSAL_LM_MAPPING,32 FLAX_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,33 FLAX_MODEL_FOR_VISION_2_SEQ_MAPPING,34)35from ..utils import ModelOutput, logging36from .configuration_utils import GenerationConfig37from .flax_logits_process import (38 FlaxForcedBOSTokenLogitsProcessor,39 FlaxForcedEOSTokenLogitsProcessor,40 FlaxForceTokensLogitsProcessor,41 FlaxLogitsProcessorList,42 FlaxMinLengthLogitsProcessor,43 FlaxNoRepeatNGramLogitsProcessor,44 FlaxSuppressTokensAtBeginLogitsProcessor,45 FlaxSuppressTokensLogitsProcessor,46 FlaxTemperatureLogitsWarper,47 FlaxTopKLogitsWarper,48 FlaxTopPLogitsWarper,49)50 51 52logger = logging.get_logger(__name__)53 54 55@flax.struct.dataclass56class FlaxGreedySearchOutput(ModelOutput):57 """58 Flax Base class for outputs of decoder-only generation models using greedy search.59 60 61 Args:62 sequences (`jnp.ndarray` of shape `(batch_size, max_length)`):63 The generated sequences.64 """65 66 sequences: Optional[jnp.ndarray] = None67 68 69@flax.struct.dataclass70class FlaxSampleOutput(ModelOutput):71 """72 Flax Base class for outputs of decoder-only generation models using sampling.73 74 75 Args:76 sequences (`jnp.ndarray` of shape `(batch_size, max_length)`):77 The generated sequences.78 """79 80 sequences: Optional[jnp.ndarray] = None81 82 83@flax.struct.dataclass84class FlaxBeamSearchOutput(ModelOutput):85 """86 Flax Base class for outputs of decoder-only generation models using greedy search.87 88 89 Args:90 sequences (`jnp.ndarray` of shape `(batch_size, max_length)`):91 The generated sequences.92 scores (`jnp.ndarray` of shape `(batch_size,)`):93 The scores (log probabilities) of the generated sequences.94 """95 96 sequences: Optional[jnp.ndarray] = None97 scores: Optional[jnp.ndarray] = None98 99 100@flax.struct.dataclass101class GreedyState:102 cur_len: jnp.ndarray103 sequences: jnp.ndarray104 running_token: jnp.ndarray105 is_sent_finished: jnp.ndarray106 model_kwargs: dict[str, jnp.ndarray]107 108 109@flax.struct.dataclass110class SampleState:111 cur_len: jnp.ndarray112 sequences: jnp.ndarray113 running_token: jnp.ndarray114 is_sent_finished: jnp.ndarray115 prng_key: jnp.ndarray116 model_kwargs: dict[str, jnp.ndarray]117 118 119@flax.struct.dataclass120class BeamSearchState:121 cur_len: jnp.ndarray122 running_sequences: jnp.ndarray123 running_scores: jnp.ndarray124 sequences: jnp.ndarray125 scores: jnp.ndarray126 is_sent_finished: jnp.ndarray127 model_kwargs: dict[str, jnp.ndarray]128 129 130class FlaxGenerationMixin:131 """132 A class containing all functions for auto-regressive text generation, to be used as a mixin in133 [`FlaxPreTrainedModel`].134 135 The class exposes [`~generation.FlaxGenerationMixin.generate`], which can be used for:136 - *greedy decoding* by calling [`~generation.FlaxGenerationMixin._greedy_search`] if `num_beams=1` and137 `do_sample=False`138 - *multinomial sampling* by calling [`~generation.FlaxGenerationMixin._sample`] if `num_beams=1` and139 `do_sample=True`140 - *beam-search decoding* by calling [`~generation.FlaxGenerationMixin._beam_search`] if `num_beams>1` and141 `do_sample=False`142 143 You do not need to call any of the above methods directly. Pass custom parameter values to 'generate' instead. To144 learn more about decoding strategies refer to the [text generation strategies guide](../generation_strategies).145 """146 147 def prepare_inputs_for_generation(self, *args, **kwargs):148 raise NotImplementedError(149 "A model class needs to define a `prepare_inputs_for_generation` method in order to use `generate`."150 )151 152 @staticmethod153 def _run_loop_in_debug(cond_fn, body_fn, init_state):154 """155 Run generation in untraced mode. This should only be used for debugging purposes.156 """157 state = init_state158 while cond_fn(state):159 state = body_fn(state)160 return state161 162 def _prepare_encoder_decoder_kwargs_for_generation(self, input_ids, params, model_kwargs):163 encoder_kwargs = {164 argument: value165 for argument, value in model_kwargs.items()166 if not (argument.startswith("decoder_") or argument.startswith("cross_attn"))167 }168 model_kwargs["encoder_outputs"] = self.encode(input_ids, params=params, return_dict=True, **encoder_kwargs)169 return model_kwargs170 171 def _prepare_decoder_input_ids_for_generation(172 self,173 batch_size: int,174 decoder_start_token_id: Optional[int] = None,175 bos_token_id: Optional[int] = None,176 model_kwargs: Optional[dict[str, jnp.ndarray]] = None,177 ) -> jnp.ndarray:178 if model_kwargs is not None and "decoder_input_ids" in model_kwargs:179 # Only use this arg if not None, otherwise just remove from model_kwargs180 decoder_input_ids = model_kwargs.pop("decoder_input_ids")181 if decoder_input_ids is not None:182 return decoder_input_ids183 decoder_start_token_id = self._get_decoder_start_token_id(decoder_start_token_id, bos_token_id)184 return jnp.array(decoder_start_token_id, dtype="i4").reshape(1, -1).repeat(batch_size, axis=0)185 186 def _get_decoder_start_token_id(187 self, decoder_start_token_id: Optional[int] = None, bos_token_id: Optional[int] = None188 ) -> int:189 # retrieve decoder_start_token_id for encoder-decoder models190 # fall back to bos_token_id if necessary191 decoder_start_token_id = (192 decoder_start_token_id193 if decoder_start_token_id is not None194 else self.generation_config.decoder_start_token_id195 )196 bos_token_id = bos_token_id if bos_token_id is not None else self.generation_config.bos_token_id197 if decoder_start_token_id is not None:198 return decoder_start_token_id199 elif (200 hasattr(self.config, "decoder")201 and hasattr(self.config.decoder, "decoder_start_token_id")202 and self.config.decoder.decoder_start_token_id is not None203 ):204 return self.config.decoder.decoder_start_token_id205 elif bos_token_id is not None:206 return bos_token_id207 elif (208 hasattr(self.config, "decoder")209 and hasattr(self.config.decoder, "bos_token_id")210 and self.config.decoder.bos_token_id is not None211 ):212 return self.config.decoder.bos_token_id213 raise ValueError(214 "`decoder_start_token_id` or `bos_token_id` has to be defined for encoder-decoder generation."215 )216 217 @staticmethod218 def _expand_to_num_beams(tensor, num_beams):219 return jnp.broadcast_to(tensor[:, None], (tensor.shape[0], num_beams) + tensor.shape[1:])220 221 def _adapt_logits_for_beam_search(self, logits):222 """223 This function can be overwritten in the specific modeling_flax_<model-name>.py classes to allow for custom beam224 search behavior. Note that the only model that overwrites this method is [`~transformers.FlaxMarianMTModel`].225 """226 return logits227 228 def _validate_model_class(self):229 """230 Confirms that the model class is compatible with generation. If not, raises an exception that points to the231 right class to use.232 """233 if not self.can_generate():234 generate_compatible_mappings = [235 FLAX_MODEL_FOR_CAUSAL_LM_MAPPING,236 FLAX_MODEL_FOR_VISION_2_SEQ_MAPPING,237 FLAX_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,238 ]239 generate_compatible_classes = set()240 for model_mapping in generate_compatible_mappings:241 supported_models = model_mapping.get(type(self.config), default=None)242 if supported_models is not None:243 generate_compatible_classes.add(supported_models.__name__)244 exception_message = (245 f"The current model class ({self.__class__.__name__}) is not compatible with `.generate()`, as "246 "it doesn't have a language model head."247 )248 if generate_compatible_classes:249 exception_message += f" Please use one of the following classes instead: {generate_compatible_classes}"250 raise TypeError(exception_message)251 252 def _validate_model_kwargs(self, model_kwargs: dict[str, Any]):253 """Validates model kwargs for generation. Generate argument typos will also be caught here."""254 unused_model_args = []255 model_args = set(inspect.signature(self.prepare_inputs_for_generation).parameters)256 # `kwargs`/`model_kwargs` is often used to handle optional forward pass inputs like `attention_mask`. If257 # `prepare_inputs_for_generation` doesn't accept them, then a stricter check can be made ;)258 if "kwargs" in model_args or "model_kwargs" in model_args:259 model_args |= set(inspect.signature(self.__call__).parameters)260 for key, value in model_kwargs.items():261 if value is not None and key not in model_args:262 unused_model_args.append(key)263 264 if unused_model_args:265 raise ValueError(266 f"The following `model_kwargs` are not used by the model: {unused_model_args} (note: typos in the"267 " generate arguments will also show up in this list)"268 )269 270 def generate(271 self,272 input_ids: jnp.ndarray,273 generation_config: Optional[GenerationConfig] = None,274 prng_key: Optional[jnp.ndarray] = None,275 trace: bool = True,276 params: Optional[dict[str, jnp.ndarray]] = None,277 logits_processor: Optional[FlaxLogitsProcessorList] = None,278 **kwargs,279 ):280 r"""281 Generates sequences of token ids for models with a language modeling head.282 283 Parameters:284 input_ids (`jnp.ndarray` of shape `(batch_size, sequence_length)`):285 The sequence used as a prompt for the generation.286 generation_config (`~generation.GenerationConfig`, *optional*):287 The generation configuration to be used as base parametrization for the generation call. `**kwargs`288 passed to generate matching the attributes of `generation_config` will override them. If289 `generation_config` is not provided, the default will be used, which had the following loading290 priority: 1) from the `generation_config.json` model file, if it exists; 2) from the model291 configuration. Please note that unspecified parameters will inherit [`~generation.GenerationConfig`]'s292 default values, whose documentation should be checked to parameterize generation.293 trace (`bool`, *optional*, defaults to `True`):294 Whether to trace generation. Setting `trace=False` should only be used for debugging and will lead to a295 considerably slower runtime.296 params (`dict[str, jnp.ndarray]`, *optional*):297 Optionally the model parameters can be passed. Can be useful for parallelized generation.298 logits_processor (`FlaxLogitsProcessorList `, *optional*):299 Custom logits processors that complement the default logits processors built from arguments and300 generation config. If a logit processor is passed that is already created with the arguments or a301 generation config an error is thrown. This feature is intended for advanced users.302 kwargs (`dict[str, Any]`, *optional*):303 Ad hoc parametrization of `generate_config` and/or additional model-specific kwargs that will be304 forwarded to the `forward` function of the model. If the model is an encoder-decoder model, encoder305 specific kwargs should not be prefixed and decoder specific kwargs should be prefixed with *decoder_*.306 307 Return:308 [`~utils.ModelOutput`].309 310 """311 # Handle `generation_config` and kwargs that might update it, and validate the `.generate()` call312 self._validate_model_class()313 314 # priority: `generation_config` argument > `model.generation_config` (the default generation config)315 if generation_config is None:316 # legacy: users may modify the model configuration to control generation. To trigger this legacy behavior,317 # two conditions must be met318 # 1) the generation config must have been created from the model config (`_from_model_config` field);319 # 2) the generation config must have seen no modification since its creation (the hash is the same).320 if self.generation_config._from_model_config and self.generation_config._original_object_hash == hash(321 self.generation_config322 ):323 new_generation_config = GenerationConfig.from_model_config(self.config)324 if new_generation_config != self.generation_config:325 warnings.warn(326 "You have modified the pretrained model configuration to control generation. This is a"327 " deprecated strategy to control generation and will be removed soon, in a future version."328 " Please use and modify the model generation configuration (see"329 " https://huggingface.co/docs/transformers/generation_strategies#default-text-generation-configuration )"330 )331 self.generation_config = new_generation_config332 generation_config = self.generation_config333 334 generation_config = copy.deepcopy(generation_config)335 model_kwargs = generation_config.update(**kwargs) # All unused kwargs must be model kwargs336 self._validate_model_kwargs(model_kwargs.copy())337 338 logits_processor = logits_processor if logits_processor is not None else FlaxLogitsProcessorList()339 340 # set init values341 prng_key = prng_key if prng_key is not None else jax.random.PRNGKey(0)342 343 if generation_config.pad_token_id is None and generation_config.eos_token_id is not None:344 if model_kwargs.get("attention_mask") is None:345 logger.warning(346 "The attention mask and the pad token id were not set. As a consequence, you may observe "347 "unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results."348 )349 eos_token_id = generation_config.eos_token_id350 if isinstance(eos_token_id, list):351 eos_token_id = eos_token_id[0]352 generation_config.pad_token_id = eos_token_id353 354 if generation_config.decoder_start_token_id is None and self.config.is_encoder_decoder:355 raise ValueError("`decoder_start_token_id` has to be defined for encoder-decoder generation.")356 357 # decoder-only models should use left-padding for generation (can't be checked with `trace=True`)358 if not self.config.is_encoder_decoder and not trace:359 if (360 generation_config.pad_token_id is not None361 and jnp.sum(input_ids[:, -1] == generation_config.pad_token_id) > 0362 ):363 logger.warning(364 "A decoder-only architecture is being used, but right-padding was detected! For correct "365 "generation results, please set `padding_side='left'` when initializing the tokenizer."366 )367 368 batch_size = input_ids.shape[0]369 370 if self.config.is_encoder_decoder:371 # add encoder_outputs to model_kwargs372 if model_kwargs.get("encoder_outputs") is None:373 model_kwargs = self._prepare_encoder_decoder_kwargs_for_generation(input_ids, params, model_kwargs)374 # prepare decoder_input_ids for generation375 input_ids = self._prepare_decoder_input_ids_for_generation(376 batch_size,377 decoder_start_token_id=generation_config.decoder_start_token_id,378 bos_token_id=generation_config.bos_token_id,379 model_kwargs=model_kwargs,380 )381 382 # Prepare `max_length` depending on other stopping criteria.383 input_ids_seq_length = input_ids.shape[-1]384 has_default_max_length = kwargs.get("max_length") is None and generation_config.max_length is not None385 if has_default_max_length and generation_config.max_new_tokens is None and generation_config.max_length == 20:386 # 20 is the default max_length of the generation config387 warnings.warn(388 f"Using the model-agnostic default `max_length` (={generation_config.max_length}) "389 "to control the generation length. recommend setting `max_new_tokens` to control the maximum length of the generation.",390 UserWarning,391 )392 elif generation_config.max_new_tokens is not None:393 if not has_default_max_length and generation_config.max_length is not None:394 logger.warning(395 f"Both `max_new_tokens` (={generation_config.max_new_tokens}) and `max_length`(="396 f"{generation_config.max_length}) seem to have been set. `max_new_tokens` will take precedence. "397 "Please refer to the documentation for more information. "398 "(https://huggingface.co/docs/transformers/main/en/main_classes/text_generation)"399 )400 generation_config.max_length = generation_config.max_new_tokens + input_ids_seq_length401 else: # by default let's always generate 20 new tokens402 if generation_config.max_length == GenerationConfig().max_length:403 generation_config.max_length = generation_config.max_length + input_ids_seq_length404 max_position_embeddings = getattr(self.config, "max_position_embeddings", None)405 if max_position_embeddings is not None:406 generation_config.max_length = min(generation_config.max_length, max_position_embeddings)407 408 if generation_config.min_length is not None and generation_config.min_length > generation_config.max_length:409 raise ValueError(410 f"Unfeasable length constraints: the minimum length ({generation_config.min_length}) is larger than"411 f" the maximum length ({generation_config.max_length})"412 )413 if input_ids_seq_length >= generation_config.max_length:414 input_ids_string = "decoder_input_ids" if self.config.is_encoder_decoder else "input_ids"415 logger.warning(416 f"Input length of {input_ids_string} is {input_ids_seq_length}, but `max_length` is set to"417 f" {generation_config.max_length}. This can lead to unexpected behavior. You should consider"418 " increasing`max_new_tokens`."419 )420 421 logits_processor = self._get_logits_processor(422 generation_config=generation_config,423 input_ids_seq_length=input_ids_seq_length,424 logits_processor=logits_processor,425 )426 427 if not generation_config.do_sample and generation_config.num_beams == 1:428 return self._greedy_search(429 input_ids,430 generation_config.max_length,431 generation_config.pad_token_id,432 generation_config.eos_token_id,433 logits_processor=logits_processor,434 trace=trace,435 params=params,436 model_kwargs=model_kwargs,437 )438 elif generation_config.do_sample and generation_config.num_beams == 1:439 logits_warper = self._get_logits_warper(generation_config=generation_config)440 return self._sample(441 input_ids,442 generation_config.max_length,443 generation_config.pad_token_id,444 generation_config.eos_token_id,445 prng_key,446 logits_warper=logits_warper,447 logits_processor=logits_processor,448 trace=trace,449 params=params,450 model_kwargs=model_kwargs,451 )452 elif not generation_config.do_sample and generation_config.num_beams > 1:453 # broadcast input_ids & encoder_outputs454 input_ids = self._expand_to_num_beams(input_ids, num_beams=generation_config.num_beams)455 456 if "encoder_outputs" in model_kwargs:457 model_kwargs["encoder_outputs"]["last_hidden_state"] = self._expand_to_num_beams(458 model_kwargs["encoder_outputs"]["last_hidden_state"], num_beams=generation_config.num_beams459 )460 461 for kwarg in ["attention_mask", "decoder_attention_mask"]:462 if kwarg in model_kwargs:463 model_kwargs[kwarg] = self._expand_to_num_beams(464 model_kwargs[kwarg], num_beams=generation_config.num_beams465 )466 467 return self._beam_search(468 input_ids,469 generation_config.max_length,470 generation_config.pad_token_id,471 generation_config.eos_token_id,472 length_penalty=generation_config.length_penalty,473 early_stopping=generation_config.early_stopping,474 logits_processor=logits_processor,475 trace=trace,476 params=params,477 num_return_sequences=generation_config.num_return_sequences,478 model_kwargs=model_kwargs,479 )480 else:481 raise NotImplementedError("`Beam sampling is currently not implemented.")482 483 def _get_logits_warper(self, generation_config: GenerationConfig) -> FlaxLogitsProcessorList:484 """485 This class returns a [`FlaxLogitsProcessorList`] list object that contains all relevant [`FlaxLogitsWarper`]486 instances used for multinomial sampling.487 """488 warpers = FlaxLogitsProcessorList()489 490 if generation_config.temperature is not None and generation_config.temperature != 1.0:491 warpers.append(FlaxTemperatureLogitsWarper(generation_config.temperature))492 if generation_config.top_k is not None and generation_config.top_k != 0:493 warpers.append(FlaxTopKLogitsWarper(top_k=generation_config.top_k, min_tokens_to_keep=1))494 if generation_config.top_p is not None and generation_config.top_p < 1.0:495 warpers.append(FlaxTopPLogitsWarper(top_p=generation_config.top_p, min_tokens_to_keep=1))496 497 return warpers498 499 def _get_logits_processor(500 self,501 generation_config: GenerationConfig,502 input_ids_seq_length: int,503 logits_processor: Optional[FlaxLogitsProcessorList],504 ) -> FlaxLogitsProcessorList:505 """506 This class returns a [`FlaxLogitsProcessorList`] list object that contains all relevant [`FlaxLogitsProcessor`]507 instances used to modify the scores of the language model head.508 """509 processors = FlaxLogitsProcessorList()510 511 if (512 generation_config.min_length is not None513 and generation_config.eos_token_id is not None514 and generation_config.min_length > -1515 ):516 processors.append(517 FlaxMinLengthLogitsProcessor(generation_config.min_length, generation_config.eos_token_id)518 )519 if generation_config.forced_bos_token_id is not None:520 processors.append(FlaxForcedBOSTokenLogitsProcessor(generation_config.forced_bos_token_id))521 if generation_config.forced_eos_token_id is not None:522 processors.append(523 FlaxForcedEOSTokenLogitsProcessor(generation_config.max_length, generation_config.forced_eos_token_id)524 )525 if generation_config.suppress_tokens is not None:526 processors.append(FlaxSuppressTokensLogitsProcessor(generation_config.suppress_tokens))527 if generation_config.begin_suppress_tokens is not None:528 begin_index = input_ids_seq_length529 begin_index = (530 begin_index531 if (input_ids_seq_length > 1 or generation_config.forced_bos_token_id is None)532 else begin_index + 1533 )534 if (535 getattr(generation_config, "forced_decoder_ids", None) is not None536 and len(generation_config.forced_decoder_ids) > 0537 ):538 # generation starts after the last token that is forced539 begin_index += generation_config.forced_decoder_ids[-1][0]540 processors.append(541 FlaxSuppressTokensAtBeginLogitsProcessor(generation_config.begin_suppress_tokens, begin_index)542 )543 if getattr(generation_config, "forced_decoder_ids", None) is not None:544 forced_decoder_ids = [545 [input_ids_seq_length + i[0] - 1, i[1]] for i in generation_config.forced_decoder_ids546 ]547 processors.append(FlaxForceTokensLogitsProcessor(forced_decoder_ids))548 if generation_config.no_repeat_ngram_size is not None and generation_config.no_repeat_ngram_size > 0:549 processors.append(FlaxNoRepeatNGramLogitsProcessor(generation_config.no_repeat_ngram_size))550 processors = self._merge_criteria_processor_list(processors, logits_processor)551 552 return processors553 554 def _merge_criteria_processor_list(555 self,556 default_list: FlaxLogitsProcessorList,557 custom_list: FlaxLogitsProcessorList,558 ) -> FlaxLogitsProcessorList:559 if len(custom_list) == 0:560 return default_list561 for default in default_list:562 for custom in custom_list:563 if type(custom) is type(default):564 object_type = "logits processor"565 raise ValueError(566 f"A custom {object_type} of type {type(custom)} with values {custom} has been passed to"567 f" `generate`, but it has already been created with the values {default}. {default} has been"568 " created by passing the corresponding arguments to generate or by the model's config default"569 f" values. If you just want to change the default values of {object_type} consider passing"570 f" them as arguments to `generate` instead of using a custom {object_type}."571 )572 default_list.extend(custom_list)573 return default_list574 575 def _greedy_search(576 self,577 input_ids: None,578 max_length: Optional[int] = None,579 pad_token_id: Optional[int] = None,580 eos_token_id: Optional[int] = None,581 logits_processor: Optional[FlaxLogitsProcessorList] = None,582 trace: bool = True,583 params: Optional[dict[str, jnp.ndarray]] = None,584 model_kwargs: Optional[dict[str, jnp.ndarray]] = None,585 ):586 # init values587 max_length = max_length if max_length is not None else self.generation_config.max_length588 pad_token_id = pad_token_id if pad_token_id is not None else self.generation_config.pad_token_id589 eos_token_id = eos_token_id if eos_token_id is not None else self.generation_config.eos_token_id590 591 batch_size, cur_len = input_ids.shape592 593 eos_token_id = jnp.array(eos_token_id, dtype=jnp.int32 if eos_token_id is not None else None)594 pad_token_id = jnp.array(pad_token_id, dtype=jnp.int32)595 cur_len = jnp.array(cur_len)596 597 # per batch-item holding current token in loop.598 sequences = jnp.full((batch_size, max_length), pad_token_id, dtype=jnp.int32)599 sequences = lax.dynamic_update_slice(sequences, input_ids, (0, 0))600 601 # per batch-item state bit indicating if sentence has finished.602 is_sent_finished = jnp.zeros((batch_size,), dtype=jnp.bool_)603 604 # For Seq2Seq generation, we only need to use the decoder instead of the whole model in generation loop605 # and pass it the `encoder_outputs`, which are part of the `model_kwargs`.606 model = self.decode if self.config.is_encoder_decoder else self607 # initialize model specific kwargs608 model_kwargs = self.prepare_inputs_for_generation(input_ids, max_length, **model_kwargs)609 610 # initialize state611 state = GreedyState(612 cur_len=cur_len,613 sequences=sequences,614 running_token=input_ids,615 is_sent_finished=is_sent_finished,616 model_kwargs=model_kwargs,617 )618 619 def greedy_search_cond_fn(state):620 """state termination condition fn."""621 has_reached_max_length = state.cur_len == max_length622 all_sequence_finished = jnp.all(state.is_sent_finished)623 finish_generation = jnp.logical_or(has_reached_max_length, all_sequence_finished)624 return ~finish_generation625 626 def greedy_search_body_fn(state):627 """state update fn."""628 model_outputs = model(state.running_token, params=params, **state.model_kwargs)629 logits = model_outputs.logits[:, -1]630 631 # apply min_length, ...632 logits = logits_processor(state.sequences, logits, state.cur_len)633 634 next_token = jnp.argmax(logits, axis=-1)635 636 next_token = next_token * ~state.is_sent_finished + pad_token_id * state.is_sent_finished637 next_is_sent_finished = state.is_sent_finished | (next_token == eos_token_id)638 next_token = next_token[:, None]639 640 next_sequences = lax.dynamic_update_slice(state.sequences, next_token, (0, state.cur_len))641 next_model_kwargs = self.update_inputs_for_generation(model_outputs, state.model_kwargs)642 return GreedyState(643 cur_len=state.cur_len + 1,644 sequences=next_sequences,645 running_token=next_token,646 is_sent_finished=next_is_sent_finished,647 model_kwargs=next_model_kwargs,648 )649 650 # The very first prompt often has sequence length > 1, so run outside of `lax.while_loop` to comply with TPU651 if input_ids.shape[1] > 1:652 state = greedy_search_body_fn(state)653 654 if not trace:655 state = self._run_loop_in_debug(greedy_search_cond_fn, greedy_search_body_fn, state)656 else:657 state = lax.while_loop(greedy_search_cond_fn, greedy_search_body_fn, state)658 659 return FlaxGreedySearchOutput(sequences=state.sequences)660 661 def _sample(662 self,663 input_ids: None,664 max_length: Optional[int] = None,665 pad_token_id: Optional[int] = None,666 eos_token_id: Optional[int] = None,667 prng_key: Optional[jnp.ndarray] = None,668 logits_processor: Optional[FlaxLogitsProcessorList] = None,669 logits_warper: Optional[FlaxLogitsProcessorList] = None,670 trace: bool = True,671 params: Optional[dict[str, jnp.ndarray]] = None,672 model_kwargs: Optional[dict[str, jnp.ndarray]] = None,673 ):674 # init values675 max_length = max_length if max_length is not None else self.generation_config.max_length676 pad_token_id = pad_token_id if pad_token_id is not None else self.generation_config.pad_token_id677 eos_token_id = eos_token_id if eos_token_id is not None else self.generation_config.eos_token_id678 prng_key = prng_key if prng_key is not None else jax.random.PRNGKey(0)679 680 batch_size, cur_len = input_ids.shape681 682 eos_token_id = jnp.array(eos_token_id, dtype=jnp.int32 if eos_token_id is not None else None)683 pad_token_id = jnp.array(pad_token_id, dtype=jnp.int32)684 cur_len = jnp.array(cur_len)685 686 # per batch-item holding current token in loop.687 sequences = jnp.full((batch_size, max_length), pad_token_id, dtype=jnp.int32)688 sequences = lax.dynamic_update_slice(sequences, input_ids, (0, 0))689 690 # per batch-item state bit indicating if sentence has finished.691 is_sent_finished = jnp.zeros((batch_size,), dtype=jnp.bool_)692 693 # For Seq2Seq generation, we only need to use the decoder instead of the whole model in generation loop694 # and pass it the `encoder_outputs`, which are part of the `model_kwargs`.695 model = self.decode if self.config.is_encoder_decoder else self696 697 # initialize model specific kwargs698 model_kwargs = self.prepare_inputs_for_generation(input_ids, max_length, **model_kwargs)699 700 # initialize state701 state = SampleState(702 cur_len=cur_len,703 sequences=sequences,704 running_token=input_ids,705 is_sent_finished=is_sent_finished,706 prng_key=prng_key,707 model_kwargs=model_kwargs,708 )709 710 def sample_search_cond_fn(state):711 """state termination condition fn."""712 has_reached_max_length = state.cur_len == max_length713 all_sequence_finished = jnp.all(state.is_sent_finished)714 finish_generation = jnp.logical_or(has_reached_max_length, all_sequence_finished)715 return ~finish_generation716 717 def sample_search_body_fn(state):718 """state update fn."""719 prng_key, prng_key_next = jax.random.split(state.prng_key)720 model_outputs = model(state.running_token, params=params, **state.model_kwargs)721 722 logits = model_outputs.logits[:, -1]723 724 # apply min_length, ...725 logits = logits_processor(state.sequences, logits, state.cur_len)726 # apply top_p, top_k, temperature727 logits = logits_warper(logits, logits, state.cur_len)728 729 next_token = jax.random.categorical(prng_key, logits, axis=-1)730 731 next_token = next_token * ~state.is_sent_finished + pad_token_id * state.is_sent_finished732 next_is_sent_finished = state.is_sent_finished | (next_token == eos_token_id)733 next_token = next_token[:, None]734 735 next_sequences = lax.dynamic_update_slice(state.sequences, next_token, (0, state.cur_len))736 next_model_kwargs = self.update_inputs_for_generation(model_outputs, state.model_kwargs)737 738 return SampleState(739 cur_len=state.cur_len + 1,740 sequences=next_sequences,741 running_token=next_token,742 is_sent_finished=next_is_sent_finished,743 model_kwargs=next_model_kwargs,744 prng_key=prng_key_next,745 )746 747 # The very first prompt often has sequence length > 1, so run outside of `lax.while_loop` to comply with TPU748 if input_ids.shape[1] > 1:749 state = sample_search_body_fn(state)750 751 if not trace:752 state = self._run_loop_in_debug(sample_search_cond_fn, sample_search_body_fn, state)753 else:754 state = lax.while_loop(sample_search_cond_fn, sample_search_body_fn, state)755 756 return FlaxSampleOutput(sequences=state.sequences)757 758 def _beam_search(759 self,760 input_ids: None,761 max_length: Optional[int] = None,762 pad_token_id: Optional[int] = None,763 eos_token_id: Optional[int] = None,764 length_penalty: Optional[float] = None,765 early_stopping: Optional[Union[bool, str]] = None,766 logits_processor: Optional[FlaxLogitsProcessorList] = None,767 trace: bool = True,768 params: Optional[dict[str, jnp.ndarray]] = None,769 num_return_sequences: Optional[int] = None,770 model_kwargs: Optional[dict[str, jnp.ndarray]] = None,771 ):772 """773 This beam search function is heavily inspired by Flax's official example:774 https://github.com/google/flax/blob/main/examples/wmt/decode.py775 """776 777 def flatten_beam_dim(tensor):778 """Flattens the first two dimensions of a non-scalar array."""779 # ignore scalars (e.g. cache index)780 if tensor.ndim == 0:781 return tensor782 return tensor.reshape((tensor.shape[0] * tensor.shape[1],) + tensor.shape[2:])783 784 def unflatten_beam_dim(tensor, batch_size, num_beams):785 """Unflattens the first, flat batch*beam dimension of a non-scalar array."""786 # ignore scalars (e.g. cache index)787 if tensor.ndim == 0:788 return tensor789 return tensor.reshape((batch_size, num_beams) + tensor.shape[1:])790 791 def gather_beams(nested, beam_indices, batch_size, new_num_beams):792 """793 Gathers the beam slices indexed by beam_indices into new beam array.794 """795 batch_indices = jnp.reshape(796 jnp.arange(batch_size * new_num_beams) // new_num_beams, (batch_size, new_num_beams)797 )798 799 def gather_fn(tensor):800 # ignore scalars (e.g. cache index)801 if tensor.ndim == 0:802 return tensor803 else:804 return tensor[batch_indices, beam_indices]805 806 return jax.tree_util.tree_map(gather_fn, nested)807 808 # init values809 max_length = max_length if max_length is not None else self.generation_config.max_length810 pad_token_id = pad_token_id if pad_token_id is not None else self.generation_config.pad_token_id811 eos_token_id = eos_token_id if eos_token_id is not None else self.generation_config.eos_token_id812 length_penalty = length_penalty if length_penalty is not None else self.generation_config.length_penalty813 early_stopping = early_stopping if early_stopping is not None else self.generation_config.early_stopping814 num_return_sequences = (815 num_return_sequences if num_return_sequences is not None else self.generation_config.num_return_sequences816 )817 818 batch_size, num_beams, cur_len = input_ids.shape819 820 eos_token_id = jnp.array(eos_token_id, dtype=jnp.int32 if eos_token_id is not None else None)821 pad_token_id = jnp.array(pad_token_id, dtype=jnp.int32)822 cur_len = jnp.array(cur_len)823 824 # record the prompt length of decoder825 decoder_prompt_len = input_ids.shape[-1]826 827 # per batch,beam-item holding current token in loop.828 sequences = jnp.full((batch_size, num_beams, max_length), pad_token_id, dtype=jnp.int32)829 running_sequences = jnp.full((batch_size, num_beams, max_length), pad_token_id, dtype=jnp.int32)830 running_sequences = lax.dynamic_update_slice(sequences, input_ids, (0, 0, 0))831 832 # per batch,beam-item state bit indicating if sentence has finished.833 is_sent_finished = jnp.zeros((batch_size, num_beams), dtype=jnp.bool_)834 835 # per batch,beam-item score, logprobs836 running_scores = jnp.tile(jnp.array([0.0] + [np.array(-1.0e7)] * (num_beams - 1)), [batch_size, 1])837 scores = jnp.ones((batch_size, num_beams)) * np.array(-1.0e7)838 839 # For Seq2Seq generation, we only need to use the decoder instead of the whole model in generation loop840 # and pass it the `encoder_outputs`, which are part of the `model_kwargs`.841 model = self.decode if self.config.is_encoder_decoder else self842 843 # flatten beam dim844 if "encoder_outputs" in model_kwargs:845 model_kwargs["encoder_outputs"]["last_hidden_state"] = flatten_beam_dim(846 model_kwargs["encoder_outputs"]["last_hidden_state"]847 )848 for kwarg in ["attention_mask", "decoder_attention_mask"]:849 if kwarg in model_kwargs:850 model_kwargs[kwarg] = flatten_beam_dim(model_kwargs[kwarg])851 852 # initialize model specific kwargs853 model_kwargs = self.prepare_inputs_for_generation(flatten_beam_dim(input_ids), max_length, **model_kwargs)854 855 # initialize state856 state = BeamSearchState(857 cur_len=cur_len,858 running_sequences=running_sequences,859 running_scores=running_scores,860 sequences=sequences,861 scores=scores,862 is_sent_finished=is_sent_finished,863 model_kwargs=model_kwargs,864 )865 866 def beam_search_cond_fn(state):867 """beam search state termination condition fn."""868 869 # 1. is less than max length?870 not_max_length_yet = state.cur_len < max_length871 872 # 2. can the new beams still improve?873 # early_stopping == False -> apply heuristic = always get the best score from `cur_len`. See the discussion874 # below for more details.875 # https://github.com/huggingface/transformers/pull/20901#issuecomment-1369845565876 # early_stopping == "never" -> compute the best score from max_length or cur_len, depending on the sign of877 # length_penalty. Positive length_penalty favors longer sequences, thus we use max_length there.878 if early_stopping == "never" and length_penalty > 0.0:879 best_running_score = state.running_scores[:, :1] / (880 (max_length - decoder_prompt_len) ** length_penalty881 )882 else:883 best_running_score = state.running_scores[:, :1] / (884 (state.cur_len - decoder_prompt_len) ** length_penalty885 )886 worst_finished_score = jnp.where(887 state.is_sent_finished, jnp.min(state.scores, axis=1, keepdims=True), np.array(-1.0e7)888 )889 improvement_still_possible = jnp.any(best_running_score > worst_finished_score)890 891 # 3. is there still a beam that has not finished?892 still_open_beam = ~(jnp.all(state.is_sent_finished) & (early_stopping is True))893 894 return not_max_length_yet & still_open_beam & improvement_still_possible895 896 def beam_search_body_fn(state, input_ids_length=1):897 """beam search state update fn."""898 # 1. Forward current tokens899 # Collect the current position slice along length to feed the fast900 # autoregressive decoder model. Flatten the beam dimension into batch901 # dimension for feeding into the model.902 # unflatten beam dimension903 # Unflatten beam dimension in attention cache arrays904 input_token = flatten_beam_dim(905 lax.dynamic_slice(906 state.running_sequences,907 (0, 0, state.cur_len - input_ids_length),908 (batch_size, num_beams, input_ids_length),909 )910 )911 model_outputs = model(input_token, params=params, **state.model_kwargs)912 913 logits = unflatten_beam_dim(model_outputs.logits[:, -1], batch_size, num_beams)914 cache = jax.tree_util.tree_map(915 lambda tensor: unflatten_beam_dim(tensor, batch_size, num_beams), model_outputs.past_key_values916 )917 918 # adapt logits for FlaxMarianMTModel919 logits = self._adapt_logits_for_beam_search(logits)920 921 # 2. Compute log probs922 # get log probabilities from logits,923 # process logits with processors (*e.g.* min_length, ...), and924 # add new logprobs to existing running logprobs scores.925 log_probs = jax.nn.log_softmax(logits)926 log_probs = logits_processor(927 flatten_beam_dim(state.running_sequences), flatten_beam_dim(log_probs), state.cur_len928 )929 log_probs = unflatten_beam_dim(log_probs, batch_size, num_beams)930 log_probs = log_probs + jnp.expand_dims(state.running_scores, axis=2)931 vocab_size = log_probs.shape[2]932 log_probs = log_probs.reshape((batch_size, num_beams * vocab_size))933 934 # 3. Retrieve top-K935 # Each item in batch has num_beams * vocab_size candidate sequences.936 # For each item, get the top 2*k candidates with the highest log-937 # probabilities. We gather the top 2*K beams here so that even if the best938 # K sequences reach EOS simultaneously, we have another K sequences939 # remaining to continue the live beam search.940 # Gather the top 2*K scores from _all_ beams.941 # Gather 2*k top beams.942 # Recover the beam index by floor division.943 # Recover token id by modulo division and expand Id array for broadcasting.944 # Update sequences for the 2*K top-k new sequences.945 beams_to_keep = 2 * num_beams946 topk_log_probs, topk_indices = lax.top_k(log_probs, k=beams_to_keep)947 topk_beam_indices = topk_indices // vocab_size948 topk_running_sequences = gather_beams(949 state.running_sequences, topk_beam_indices, batch_size, beams_to_keep950 )951 topk_ids = jnp.expand_dims(topk_indices % vocab_size, axis=2)952 topk_sequences = lax.dynamic_update_slice(topk_running_sequences, topk_ids, (0, 0, state.cur_len))953 954 # 4. Check which sequences have ended955 # Update current sequences:956 # Did any of these sequences reach an end marker?957 # To prevent these just finished sequences from being added to the current sequences958 # set of active beam search sequences, set their log probs to a very large959 # negative value.960 did_topk_just_finished = topk_sequences[:, :, state.cur_len] == eos_token_id961 running_topk_log_probs = topk_log_probs + did_topk_just_finished * np.array(-1.0e7)962 # 5. Get running sequences scores for next963 # Determine the top k beam indices (from top 2*k beams) from log probs964 # and gather top k beams (from top 2*k beams).965 next_topk_indices = lax.top_k(running_topk_log_probs, k=num_beams)[1]966 next_running_sequences, next_running_scores = gather_beams(967 [topk_sequences, running_topk_log_probs], next_topk_indices, batch_size, num_beams968 )969 970 # 6. Process topk logits971 # Further process log probs:972 # - add length penalty973 # - make sure no scores can be added anymore if beam is full974 # - make sure still running sequences cannot be chosen as finalized beam975 topk_log_probs = topk_log_probs / ((state.cur_len + 1 - decoder_prompt_len) ** length_penalty)976 beams_in_batch_are_full = jnp.broadcast_to(977 state.is_sent_finished.all(axis=-1, keepdims=True), did_topk_just_finished.shape978 ) & (early_stopping is True)979 add_penalty = ~did_topk_just_finished | beams_in_batch_are_full980 topk_log_probs += add_penalty * np.array(-1.0e7)981 982 # 7. Get scores, sequences, is sentence finished for next.983 # Combine sequences, scores, and flags along the beam dimension and compare984 # new finished sequence scores to existing finished scores and select the985 # best from the new set of beams986 merged_sequences = jnp.concatenate([state.sequences, topk_sequences], axis=1)987 merged_scores = jnp.concatenate([state.scores, topk_log_probs], axis=1)988 merged_is_sent_finished = jnp.concatenate([state.is_sent_finished, did_topk_just_finished], axis=1)989 topk_merged_indices = lax.top_k(merged_scores, k=num_beams)[1]990 next_sequences, next_scores, next_is_sent_finished = gather_beams(991 [merged_sequences, merged_scores, merged_is_sent_finished], topk_merged_indices, batch_size, num_beams992 )993 994 # 8. Update model kwargs.995 # Determine the top k beam indices from the original set of all beams.996 # With these, gather the top k beam-associated caches.997 next_running_indices = gather_beams(topk_beam_indices, next_topk_indices, batch_size, num_beams)998 next_cache = gather_beams(cache, next_running_indices, batch_size, num_beams)999 model_outputs["past_key_values"] = jax.tree_util.tree_map(lambda x: flatten_beam_dim(x), next_cache)1000 next_model_kwargs = self.update_inputs_for_generation(model_outputs, state.model_kwargs)1001 1002 return BeamSearchState(1003 cur_len=state.cur_len + 1,1004 running_scores=next_running_scores,1005 running_sequences=next_running_sequences,1006 scores=next_scores,1007 sequences=next_sequences,1008 is_sent_finished=next_is_sent_finished,1009 model_kwargs=next_model_kwargs,1010 )1011 1012 # Always run first iteration outside of `lax.while_loop` to avoid calling `beam_search_cond_fn`1013 # when `state.cur_len` equals `decoder_prompt_len`. This also helps to comply with TPU when1014 # the very first prompt has sequence length > 1.1015 state = partial(beam_search_body_fn, input_ids_length=input_ids.shape[-1])(state)1016 1017 if not trace:1018 state = self._run_loop_in_debug(beam_search_cond_fn, beam_search_body_fn, state)1019 else:1020 state = lax.while_loop(beam_search_cond_fn, beam_search_body_fn, state)1021 1022 # Account for the edge-case where there are no finished sequences for a1023 # particular batch item. If so, return running sequences for that batch item.1024 none_finished = jnp.any(state.is_sent_finished, axis=1)1025 sequences = jnp.where(none_finished[:, None, None], state.sequences, state.running_sequences)1026 scores = jnp.where(none_finished[:, None], state.scores, state.running_scores)1027 1028 # Take best beams for each batch (the score is sorted in descending order)1029 sequences = flatten_beam_dim(sequences[:, :num_return_sequences, :])1030 scores = flatten_beam_dim(scores[:, :num_return_sequences])1031 1032 return FlaxBeamSearchOutput(sequences=sequences, scores=scores)1033