




from collections import OrderedDict
from typing import Any, Mapping, Optional

from ... import PreTrainedTokenizer
from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig, OnnxConfigWithPast, OnnxSeqaSeqConfigWithPast
from ...onnx.utils import compute_effective_axis_dimension
from ...utils import TensorType, is_torch_available, logging


__UpperCamelCase   :			Union[str, Any]          = logging.get_logger(__name__)

__UpperCamelCase   :			Optional[Any]          = {
    'Helsinki-NLP/opus-mt-en-de': 'https://huggingface.co/Helsinki-NLP/opus-mt-en-de/resolve/main/config.json',
    # See all Marian models at https://huggingface.co/models?filter=marian
}

class 						__lowerCAmelCase      (      _lowerCamelCase					):
							UpperCamelCase__   	=    '''marian'''
							UpperCamelCase__   	=    ['''past_key_values''']
							UpperCamelCase__   	=    {'''num_attention_heads''': '''encoder_attention_heads''', '''hidden_size''': '''d_model'''}
							def __init__(		self		:List[str]						, __magic_name__		:List[Any]=5_8101						, __magic_name__		:str=None						, __magic_name__		:str=1024						, __magic_name__		:Union[str, Any]=12						, __magic_name__		:Dict=4096						, __magic_name__		:Union[str, Any]=16						, __magic_name__		:int=12						, __magic_name__		:Any=4096						, __magic_name__		:Dict=16						, __magic_name__		:List[str]=0.0						, __magic_name__		:List[Any]=0.0						, __magic_name__		:Optional[Any]=True						, __magic_name__		:Optional[Any]=True						, __magic_name__		:List[str]="gelu"						, __magic_name__		:Tuple=1024						, __magic_name__		:Dict=0.1						, __magic_name__		:Any=0.0						, __magic_name__		:List[str]=0.0						, __magic_name__		:List[Any]=0.02						, __magic_name__		:Optional[int]=5_8100						, __magic_name__		:Dict=False						, __magic_name__		:str=5_8100						, __magic_name__		:Tuple=0						, __magic_name__		:Tuple=0						, __magic_name__		:Optional[Any]=True						, **__magic_name__		:int						, ):







								'''simple docstring'''







								a									=				vocab_size
								a									=				decoder_vocab_size or vocab_size
								a									=				max_position_embeddings
								a									=				d_model
								a									=				encoder_ffn_dim
								a									=				encoder_layers
								a									=				encoder_attention_heads
								a									=				decoder_ffn_dim
								a									=				decoder_layers
								a									=				decoder_attention_heads
								a									=				dropout
								a									=				attention_dropout
								a									=				activation_dropout
								a									=				activation_function
								a									=				init_std
								a									=				encoder_layerdrop
								a									=				decoder_layerdrop
								a									=				use_cache
								a									=				encoder_layers
								a									=				scale_embedding  # scale factor will be sqrt(d_model) if True
								a									=				share_encoder_decoder_embeddings
								super().__init__(
								    pad_token_id=_SCREAMING_SNAKE_CASE						, eos_token_id=_SCREAMING_SNAKE_CASE						, is_encoder_decoder=_SCREAMING_SNAKE_CASE						, decoder_start_token_id=_SCREAMING_SNAKE_CASE						, forced_eos_token_id=_SCREAMING_SNAKE_CASE						, **_SCREAMING_SNAKE_CASE						, )







