Aluode/PerceptionLabPortable
0
1# coding=utf-82# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.3# Copyright (c) 2018, 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 17import copy18import inspect19import warnings20from dataclasses import dataclass21from typing import Any, Optional, Union22 23import numpy as np24import tensorflow as tf25from tensorflow.compiler.tf2xla.python.xla import dynamic_update_slice26 27from ..modeling_tf_outputs import TFCausalLMOutputWithPast, TFSeq2SeqLMOutput28from ..models.auto import (29 TF_MODEL_FOR_CAUSAL_LM_MAPPING,30 TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,31 TF_MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING,32 TF_MODEL_FOR_VISION_2_SEQ_MAPPING,33)34from ..tf_utils import shape_list, stable_softmax35from ..utils import ModelOutput, logging36from .configuration_utils import GenerationConfig37from .tf_logits_process import (38 TFForcedBOSTokenLogitsProcessor,39 TFForcedEOSTokenLogitsProcessor,40 TFForceTokensLogitsProcessor,41 TFLogitsProcessorList,42 TFMinLengthLogitsProcessor,43 TFNoBadWordsLogitsProcessor,44 TFNoRepeatNGramLogitsProcessor,45 TFRepetitionPenaltyLogitsProcessor,46 TFSuppressTokensAtBeginLogitsProcessor,47 TFSuppressTokensLogitsProcessor,48 TFTemperatureLogitsWarper,49 TFTopKLogitsWarper,50 TFTopPLogitsWarper,51)52 53 54logger = logging.get_logger(__name__)55 56 57@dataclass58class TFGreedySearchDecoderOnlyOutput(ModelOutput):59 """60 Base class for outputs of decoder-only generation models using greedy search.61 62 63 Args:64 sequences (`tf.Tensor` of shape `(batch_size, sequence_length)`):65 The generated sequences. The second dimension (sequence_length) is either equal to `max_length` or shorter66 if all batches finished early due to the `eos_token_id`.67 scores (`tuple(tf.Tensor)` *optional*, returned when `output_scores=True` is passed or when `config.output_scores=True`):68 Processed prediction scores of the language modeling head (scores for each vocabulary token before SoftMax)69 at each generation step. Tuple of `tf.Tensor` with up to `max_new_tokens` elements (one element for each70 generated token), with each tensor of shape `(batch_size, config.vocab_size)`.71 attentions (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_attentions=True` is passed or `config.output_attentions=True`):72 Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of73 `tf.Tensor` of shape `(batch_size, num_heads, generated_length, sequence_length)`.74 hidden_states (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):75 Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of76 `tf.Tensor` of shape `(batch_size, generated_length, hidden_size)`.77 """78 79 sequences: Optional[tf.Tensor] = None80 scores: Optional[tuple[tf.Tensor]] = None81 attentions: Optional[tuple[tuple[tf.Tensor]]] = None82 hidden_states: Optional[tuple[tuple[tf.Tensor]]] = None83 84 85@dataclass86class TFGreedySearchEncoderDecoderOutput(ModelOutput):87 """88 Base class for outputs of encoder-decoder generation models using greedy search. Hidden states and attention89 weights of the decoder (respectively the encoder) can be accessed via the encoder_attentions and the90 encoder_hidden_states attributes (respectively the decoder_attentions and the decoder_hidden_states attributes)91 92 93 Args:94 sequences (`tf.Tensor` of shape `(batch_size, sequence_length)`):95 The generated sequences. The second dimension (sequence_length) is either equal to `max_length` or shorter96 if all batches finished early due to the `eos_token_id`.97 scores (`tuple(tf.Tensor)` *optional*, returned when `output_scores=True` is passed or when `config.output_scores=True`):98 Processed prediction scores of the language modeling head (scores for each vocabulary token before SoftMax)99 at each generation step. Tuple of `tf.Tensor` with up to `max_new_tokens` elements (one element for each100 generated token), with each tensor of shape `(batch_size, config.vocab_size)`.101 encoder_attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or `config.output_attentions=True`):102 Tuple of `tf.Tensor` (one for each layer of the decoder) of shape `(batch_size, num_heads, sequence_length,103 sequence_length)`.104 encoder_hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):105 Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape106 `(batch_size, sequence_length, hidden_size)`.107 decoder_attentions (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_attentions=True` is passed or `config.output_attentions=True`):108 Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of109 `tf.Tensor` of shape `(batch_size, num_heads, generated_length, sequence_length)`.110 cross_attentions (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_attentions=True` is passed or `config.output_attentions=True`):111 Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of112 `tf.Tensor` of shape `(batch_size, num_heads, generated_length, sequence_length)`.113 decoder_hidden_states (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):114 Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of115 `tf.Tensor` of shape `(batch_size, generated_length, hidden_size)`.116 """117 118 sequences: Optional[tf.Tensor] = None119 scores: Optional[tuple[tf.Tensor]] = None120 encoder_attentions: Optional[tuple[tf.Tensor]] = None121 encoder_hidden_states: Optional[tuple[tf.Tensor]] = None122 decoder_attentions: Optional[tuple[tuple[tf.Tensor]]] = None123 cross_attentions: Optional[tuple[tuple[tf.Tensor]]] = None124 decoder_hidden_states: Optional[tuple[tuple[tf.Tensor]]] = None125 126 127@dataclass128class TFSampleDecoderOnlyOutput(ModelOutput):129 """130 Base class for outputs of decoder-only generation models using sampling.131 132 133 Args:134 sequences (`tf.Tensor` of shape `(batch_size*num_return_sequences, sequence_length)`):135 The generated sequences. The second dimension (sequence_length) is either equal to `max_length` or shorter136 if all batches finished early due to the `eos_token_id`.137 scores (`tuple(tf.Tensor)` *optional*, returned when `output_scores=True` is passed or when `config.output_scores=True`):138 Processed prediction scores of the language modeling head (scores for each vocabulary token before SoftMax)139 at each generation step. Tuple of `tf.Tensor` with up to `max_new_tokens` elements (one element for each140 generated token), with each tensor of shape `(batch_size*num_return_sequences, config.vocab_size)`.141 attentions (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_attentions=True` is passed or `config.output_attentions=True`):142 Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of143 `tf.Tensor` of shape `(num_return_sequences*batch_size, num_heads, generated_length, sequence_length)`.144 hidden_states (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):145 Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of146 `tf.Tensor` of shape `(num_return_sequences*batch_size, generated_length, hidden_size)`.147 """148 149 sequences: Optional[tf.Tensor] = None150 scores: Optional[tuple[tf.Tensor]] = None151 attentions: Optional[tuple[tuple[tf.Tensor]]] = None152 hidden_states: Optional[tuple[tuple[tf.Tensor]]] = None153 154 155@dataclass156class TFSampleEncoderDecoderOutput(ModelOutput):157 """158 Base class for outputs of encoder-decoder generation models using sampling. Hidden states and attention weights of159 the decoder (respectively the encoder) can be accessed via the encoder_attentions and the encoder_hidden_states160 attributes (respectively the decoder_attentions and the decoder_hidden_states attributes)161 162 163 Args:164 sequences (`tf.Tensor` of shape `(batch_size*num_return_sequences, sequence_length)`):165 The generated sequences. The second dimension (sequence_length) is either equal to `max_length` or shorter166 if all batches finished early due to the `eos_token_id`.167 scores (`tuple(tf.Tensor)` *optional*, returned when `output_scores=True` is passed or when `config.output_scores=True`):168 Processed prediction scores of the language modeling head (scores for each vocabulary token before SoftMax)169 at each generation step. Tuple of `tf.Tensor` with up to `max_new_tokens` elements (one element for each170 generated token), with each tensor of shape `(batch_size*num_return_sequences, config.vocab_size)`.171 encoder_attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or `config.output_attentions=True`):172 Tuple of `tf.Tensor` (one for each layer of the decoder) of shape `(batch_size*num_return_sequences,173 num_heads, sequence_length, sequence_length)`.174 encoder_hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):175 Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape176 `(batch_size*num_return_sequences, sequence_length, hidden_size)`.177 decoder_attentions (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_attentions=True` is passed or `config.output_attentions=True`):178 Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of179 `tf.Tensor` of shape `(batch_size*num_return_sequences, num_heads, generated_length, sequence_length)`.180 cross_attentions (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_attentions=True` is passed or `config.output_attentions=True`):181 Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of182 `tf.Tensor` of shape `(batch_size, num_heads, generated_length, sequence_length)`.183 decoder_hidden_states (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):184 Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of185 `tf.Tensor` of shape `(batch_size*num_return_sequences, generated_length, hidden_size)`.186 """187 188 sequences: Optional[tf.Tensor] = None189 scores: Optional[tuple[tf.Tensor]] = None190 encoder_attentions: Optional[tuple[tf.Tensor]] = None191 encoder_hidden_states: Optional[tuple[tf.Tensor]] = None192 decoder_attentions: Optional[tuple[tuple[tf.Tensor]]] = None193 cross_attentions: Optional[tuple[tuple[tf.Tensor]]] = None194 decoder_hidden_states: Optional[tuple[tuple[tf.Tensor]]] = None195 196 197@dataclass198class TFBeamSearchDecoderOnlyOutput(ModelOutput):199 """200 Base class for outputs of decoder-only generation models using beam search.201 202 Args:203 sequences (`tf.Tensor` of shape `(batch_size*num_return_sequences, sequence_length)`):204 The generated sequences. The second dimension (sequence_length) is either equal to `max_length` or shorter205 if all batches finished early due to the `eos_token_id`.206 sequences_scores (`tf.Tensor` of shape `(batch_size*num_return_sequences)`, *optional*, returned when `output_scores=True` is passed or when `config.output_scores=True`):207 Final beam scores of the generated `sequences`.208 scores (`tuple(tf.Tensor)` *optional*, returned when `output_scores=True` is passed or when `config.output_scores=True`):209 Processed beam scores for each vocabulary token at each generation step. Beam scores consisting of log210 softmax scores for each vocabulary token and sum of log softmax of previously generated tokens in this211 beam. Tuple of `tf.Tensor` with up to `max_new_tokens` elements (one element for each generated token),212 with each tensor of shape `(batch_size*num_beams*num_return_sequences, config.vocab_size)`.213 beam_indices (`tf.Tensor`, *optional*, returned when `output_scores=True` is passed or when `config.output_scores=True`):214 Beam indices of generated token id at each generation step. `tf.Tensor` of shape215 `(batch_size*num_return_sequences, sequence_length)`.216 attentions (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_attentions=True` is passed or `config.output_attentions=True`):217 Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of218 `tf.Tensor` of shape `(batch_size*num_beams, num_heads, generated_length, sequence_length)`.219 hidden_states (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):220 Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of221 `tf.Tensor` of shape `(batch_size*num_beams*num_return_sequences, generated_length, hidden_size)`.222 """223 224 sequences: Optional[tf.Tensor] = None225 sequences_scores: Optional[tf.Tensor] = None226 scores: Optional[tuple[tf.Tensor]] = None227 beam_indices: Optional[tf.Tensor] = None228 attentions: Optional[tuple[tuple[tf.Tensor]]] = None229 hidden_states: Optional[tuple[tuple[tf.Tensor]]] = None230 231 232@dataclass233class TFBeamSearchEncoderDecoderOutput(ModelOutput):234 """235 Base class for outputs of encoder-decoder generation models using beam search. Hidden states and attention weights236 of the decoder (respectively the encoder) can be accessed via the encoder_attentions and the encoder_hidden_states237 attributes (respectively the decoder_attentions and the decoder_hidden_states attributes)238 239 Args:240 sequences (`tf.Tensor` of shape `(batch_size*num_return_sequences, sequence_length)`):241 The generated sequences. The second dimension (sequence_length) is either equal to `max_length` or shorter242 if all batches finished early due to the `eos_token_id`.243 sequences_scores (`tf.Tensor` of shape `(batch_size*num_return_sequences)`, *optional*, returned when `output_scores=True` is passed or when `config.output_scores=True`):244 Final beam scores of the generated `sequences`.245 scores (`tuple(tf.Tensor)` *optional*, returned when `output_scores=True` is passed or when `config.output_scores=True`):246 Processed beam scores for each vocabulary token at each generation step. Beam scores consisting of log247 softmax scores for each vocabulary token and sum of log softmax of previously generated tokens in this248 beam. `Tuple of `tf.Tensor` with up to `max_new_tokens` elements (one element for each generated token),249 with each tensor of shape `(batch_size*num_beams, config.vocab_size)`.250 beam_indices (`tf.Tensor`, *optional*, returned when `output_scores=True` is passed or when `config.output_scores=True`):251 Beam indices of generated token id at each generation step. `tf.Tensor` of shape252 `(batch_size*num_return_sequences, sequence_length)`.253 encoder_attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or `config.output_attentions=True`):254 Tuple of `tf.Tensor` (one for each layer of the decoder) of shape `(batch_size, num_heads, sequence_length,255 sequence_length)`.256 encoder_hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):257 Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape258 `(batch_size*num_beams*num_return_sequences, sequence_length, hidden_size)`.259 decoder_attentions (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_attentions=True` is passed or `config.output_attentions=True`):260 Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of261 `tf.Tensor` of shape `(batch_size*num_beams*num_return_sequences, num_heads, generated_length,262 sequence_length)`.263 cross_attentions (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_attentions=True` is passed or `config.output_attentions=True`):264 Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of265 `tf.Tensor` of shape `(batch_size, num_heads, generated_length, sequence_length)`.266 decoder_hidden_states (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):267 Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of268 `tf.Tensor` of shape `(batch_size*num_beams*num_return_sequences, generated_length, hidden_size)`.269 """270 271 sequences: Optional[tf.Tensor] = None272 sequences_scores: Optional[tf.Tensor] = None273 scores: Optional[tuple[tf.Tensor]] = None274 beam_indices: Optional[tf.Tensor] = None275 encoder_attentions: Optional[tuple[tf.Tensor]] = None276 encoder_hidden_states: Optional[tuple[tf.Tensor]] = None277 decoder_attentions: Optional[tuple[tuple[tf.Tensor]]] = None278 cross_attentions: Optional[tuple[tuple[tf.Tensor]]] = None279 decoder_hidden_states: Optional[tuple[tuple[tf.Tensor]]] = None280 281 282@dataclass283class TFBeamSampleDecoderOnlyOutput(ModelOutput):284 """285 Base class for outputs of decoder-only generation models using beam sample.286 287 Args:288 sequences (`tf.Tensor` of shape `(batch_size*num_return_sequences, sequence_length)`):289 The generated sequences. The second dimension (sequence_length) is either equal to `max_length` or shorter290 if all batches finished early due to the `eos_token_id`.291 sequences_scores (`tf.Tensor` of shape `(batch_size * num_return_sequence)`, *optional*, returned when `output_scores=True` is passed or when `config.output_scores=True`):292 Final beam scores of the generated `sequences`.293 scores (`tuple(tf.Tensor)` *optional*, returned when `output_scores=True` is passed or when `config.output_scores=True`):294 Processed beam scores for each vocabulary token at each generation step. Beam scores consisting of log295 softmax scores for each vocabulary token and sum of log softmax of previously generated tokens in this296 beam. Tuple of `tf.Tensor` with up to `max_new_tokens` elements (one element for each generated token),297 with each tensor of shape `(batch_size*num_beams*num_return_sequences, config.vocab_size)`.298 beam_indices (`tf.Tensor`, *optional*, returned when `output_scores=True` is passed or when `config.output_scores=True`):299 Beam indices of generated token id at each generation step. `tf.Tensor` of shape300 `(batch_size*num_return_sequences, sequence_length)`.301 attentions (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_attentions=True` is passed or `config.output_attentions=True`):302 Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of303 `tf.Tensor` of shape `(batch_size*num_beams, num_heads, generated_length, sequence_length)`.304 hidden_states (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):305 Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of306 `tf.Tensor` of shape `(batch_size*num_beams, generated_length, hidden_size)`.307 """308 309 sequences: Optional[tf.Tensor] = None310 sequences_scores: Optional[tf.Tensor] = None311 scores: Optional[tuple[tf.Tensor]] = None312 beam_indices: Optional[tf.Tensor] = None313 attentions: Optional[tuple[tuple[tf.Tensor]]] = None314 hidden_states: Optional[tuple[tuple[tf.Tensor]]] = None315 316 317@dataclass318class TFBeamSampleEncoderDecoderOutput(ModelOutput):319 """320 Base class for outputs of encoder-decoder generation models using beam sampling. Hidden states and attention321 weights of the decoder (respectively the encoder) can be accessed via the encoder_attentions and the322 encoder_hidden_states attributes (respectively the decoder_attentions and the decoder_hidden_states attributes)323 324 Args:325 sequences (`tf.Tensor` of shape `(batch_size*num_beams, sequence_length)`):326 The generated sequences. The second dimension (sequence_length) is either equal to `max_length` or shorter327 if all batches finished early due to the `eos_token_id`.328 sequences_scores (`tf.Tensor` of shape `(batch_size * num_return_sequence)`, *optional*, returned when `output_scores=True` is passed or when `config.output_scores=True`):329 Final beam scores of the generated `sequences`.330 scores (`tuple(tf.Tensor)` *optional*, returned when `output_scores=True` is passed or when `config.output_scores=True`):331 Processed beam scores for each vocabulary token at each generation step. Beam scores consisting of log332 softmax scores for each vocabulary token and sum of log softmax of previously generated tokens in this333 beam. Tuple of `tf.Tensor` with up to `max_new_tokens` elements (one element for each generated token),334 with each tensor of shape `(batch_size*num_beams, config.vocab_size)`.335 beam_indices (`tf.Tensor`, *optional*, returned when `output_scores=True` is passed or when `config.output_scores=True`):336 Beam indices of generated token id at each generation step. `tf.Tensor` of shape337 `(batch_size*num_return_sequences, sequence_length)`.338 encoder_attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or `config.output_attentions=True`):339 Tuple of `tf.Tensor` (one for each layer of the decoder) of shape `(batch_size, num_heads, sequence_length,340 sequence_length)`.341 encoder_hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):342 Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape343 `(batch_size*num_beams, sequence_length, hidden_size)`.344 decoder_attentions (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_attentions=True` is passed or `config.output_attentions=True`):345 Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of346 `tf.Tensor` of shape `(batch_size*num_beams, num_heads, generated_length, sequence_length)`.347 cross_attentions (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_attentions=True` is passed or `config.output_attentions=True`):348 Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of349 `tf.Tensor` of shape `(batch_size, num_heads, generated_length, sequence_length)`.350 decoder_hidden_states (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):351 Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of352 `tf.Tensor` of shape `(batch_size*num_beams, generated_length, hidden_size)`.353 """354 355 sequences: Optional[tf.Tensor] = None356 sequences_scores: Optional[tf.Tensor] = None357 scores: Optional[tuple[tf.Tensor]] = None358 beam_indices: Optional[tf.Tensor] = None359 encoder_attentions: Optional[tuple[tf.Tensor]] = None360 encoder_hidden_states: Optional[tuple[tf.Tensor]] = None361 decoder_attentions: Optional[tuple[tuple[tf.Tensor]]] = None362 cross_attentions: Optional[tuple[tuple[tf.Tensor]]] = None363 decoder_hidden_states: Optional[tuple[tuple[tf.Tensor]]] = None364 365 366@dataclass367class TFContrastiveSearchDecoderOnlyOutput(ModelOutput):368 """369 Base class for outputs of decoder-only generation models using contrastive search.370 371 Args:372 sequences (`tf.Tensor` of shape `(batch_size, sequence_length)`):373 The generated sequences. The second dimension (sequence_length) is either equal to `max_length` or shorter374 if all batches finished early due to the `eos_token_id`.375 scores (`tuple(tf.Tensor)` *optional*, returned when `output_scores=True` is passed or when `config.output_scores=True`):376 Processed prediction scores of the language modeling head (scores for each vocabulary token before SoftMax)377 at each generation step. Tuple of `tf.Tensor` with up to `max_new_tokens` elements (one element for each378 generated token), with each tensor of shape `(batch_size, config.vocab_size)`.379 attentions (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_attentions=True` is passed or `config.output_attentions=True`):380 Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of381 `tf.Tensor` of shape `(batch_size, num_heads, generated_length, sequence_length)`.382 hidden_states (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):383 Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of384 `tf.Tensor` of shape `(batch_size, generated_length, hidden_size)`.385 """386 387 sequences: Optional[tf.Tensor] = None388 scores: Optional[tuple[tf.Tensor]] = None389 attentions: Optional[tuple[tuple[tf.Tensor]]] = None390 hidden_states: Optional[tuple[tuple[tf.Tensor]]] = None391 392 393@dataclass394class TFContrastiveSearchEncoderDecoderOutput(ModelOutput):395 """396 Base class for outputs of encoder-decoder generation models using contrastive search. Hidden states and attention397 weights of the decoder (respectively the encoder) can be accessed via the encoder_attentions and the398 encoder_hidden_states attributes (respectively the decoder_attentions and the decoder_hidden_states attributes)399 400 Args:401 sequences (`tf.Tensor` of shape `(batch_size, sequence_length)`):402 The generated sequences. The second dimension (sequence_length) is either equal to `max_length` or shorter403 if all batches finished early due to the `eos_token_id`.404 scores (`tuple(tf.Tensor)` *optional*, returned when `output_scores=True` is passed or when `config.output_scores=True`):405 Processed prediction scores of the language modeling head (scores for each vocabulary token before SoftMax)406 at each generation step. Tuple of `tf.Tensor` with up to `max_new_tokens` elements (one element for each407 generated token), with each tensor of shape `(batch_size, config.vocab_size)`.408 encoder_attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or `config.output_attentions=True`):409 Tuple of `tf.Tensor` (one for each layer of the decoder) of shape `(batch_size, num_heads, sequence_length,410 sequence_length)`.411 encoder_hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):412 Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape413 `(batch_size, sequence_length, hidden_size)`.414 decoder_attentions (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_attentions=True` is passed or `config.output_attentions=True`):415 Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of416 `tf.Tensor` of shape `(batch_size, num_heads, generated_length, sequence_length)`.417 cross_attentions (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_attentions=True` is passed or `config.output_attentions=True`):418 Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of419 `tf.Tensor` of shape `(batch_size, num_heads, generated_length, sequence_length)`.420 decoder_hidden_states (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):421 Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of422 `tf.Tensor` of shape `(batch_size, generated_length, hidden_size)`.423 """424 425 sequences: Optional[tf.Tensor] = None426 scores: Optional[tuple[tf.Tensor]] = None427 encoder_attentions: Optional[tuple[tf.Tensor]] = None428 encoder_hidden_states: Optional[tuple[tf.Tensor]] = None429 decoder_attentions: Optional[tuple[tuple[tf.Tensor]]] = None430 cross_attentions: Optional[tuple[tuple[tf.Tensor]]] = None431 decoder_hidden_states: Optional[tuple[tuple[tf.Tensor]]] = None432 433 434TFGreedySearchOutput = Union[TFGreedySearchEncoderDecoderOutput, TFGreedySearchDecoderOnlyOutput]435TFSampleOutput = Union[TFSampleEncoderDecoderOutput, TFSampleDecoderOnlyOutput]436TFBeamSearchOutput = Union[TFBeamSearchEncoderDecoderOutput, TFBeamSearchDecoderOnlyOutput]437TFBeamSampleOutput = Union[TFBeamSampleEncoderDecoderOutput, TFBeamSampleDecoderOnlyOutput]438TFContrastiveSearchOutput = Union[TFContrastiveSearchEncoderDecoderOutput, TFContrastiveSearchDecoderOnlyOutput]439TFGenerateOutput = Union[440 TFGreedySearchOutput, TFSampleOutput, TFBeamSearchOutput, TFBeamSampleOutput, TFContrastiveSearchOutput441]442 443 444class TFGenerationMixin:445 """446 A class containing all of the functions supporting generation, to be used as a mixin in [`TFPreTrainedModel`].447 448 The class exposes [`~generation.TFGenerationMixin.generate`], which can be used for:449 - *greedy decoding* by calling [`~generation.TFGenerationMixin.greedy_search`] if `num_beams=1` and450 `do_sample=False`451 - *contrastive search* by calling [`~generation.TFGenerationMixin.contrastive_search`] if `penalty_alpha>0` and452 `top_k>1`453 - *multinomial sampling* by calling [`~generation.TFGenerationMixin.sample`] if `num_beams=1` and454 `do_sample=True`455 - *beam-search decoding* by calling [`~generation.TFGenerationMixin.beam_search`] if `num_beams>1`456 457 You do not need to call any of the above methods directly. Pass custom parameter values to 'generate' instead. To458 learn more about decoding strategies refer to the [text generation strategies guide](../generation_strategies).459 """460 461 _seed_generator = None462 463 @property464 def seed_generator(self):465 warnings.warn("`seed_generator` is deprecated and will be removed in a future version.", UserWarning)466 if self._seed_generator is None:467 self._seed_generator = tf.random.Generator.from_non_deterministic_state()468 return self._seed_generator469 470 supports_xla_generation = True471 472 def prepare_inputs_for_generation(self, *args, **kwargs):473 raise NotImplementedError(474 "A model class needs to define a `prepare_inputs_for_generation` method in order to use `generate`."475 )476 477 def compute_transition_scores(478 self,479 sequences: tf.Tensor,480 scores: tuple[tf.Tensor],481 beam_indices: Optional[tf.Tensor] = None,482 normalize_logits: bool = False,483 ) -> tf.Tensor:484 """485 Computes the transition scores of sequences given the generation scores (and beam indices, if beam search was486 used). This is a convenient method to quickly obtain the scores of the selected tokens at generation time.487 488 Parameters:489 sequences (`tf.Tensor`):490 The generated sequences. The second dimension (sequence_length) is either equal to `max_length` or491 shorter if all batches finished early due to the `eos_token_id`.492 scores (`tuple(tf.Tensor)`):493 Transition scores for each vocabulary token at each generation step. Beam transition scores consisting494 of log probabilities of tokens conditioned on log softmax of previously generated tokens Tuple of495 `tf.Tensor` with up to `max_new_tokens` elements (one element for each generated token), with each496 tensor of shape `(batch_size*num_beams, config.vocab_size)`.497 beam_indices (`tf.Tensor`, *optional*):498 Beam indices of generated token id at each generation step. `tf.Tensor` of shape499 `(batch_size*num_return_sequences, sequence_length)`. Only required if a `num_beams>1` at500 generate-time.501 normalize_logits (`bool`, *optional*, defaults to `False`):502 Whether to normalize the logits (which, for legacy reasons, may be unnormalized).503 504 Return:505 `tf.Tensor`: A `tf.Tensor` of shape `(batch_size*num_return_sequences, sequence_length)` containing506 the transition scores (logits)507 508 Examples:509 510 ```python511 >>> from transformers import GPT2Tokenizer, TFAutoModelForCausalLM512 >>> import numpy as np513 514 >>> tokenizer = GPT2Tokenizer.from_pretrained("openai-community/gpt2")515 >>> model = TFAutoModelForCausalLM.from_pretrained("openai-community/gpt2")516 >>> tokenizer.pad_token_id = tokenizer.eos_token_id517 >>> inputs = tokenizer(["Today is"], return_tensors="tf")518 519 >>> # Example 1: Print the scores for each token generated with Greedy Search520 >>> outputs = model.generate(**inputs, max_new_tokens=5, return_dict_in_generate=True, output_scores=True)521 >>> transition_scores = model.compute_transition_scores(522 ... outputs.sequences, outputs.scores, normalize_logits=True523 ... )524 >>> # input_length is the length of the input prompt for decoder-only models, like the GPT family, and 1 for525 >>> # encoder-decoder models, like BART or T5.526 >>> input_length = 1 if model.config.is_encoder_decoder else inputs.input_ids.shape[1]527 >>> generated_tokens = outputs.sequences[:, input_length:]528 >>> for tok, score in zip(generated_tokens[0], transition_scores[0]):529 ... # | token | token string | logits | probability530 ... print(f"| {tok:5d} | {tokenizer.decode(tok):8s} | {score.numpy():.3f} | {np.exp(score.numpy()):.2%}")531 | 262 | the | -1.414 | 24.33%532 | 1110 | day | -2.609 | 7.36%533 | 618 | when | -2.010 | 13.40%534 | 356 | we | -1.859 | 15.58%535 | 460 | can | -2.508 | 8.14%536 537 >>> # Example 2: Reconstruct the sequence scores from Beam Search538 >>> outputs = model.generate(539 ... **inputs,540 ... max_new_tokens=5,541 ... num_beams=4,542 ... num_return_sequences=4,543 ... return_dict_in_generate=True,544 ... output_scores=True,545 ... )546 >>> transition_scores = model.compute_transition_scores(547 ... outputs.sequences, outputs.scores, outputs.beam_indices, normalize_logits=False548 ... )549 >>> # If you sum the generated tokens' scores and apply the length penalty, you'll get the sequence scores.550 >>> # Tip: recomputing the scores is only guaranteed to match with `normalize_logits=False`. Depending on the551 >>> # use case, you might want to recompute it with `normalize_logits=True`.552 >>> output_length = np.sum(transition_scores.numpy() < 0, axis=1)553 >>> length_penalty = model.generation_config.length_penalty554 >>> reconstructed_scores = np.sum(transition_scores, axis=1) / (output_length**length_penalty)555 >>> print(np.allclose(outputs.sequences_scores, reconstructed_scores))556 True557 ```"""558 # 1. In absence of `beam_indices`, we can assume that we come from e.g. greedy search, which is equivalent559 # to a beam search approach were the first (and only) beam is always selected560 if beam_indices is None:561 beam_indices = tf.tile(tf.expand_dims(tf.range(scores[0].shape[0]), axis=1), [1, len(scores)])562 563 # 2. reshape scores as [batch_size, vocab_size, # generation steps] with # generation steps being564 # seq_len - input_length565 scores = tf.transpose(tf.reshape(tf.stack(scores), (len(scores), -1)), (1, 0))566 scores = tf.reshape(scores, (-1, self.config.vocab_size, scores.shape[-1]))567 568 # 3. Optionally normalize the logits (across the vocab dimension)569 if normalize_logits:570 scores = tf.nn.log_softmax(scores, axis=1)571 572 # 4. cut beam_indices to longest beam length573 beam_indices_mask = beam_indices < 0574 max_beam_length = tf.math.reduce_max(575 tf.math.reduce_sum((1 - tf.cast(beam_indices_mask, dtype=tf.int32)), axis=-1)576 )577 beam_indices = beam_indices[:, -max_beam_length:]578 beam_indices_mask = beam_indices_mask[:, -max_beam_length:]579 580 # 5. Set indices of beams that finished early to 0; such indices will be masked correctly afterwards581 beam_indices = tf.where(beam_indices_mask, 0, beam_indices)582 583 # 6. Define which indices contributed to scores584 cut_idx = sequences.shape[-1] - max_beam_length585 token_indices = sequences[:, cut_idx:]586 gen_step_idx = tf.broadcast_to(tf.range(scores.shape[-1]), token_indices.shape)587 indices = tf.stack([beam_indices, token_indices, gen_step_idx], axis=-1)588 589 # 7. Compute scores590 transition_scores = tf.gather_nd(scores, indices)591 592 # 8. Mask out transition_scores of beams that stopped early593 transition_scores = tf.where(beam_indices_mask, 0, transition_scores)594 595 return transition_scores596 597 def _validate_model_class(self):598 """599 Confirms that the model class is compatible with generation. If not, raises an exception that points to the600 right class to use.601 """602 if not self.can_generate():603 generate_compatible_mappings = [604 TF_MODEL_FOR_CAUSAL_LM_MAPPING,605 TF_MODEL_FOR_VISION_2_SEQ_MAPPING,606 TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,607 TF_MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING,608 ]609 generate_compatible_classes = set()610 for model_mapping in generate_compatible_mappings:611 supported_models = model_mapping.get(type(self.config), default=None)612 if supported_models is not None:613 generate_compatible_classes.add(supported_models.__name__)614 exception_message = (615 f"The current model class ({self.__class__.__name__}) is not compatible with `.generate()`, as "616 "it doesn't have a language model head."617 )618 if generate_compatible_classes:619 exception_message += f" Please use one of the following classes instead: {generate_compatible_classes}"620 raise TypeError(exception_message)621 622 def _validate_model_kwargs(self, model_kwargs: dict[str, Any]):623 """Validates model kwargs for generation. Generate argument typos will also be caught here."""624 # Excludes arguments that are handled before calling any model function625 if self.config.is_encoder_decoder:626 for key in ["decoder_input_ids"]:627 model_kwargs.pop(key, None)628 629 unused_model_args = []630 model_args = set(inspect.signature(self.prepare_inputs_for_generation).parameters)631 # `kwargs`/`model_kwargs` is often used to handle optional forward pass inputs like `attention_mask`. If632 # `prepare_inputs_for_generation` doesn't accept them, then a stricter check can be made ;)633 if "kwargs" in model_args or "model_kwargs" in model_args:634 model_args |= set(inspect.signature(self.call).parameters)635 for key, value in model_kwargs.items():636 if value is not None and key not in model_args:637 unused_model_args.append(key)638 639 if unused_model_args:640 raise ValueError(641 f"The following `model_kwargs` are not used by the model: {unused_model_args} (note: typos in the"642 " generate arguments will also show up in this list)"643 )644 645 def generate(646 self,647 inputs: Optional[tf.Tensor] = None,648 generation_config: Optional[GenerationConfig] = None,649 logits_processor: Optional[TFLogitsProcessorList] = None,650 seed=None,651 **kwargs,652 ) -> Union[TFGenerateOutput, tf.Tensor]:653 r"""654 Generates sequences of token ids for models with a language modeling head.655 656 <Tip warning={true}>657 658 Most generation-controlling parameters are set in `generation_config` which, if not passed, will be set to the659 model's default generation configuration. You can override any `generation_config` by passing the corresponding660 parameters to generate, e.g. `.generate(inputs, num_beams=4, do_sample=True)`.661 662 For an overview of generation strategies and code examples, check out the [following663 guide](../generation_strategies).664 665 </Tip>666 667 Parameters:668 inputs (`tf.Tensor` of varying shape depending on the modality, *optional*):669 The sequence used as a prompt for the generation or as model inputs to the encoder. If `None` the670 method initializes it with `bos_token_id` and a batch size of 1. For decoder-only models `inputs`671 should of in the format of `input_ids`. For encoder-decoder models *inputs* can represent any of672 `input_ids`, `input_values`, `input_features`, or `pixel_values`.673 generation_config (`~generation.GenerationConfig`, *optional*):674 The generation configuration to be used as base parametrization for the generation call. `**kwargs`675 passed to generate matching the attributes of `generation_config` will override them. If676 `generation_config` is not provided, the default will be used, which had the following loading677 priority: 1) from the `generation_config.json` model file, if it exists; 2) from the model678 configuration. Please note that unspecified parameters will inherit [`~generation.GenerationConfig`]'s679 default values, whose documentation should be checked to parameterize generation.680 logits_processor (`LogitsProcessorList`, *optional*):681 Custom logits processors that complement the default logits processors built from arguments and682 generation config. If a logit processor is passed that is already created with the arguments or a683 generation config an error is thrown. This feature is intended for advanced users.684 seed (`list[int]`, *optional*):685 Random seed to control sampling, containing two integers, used when `do_sample` is `True`. See the686 `seed` argument from stateless functions in `tf.random`.687 kwargs (`dict[str, Any]`, *optional*):688 Ad hoc parametrization of `generate_config` and/or additional model-specific kwargs that will be689 forwarded to the `forward` function of the model. If the model is an encoder-decoder model, encoder690 specific kwargs should not be prefixed and decoder specific kwargs should be prefixed with *decoder_*.691 692 Return:693 [`~utils.ModelOutput`] or `tf.Tensor`: A [`~utils.ModelOutput`] (if `return_dict_in_generate=True` or when694 `config.return_dict_in_generate=True`) or a `tf.Tensor`.695 696 If the model is *not* an encoder-decoder model (`model.config.is_encoder_decoder=False`), the possible697 [`~utils.ModelOutput`] types are:698 699 - [`~generation.TFGreedySearchDecoderOnlyOutput`],700 - [`~generation.TFSampleDecoderOnlyOutput`],701 - [`~generation.TFBeamSearchDecoderOnlyOutput`],702 - [`~generation.TFBeamSampleDecoderOnlyOutput`]703 704 If the model is an encoder-decoder model (`model.config.is_encoder_decoder=True`), the possible705 [`~utils.ModelOutput`] types are:706 707 - [`~generation.TFGreedySearchEncoderDecoderOutput`],708 - [`~generation.TFSampleEncoderDecoderOutput`],709 - [`~generation.TFBeamSearchEncoderDecoderOutput`],710 - [`~generation.TFBeamSampleEncoderDecoderOutput`]711 712 """713 714 # 1. Handle `generation_config` and kwargs that might update it, and validate the `.generate()` call715 self._validate_model_class()716 717 # priority: `generation_config` argument > `model.generation_config` (the default generation config)718 if generation_config is None:719 # legacy: users may modify the model configuration to control generation. To trigger this legacy behavior,720 # two conditions must be met721 # 1) the generation config must have been created from the model config (`_from_model_config` field);722 # 2) the generation config must have seen no modification since its creation (the hash is the same).723 if self.generation_config._from_model_config and self.generation_config._original_object_hash == hash(724 self.generation_config725 ):726 new_generation_config = GenerationConfig.from_model_config(self.config)727 if new_generation_config != self.generation_config:728 warnings.warn(729 "You have modified the pretrained model configuration to control generation. This is a"730 " deprecated strategy to control generation and will be removed soon, in a future version."731 " Please use and modify the model generation configuration (see"732 " https://huggingface.co/docs/transformers/generation_strategies#default-text-generation-configuration )"733 )734 self.generation_config = new_generation_config735 generation_config = self.generation_config736 737 generation_config = copy.deepcopy(generation_config)738 model_kwargs = generation_config.update(**kwargs) # All unused kwargs must be model kwargs739 self._validate_model_kwargs(model_kwargs.copy())740 741 # 2. Cast input dtypes to tf.int32 unless they're floats (which happens for some image models)742 if inputs is not None:743 if isinstance(inputs, tf.Tensor) and inputs.dtype.is_floating:744 pass745 elif isinstance(inputs, np.ndarray) and np.issubdtype(inputs.dtype, np.floating):746 pass747 else:748 inputs = tf.cast(inputs, tf.int32)749 if model_kwargs.get("attention_mask") is not None:750 model_kwargs["attention_mask"] = tf.cast(model_kwargs["attention_mask"], tf.int32)751 if "decoder_input_ids" in model_kwargs:752 if (753 isinstance(model_kwargs["decoder_input_ids"], tf.Tensor)754 and model_kwargs["decoder_input_ids"].dtype.is_floating755 ):756 pass757 elif isinstance(model_kwargs["decoder_input_ids"], np.ndarray) and np.issubdtype(758 model_kwargs["decoder_input_ids"].dtype, np.floating759 ):760 pass761 else:762 model_kwargs["decoder_input_ids"] = tf.cast(model_kwargs["decoder_input_ids"], tf.int32)763 764 # 3. Set generation parameters if not already defined765 logits_processor = logits_processor if logits_processor is not None else TFLogitsProcessorList()766 767 if generation_config.pad_token_id is None and generation_config.eos_token_id is not None:768 if model_kwargs.get("attention_mask") is None:769 logger.warning(770 "The attention mask and the pad token id were not set. As a consequence, you may observe "771 "unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results."772 )773 eos_token_id = generation_config.eos_token_id774 if isinstance(eos_token_id, list):775 eos_token_id = eos_token_id[0]776 generation_config.pad_token_id = eos_token_id777 778 use_xla = not tf.executing_eagerly()779 if use_xla and not self.supports_xla_generation:780 raise ValueError(781 "The selected model does not support Graph mode nor XLA generation (e.g. from tf.function())"782 )783 784 # 4. Define model inputs785 inputs_tensor, model_input_name, model_kwargs = self._prepare_model_inputs(786 inputs, generation_config.bos_token_id, model_kwargs787 )788 # inputs_ids now has to be defined and cannot be None anymore789 batch_size = shape_list(inputs_tensor)[0]790 791 # 5. Prepare other model kwargs792 model_kwargs["output_attentions"] = generation_config.output_attentions793 model_kwargs["output_hidden_states"] = generation_config.output_hidden_states794 model_kwargs["use_cache"] = generation_config.use_cache795 796 accepts_attention_mask = "attention_mask" in set(inspect.signature(self.call).parameters.keys())797 requires_attention_mask = "encoder_outputs" not in model_kwargs798 799 if model_kwargs.get("attention_mask", None) is None and requires_attention_mask and accepts_attention_mask:800 model_kwargs["attention_mask"] = self._prepare_attention_mask_for_generation(801 inputs_tensor, generation_config.pad_token_id, generation_config.eos_token_id802 )803 804 # decoder-only models should use left-padding for generation805 if not self.config.is_encoder_decoder:806 if generation_config.pad_token_id is not None and tf.math.reduce_any(807 inputs_tensor[:, -1] == generation_config.pad_token_id808 ):809 logger.warning(810 "A decoder-only architecture is being used, but right-padding was detected! For correct "811 "generation results, please set `padding_side='left'` when initializing the tokenizer."812 )813 if self.config.is_encoder_decoder and "encoder_outputs" not in model_kwargs:814 # if model is encoder decoder encoder_outputs are created and added to `model_kwargs`815 model_kwargs = self._prepare_encoder_decoder_kwargs_for_generation(816 inputs_tensor, model_kwargs, model_input_name817 )818 819 # 6. Prepare model inputs which will be used for auto-regressive generation820 if self.config.is_encoder_decoder:821 input_ids, model_kwargs = self._prepare_decoder_input_ids_for_generation(822 batch_size=batch_size,823 model_input_name=model_input_name,824 model_kwargs=model_kwargs,825 decoder_start_token_id=generation_config.decoder_start_token_id,826 bos_token_id=generation_config.bos_token_id,827 )828 else:829 input_ids = inputs_tensor if model_input_name == "input_ids" else model_kwargs.pop("input_ids")830 831 # 7. Prepare `max_length` depending on other stopping criteria.832 input_ids_seq_length = shape_list(input_ids)[-1]833 has_default_max_length = kwargs.get("max_length") is None and generation_config.max_length is not None834 if has_default_max_length and generation_config.max_new_tokens is None and generation_config.max_length == 20:835 # 20 is the default max_length of the generation config836 warnings.warn(837 f"Using the model-agnostic default `max_length` (={generation_config.max_length}) "838 "to control the generation length. recommend setting `max_new_tokens` to control the maximum length of the generation.",839 UserWarning,840 )841 elif generation_config.max_new_tokens is not None:842 if not has_default_max_length and generation_config.max_length is not None:843 logger.warning(844 f"Both `max_new_tokens` (={generation_config.max_new_tokens}) and `max_length`(="845 f"{generation_config.max_length}) seem to have been set. `max_new_tokens` will take precedence. "846 "Please refer to the documentation for more information. "847 "(https://huggingface.co/docs/transformers/main/en/main_classes/text_generation)"848 )849 generation_config.max_length = generation_config.max_new_tokens + input_ids_seq_length850 851 # If the input length is a tensor (i.e. dynamic length), skip length checks852 if not isinstance(input_ids_seq_length, tf.Tensor):853 if (854 generation_config.min_length is not None855 and generation_config.min_length > generation_config.max_length856 ):857 raise ValueError(858 f"Unfeasable length constraints: the minimum length ({generation_config.min_length}) is larger"859 f" than the maximum length ({generation_config.max_length})"860 )861 if input_ids_seq_length >= generation_config.max_length:862 input_ids_string = "decoder_input_ids" if self.config.is_encoder_decoder else "input_ids"863 logger.warning(864 f"Input length of {input_ids_string} is {input_ids_seq_length}, but `max_length` is set to"865 f" {generation_config.max_length}. This can lead to unexpected behavior. You should consider"866 " increasing`max_new_tokens`."867 )868 869 # 8. determine generation mode870 is_contrastive_search_gen_mode = (871 generation_config.top_k is not None872 and generation_config.top_k > 1873 and generation_config.do_sample is False874 and generation_config.penalty_alpha is not None875 and generation_config.penalty_alpha > 0876 )877 is_greedy_gen_mode = (878 not is_contrastive_search_gen_mode879 and (generation_config.num_beams == 1)880 and generation_config.do_sample is False881 )882 is_beam_gen_mode = (883 not is_contrastive_search_gen_mode884 and (generation_config.num_beams > 1)885 and generation_config.do_sample is False886 )887 is_sample_gen_mode = (generation_config.num_beams == 1) and generation_config.do_sample is True888 is_beam_sample_gen_mode = (generation_config.num_beams > 1) and generation_config.do_sample is True889 890 # 9. prepare distribution pre_processing samplers891 logits_processor = self._get_logits_processor(892 generation_config=generation_config,893 input_ids_seq_length=input_ids_seq_length,894 logits_processor=logits_processor,895 )896 897 # 10. go into different generation modes898 if is_greedy_gen_mode:899 if generation_config.num_return_sequences > 1:900 raise ValueError(901 f"num_return_sequences has to be 1, but is {generation_config.num_return_sequences} when doing"902 " greedy search."903 )904 # 11. run greedy search905 return self.greedy_search(906 input_ids,907 max_length=generation_config.max_length,908 pad_token_id=generation_config.pad_token_id,909 eos_token_id=generation_config.eos_token_id,910 logits_processor=logits_processor,911 output_scores=generation_config.output_scores,912 return_dict_in_generate=generation_config.return_dict_in_generate,913 **model_kwargs,914 )915 elif is_contrastive_search_gen_mode:916 if generation_config.num_return_sequences > 1:917 raise ValueError(918 f"num_return_sequences has to be 1, but is {generation_config.num_return_sequences} when doing"919 " contrastive search."920 )921 # 11. run contrastive search922 return self.contrastive_search(923 input_ids,924 top_k=generation_config.top_k,925 penalty_alpha=generation_config.penalty_alpha,926 logits_processor=logits_processor,927 max_length=generation_config.max_length,928 pad_token_id=generation_config.pad_token_id,929 eos_token_id=generation_config.eos_token_id,930 output_scores=generation_config.output_scores,931 return_dict_in_generate=generation_config.return_dict_in_generate,932 **model_kwargs,933 )934 elif is_sample_gen_mode:935 # 11. prepare logits warper936 logits_warper = self._get_logits_warper(generation_config=generation_config)937 938 # 12. expand input_ids with `num_return_sequences` additional sequences per batch939 input_ids, model_kwargs = self._expand_inputs_for_generation(940 input_ids=input_ids,941 expand_size=generation_config.num_return_sequences,942 is_encoder_decoder=self.config.is_encoder_decoder,943 **model_kwargs,944 )945 946 # 13. run sample947 return self.sample(948 input_ids,949 logits_processor=logits_processor,950 logits_warper=logits_warper,951 max_length=generation_config.max_length,952 pad_token_id=generation_config.pad_token_id,953 eos_token_id=generation_config.eos_token_id,954 seed=seed,955 output_scores=generation_config.output_scores,956 return_dict_in_generate=generation_config.return_dict_in_generate,957 **model_kwargs,958 )959 960 elif is_beam_gen_mode:961 if generation_config.num_beams < generation_config.num_return_sequences:962 raise ValueError(963 "Beam search decoding cannot return more sequences than it has beams. Please set num_beams >="964 f" num_return_sequences, got {generation_config.num_beams} and"965 f" {generation_config.num_return_sequences} (respectively)"966 )967 968 # 11. broadcast inputs to the desired number of beams969 input_ids, model_kwargs = self._expand_inputs_for_generation(970 input_ids=input_ids,971 expand_size=generation_config.num_beams,972 is_encoder_decoder=self.config.is_encoder_decoder,973 expand_in_new_axis=True,974 **model_kwargs,975 )976 977 # 12. run beam search978 return self.beam_search(979 input_ids,980 max_length=generation_config.max_length,981 pad_token_id=generation_config.pad_token_id,982 eos_token_id=generation_config.eos_token_id,983 length_penalty=generation_config.length_penalty,984 early_stopping=generation_config.early_stopping,985 logits_processor=logits_processor,986 output_scores=generation_config.output_scores,987 return_dict_in_generate=generation_config.return_dict_in_generate,988 num_return_sequences=generation_config.num_return_sequences,989 **model_kwargs,990 )991 992 elif is_beam_sample_gen_mode:993 if generation_config.num_beams < generation_config.num_return_sequences:994 raise ValueError(995 "Beam search decoding cannot return more sequences than it has beams. Please set num_beams >="996 f" num_return_sequences, got {generation_config.num_beams} and"997 f" {generation_config.num_return_sequences} (respectively)"998 )999 1000 # 11. prepare logits warper1001 logits_warper = self._get_logits_warper(generation_config=generation_config)1002 1003 # 12. broadcast inputs to the desired number of beams1004 input_ids, model_kwargs = self._expand_inputs_for_generation(1005 input_ids=input_ids,1006 expand_size=generation_config.num_beams,1007 is_encoder_decoder=self.config.is_encoder_decoder,1008 expand_in_new_axis=True,1009 **model_kwargs,1010 )1011 1012 # 13. run beam sample (beam search with sampling)1013 return self.beam_search(1014 input_ids,1015 do_sample=True,1016 max_length=generation_config.max_length,1017 pad_token_id=generation_config.pad_token_id,1018 eos_token_id=generation_config.eos_token_id,1019 length_penalty=generation_config.length_penalty,1020 early_stopping=generation_config.early_stopping,1021 logits_processor=logits_processor,1022 logits_warper=logits_warper,1023 output_scores=generation_config.output_scores,1024 return_dict_in_generate=generation_config.return_dict_in_generate,1025 num_return_sequences=generation_config.num_return_sequences,1026 **model_kwargs,1027 )1028 1029 def _prepare_attention_mask_for_generation(1030 self,1031 inputs: tf.Tensor,1032 pad_token_id: Optional[int],1033 eos_token_id: Optional[int],1034 ) -> tf.Tensor:1035 is_input_ids = len(inputs.shape) == 2 and inputs.dtype in (tf.int32, tf.int64)1036 is_pad_token_in_inputs = (pad_token_id is not None) and tf.math.reduce_any(inputs == pad_token_id)1037 is_pad_token_not_equal_to_eos_token_id = (eos_token_id is None) or (pad_token_id != eos_token_id)1038 1039 # Check if input is input_ids and padded -> only then is attention_mask defined1040 if is_input_ids and is_pad_token_in_inputs and is_pad_token_not_equal_to_eos_token_id:1041 return tf.cast(tf.math.not_equal(inputs, pad_token_id), dtype=tf.int32)1042 else:1043 return tf.ones(inputs.shape[:2], dtype=tf.int32)1044 1045 def _prepare_encoder_decoder_kwargs_for_generation(1046 self, inputs_tensor: tf.Tensor, model_kwargs, model_input_name: Optional[str] = None1047 ) -> dict[str, Any]:1048 # 1. get encoder and store encoder outputs1049 encoder = self.get_encoder()1050 1051 # 2. prepare encoder args and encoder kwargs from model kwargs1052 irrelevant_prefix = ["decoder_", "cross_attn", "use_cache"]1053 encoder_kwargs = {1054 argument: value1055 for argument, value in model_kwargs.items()1056 if not any(argument.startswith(p) for p in irrelevant_prefix)1057 }1058 encoder_signature = set(inspect.signature(encoder.call).parameters)1059 encoder_accepts_wildcard = "kwargs" in encoder_signature or "model_kwargs" in encoder_signature1060 if not encoder_accepts_wildcard:1061 encoder_kwargs = {1062 argument: value for argument, value in encoder_kwargs.items() if argument in encoder_signature1063 }1064 1065 # 3. vision models don't use `attention_mask`.1066 encoder_kwargs["return_dict"] = True1067 encoder_kwargs[model_input_name] = inputs_tensor1068 if model_input_name != self.main_input_name: # in Keras, the first input must always be passed1069 encoder_kwargs[self.main_input_name] = None1070 encoder_outputs = encoder(**encoder_kwargs)1071 model_kwargs["encoder_outputs"] = encoder_outputs1072 1073 return model_kwargs1074 1075 def _prepare_decoder_input_ids_for_generation(1076 self,1077 batch_size: int,1078 model_input_name: str,1079 model_kwargs: dict[str, tf.Tensor],1080 decoder_start_token_id: Optional[int] = None,1081 bos_token_id: Optional[int] = None,1082 ) -> tuple[tf.Tensor, dict[str, tf.Tensor]]:1083 """Prepares `decoder_input_ids` for generation with encoder-decoder models"""1084 # 1. Check whether the user has defined `decoder_input_ids` manually. To facilitate in terms of input naming,1085 # we also allow the user to pass it under `input_ids`, if the encoder does not use it as the main input.1086 if model_kwargs is not None and "decoder_input_ids" in model_kwargs:1087 decoder_input_ids = model_kwargs.pop("decoder_input_ids")1088 elif "input_ids" in model_kwargs and model_input_name != "input_ids":1089 decoder_input_ids = model_kwargs.pop("input_ids")1090 else:1091 decoder_input_ids = None1092 1093 # 2. Encoder-decoder models expect the `decoder_input_ids` to start with a special token. Let's ensure that.1094 decoder_start_token_id = self._get_decoder_start_token_id(decoder_start_token_id, bos_token_id)1095 decoder_input_ids_start = tf.ones((batch_size, 1), dtype=tf.int32) * decoder_start_token_id1096 1097 # no user input -> use decoder_start_token_id as decoder_input_ids1098 if decoder_input_ids is None:1099 decoder_input_ids = decoder_input_ids_start1100 # user input but doesn't start with decoder_start_token_id -> prepend decoder_start_token_id (and adjust1101 # decoder_attention_mask if provided)1102 elif tf.reduce_all(decoder_input_ids[:, 0] != decoder_start_token_id):1103 decoder_input_ids = tf.concat([decoder_input_ids_start, decoder_input_ids], axis=-1)1104 if "decoder_attention_mask" in model_kwargs:1105 decoder_attention_mask = model_kwargs["decoder_attention_mask"]1106 decoder_attention_mask = tf.concat(1107 (tf.ones_like(decoder_attention_mask)[:, :1], decoder_attention_mask),1108 axis=-1,1109 )1110 model_kwargs["decoder_attention_mask"] = decoder_attention_mask1111 1112 return decoder_input_ids, model_kwargs1113 1114 def _get_decoder_start_token_id(1115 self, decoder_start_token_id: Optional[int] = None, bos_token_id: Optional[int] = None1116 ) -> int:1117 # retrieve decoder_start_token_id for encoder-decoder models1118 # fall back to bos_token_id if necessary1119 decoder_start_token_id = (1120 decoder_start_token_id1121 if decoder_start_token_id is not None1122 else self.generation_config.decoder_start_token_id1123 )1124 bos_token_id = bos_token_id if bos_token_id is not None else self.generation_config.bos_token_id1125 1126 if decoder_start_token_id is not None:1127 return decoder_start_token_id1128 elif bos_token_id is not None:1129 return bos_token_id1130 raise ValueError(1131 "`decoder_start_token_id` or `bos_token_id` has to be defined for encoder-decoder generation."1132 )1133 1134 @staticmethod1135 def _expand_inputs_for_generation(1136 expand_size: int = 1,1137 is_encoder_decoder: bool = False,1138 input_ids: Optional[tf.Tensor] = None,1139 expand_in_new_axis: bool = False,1140 **model_kwargs,1141 ) -> tuple[tf.Tensor, dict[str, Any]]:1142 """1143 Expands tensors from [batch_size, ...] to [batch_size * expand_size, ...] or [batch_size, expand_size, ...],1144 depending on `expand_in_new_axis`. Beam-based approaches expect this function to be used with1145 `expand_in_new_axis=True`1146 """1147 1148 def _expand_tensor(tensor: tf.Tensor):1149 if expand_in_new_axis:1150 shape = shape_list(tensor)1151 return tf.broadcast_to(tensor[:, None], (shape[0], expand_size) + tuple(shape[1:]))1152 else:1153 return tf.repeat(tensor, expand_size, axis=0)1154 1155 def _expand_dict_for_generation(dict_to_expand):1156 for key in dict_to_expand:1157 if dict_to_expand[key] is not None and isinstance(dict_to_expand[key], tf.Tensor):1158 dict_to_expand[key] = _expand_tensor(dict_to_expand[key])1159 return dict_to_expand1160 1161 if input_ids is not None:1162 input_ids = _expand_tensor(input_ids)1163 1164 model_kwargs = _expand_dict_for_generation(model_kwargs)1165 1166 if is_encoder_decoder:1167 if model_kwargs.get("encoder_outputs") is None:1168 raise ValueError("If `is_encoder_decoder` is True, make sure that `encoder_outputs` is defined.")1169 model_kwargs["encoder_outputs"] = _expand_dict_for_generation(model_kwargs["encoder_outputs"])1170 1171 return input_ids, model_kwargs1172 1173 def _prepare_model_inputs(1174 self,1175 inputs: Optional[tf.Tensor] = None,1176 bos_token_id: Optional[int] = None,1177 model_kwargs: Optional[dict[str, tf.Tensor]] = None,1178 ) -> tuple[tf.Tensor, Optional[str], dict[str, tf.Tensor]]:1179 """1180 This function extracts the model-specific `inputs` for generation.1181 """1182 # 1. retrieve all kwargs that are non-None or non-model input related.1183 # some encoder-decoder models have different names for model and encoder1184 if (1185 self.config.is_encoder_decoder1186 and hasattr(self, "encoder")1187 and hasattr(self.encoder, "main_input_name")1188 and self.encoder.main_input_name != self.main_input_name1189 ):1190 input_name = self.encoder.main_input_name1191 else:1192 input_name = self.main_input_name1193 1194 model_kwargs = {k: v for k, v in model_kwargs.items() if v is not None or k != input_name}1195 1196 # 2. check whether model_input_name is passed as kwarg1197 # if yes and `inputs` is None use kwarg inputs1198 inputs_kwarg = model_kwargs.pop(input_name, None)1199 if inputs_kwarg is not None and inputs is not None:1200 raise ValueError(