



"""simple docstring"""
from __future__ import annotations

import unittest

import numpy as np

from transformers import OPTConfig, is_tf_available
from transformers.testing_utils import require_sentencepiece, require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
       import tensorflow as tf

       from transformers import GPTaTokenizer, TFOPTForCausalLM, TFOPTModel


def 			_SCREAMING_SNAKE_CASE						(__lowerCAmelCase ,  __lowerCAmelCase ,  __lowerCAmelCase=None ,  __lowerCAmelCase=None							) ->						Tuple:







  '''simple docstring'''


  if attention_mask is None:
    lowercase_								= tf.cast(tf.math.not_equal(__lowerCAmelCase ,  config.pad_token_id							) ,  tf.inta							)
  return {"input_ids": input_ids, "attention_mask": attention_mask}


@require_tf
class    SCREAMING_SNAKE_CASE__ :
    lowercase__      =   OPTConfig
    lowercase__      =   {}
    lowercase__      =   "gelu"

    def __init__(				self			:				Tuple ,  lowerCAmelCase_			:				Tuple ,  lowerCAmelCase_			:				int=1_3 ,  lowerCAmelCase_			:				Optional[Any]=7 ,  lowerCAmelCase_			:				Optional[Any]=True ,  lowerCAmelCase_			:				Dict=False ,  lowerCAmelCase_			:				str=9_9 ,  lowerCAmelCase_			:				Dict=1_6 ,  lowerCAmelCase_			:				Optional[int]=2 ,  lowerCAmelCase_			:				int=4 ,  lowerCAmelCase_			:				List[Any]=4 ,  lowerCAmelCase_			:				Union[str, Any]="gelu" ,  lowerCAmelCase_			:				Optional[Any]=0.1 ,  lowerCAmelCase_			:				int=0.1 ,  lowerCAmelCase_			:				Union[str, Any]=2_0 ,  lowerCAmelCase_			:				Optional[Any]=2 ,  lowerCAmelCase_			:				Optional[Any]=1 ,  lowerCAmelCase_			:				Optional[Any]=0 ,  lowerCAmelCase_			:				str=1_6 ,  lowerCAmelCase_			:				Union[str, Any]=1_6 ,  ):







      """simple docstring"""







      lowercase_								= parent
      lowercase_								= batch_size
      lowercase_								= seq_length
      lowercase_								= is_training
      lowercase_								= use_labels
      lowercase_								= vocab_size
      lowercase_								= hidden_size
      lowercase_								= num_hidden_layers
      lowercase_								= num_attention_heads
      lowercase_								= intermediate_size
      lowercase_								= hidden_act
      lowercase_								= hidden_dropout_prob
      lowercase_								= attention_probs_dropout_prob
      lowercase_								= max_position_embeddings
      lowercase_								= eos_token_id
      lowercase_								= pad_token_id
      lowercase_								= bos_token_id
      lowercase_								= embed_dim
      lowercase_								= word_embed_proj_dim
      lowercase_								= False

    def 					_UpperCAmelCase	(				self			:				Union[str, Any]):







      """simple docstring"""







      lowercase_								= ids_tensor([self.batch_size, self.seq_length - 1] ,  self.vocab_size)
      lowercase_								= tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size) ,  1)
      lowercase_								= tf.concat([input_ids, eos_tensor] ,  axis=1)

      lowercase_								= self.config_cls(
          vocab_size=self.vocab_size ,  hidden_size=self.hidden_size ,  num_hidden_layers=self.num_hidden_layers ,  num_attention_heads=self.num_attention_heads ,  ffn_dim=self.intermediate_size ,  dropout=self.hidden_dropout_prob ,  attention_dropout=self.attention_probs_dropout_prob ,  max_position_embeddings=self.max_position_embeddings ,  eos_token_id=self.eos_token_id ,  bos_token_id=self.bos_token_id ,  pad_token_id=self.pad_token_id ,  embed_dim=self.embed_dim ,  word_embed_proj_dim=self.word_embed_proj_dim ,  is_encoder_decoder=lowerCAmelCase_ ,  **self.config_updates ,  )
      lowercase_								= prepare_opt_inputs_dict(lowerCAmelCase_ ,  lowerCAmelCase_)
      return config, inputs_dict




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







      """simple docstring"""







      lowercase_								= TFOPTModel(config=lowerCAmelCase_)
      lowercase_								= inputs_dict["""input_ids"""]

      lowercase_								= input_ids[:1, :]
      lowercase_								= inputs_dict["""attention_mask"""][:1, :]
      lowercase_								= 1

      # first forward pass
      lowercase_								= model(lowerCAmelCase_ ,  attention_mask=lowerCAmelCase_ ,  use_cache=lowerCAmelCase_)

      lowercase_					,    lowercase_								= outputs.to_tuple()

      # create hypothetical next token and extent to next_input_ids
      lowercase_								= ids_tensor((self.batch_size, 3) ,  config.vocab_size)
      lowercase_								= tf.cast(ids_tensor((self.batch_size, 3) ,  2) ,  tf.inta)

      # append to next input_ids and
      lowercase_								= tf.concat([input_ids, next_tokens] ,  axis=-1)
      lowercase_								= tf.concat([attention_mask, next_attn_mask] ,  axis=-1)

      lowercase_								= model(lowerCAmelCase_ ,  attention_mask=lowerCAmelCase_)[0]
      lowercase_								= model(lowerCAmelCase_ ,  attention_mask=lowerCAmelCase_ ,  past_key_values=lowerCAmelCase_)[0]

      self.parent.assertEqual(next_tokens.shape[1] ,  output_from_past.shape[1])

      # select random slice
      lowercase_								= int(ids_tensor((1,) ,  output_from_past.shape[-1]))
      lowercase_								= output_from_no_past[:, -3:, random_slice_idx]
      lowercase_								= output_from_past[:, :, random_slice_idx]

      # test that outputs are equal for slice
      tf.debugging.assert_near(lowerCAmelCase_ ,  lowerCAmelCase_ ,  rtol=1E-3)