class 						__lowerCAmelCase      (      _lowerCamelCase					):
							@property
							# Copied from transformers.models.bart.configuration_bart.BartOnnxConfig.inputs
							def  lowerCamelCase__				(		self		:List[Any]							):







								'''simple docstring'''







								if self.task in ["default", "seq2seq-lm"]:
									a									=				OrderedDict(
									    [
									        ("""input_ids""", {0: """batch""", 1: """encoder_sequence"""}),
									        ("""attention_mask""", {0: """batch""", 1: """encoder_sequence"""}),
									    ]							)

									if self.use_past:
										a									=				{0: 'batch'}
										a									=				{0: 'batch', 1: 'past_decoder_sequence + sequence'}
									else:
										a									=				{0: 'batch', 1: 'decoder_sequence'}
										a									=				{0: 'batch', 1: 'decoder_sequence'}

									if self.use_past:
										self.fill_with_past_key_values_(_SCREAMING_SNAKE_CASE						, direction="""inputs"""							)
								elif self.task == "causal-lm":
									# TODO: figure this case out.
									a									=				OrderedDict(
									    [
									        ("""input_ids""", {0: """batch""", 1: """encoder_sequence"""}),
									        ("""attention_mask""", {0: """batch""", 1: """encoder_sequence"""}),
									    ]							)
									if self.use_past:
										a									=				self.num_layers
										for i in range(_SCREAMING_SNAKE_CASE							):
											a									=				{0: 'batch', 2: 'past_sequence + sequence'}
											a									=				{0: 'batch', 2: 'past_sequence + sequence'}
								else:
									a									=				OrderedDict(
									    [
									        ("""input_ids""", {0: """batch""", 1: """encoder_sequence"""}),
									        ("""attention_mask""", {0: """batch""", 1: """encoder_sequence"""}),
									        ("""decoder_input_ids""", {0: """batch""", 1: """decoder_sequence"""}),
									        ("""decoder_attention_mask""", {0: """batch""", 1: """decoder_sequence"""}),
									    ]							)

								return common_inputs
							@property
							# Copied from transformers.models.bart.configuration_bart.BartOnnxConfig.outputs
							def  lowerCamelCase__				(		self		:Dict							):







								'''simple docstring'''







								if self.task in ["default", "seq2seq-lm"]:
									a									=				super().outputs
								else:
									a									=				super(_SCREAMING_SNAKE_CASE						, self							).outputs
									if self.use_past:
										a									=				self.num_layers
										for i in range(_SCREAMING_SNAKE_CASE							):
											a									=				{0: 'batch', 2: 'past_sequence + sequence'}
											a									=				{0: 'batch', 2: 'past_sequence + sequence'}
								return common_outputs
							def  lowerCamelCase__				(		self		:Tuple						, __magic_name__		:PreTrainedTokenizer						, __magic_name__		:int = -1						, __magic_name__		:int = -1						, __magic_name__		:bool = False						, __magic_name__		:Optional[TensorType] = None						, ):







								'''simple docstring'''







								a									=				self._generate_dummy_inputs_for_encoder_and_decoder(
								    _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE							)

								# Generate decoder inputs
								a									=				seq_length if not self.use_past else 1
								a									=				self._generate_dummy_inputs_for_encoder_and_decoder(
								    _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE							)
								a									=				{F'decoder_{name}': tensor for name, tensor in decoder_inputs.items()}
								a									=				dict(**_SCREAMING_SNAKE_CASE						, **_SCREAMING_SNAKE_CASE							)

								if self.use_past:
									if not is_torch_available():
										raise ValueError("""Cannot generate dummy past_keys inputs without PyTorch installed."""							)
									else:
										import torch
									a									=				common_inputs['input_ids'].shape
									a									=				common_inputs['decoder_input_ids'].shape[1]
									a									=				self.num_attention_heads
									a									=				(
									    batch,
									    num_encoder_attention_heads,
									    encoder_seq_length,
									    self._config.hidden_size // num_encoder_attention_heads,
									)
									a									=				decoder_seq_length + 3
									a									=				(
									    batch,
									    num_decoder_attention_heads,
									    decoder_past_length,
									    self._config.hidden_size // num_decoder_attention_heads,
									)

									a									=				torch.cat(
									    [common_inputs["""decoder_attention_mask"""], torch.ones(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE							)]						, dim=1							)

									a									=				[]
									# If the number of encoder and decoder layers are present in the model configuration, both are considered
									a									=				self.num_layers
									a									=				min(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE							)
									a									=				max(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE							) - min_num_layers
									a									=				'encoder' if num_encoder_layers > num_decoder_layers else 'decoder'

									for _ in range(_SCREAMING_SNAKE_CASE							):
										common_inputs["past_key_values"].append(
										    (
										        torch.zeros(_SCREAMING_SNAKE_CASE							),
										        torch.zeros(_SCREAMING_SNAKE_CASE							),
										        torch.zeros(_SCREAMING_SNAKE_CASE							),
										        torch.zeros(_SCREAMING_SNAKE_CASE							),
										    )							)
									# TODO: test this.
									a									=				encoder_shape if remaining_side_name == 'encoder' else decoder_shape
									for _ in range(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE							):
										common_inputs["past_key_values"].append((torch.zeros(_SCREAMING_SNAKE_CASE							), torch.zeros(_SCREAMING_SNAKE_CASE							))							)
								return common_inputs
							def  lowerCamelCase__				(		self		:int						, __magic_name__		:PreTrainedTokenizer						, __magic_name__		:int = -1						, __magic_name__		:int = -1						, __magic_name__		:bool = False						, __magic_name__		:Optional[TensorType] = None						, ):







								'''simple docstring'''







								a									=				self._generate_dummy_inputs_for_encoder_and_decoder(
								    _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE							)

								if self.use_past:
									if not is_torch_available():
										raise ValueError("""Cannot generate dummy past_keys inputs without PyTorch installed."""							)
									else:
										import torch
									a									=				common_inputs['input_ids'].shape
									# Not using the same length for past_key_values
									a									=				seqlen + 2
									a									=				self.num_layers
									a									=				self.num_attention_heads
									a									=				(
									    batch,
									    num_encoder_attention_heads,
									    past_key_values_length,
									    self._config.hidden_size // num_encoder_attention_heads,
									)

									a									=				common_inputs['attention_mask'].dtype
									a									=				torch.cat(
									    [common_inputs["""attention_mask"""], torch.ones(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, dtype=_SCREAMING_SNAKE_CASE							)]						, dim=1							)
									a									=				[
									    (torch.zeros(_SCREAMING_SNAKE_CASE							), torch.zeros(_SCREAMING_SNAKE_CASE							)) for _ in range(_SCREAMING_SNAKE_CASE							)
									]
								return common_inputs
							def  lowerCamelCase__				(		self		:Optional[Any]						, __magic_name__		:PreTrainedTokenizer						, __magic_name__		:int = -1						, __magic_name__		:int = -1						, __magic_name__		:bool = False						, __magic_name__		:Optional[TensorType] = None						, ):







								'''simple docstring'''







								a									=				compute_effective_axis_dimension(
								    _SCREAMING_SNAKE_CASE						, fixed_dimension=OnnxConfig.default_fixed_batch						, num_token_to_add=0							)

								# If dynamic axis (-1) we forward with a fixed dimension of 8 tokens to avoid optimizations made by ONNX
								a									=				tokenizer.num_special_tokens_to_add(_SCREAMING_SNAKE_CASE							)
								a									=				compute_effective_axis_dimension(
								    _SCREAMING_SNAKE_CASE						, fixed_dimension=OnnxConfig.default_fixed_sequence						, num_token_to_add=_SCREAMING_SNAKE_CASE							)

								# Generate dummy inputs according to compute batch and sequence
								a									=				[' '.join([tokenizer.unk_token]							) * seq_length] * batch_size
								a									=				dict(tokenizer(_SCREAMING_SNAKE_CASE						, return_tensors=_SCREAMING_SNAKE_CASE							)							)
								return common_inputs
							def  lowerCamelCase__				(		self		:Union[str, Any]						, __magic_name__		:PreTrainedTokenizer						, __magic_name__		:int = -1						, __magic_name__		:int = -1						, __magic_name__		:bool = False						, __magic_name__		:Optional[TensorType] = None						, ):







								'''simple docstring'''







								if self.task in ["default", "seq2seq-lm"]:
									a									=				self._generate_dummy_inputs_for_default_and_seqaseq_lm(
									    _SCREAMING_SNAKE_CASE						, batch_size=_SCREAMING_SNAKE_CASE						, seq_length=_SCREAMING_SNAKE_CASE						, is_pair=_SCREAMING_SNAKE_CASE						, framework=_SCREAMING_SNAKE_CASE							)

								else:
									a									=				self._generate_dummy_inputs_for_causal_lm(
									    _SCREAMING_SNAKE_CASE						, batch_size=_SCREAMING_SNAKE_CASE						, seq_length=_SCREAMING_SNAKE_CASE						, is_pair=_SCREAMING_SNAKE_CASE						, framework=_SCREAMING_SNAKE_CASE							)

								return common_inputs
							def  lowerCamelCase__				(		self		:List[Any]						, __magic_name__		:Optional[int]						, __magic_name__		:List[Any]						, __magic_name__		:Optional[Any]						, __magic_name__		:int							):







								'''simple docstring'''







								if self.task in ["default", "seq2seq-lm"]:
									a									=				super()._flatten_past_key_values_(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE							)
								else:
									a									=				super(_SCREAMING_SNAKE_CASE						, self							)._flatten_past_key_values_(
									    _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE							)





							@property
							def  lowerCamelCase__				(		self		:List[Any]							):







								'''simple docstring'''







								return 1E-4






