



"""simple docstring"""
import unittest

import numpy as np

from transformers import is_flax_available
from transformers.testing_utils import require_flax

from ..test_modeling_flax_common import ids_tensor


if is_flax_available():
       import jax
       import jax.numpy as jnp

       from transformers.generation import (
           FlaxForcedBOSTokenLogitsProcessor,
           FlaxForcedEOSTokenLogitsProcessor,
           FlaxLogitsProcessorList,
           FlaxMinLengthLogitsProcessor,
           FlaxTemperatureLogitsWarper,
           FlaxTopKLogitsWarper,
           FlaxTopPLogitsWarper,
       )


@require_flax
class    SCREAMING_SNAKE_CASE__ (							unittest.TestCase   ):

    def 					_UpperCAmelCase	(				self			:				Tuple ,  lowerCAmelCase_			:				int ,  lowerCAmelCase_			:				int):







      """simple docstring"""







      lowercase_								= jnp.ones((batch_size, length)) / length
      return scores

    def 					_UpperCAmelCase	(				self			:				Tuple):







      """simple docstring"""







      lowercase_								= None
      lowercase_								= 2_0

      lowercase_								= self._get_uniform_logits(batch_size=2 ,  length=lowerCAmelCase_)

      # tweak scores to not be uniform anymore
      lowercase_								= scores.at[1, 5].set((1 / length) + 0.1)  # peak, 1st batch
      lowercase_								= scores.at[1, 1_0].set((1 / length) - 0.4)  # valley, 1st batch

      # compute softmax
      lowercase_								= jax.nn.softmax(lowerCAmelCase_ ,  axis=-1)

      lowercase_								= FlaxTemperatureLogitsWarper(temperature=0.5)
      lowercase_								= FlaxTemperatureLogitsWarper(temperature=1.3)

      lowercase_								= jax.nn.softmax(temp_dist_warper_sharper(lowerCAmelCase_ ,  scores.copy() ,  cur_len=lowerCAmelCase_) ,  axis=-1)
      lowercase_								= jax.nn.softmax(temp_dist_warper_smoother(lowerCAmelCase_ ,  scores.copy() ,  cur_len=lowerCAmelCase_) ,  axis=-1)

      # uniform distribution stays uniform
      self.assertTrue(jnp.allclose(probs[0, :] ,  warped_prob_sharp[0, :] ,  atol=1E-3))
      self.assertTrue(jnp.allclose(probs[0, :] ,  warped_prob_smooth[0, :] ,  atol=1E-3))

      # sharp peaks get higher, valleys get lower
      self.assertLess(probs[1, :].max() ,  warped_prob_sharp[1, :].max())
      self.assertGreater(probs[1, :].min() ,  warped_prob_sharp[1, :].min())

      # smooth peaks get lower, valleys get higher
      self.assertGreater(probs[1, :].max() ,  warped_prob_smooth[1, :].max())
      self.assertLess(probs[1, :].min() ,  warped_prob_smooth[1, :].min())

    def 					_UpperCAmelCase	(				self			:				List[Any]):







      """simple docstring"""







      lowercase_								= None
      lowercase_								= 1_0
      lowercase_								= 2

      # create ramp distribution
      lowercase_								= np.broadcast_to(np.arange(lowerCAmelCase_)[None, :] ,  (batch_size, vocab_size)).copy()
      lowercase_								= ramp_logits[1:, : vocab_size // 2] + vocab_size

      lowercase_								= FlaxTopKLogitsWarper(3)

      lowercase_								= top_k_warp(lowerCAmelCase_ ,  lowerCAmelCase_ ,  cur_len=lowerCAmelCase_)

      # check that correct tokens are filtered
      self.assertListEqual(jnp.isinf(scores[0]).tolist() ,  7 * [True] + 3 * [False])
      self.assertListEqual(jnp.isinf(scores[1]).tolist() ,  2 * [True] + 3 * [False] + 5 * [True])

      # check special case
      lowercase_								= 5
      lowercase_								= FlaxTopKLogitsWarper(top_k=1 ,  filter_value=0.0 ,  min_tokens_to_keep=3)

      lowercase_								= np.broadcast_to(np.arange(lowerCAmelCase_)[None, :] ,  (batch_size, length)).copy()
      lowercase_								= top_k_warp_safety_check(lowerCAmelCase_ ,  lowerCAmelCase_ ,  cur_len=lowerCAmelCase_)

      # min_tokens overwrites k: 3 tokens are kept => 2 tokens are nullified
      self.assertListEqual((scores == 0.0).sum(axis=-1).tolist() ,  [2, 2])

    def 					_UpperCAmelCase	(				self			:				int):







      """simple docstring"""







      lowercase_								= None
      lowercase_								= 1_0
      lowercase_								= 2

      # create distribution and take log (inverse to Softmax as taken in TopPLogitsWarper)
      lowercase_								= np.log(np.array([[0.3, 0.1, 0.1, 0.5], [0.15, 0.3, 0.3, 0.25]]))

      lowercase_								= FlaxTopPLogitsWarper(0.8)
      lowercase_								= np.exp(top_p_warp(lowerCAmelCase_ ,  lowerCAmelCase_ ,  cur_len=lowerCAmelCase_))

      # dist should be filtered to keep min num values so that sum is >= top_p
      # exp (-inf) => 0
      lowercase_								= np.array([[0.3, 0.0, 0.0, 0.5], [0.0, 0.3, 0.3, 0.25]])
      self.assertTrue(np.allclose(lowerCAmelCase_ ,  lowerCAmelCase_ ,  atol=1E-3))

      # check edge cases with negative and extreme logits
      lowercase_								= np.broadcast_to(np.arange(lowerCAmelCase_)[None, :] ,  (batch_size, vocab_size)).copy() - (
          vocab_size // 2
      )

      # make ramp_logits more extreme
      lowercase_								= ramp_logits[1] * 100.0

      # make sure at least 2 tokens are kept
      lowercase_								= FlaxTopPLogitsWarper(0.9 ,  min_tokens_to_keep=2 ,  filter_value=0.0)
      lowercase_								= top_p_warp(lowerCAmelCase_ ,  lowerCAmelCase_ ,  cur_len=lowerCAmelCase_)

      # first batch should keep three tokens, second batch would keep only 1, but due to `min_tokens_to_keep=2` keeps 2.
      self.assertListEqual((filtered_dist != 0.0).sum(axis=-1).tolist() ,  [3, 2])

    def 					_UpperCAmelCase	(				self			:				Dict):







      """simple docstring"""







      lowercase_								= 2_0
      lowercase_								= 4
      lowercase_								= 0

      lowercase_								= FlaxMinLengthLogitsProcessor(min_length=1_0 ,  eos_token_id=lowerCAmelCase_)

      # check that min length is applied at length 5
      lowercase_								= ids_tensor((batch_size, 2_0) ,  vocab_size=2_0)
      lowercase_								= 5
      lowercase_								= self._get_uniform_logits(lowerCAmelCase_ ,  lowerCAmelCase_)
      lowercase_								= min_dist_processor(lowerCAmelCase_ ,  lowerCAmelCase_ ,  cur_len=lowerCAmelCase_)
      self.assertListEqual(scores_before_min_length[:, eos_token_id].tolist() ,  4 * [-float("""inf""")])

      # check that min length is not applied anymore at length 15
      lowercase_								= self._get_uniform_logits(lowerCAmelCase_ ,  lowerCAmelCase_)
      lowercase_								= 1_5
      lowercase_								= min_dist_processor(lowerCAmelCase_ ,  lowerCAmelCase_ ,  cur_len=lowerCAmelCase_)
      self.assertFalse(jnp.isinf(lowerCAmelCase_).any())

    def 					_UpperCAmelCase	(				self			:				str):







      """simple docstring"""







      lowercase_								= 2_0
      lowercase_								= 4
      lowercase_								= 0

      lowercase_								= FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowerCAmelCase_)

      # check that all scores are -inf except the bos_token_id score
      lowercase_								= ids_tensor((batch_size, 1) ,  vocab_size=2_0)
      lowercase_								= 1
      lowercase_								= self._get_uniform_logits(lowerCAmelCase_ ,  lowerCAmelCase_)
      lowercase_								= logits_processor(lowerCAmelCase_ ,  lowerCAmelCase_ ,  cur_len=lowerCAmelCase_)
      self.assertTrue(jnp.isneginf(scores[:, bos_token_id + 1 :]).all())
      self.assertListEqual(scores[:, bos_token_id].tolist() ,  4 * [0])  # score for bos_token_id shold be zero

      # check that bos_token_id is not forced if current length is greater than 1
      lowercase_								= 3
      lowercase_								= self._get_uniform_logits(lowerCAmelCase_ ,  lowerCAmelCase_)
      lowercase_								= logits_processor(lowerCAmelCase_ ,  lowerCAmelCase_ ,  cur_len=lowerCAmelCase_)
      self.assertFalse(jnp.isinf(lowerCAmelCase_).any())

    def 					_UpperCAmelCase	(				self			:				Optional[Any]):







      """simple docstring"""







      lowercase_								= 2_0
      lowercase_								= 4
      lowercase_								= 0
      lowercase_								= 5

      lowercase_								= FlaxForcedEOSTokenLogitsProcessor(max_length=lowerCAmelCase_ ,  eos_token_id=lowerCAmelCase_)

      # check that all scores are -inf except the eos_token_id when max_length is reached
      lowercase_								= ids_tensor((batch_size, 4) ,  vocab_size=2_0)
      lowercase_								= 4
      lowercase_								= self._get_uniform_logits(lowerCAmelCase_ ,  lowerCAmelCase_)
      lowercase_								= logits_processor(lowerCAmelCase_ ,  lowerCAmelCase_ ,  cur_len=lowerCAmelCase_)
      self.assertTrue(jnp.isneginf(scores[:, eos_token_id + 1 :]).all())
      self.assertListEqual(scores[:, eos_token_id].tolist() ,  4 * [0])  # score for eos_token_id should be zero

      # check that eos_token_id is not forced if max_length is not reached
      lowercase_								= 3
      lowercase_								= self._get_uniform_logits(lowerCAmelCase_ ,  lowerCAmelCase_)
      lowercase_								= logits_processor(lowerCAmelCase_ ,  lowerCAmelCase_ ,  cur_len=lowerCAmelCase_)
      self.assertFalse(jnp.isinf(lowerCAmelCase_).any())

    def 					_UpperCAmelCase	(				self			:				List[Any]):







      """simple docstring"""







      lowercase_								= 4
      lowercase_								= 1_0
      lowercase_								= 1_5
      lowercase_								= 2
      lowercase_								= 1
      lowercase_								= 1_5

      # dummy input_ids and scores
      lowercase_								= ids_tensor((batch_size, sequence_length) ,  lowerCAmelCase_)
      lowercase_								= input_ids.copy()

      lowercase_								= self._get_uniform_logits(lowerCAmelCase_ ,  lowerCAmelCase_)
      lowercase_								= scores.copy()

      # instantiate all dist processors
      lowercase_								= FlaxTemperatureLogitsWarper(temperature=0.5)
      lowercase_								= FlaxTopKLogitsWarper(3)
      lowercase_								= FlaxTopPLogitsWarper(0.8)

      # instantiate all logits processors
      lowercase_								= FlaxMinLengthLogitsProcessor(min_length=1_0 ,  eos_token_id=lowerCAmelCase_)
      lowercase_								= FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowerCAmelCase_)
      lowercase_								= FlaxForcedEOSTokenLogitsProcessor(max_length=lowerCAmelCase_ ,  eos_token_id=lowerCAmelCase_)

      lowercase_								= 1_0

      # no processor list
      lowercase_								= temp_dist_warp(lowerCAmelCase_ ,  lowerCAmelCase_ ,  cur_len=lowerCAmelCase_)
      lowercase_								= top_k_warp(lowerCAmelCase_ ,  lowerCAmelCase_ ,  cur_len=lowerCAmelCase_)
      lowercase_								= top_p_warp(lowerCAmelCase_ ,  lowerCAmelCase_ ,  cur_len=lowerCAmelCase_)
      lowercase_								= min_dist_proc(lowerCAmelCase_ ,  lowerCAmelCase_ ,  cur_len=lowerCAmelCase_)
      lowercase_								= bos_dist_proc(lowerCAmelCase_ ,  lowerCAmelCase_ ,  cur_len=lowerCAmelCase_)
      lowercase_								= eos_dist_proc(lowerCAmelCase_ ,  lowerCAmelCase_ ,  cur_len=lowerCAmelCase_)

      # with processor list
      lowercase_								= FlaxLogitsProcessorList(
          [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc])
      lowercase_								= processor(lowerCAmelCase_ ,  lowerCAmelCase_ ,  cur_len=lowerCAmelCase_)

      # scores should be equal
      self.assertTrue(jnp.allclose(lowerCAmelCase_ ,  lowerCAmelCase_ ,  atol=1E-3))

      # input_ids should never be changed
      self.assertListEqual(input_ids.tolist() ,  input_ids_comp.tolist())




    def 					_UpperCAmelCase	(				self			:				int):







      """simple docstring"""







      lowercase_								= 4
      lowercase_								= 1_0
      lowercase_								= 1_5
      lowercase_								= 2
      lowercase_								= 1
      lowercase_								= 1_5

      # dummy input_ids and scores
      lowercase_								= ids_tensor((batch_size, sequence_length) ,  lowerCAmelCase_)
      lowercase_								= input_ids.copy()

      lowercase_								= self._get_uniform_logits(lowerCAmelCase_ ,  lowerCAmelCase_)
      lowercase_								= scores.copy()

      # instantiate all dist processors
      lowercase_								= FlaxTemperatureLogitsWarper(temperature=0.5)
      lowercase_								= FlaxTopKLogitsWarper(3)
      lowercase_								= FlaxTopPLogitsWarper(0.8)

      # instantiate all logits processors
      lowercase_								= FlaxMinLengthLogitsProcessor(min_length=1_0 ,  eos_token_id=lowerCAmelCase_)
      lowercase_								= FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowerCAmelCase_)
      lowercase_								= FlaxForcedEOSTokenLogitsProcessor(max_length=lowerCAmelCase_ ,  eos_token_id=lowerCAmelCase_)

      lowercase_								= 1_0

      # no processor list
      def run_no_processor_list(lowerCAmelCase_			:				int ,  lowerCAmelCase_			:				Dict ,  lowerCAmelCase_			:				List[Any]):
        lowercase_								= temp_dist_warp(lowerCAmelCase_ ,  lowerCAmelCase_ ,  cur_len=lowerCAmelCase_)
        lowercase_								= top_k_warp(lowerCAmelCase_ ,  lowerCAmelCase_ ,  cur_len=lowerCAmelCase_)
        lowercase_								= top_p_warp(lowerCAmelCase_ ,  lowerCAmelCase_ ,  cur_len=lowerCAmelCase_)
        lowercase_								= min_dist_proc(lowerCAmelCase_ ,  lowerCAmelCase_ ,  cur_len=lowerCAmelCase_)
        lowercase_								= bos_dist_proc(lowerCAmelCase_ ,  lowerCAmelCase_ ,  cur_len=lowerCAmelCase_)
        lowercase_								= eos_dist_proc(lowerCAmelCase_ ,  lowerCAmelCase_ ,  cur_len=lowerCAmelCase_)
        return scores

      # with processor list
      def run_processor_list(lowerCAmelCase_			:				Union[str, Any] ,  lowerCAmelCase_			:				int ,  lowerCAmelCase_			:				Optional[Any]):
        lowercase_								= FlaxLogitsProcessorList(
            [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc])
        lowercase_								= processor(lowerCAmelCase_ ,  lowerCAmelCase_ ,  cur_len=lowerCAmelCase_)
        return scores

      lowercase_								= jax.jit(lowerCAmelCase_)
      lowercase_								= jax.jit(lowerCAmelCase_)

      lowercase_								= jitted_run_no_processor_list(lowerCAmelCase_ ,  lowerCAmelCase_ ,  lowerCAmelCase_)
      lowercase_								= jitted_run_processor_list(lowerCAmelCase_ ,  lowerCAmelCase_ ,  lowerCAmelCase_)

      # scores should be equal
      self.assertTrue(jnp.allclose(lowerCAmelCase_ ,  lowerCAmelCase_ ,  atol=1E-3))

      # input_ids should never be changed
      self.assertListEqual(input_ids.tolist() ,  input_ids_comp.tolist())