@require_tf
class    SCREAMING_SNAKE_CASE__ (							__UpperCAmelCase      ,  __UpperCAmelCase      ,  unittest.TestCase   ):
    lowercase__      =   (TFOPTModel, TFOPTForCausalLM) if is_tf_available() else ()
    lowercase__      =   (TFOPTForCausalLM,) if is_tf_available() else ()
    lowercase__      =   (
        {"feature-extraction": TFOPTModel, "text-generation": TFOPTForCausalLM} if is_tf_available() else {}
    )
    lowercase__      =   False
    lowercase__      =   False
    lowercase__      =   False
    lowercase__      =   10

    def 					_UpperCAmelCase	(				self			:				Any):







      """simple docstring"""







      lowercase_								= TFOPTModelTester(self)
      lowercase_								= ConfigTester(self ,  config_class=lowerCAmelCase_)

    def 					_UpperCAmelCase	(				self			:				List[str]):







      """simple docstring"""







      self.config_tester.run_common_tests()

    def 					_UpperCAmelCase	(				self			:				List[str]):







      """simple docstring"""







      lowercase_								= self.model_tester.prepare_config_and_inputs_for_common()
      self.model_tester.check_decoder_model_past_large_inputs(*lowerCAmelCase_)




    def 					_UpperCAmelCase	(				self			:				Dict):







      """simple docstring"""







      lowercase_					,    lowercase_								= self.model_tester.prepare_config_and_inputs_for_common()

      def _get_word_embedding_weight(lowerCAmelCase_			:				List[str] ,  lowerCAmelCase_			:				List[Any]):
        if hasattr(lowerCAmelCase_ ,  """weight"""):
          return embedding_layer.weight
        else:
          # Here we build the word embeddings weights if not exists.
          # And then we retry to get the attribute once built.
          model.build()
          if hasattr(lowerCAmelCase_ ,  """weight"""):
            return embedding_layer.weight
          else:
            return None

      for model_class in self.all_model_classes:
        for size in [config.vocab_size - 1_0, config.vocab_size + 1_0]:
          # build the embeddings
          lowercase_								= model_class(config=lowerCAmelCase_)
          lowercase_								= _get_word_embedding_weight(lowerCAmelCase_ ,  model.get_input_embeddings())
          lowercase_								= _get_word_embedding_weight(lowerCAmelCase_ ,  model.get_output_embeddings())

          # reshape the embeddings
          model.resize_token_embeddings(lowerCAmelCase_)
          lowercase_								= _get_word_embedding_weight(lowerCAmelCase_ ,  model.get_input_embeddings())
          lowercase_								= _get_word_embedding_weight(lowerCAmelCase_ ,  model.get_output_embeddings())

          # check that the resized embeddings size matches the desired size.
          lowercase_								= size if size is not None else config.vocab_size

          self.assertEqual(new_input_embeddings.shape[0] ,  lowerCAmelCase_)

          # check that weights remain the same after resizing
          lowercase_								= True
          for pa, pa in zip(old_input_embeddings.value() ,  new_input_embeddings.value()):
            if tf.math.reduce_sum(tf.math.abs(pa - pa)) > 0:
              lowercase_								= False
          self.assertTrue(lowerCAmelCase_)

          if old_output_embeddings is not None and new_output_embeddings is not None:
            self.assertEqual(new_output_embeddings.shape[0] ,  lowerCAmelCase_)

            lowercase_								= True
            for pa, pa in zip(old_output_embeddings.value() ,  new_output_embeddings.value()):
              if tf.math.reduce_sum(tf.math.abs(pa - pa)) > 0:
                lowercase_								= False
            self.assertTrue(lowerCAmelCase_)