"""simple docstring"""

from typing import Optional

import numpy as np
import torch
from torch import nn
from transformers import GPTaConfig, GPTaLMHeadModel
from transformers.modeling_utils import ModuleUtilsMixin

from ...configuration_utils import ConfigMixin, register_to_config
from ...models import ModelMixin
class     lowerCAmelCase__						(	_lowerCamelCase							,							_lowerCamelCase							,							_lowerCamelCase ):


	'''simple docstring'''







	_SCREAMING_SNAKE_CASE	:		Dict  						= [r'''h\.\d+\.attn\.bias''', r'''h\.\d+\.attn\.masked_bias''']





	@register_to_config
	def __init__(      self				:       Any     ,       _SCREAMING_SNAKE_CASE				:       int     ,       _SCREAMING_SNAKE_CASE				:       int     ,       _SCREAMING_SNAKE_CASE				:       Optional[int] = None     ,       _SCREAMING_SNAKE_CASE				:       int = 50_257     ,       _SCREAMING_SNAKE_CASE				:       int = 1_024     ,       _SCREAMING_SNAKE_CASE				:       int = 768     ,       _SCREAMING_SNAKE_CASE				:       int = 12     ,       _SCREAMING_SNAKE_CASE				:       int = 12     ,       _SCREAMING_SNAKE_CASE				:       Optional[int] = None     ,       _SCREAMING_SNAKE_CASE				:       str = "gelu_new"     ,       _SCREAMING_SNAKE_CASE				:       float = 0.1     ,       _SCREAMING_SNAKE_CASE				:       float = 0.1     ,       _SCREAMING_SNAKE_CASE				:       float = 0.1     ,       _SCREAMING_SNAKE_CASE				:       float = 1E-5     ,       _SCREAMING_SNAKE_CASE				:       float = 0.0_2     ,       _SCREAMING_SNAKE_CASE				:       bool = True     ,       _SCREAMING_SNAKE_CASE				:       bool = True     ,       _SCREAMING_SNAKE_CASE				:       bool = False     ,       _SCREAMING_SNAKE_CASE				:       bool = False     ,       )						->	Tuple:





					"""simple docstring"""


					super().__init__()

					SCREAMING_SNAKE_CASE						:							Optional[int]							=     prefix_length

					if prefix_inner_dim != n_embd and prefix_hidden_dim is None:
									raise ValueError(
									    f"""`prefix_hidden_dim` cannot be `None` when `prefix_inner_dim`: {prefix_hidden_dim} and"""
									    f""" `n_embd`: {n_embd} are not equal."""    )

					SCREAMING_SNAKE_CASE						:							int							=     prefix_inner_dim
					SCREAMING_SNAKE_CASE						:							Dict							=     prefix_hidden_dim

					SCREAMING_SNAKE_CASE						:							int							=     (
					    nn.Linear(self.prefix_inner_dim     ,       self.prefix_hidden_dim    )
					    if self.prefix_hidden_dim is not None
					    else nn.Identity()
					)
					SCREAMING_SNAKE_CASE						:							List[Any]							=     (
					    nn.Linear(self.prefix_hidden_dim     ,       _SCREAMING_SNAKE_CASE    ) if self.prefix_hidden_dim is not None else nn.Identity()
					)

					SCREAMING_SNAKE_CASE						:							Optional[Any]							=     GPTaConfig(
					    vocab_size=_SCREAMING_SNAKE_CASE     ,       n_positions=_SCREAMING_SNAKE_CASE     ,       n_embd=_SCREAMING_SNAKE_CASE     ,       n_layer=_SCREAMING_SNAKE_CASE     ,       n_head=_SCREAMING_SNAKE_CASE     ,       n_inner=_SCREAMING_SNAKE_CASE     ,       activation_function=_SCREAMING_SNAKE_CASE     ,       resid_pdrop=_SCREAMING_SNAKE_CASE     ,       embd_pdrop=_SCREAMING_SNAKE_CASE     ,       attn_pdrop=_SCREAMING_SNAKE_CASE     ,       layer_norm_epsilon=_SCREAMING_SNAKE_CASE     ,       initializer_range=_SCREAMING_SNAKE_CASE     ,       scale_attn_weights=_SCREAMING_SNAKE_CASE     ,       use_cache=_SCREAMING_SNAKE_CASE     ,       scale_attn_by_inverse_layer_idx=_SCREAMING_SNAKE_CASE     ,       reorder_and_upcast_attn=_SCREAMING_SNAKE_CASE     ,       )
					SCREAMING_SNAKE_CASE						:							Union[str, Any]							=     GPTaLMHeadModel(_SCREAMING_SNAKE_CASE    )





	def 			_lowerCAmelCase     (      self				:       Any     ,       _SCREAMING_SNAKE_CASE				:       torch.Tensor     ,       _SCREAMING_SNAKE_CASE				:       torch.Tensor     ,       _SCREAMING_SNAKE_CASE				:       Optional[torch.Tensor] = None     ,       _SCREAMING_SNAKE_CASE				:       Optional[torch.Tensor] = None     ,       )						->	Union[str, Any]:





					"""simple docstring"""


					SCREAMING_SNAKE_CASE						:							Tuple							=     self.transformer.transformer.wte(_SCREAMING_SNAKE_CASE    )
					SCREAMING_SNAKE_CASE						:							str							=     self.encode_prefix(_SCREAMING_SNAKE_CASE    )
					SCREAMING_SNAKE_CASE						:							Dict							=     self.decode_prefix(_SCREAMING_SNAKE_CASE    )
					SCREAMING_SNAKE_CASE						:							Optional[int]							=     torch.cat((prefix_embeds, embedding_text)     ,       dim=1    )

					if labels is not None:
									SCREAMING_SNAKE_CASE						:							Dict							=     self.get_dummy_token(input_ids.shape[0]     ,       input_ids.device    )
									SCREAMING_SNAKE_CASE						:							Union[str, Any]							=     torch.cat((dummy_token, input_ids)     ,       dim=1    )
					SCREAMING_SNAKE_CASE						:							List[Any]							=     self.transformer(inputs_embeds=_SCREAMING_SNAKE_CASE     ,       labels=_SCREAMING_SNAKE_CASE     ,       attention_mask=_SCREAMING_SNAKE_CASE    )
					if self.prefix_hidden_dim is not None:
									return out, hidden
					else:
									return out





	def 			_lowerCAmelCase     (      self				:       List[Any]     ,       _SCREAMING_SNAKE_CASE				:       int     ,       _SCREAMING_SNAKE_CASE				:       torch.device    )						->	torch.Tensor:





					"""simple docstring"""


					return torch.zeros(_SCREAMING_SNAKE_CASE     ,       self.prefix_length     ,       dtype=torch.intaa     ,       device=_SCREAMING_SNAKE_CASE    )





	def 			_lowerCAmelCase     (      self				:       Optional[int]     ,       _SCREAMING_SNAKE_CASE				:       Union[str, Any]    )						->	str:





					"""simple docstring"""


					return self.encode_prefix(_SCREAMING_SNAKE_CASE    )





	@torch.no_grad()
	def 			_lowerCAmelCase     (      self				:       Dict     ,       _SCREAMING_SNAKE_CASE				:       Dict     ,       _SCREAMING_SNAKE_CASE				:       Union[str, Any]     ,       _SCREAMING_SNAKE_CASE				:       List[str]    )						->	Dict:





					"""simple docstring"""


					SCREAMING_SNAKE_CASE						:							Optional[int]							=     torch.split(_SCREAMING_SNAKE_CASE     ,       1     ,       dim=0    )
					SCREAMING_SNAKE_CASE						:							int							=     []
					SCREAMING_SNAKE_CASE						:							str							=     []
					for feature in features:
									SCREAMING_SNAKE_CASE						:							Optional[Any]							=     self.decode_prefix(feature.to(_SCREAMING_SNAKE_CASE    )    )  # back to the clip feature
									# Only support beam search for now
									SCREAMING_SNAKE_CASE				,						SCREAMING_SNAKE_CASE						:							List[str]							=     self.generate_beam(
									    input_embeds=_SCREAMING_SNAKE_CASE     ,       device=_SCREAMING_SNAKE_CASE     ,       eos_token_id=_SCREAMING_SNAKE_CASE    )
									generated_tokens.append(output_tokens[0]    )
									generated_seq_lengths.append(seq_lengths[0]    )
					SCREAMING_SNAKE_CASE						:							List[str]							=     torch.stack(_SCREAMING_SNAKE_CASE    )
					SCREAMING_SNAKE_CASE						:							Dict							=     torch.stack(_SCREAMING_SNAKE_CASE    )
					return generated_tokens, generated_seq_lengths







	@torch.no_grad()
	def 			_lowerCAmelCase     (      self				:       List[str]     ,       _SCREAMING_SNAKE_CASE				:       int=None     ,       _SCREAMING_SNAKE_CASE				:       Optional[Any]=None     ,       _SCREAMING_SNAKE_CASE				:       List[Any]=None     ,       _SCREAMING_SNAKE_CASE				:       int = 5     ,       _SCREAMING_SNAKE_CASE				:       int = 67     ,       _SCREAMING_SNAKE_CASE				:       float = 1.0     ,       _SCREAMING_SNAKE_CASE				:       Optional[int] = None     ,       )						->	Optional[int]:





					"""simple docstring"""


					SCREAMING_SNAKE_CASE						:							str							=     eos_token_id
					SCREAMING_SNAKE_CASE						:							Tuple							=     None
					SCREAMING_SNAKE_CASE						:							int							=     None
					SCREAMING_SNAKE_CASE						:							Optional[Any]							=     torch.ones(_SCREAMING_SNAKE_CASE     ,       device=_SCREAMING_SNAKE_CASE     ,       dtype=torch.int    )
					SCREAMING_SNAKE_CASE						:							Dict							=     torch.zeros(_SCREAMING_SNAKE_CASE     ,       device=_SCREAMING_SNAKE_CASE     ,       dtype=torch.bool    )

					if input_embeds is not None:
									SCREAMING_SNAKE_CASE						:							List[Any]							=     input_embeds
					else:
									SCREAMING_SNAKE_CASE						:							Optional[Any]							=     self.transformer.transformer.wte(_SCREAMING_SNAKE_CASE    )

					for i in range(_SCREAMING_SNAKE_CASE    ):
									SCREAMING_SNAKE_CASE						:							Tuple							=     self.transformer(inputs_embeds=_SCREAMING_SNAKE_CASE    )
									SCREAMING_SNAKE_CASE						:							int							=     outputs.logits
									SCREAMING_SNAKE_CASE						:							Optional[int]							=     logits[:, -1, :] / (temperature if temperature > 0 else 1.0)
									SCREAMING_SNAKE_CASE						:							Optional[Any]							=     logits.softmax(-1    ).log()

									if scores is None:
													SCREAMING_SNAKE_CASE				,						SCREAMING_SNAKE_CASE						:							Tuple							=     logits.topk(_SCREAMING_SNAKE_CASE     ,       -1    )
													SCREAMING_SNAKE_CASE						:							Optional[Any]							=     generated.expand(_SCREAMING_SNAKE_CASE     ,       *generated.shape[1:]    )
													SCREAMING_SNAKE_CASE				,						SCREAMING_SNAKE_CASE						:							List[str]							=     next_tokens.permute(1     ,       0    ), scores.squeeze(0    )
													if tokens is None:
																	SCREAMING_SNAKE_CASE						:							Optional[Any]							=     next_tokens
													else:
																	SCREAMING_SNAKE_CASE						:							List[Any]							=     tokens.expand(_SCREAMING_SNAKE_CASE     ,       *tokens.shape[1:]    )
																	SCREAMING_SNAKE_CASE						:							List[str]							=     torch.cat((tokens, next_tokens)     ,       dim=1    )
									else:
													SCREAMING_SNAKE_CASE						:							List[str]							=     -float(np.inf    )
													SCREAMING_SNAKE_CASE						:							int							=     0
													SCREAMING_SNAKE_CASE						:							Tuple							=     scores[:, None] + logits
													seq_lengths[~is_stopped] += 1
													SCREAMING_SNAKE_CASE						:							Any							=     scores_sum / seq_lengths[:, None]
													SCREAMING_SNAKE_CASE				,						SCREAMING_SNAKE_CASE						:							Dict							=     scores_sum_average.view(-1    ).topk(_SCREAMING_SNAKE_CASE     ,       -1    )
													SCREAMING_SNAKE_CASE						:							Optional[Any]							=     next_tokens // scores_sum.shape[1]
													SCREAMING_SNAKE_CASE						:							Dict							=     seq_lengths[next_tokens_source]
													SCREAMING_SNAKE_CASE						:							List[str]							=     next_tokens % scores_sum.shape[1]
													SCREAMING_SNAKE_CASE						:							int							=     next_tokens.unsqueeze(1    )
													SCREAMING_SNAKE_CASE						:							Union[str, Any]							=     tokens[next_tokens_source]
													SCREAMING_SNAKE_CASE						:							List[Any]							=     torch.cat((tokens, next_tokens)     ,       dim=1    )
													SCREAMING_SNAKE_CASE						:							Any							=     generated[next_tokens_source]
													SCREAMING_SNAKE_CASE						:							Any							=     scores_sum_average * seq_lengths
													SCREAMING_SNAKE_CASE						:							Tuple							=     is_stopped[next_tokens_source]

									SCREAMING_SNAKE_CASE						:							str							=     self.transformer.transformer.wte(next_tokens.squeeze()    ).view(generated.shape[0]     ,       1     ,       -1    )
									SCREAMING_SNAKE_CASE						:							Tuple							=     torch.cat((generated, next_token_embed)     ,       dim=1    )
									SCREAMING_SNAKE_CASE						:							str							=     is_stopped + next_tokens.eq(_SCREAMING_SNAKE_CASE    ).squeeze()
									if is_stopped.all():
													break

					SCREAMING_SNAKE_CASE						:							Tuple							=     scores / seq_lengths
					SCREAMING_SNAKE_CASE						:							Tuple							=     scores.argsort(descending=_SCREAMING_SNAKE_CASE    )
					# tokens tensors are already padded to max_seq_length
					SCREAMING_SNAKE_CASE						:							Union[str, Any]							=     [tokens[i] for i in order]
					SCREAMING_SNAKE_CASE						:							List[str]							=     torch.stack(_SCREAMING_SNAKE_CASE     ,       dim=0    )
					SCREAMING_SNAKE_CASE						:							int							=     torch.tensor([seq_lengths[i] for i in order]     ,       dtype=seq_lengths.dtype    )
					return output_texts, seq_lengths





