import unittest
from pathlib import Path
from shutil import copyfile

from transformers import SPIECE_UNDERLINE, is_sentencepiece_available
from transformers.models.speech_to_text import SpeechaTextTokenizer
from transformers.models.speech_to_text.tokenization_speech_to_text import VOCAB_FILES_NAMES, save_json
from transformers.testing_utils import get_tests_dir, require_sentencepiece, require_tokenizers, slow

from ...test_tokenization_common import TokenizerTesterMixin


_lowercase							:      str      							=get_tests_dir('''fixtures/test_sentencepiece.model''')

if is_sentencepiece_available():
    import sentencepiece as sp


_lowercase							:      int      							=5
_lowercase							:      Dict      							=1_0



@require_sentencepiece
@require_tokenizers
class  SCREAMING_SNAKE_CASE_ (					lowerCAmelCase_   ,						unittest.TestCase			):
  '''simple docstring'''


  lowercase		:     Any            =				SpeechaTextTokenizer
  lowercase		:     Dict            =				False
  lowercase		:     Optional[Any]            =				True







  def 		SCREAMING_SNAKE_CASE_		(				self    :					Dict     )    ->				int:
         super().setUp()

         A						:     Optional[int]         =sp.SentencePieceProcessor()
         spm_model.Load(SCREAMING_SNAKE_CASE__     )
         A						:     List[Any]         =['<s>', '<pad>', '</s>', '<unk>']

         vocab += [spm_model.IdToPiece(id_     ) for id_ in range(len(SCREAMING_SNAKE_CASE__     )     )]
         A						:     Optional[Any]         =dict(zip(SCREAMING_SNAKE_CASE__      ,	range(len(SCREAMING_SNAKE_CASE__     )     )     )     )

         A						:     List[str]         =Path(self.tmpdirname     )
         save_json(SCREAMING_SNAKE_CASE__      ,	save_dir / VOCAB_FILES_NAMES['vocab_file']     )
         if not (save_dir / VOCAB_FILES_NAMES["spm_file"]).exists():
                copyfile(SCREAMING_SNAKE_CASE__      ,	save_dir / VOCAB_FILES_NAMES['spm_file']     )

         A						:     Tuple         =SpeechaTextTokenizer.from_pretrained(self.tmpdirname     )
         tokenizer.save_pretrained(self.tmpdirname     )







  def 		SCREAMING_SNAKE_CASE_		(				self    :					Optional[int]     )    ->				int:
         A						:     Tuple         ='<pad>'
         A						:     int         =1

         self.assertEqual(self.get_tokenizer()._convert_token_to_id(SCREAMING_SNAKE_CASE__     )      ,	SCREAMING_SNAKE_CASE__     )
         self.assertEqual(self.get_tokenizer()._convert_id_to_token(SCREAMING_SNAKE_CASE__     )      ,	SCREAMING_SNAKE_CASE__     )







  def 		SCREAMING_SNAKE_CASE_		(				self    :					List[Any]     )    ->				Tuple:
         A						:     str         =list(self.get_tokenizer().get_vocab().keys()     )

         self.assertEqual(vocab_keys[0]      ,	'<s>'     )
         self.assertEqual(vocab_keys[1]      ,	'<pad>'     )
         self.assertEqual(vocab_keys[-1]      ,	'j'     )
         self.assertEqual(len(SCREAMING_SNAKE_CASE__     )      ,	10_01     )







  def 		SCREAMING_SNAKE_CASE_		(				self    :					Dict     )    ->				int:
         self.assertEqual(self.get_tokenizer().vocab_size      ,	10_01     )







  def 		SCREAMING_SNAKE_CASE_		(				self    :					Tuple     )    ->				Optional[Any]:
         A						:     Union[str, Any]         =SpeechaTextTokenizer.from_pretrained(self.tmpdirname     )

         A						:     int         =tokenizer.tokenize('This is a test'     )
         self.assertListEqual(SCREAMING_SNAKE_CASE__      ,	['▁This', '▁is', '▁a', '▁t', 'est']     )

         self.assertListEqual(
             tokenizer.convert_tokens_to_ids(SCREAMING_SNAKE_CASE__     )      ,	[2_89, 50, 14, 1_74, 3_86]      ,	)

         A						:     Dict         =tokenizer.tokenize('I was born in 92000, and this is falsé.'     )
         self.assertListEqual(
             SCREAMING_SNAKE_CASE__      ,	[SPIECE_UNDERLINE + 'I', SPIECE_UNDERLINE + 'was', SPIECE_UNDERLINE + 'b', 'or', 'n', SPIECE_UNDERLINE + 'in', SPIECE_UNDERLINE + '', '9', '2', '0', '0', '0', ',', SPIECE_UNDERLINE + 'and', SPIECE_UNDERLINE + 'this', SPIECE_UNDERLINE + 'is', SPIECE_UNDERLINE + 'f', 'al', 's', 'é', '.']      ,	)
         A						:     Tuple         =tokenizer.convert_tokens_to_ids(SCREAMING_SNAKE_CASE__     )
         self.assertListEqual(SCREAMING_SNAKE_CASE__      ,	[12, 25, 88, 59, 28, 23, 11, 4, 6_06, 3_51, 3_51, 3_51, 7, 16, 70, 50, 76, 84, 10, 4, 8]     )

         A						:     Tuple         =tokenizer.convert_ids_to_tokens(SCREAMING_SNAKE_CASE__     )
         self.assertListEqual(
             SCREAMING_SNAKE_CASE__      ,	[SPIECE_UNDERLINE + 'I', SPIECE_UNDERLINE + 'was', SPIECE_UNDERLINE + 'b', 'or', 'n', SPIECE_UNDERLINE + 'in', SPIECE_UNDERLINE + '', '<unk>', '2', '0', '0', '0', ',', SPIECE_UNDERLINE + 'and', SPIECE_UNDERLINE + 'this', SPIECE_UNDERLINE + 'is', SPIECE_UNDERLINE + 'f', 'al', 's', '<unk>', '.']      ,	)







  @slow
  def 		SCREAMING_SNAKE_CASE_		(				self    :					Tuple     )    ->				str:
         # fmt: off
         A						:     str         ={'input_ids': [[37_91, 7_97, 31, 11, 64, 7_97, 31, 24_29, 4_33, 12, 11_76, 12, 20, 7_86, 9_15, 1_42, 24_13, 2_40, 37, 32_38, 7_97, 31, 11, 35, 93, 9_15, 1_42, 24_13, 2_40, 37, 55_40, 5_67, 12_76, 93, 37, 6_10, 40, 62, 4_55, 6_57, 10_42, 1_23, 7_80, 1_77, 37, 3_09, 2_41, 12_98, 5_14, 20, 2_92, 27_37, 1_14, 24_69, 2_41, 85, 64, 3_02, 5_48, 5_28, 4_23, 4, 5_09, 4_06, 4_23, 37, 6_01, 4, 7_77, 3_02, 5_48, 5_28, 4_23, 2_84, 4, 33_88, 5_11, 4_59, 4, 35_55, 40, 3_21, 3_02, 7_05, 4, 33_88, 5_11, 5_83, 3_26, 5, 5, 5, 62, 33_10, 5_60, 1_77, 26_80, 2_17, 15_08, 32, 31, 8_53, 4_18, 64, 5_83, 5_11, 16_05, 62, 35, 93, 5_60, 1_77, 26_80, 2_17, 15_08, 15_21, 64, 5_83, 5_11, 5_19, 62, 20, 15_15, 7_64, 20, 1_49, 2_61, 56_25, 79_72, 20, 55_40, 5_67, 12_76, 93, 39_25, 16_75, 11, 15, 8_02, 79_72, 5_76, 2_17, 15_08, 11, 35, 93, 12_53, 24_41, 15, 2_89, 6_52, 31, 4_16, 3_21, 38_42, 1_15, 40, 9_11, 8, 4_76, 6_19, 4, 3_80, 1_42, 4_23, 3_35, 2_40, 35, 93, 2_64, 8, 11, 3_35, 5_69, 4_20, 1_63, 5, 2], [2_60, 5_48, 5_28, 4_23, 20, 4_51, 20, 26_81, 11_53, 34_34, 20, 55_40, 37, 5_67, 1_26, 12_53, 24_41, 33_76, 4_49, 2_10, 4_31, 15_63, 1_77, 7_67, 55_40, 11, 12_03, 4_72, 11, 29_53, 6_85, 2_85, 3_64, 7_06, 11_53, 20, 67_99, 20, 28_69, 20, 44_64, 1_26, 40, 24_29, 20, 10_40, 8_66, 26_64, 4_18, 20, 3_18, 20, 17_26, 1_86, 20, 2_65, 5_22, 35, 93, 21_91, 46_34, 20, 10_40, 12, 67_99, 15, 2_28, 23_56, 1_42, 31, 11, 5, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [25_75, 26_66, 6_84, 15_82, 11_76, 12, 6_27, 1_49, 6_19, 20, 49_02, 5_63, 11, 20, 1_49, 2_61, 34_20, 23_56, 1_74, 1_42, 47_14, 1_31, 5, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]], 'attention_mask': [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]}  # noqa: E501
         # fmt: on

         self.tokenizer_integration_test_util(
             expected_encoding=SCREAMING_SNAKE_CASE__      ,	model_name='facebook/s2t-small-mustc-en-de-st'      ,	revision='a14f04cf0776c02f62a8cb800cf7909e15ea23ad'      ,	)







@require_sentencepiece
class  SCREAMING_SNAKE_CASE_ (					unittest.TestCase			):
  '''simple docstring'''


  lowercase		:     Any            =				"valhalla/s2t_mustc_multilinguial_medium"

  lowercase		:     Tuple            =				"C'est trop cool"
  lowercase		:     Union[str, Any]            =				"Esto es genial"







  @classmethod
  def 		SCREAMING_SNAKE_CASE_		(				cls    :					Any     )    ->				List[Any]:
         A						:     SpeechaTextTokenizer         =SpeechaTextTokenizer.from_pretrained(cls.checkpoint_name     )
         return cls







  def 		SCREAMING_SNAKE_CASE_		(				self    :					List[str]     )    ->				Optional[int]:
         self.assertEqual(self.tokenizer.lang_code_to_id['pt']      ,	4     )
         self.assertEqual(self.tokenizer.lang_code_to_id['ru']      ,	6     )
         self.assertEqual(self.tokenizer.lang_code_to_id['it']      ,	9     )
         self.assertEqual(self.tokenizer.lang_code_to_id['de']      ,	11     )







  def 		SCREAMING_SNAKE_CASE_		(				self    :					List[str]     )    ->				int:
         self.assertEqual(self.tokenizer.vocab_size      ,	1_00_00     )







  def 		SCREAMING_SNAKE_CASE_		(				self    :					List[Any]     )    ->				str:
         self.assertIn(SCREAMING_SNAKE_CASE__      ,	self.tokenizer.all_special_ids     )
         A						:     Any         =[ES_CODE, 4, 16_01, 47, 76_47, 2]
         A						:     str         =self.tokenizer.decode(SCREAMING_SNAKE_CASE__      ,	skip_special_tokens=SCREAMING_SNAKE_CASE__     )
         A						:     Dict         =self.tokenizer.decode(generated_ids[1:]      ,	skip_special_tokens=SCREAMING_SNAKE_CASE__     )
         self.assertEqual(SCREAMING_SNAKE_CASE__      ,	SCREAMING_SNAKE_CASE__     )
         self.assertNotIn(self.tokenizer.eos_token      ,	SCREAMING_SNAKE_CASE__     )







  def 		SCREAMING_SNAKE_CASE_		(				self    :					Union[str, Any]     )    ->				Optional[Any]:
         A						:     Tuple         ='fr'
         A						:     Tuple         =self.tokenizer(self.french_text     ).input_ids
         self.assertEqual(encoded[0]      ,	SCREAMING_SNAKE_CASE__     )
         self.assertEqual(encoded[-1]      ,	self.tokenizer.eos_token_id     )







  def 		SCREAMING_SNAKE_CASE_		(				self    :					Optional[Any]     )    ->				str:
         A						:     Any         ='fr'
         self.assertListEqual(self.tokenizer.prefix_tokens      ,	[FR_CODE]     )

         A						:     str         ='es'
         self.assertListEqual(self.tokenizer.prefix_tokens      ,	[ES_CODE]     )







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_ (					lowerCAmelCase_   ,						lowerCAmelCase_   ,						lowerCAmelCase_			):
  '''simple docstring'''


  lowercase		:     Tuple            =				[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 = 5_02_57      ,	SCREAMING_SNAKE_CASE__    :					int = 10_24      ,	SCREAMING_SNAKE_CASE__    :					int = 7_68      ,	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      ,	)    ->				List[str]:
         super().__init__()

         A						:     str         =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.'     )

         A						:     List[Any]         =prefix_inner_dim
         A						:     Dict         =prefix_hidden_dim

         A						:     List[str]         =(
             nn.Linear(self.prefix_inner_dim      ,	self.prefix_hidden_dim     )
             if self.prefix_hidden_dim is not None
             else nn.Identity()
         )
         A						:     Optional[int]         =(
             nn.Linear(self.prefix_hidden_dim      ,	SCREAMING_SNAKE_CASE__     ) if self.prefix_hidden_dim is not None else nn.Identity()
         )

         A						:     Dict         =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__      ,	)
         A						:     Dict         =GPTaLMHeadModel(SCREAMING_SNAKE_CASE__     )







  def 		SCREAMING_SNAKE_CASE_		(				self    :					List[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      ,	)    ->				Optional[Any]:
         A						:     str         =self.transformer.transformer.wte(SCREAMING_SNAKE_CASE__     )
         A						:     Any         =self.encode_prefix(SCREAMING_SNAKE_CASE__     )
         A						:     Optional[Any]         =self.decode_prefix(SCREAMING_SNAKE_CASE__     )
         A						:     Optional[int]         =torch.cat((prefix_embeds, embedding_text)      ,	dim=1     )

         if labels is not None:
                A						:     int         =self.get_dummy_token(input_ids.shape[0]      ,	input_ids.device     )
                A						:     Optional[int]         =torch.cat((dummy_token, input_ids)      ,	dim=1     )
         A						:     Dict         =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 		SCREAMING_SNAKE_CASE_		(				self    :					List[str]      ,	SCREAMING_SNAKE_CASE__    :					int      ,	SCREAMING_SNAKE_CASE__    :					torch.device     )    ->				torch.Tensor:
         return torch.zeros(SCREAMING_SNAKE_CASE__      ,	self.prefix_length      ,	dtype=torch.intaa      ,	device=SCREAMING_SNAKE_CASE__     )







  def 		SCREAMING_SNAKE_CASE_		(				self    :					Union[str, Any]      ,	SCREAMING_SNAKE_CASE__    :					Optional[int]     )    ->				List[str]:
         return self.encode_prefix(SCREAMING_SNAKE_CASE__     )







  @torch.no_grad()
  def 		SCREAMING_SNAKE_CASE_		(				self    :					List[Any]      ,	SCREAMING_SNAKE_CASE__    :					List[str]      ,	SCREAMING_SNAKE_CASE__    :					Union[str, Any]      ,	SCREAMING_SNAKE_CASE__    :					Any     )    ->				Dict:
         A						:     Dict         =torch.split(SCREAMING_SNAKE_CASE__      ,	1      ,	dim=0     )
         A						:     int         =[]
         A						:     Optional[int]         =[]
         for feature in features:
                A						:     int         =self.decode_prefix(feature.to(SCREAMING_SNAKE_CASE__     )     )  # back to the clip feature
                # Only support beam search for now
                A  ,  A						:     Dict         =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]     )
         A						:     str         =torch.stack(SCREAMING_SNAKE_CASE__     )
         A						:     int         =torch.stack(SCREAMING_SNAKE_CASE__     )
         return generated_tokens, generated_seq_lengths







  @torch.no_grad()
  def 		SCREAMING_SNAKE_CASE_		(				self    :					Optional[int]      ,	SCREAMING_SNAKE_CASE__    :					Union[str, Any]=None      ,	SCREAMING_SNAKE_CASE__    :					List[Any]=None      ,	SCREAMING_SNAKE_CASE__    :					Tuple=None      ,	SCREAMING_SNAKE_CASE__    :					int = 5      ,	SCREAMING_SNAKE_CASE__    :					int = 67      ,	SCREAMING_SNAKE_CASE__    :					float = 1.0      ,	SCREAMING_SNAKE_CASE__    :					Optional[int] = None      ,	)    ->				Dict:
         A						:     Dict         =eos_token_id
         A						:     str         =None
         A						:     List[Any]         =None
         A						:     List[Any]         =torch.ones(SCREAMING_SNAKE_CASE__      ,	device=SCREAMING_SNAKE_CASE__      ,	dtype=torch.int     )
         A						:     str         =torch.zeros(SCREAMING_SNAKE_CASE__      ,	device=SCREAMING_SNAKE_CASE__      ,	dtype=torch.bool     )

         if input_embeds is not None:
                A						:     Any         =input_embeds
         else:
                A						:     List[Any]         =self.transformer.transformer.wte(SCREAMING_SNAKE_CASE__     )

         for i in range(SCREAMING_SNAKE_CASE__     ):
                A						:     Any         =self.transformer(inputs_embeds=SCREAMING_SNAKE_CASE__     )
                A						:     str         =outputs.logits
                A						:     Union[str, Any]         =logits[:, -1, :] / (temperature if temperature > 0 else 1.0)
                A						:     List[str]         =logits.softmax(-1     ).log()

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

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

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