def 			_SCREAMING_SNAKE_CASE						(__lowerCAmelCase							) ->						str:







  '''simple docstring'''


  return tf.constant(__lowerCAmelCase ,  dtype=tf.intaa							)


@require_tf
class    SCREAMING_SNAKE_CASE__ (							unittest.TestCase   ):
    lowercase__      =   99

    def 					_UpperCAmelCase	(				self			:				Tuple):







      """simple docstring"""







      lowercase_								= tf.ones((4, 1) ,  dtype=tf.intaa) * 2
      lowercase_								= tf.concat([ids_tensor((4, 6) ,  self.vocab_size - 3) + 3, eos_column_vector] ,  axis=1)
      lowercase_								= input_ids.shape[0]
      lowercase_								= OPTConfig(
          vocab_size=self.vocab_size ,  hidden_size=2_4 ,  num_hidden_layers=2 ,  num_attention_heads=2 ,  ffn_dim=3_2 ,  max_position_embeddings=4_8 ,  eos_token_id=2 ,  pad_token_id=1 ,  bos_token_id=0 ,  )
      return config, input_ids, batch_size







@require_sentencepiece
@require_tf
class    SCREAMING_SNAKE_CASE__ (							unittest.TestCase   ):

    @slow
    def 					_UpperCAmelCase	(				self			:				Dict):







      """simple docstring"""







      lowercase_								= TFOPTModel.from_pretrained("""facebook/opt-350m""")
      lowercase_								= _long_tensor([[0, 3_1_4_1_4, 2_3_2, 3_2_8, 7_4_0, 1_1_4_0, 1_2_6_9_5, 6_9, 4_6_0_7_8, 1_5_8_8, 2]])
      lowercase_								= tf.not_equal(lowerCAmelCase_ ,  model.config.pad_token_id)
      with tf.GradientTape():
        lowercase_								= model(input_ids=lowerCAmelCase_ ,  attention_mask=lowerCAmelCase_).last_hidden_state
      lowercase_								= (1, 1_1, 5_1_2)
      self.assertEqual(output.shape ,  lowerCAmelCase_)
      lowercase_								= tf.constant(
          [[-0.2_873, -1.9_218, -0.3_033], [-1.2_710, -0.1_338, -0.1_902], [0.4_095, 0.1_214, -1.3_121]])
      self.assertTrue(np.allclose(output[:, :3, :3] ,  lowerCAmelCase_ ,  atol=4E-3))

      lowercase_								= tf.function(lowerCAmelCase_ ,  jit_compile=lowerCAmelCase_)
      lowercase_								= xla_generate(lowerCAmelCase_ ,  lowerCAmelCase_)[0]
      self.assertTrue(np.allclose(output[:, :3, :3] ,  lowerCAmelCase_ ,  atol=4E-2))