"""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    SCREAMING_SNAKE_CASE__ (							__UpperCAmelCase      ,  __UpperCAmelCase      ,  __UpperCAmelCase   ):
    lowercase__      =   [r"h\.\d+\.attn\.bias", r"h\.\d+\.attn\.masked_bias"]

    @register_to_config
    def __init__(				self			:				int ,  lowerCAmelCase_			:				int ,  lowerCAmelCase_			:				int ,  lowerCAmelCase_			:				Optional[int] = None ,  lowerCAmelCase_			:				int = 5_0_2_5_7 ,  lowerCAmelCase_			:				int = 1_0_2_4 ,  lowerCAmelCase_			:				int = 7_6_8 ,  lowerCAmelCase_			:				int = 1_2 ,  lowerCAmelCase_			:				int = 1_2 ,  lowerCAmelCase_			:				Optional[int] = None ,  lowerCAmelCase_			:				str = "gelu_new" ,  lowerCAmelCase_			:				float = 0.1 ,  lowerCAmelCase_			:				float = 0.1 ,  lowerCAmelCase_			:				float = 0.1 ,  lowerCAmelCase_			:				float = 1E-5 ,  lowerCAmelCase_			:				float = 0.02 ,  lowerCAmelCase_			:				bool = True ,  lowerCAmelCase_			:				bool = True ,  lowerCAmelCase_			:				bool = False ,  lowerCAmelCase_			:				bool = False ,  ):







      """simple docstring"""







      super().__init__()

      lowercase_								= 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.''')

      lowercase_								= prefix_inner_dim
      lowercase_								= prefix_hidden_dim

      lowercase_								= (
          nn.Linear(self.prefix_inner_dim ,  self.prefix_hidden_dim)
          if self.prefix_hidden_dim is not None
          else nn.Identity()
      )
      lowercase_								= (
          nn.Linear(self.prefix_hidden_dim ,  lowerCAmelCase_) if self.prefix_hidden_dim is not None else nn.Identity()
      )

      lowercase_								= GPTaConfig(
          vocab_size=lowerCAmelCase_ ,  n_positions=lowerCAmelCase_ ,  n_embd=lowerCAmelCase_ ,  n_layer=lowerCAmelCase_ ,  n_head=lowerCAmelCase_ ,  n_inner=lowerCAmelCase_ ,  activation_function=lowerCAmelCase_ ,  resid_pdrop=lowerCAmelCase_ ,  embd_pdrop=lowerCAmelCase_ ,  attn_pdrop=lowerCAmelCase_ ,  layer_norm_epsilon=lowerCAmelCase_ ,  initializer_range=lowerCAmelCase_ ,  scale_attn_weights=lowerCAmelCase_ ,  use_cache=lowerCAmelCase_ ,  scale_attn_by_inverse_layer_idx=lowerCAmelCase_ ,  reorder_and_upcast_attn=lowerCAmelCase_ ,  )
      lowercase_								= GPTaLMHeadModel(lowerCAmelCase_)

    def 					_UpperCAmelCase	(				self			:				int ,  lowerCAmelCase_			:				torch.Tensor ,  lowerCAmelCase_			:				torch.Tensor ,  lowerCAmelCase_			:				Optional[torch.Tensor] = None ,  lowerCAmelCase_			:				Optional[torch.Tensor] = None ,  ):







      """simple docstring"""







      lowercase_								= self.transformer.transformer.wte(lowerCAmelCase_)
      lowercase_								= self.encode_prefix(lowerCAmelCase_)
      lowercase_								= self.decode_prefix(lowerCAmelCase_)
      lowercase_								= torch.cat((prefix_embeds, embedding_text) ,  dim=1)

      if labels is not None:
        lowercase_								= self.get_dummy_token(input_ids.shape[0] ,  input_ids.device)
        lowercase_								= torch.cat((dummy_token, input_ids) ,  dim=1)
      lowercase_								= self.transformer(inputs_embeds=lowerCAmelCase_ ,  labels=lowerCAmelCase_ ,  attention_mask=lowerCAmelCase_)
      if self.prefix_hidden_dim is not None:
        return out, hidden
      else:
        return out

    def 					_UpperCAmelCase	(				self			:				Optional[int] ,  lowerCAmelCase_			:				int ,  lowerCAmelCase_			:				torch.device):







      """simple docstring"""







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

    def 					_UpperCAmelCase	(				self			:				List[str] ,  lowerCAmelCase_			:				List[Any]):







      """simple docstring"""







      return self.encode_prefix(lowerCAmelCase_)

    @torch.no_grad()
    def 					_UpperCAmelCase	(				self			:				Tuple ,  lowerCAmelCase_			:				str ,  lowerCAmelCase_			:				Dict ,  lowerCAmelCase_			:				Union[str, Any]):







      """simple docstring"""







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




    @torch.no_grad()
    def 					_UpperCAmelCase	(				self			:				int ,  lowerCAmelCase_			:				List[str]=None ,  lowerCAmelCase_			:				Tuple=None ,  lowerCAmelCase_			:				str=None ,  lowerCAmelCase_			:				int = 5 ,  lowerCAmelCase_			:				int = 6_7 ,  lowerCAmelCase_			:				float = 1.0 ,  lowerCAmelCase_			:				Optional[int] = None ,  ):







      """simple docstring"""







      lowercase_								= eos_token_id
      lowercase_								= None
      lowercase_								= None
      lowercase_								= torch.ones(lowerCAmelCase_ ,  device=lowerCAmelCase_ ,  dtype=torch.int)
      lowercase_								= torch.zeros(lowerCAmelCase_ ,  device=lowerCAmelCase_ ,  dtype=torch.bool)

      if input_embeds is not None:
        lowercase_								= input_embeds
      else:
        lowercase_								= self.transformer.transformer.wte(lowerCAmelCase_)

      for i in range(lowerCAmelCase_):
        lowercase_								= self.transformer(inputs_embeds=lowerCAmelCase_)
        lowercase_								= outputs.logits
        lowercase_								= logits[:, -1, :] / (temperature if temperature > 0 else 1.0)
        lowercase_								= logits.softmax(-1).log()

        if scores is None:
          lowercase_					,    lowercase_								= logits.topk(lowerCAmelCase_ ,  -1)
          lowercase_								= generated.expand(lowerCAmelCase_ ,  *generated.shape[1:])
          lowercase_					,    lowercase_								= next_tokens.permute(1 ,  0), scores.squeeze(0)
          if tokens is None:
            lowercase_								= next_tokens
          else:
            lowercase_								= tokens.expand(lowerCAmelCase_ ,  *tokens.shape[1:])
            lowercase_								= torch.cat((tokens, next_tokens) ,  dim=1)
        else:
          lowercase_								= -float(np.inf)
          lowercase_								= 0
          lowercase_								= scores[:, None] + logits
          seq_lengths[~is_stopped] += 1
          lowercase_								= scores_sum / seq_lengths[:, None]
          lowercase_					,    lowercase_								= scores_sum_average.view(-1).topk(lowerCAmelCase_ ,  -1)
          lowercase_								= next_tokens // scores_sum.shape[1]
          lowercase_								= seq_lengths[next_tokens_source]
          lowercase_								= next_tokens % scores_sum.shape[1]
          lowercase_								= next_tokens.unsqueeze(1)
          lowercase_								= tokens[next_tokens_source]
          lowercase_								= torch.cat((tokens, next_tokens) ,  dim=1)
          lowercase_								= generated[next_tokens_source]
          lowercase_								= scores_sum_average * seq_lengths
          lowercase_								= is_stopped[next_tokens_source]

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

      lowercase_								= scores / seq_lengths
      lowercase_								= scores.argsort(descending=lowerCAmelCase_)
      # tokens tensors are already padded to max_seq_length
      lowercase_								= [tokens[i] for i in order]
      lowercase_								= torch.stack(lowerCAmelCase_ ,  dim=0)
      lowercase_								= torch.tensor([seq_lengths[i] for i in order] ,  dtype=seq_lengths.dtype)
      return output_texts, seq_lengths