@require_tf
@slow
class    SCREAMING_SNAKE_CASE__ (							unittest.TestCase   ):

    def 					_UpperCAmelCase	(				self			:				Any):







      """simple docstring"""







      super().setUp()
      lowercase_								= """facebook/opt-350m"""




    def 					_UpperCAmelCase	(				self			:				int):







      """simple docstring"""







      lowercase_								= TFOPTForCausalLM.from_pretrained(self.path_model)
      lowercase_								= GPTaTokenizer.from_pretrained(self.path_model)

      lowercase_								= [
          """Today is a beautiful day and I want to""",
          """In the city of""",
          """Paris is the capital of France and""",
          """Computers and mobile phones have taken""",
      ]
      # verify that prompt without BOS token is identical to Metaseq -> add_special_tokens=False
      lowercase_								= tokenizer(lowerCAmelCase_ ,  return_tensors="""tf""" ,  padding=lowerCAmelCase_ ,  add_special_tokens=lowerCAmelCase_)
      lowercase_								= tf.math.reduce_mean(model(inputs.input_ids ,  attention_mask=inputs.attention_mask)[0] ,  axis=-1)
      lowercase_								= tf.constant(
          [
              [1.3_851, -13.8_923, -10.5_229, -10.7_533, -0.2_309, -10.2_384, -0.5_365, -9.0_947, -5.1_670],
              [-4.7_073, -10.6_276, -3.9_415, -21.5_242, -0.2_822, -0.2_822, -0.2_822, -0.2_822, -0.2_822],
              [0.6_247, -3.4_229, -8.9_179, -1.4_297, -14.1_650, 1.4_146, -9.0_218, -0.2_703, -0.2_703],
              [6.4_783, -1.9_913, -10.7_926, -2.3_336, 1.5_092, -0.9_974, -6.8_213, 1.3_477, 1.3_477],
          ])
      self.assertTrue(np.allclose(lowerCAmelCase_ ,  lowerCAmelCase_ ,  atol=1E-4))

      lowercase_								= tf.function(lowerCAmelCase_ ,  jit_compile=lowerCAmelCase_)
      lowercase_								= tf.math.reduce_mean(xla_generate(inputs.input_ids ,  attention_mask=inputs.attention_mask)[0] ,  axis=-1)
      self.assertTrue(np.allclose(lowerCAmelCase_ ,  lowerCAmelCase_ ,  atol=1E-4))







@require_tf
@slow
class    SCREAMING_SNAKE_CASE__ (							unittest.TestCase   ):

    @property
    def 					_UpperCAmelCase	(				self			:				Optional[Any]):







      """simple docstring"""







      return [
          "Today is a beautiful day and I want",
          "In the city of",
          "Paris is the capital of France and",
          "Computers and mobile phones have taken",
      ]

    def 					_UpperCAmelCase	(				self			:				Tuple):







      """simple docstring"""







      lowercase_								= """facebook/opt-125m"""

      lowercase_								= [
          """Today is a beautiful day and I want to""",
          """In the city of New York, the city""",
          """Paris is the capital of France and the capital""",
          """Computers and mobile phones have taken over the""",
      ]

      lowercase_								= []
      lowercase_								= GPTaTokenizer.from_pretrained(lowerCAmelCase_)
      lowercase_								= TFOPTForCausalLM.from_pretrained(lowerCAmelCase_)

      for prompt in self.prompts:
        lowercase_								= tokenizer(lowerCAmelCase_ ,  return_tensors="""tf""").input_ids

        lowercase_								= model.generate(lowerCAmelCase_ ,  max_length=1_0)

        lowercase_								= tokenizer.batch_decode(lowerCAmelCase_ ,  skip_special_tokens=lowerCAmelCase_)
        predicted_outputs += generated_string

      self.assertListEqual(lowerCAmelCase_ ,  lowerCAmelCase_)

    def 					_UpperCAmelCase	(				self			:				Dict):







      """simple docstring"""







      lowercase_								= """facebook/opt-350m"""

      lowercase_								= GPTaTokenizer.from_pretrained(lowerCAmelCase_)
      lowercase_								= TFOPTForCausalLM.from_pretrained(lowerCAmelCase_)

      lowercase_								= """left"""

      # use different length sentences to test batching
      lowercase_								= [
          """Hello, my dog is a little""",
          """Today, I""",
      ]

      lowercase_								= tokenizer(lowerCAmelCase_ ,  return_tensors="""tf""" ,  padding=lowerCAmelCase_)
      lowercase_								= inputs["""input_ids"""]

      lowercase_								= model.generate(input_ids=lowerCAmelCase_ ,  attention_mask=inputs["""attention_mask"""])

      lowercase_								= tokenizer(sentences[0] ,  return_tensors="""tf""").input_ids
      lowercase_								= model.generate(input_ids=lowerCAmelCase_)

      lowercase_								= inputs_non_padded.shape[-1] - tf.math.reduce_sum(
          tf.cast(inputs["""attention_mask"""][-1] ,  tf.intaa))
      lowercase_								= tokenizer(sentences[1] ,  return_tensors="""tf""").input_ids
      lowercase_								= model.generate(input_ids=lowerCAmelCase_ ,  max_length=model.config.max_length - num_paddings)

      lowercase_								= tokenizer.batch_decode(lowerCAmelCase_ ,  skip_special_tokens=lowerCAmelCase_)
      lowercase_								= tokenizer.decode(output_non_padded[0] ,  skip_special_tokens=lowerCAmelCase_)
      lowercase_								= tokenizer.decode(output_padded[0] ,  skip_special_tokens=lowerCAmelCase_)

      lowercase_								= [
          """Hello, my dog is a little bit of a dork.\nI'm a little bit""",
          """Today, I was in the middle of a conversation with a friend about the""",
      ]
      self.assertListEqual(lowerCAmelCase_ ,  lowerCAmelCase_)
      self.assertListEqual(lowerCAmelCase_ ,  [non_padded_sentence, padded_sentence])




    def 					_UpperCAmelCase	(				self			:				int):







      """simple docstring"""







      lowercase_								= """facebook/opt-350m"""

      lowercase_								= [
          """Today is a beautiful day and I want to""",
          """In the city of San Francisco, the city""",
          """Paris is the capital of France and the capital""",
          """Computers and mobile phones have taken over the""",
      ]

      lowercase_								= []
      lowercase_								= GPTaTokenizer.from_pretrained(lowerCAmelCase_)
      lowercase_								= TFOPTForCausalLM.from_pretrained(lowerCAmelCase_)

      for prompt in self.prompts:
        lowercase_								= tokenizer(lowerCAmelCase_ ,  return_tensors="""tf""").input_ids

        lowercase_								= model.generate(lowerCAmelCase_ ,  max_length=1_0)

        lowercase_								= tokenizer.batch_decode(lowerCAmelCase_ ,  skip_special_tokens=lowerCAmelCase_)
        predicted_outputs += generated_string

      self.assertListEqual(lowerCAmelCase_ ,  lowerCAmelCase_)






"""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

