


"""simple docstring"""



import shutil
import tempfile
import unittest

from transformers import (
    SPIECE_UNDERLINE,
    AddedToken,
    BatchEncoding,
    NllbTokenizer,
    NllbTokenizerFast,
    is_torch_available,
)
from transformers.testing_utils import (
    get_tests_dir,
    nested_simplify,
    require_sentencepiece,
    require_tokenizers,
    require_torch,
)

from ...test_tokenization_common import TokenizerTesterMixin


lowercase_							=							get_tests_dir('fixtures/test_sentencepiece.model')


if is_torch_available():
		from transformers.models.mam_aaa.modeling_mam_aaa import shift_tokens_right

lowercase_							=							25_60_47
lowercase_							=							25_61_45







@require_sentencepiece
@require_tokenizers
class    __a						( __snake_case					,			unittest.TestCase							):
				lowerCamelCase     :       Any			       =NllbTokenizer
				lowerCamelCase     :       Any			       =NllbTokenizerFast
				lowerCamelCase     :       Tuple			       =True
				lowerCamelCase     :       Dict			       =True
				lowerCamelCase     :       Union[str, Any]			       ={}







				def      lowerCamelCase_     ( self   ):

					'''simple docstring'''






					super().setUp()

					# We have a SentencePiece fixture for testing
					lowerCAmelCase_      							=			NllbTokenizer(UpperCAmelCase      ,	keep_accents=UpperCAmelCase   )
					tokenizer.save_pretrained(self.tmpdirname   )







				def      lowerCamelCase_     ( self   ):

					'''simple docstring'''






					lowerCAmelCase_      							=			NllbTokenizer(UpperCAmelCase      ,	keep_accents=UpperCAmelCase   )

					lowerCAmelCase_      							=			tokenizer.tokenize('''This is a test'''   )
					self.assertListEqual(UpperCAmelCase      ,	['''▁This''', '''▁is''', '''▁a''', '''▁t''', '''est''']   )

					self.assertListEqual(
					    tokenizer.convert_tokens_to_ids(UpperCAmelCase   )      ,	[value + tokenizer.fairseq_offset for value in [285, 46, 10, 170, 382]]      ,	)

					lowerCAmelCase_      							=			tokenizer.tokenize('''I was born in 92000, and this is falsé.'''   )
					self.assertListEqual(
					    UpperCAmelCase      ,	[
					        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''',
					        '''é''',
					        '''.''',
					    ]      ,	)
					lowerCAmelCase_      							=			tokenizer.convert_tokens_to_ids(UpperCAmelCase   )
					self.assertListEqual(
					    UpperCAmelCase      ,	[
					        value + tokenizer.fairseq_offset
					        for value in [8, 21, 84, 55, 24, 19, 7, 2, 602, 347, 347, 347, 3, 12, 66, 46, 72, 80, 6, 2, 4]
					    ]      ,	)

					lowerCAmelCase_      							=			tokenizer.convert_ids_to_tokens(UpperCAmelCase   )
					self.assertListEqual(
					    UpperCAmelCase      ,	[
					        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>''',
					        '''.''',
					    ]      ,	)







				def      lowerCamelCase_     ( self   ):

					'''simple docstring'''






					lowerCAmelCase_      							=			(self.rust_tokenizer_class, '''hf-internal-testing/tiny-random-nllb''', {})
					for tokenizer, pretrained_name, kwargs in self.tokenizers_list:
						with self.subTest(F"""{tokenizer.__class__.__name__} ({pretrained_name})"""   ):
							lowerCAmelCase_      							=			self.rust_tokenizer_class.from_pretrained(UpperCAmelCase      ,	**UpperCAmelCase   )
							lowerCAmelCase_      							=			self.tokenizer_class.from_pretrained(UpperCAmelCase      ,	**UpperCAmelCase   )

							lowerCAmelCase_      							=			tempfile.mkdtemp()

							lowerCAmelCase_      							=			tokenizer_r.save_pretrained(UpperCAmelCase   )
							lowerCAmelCase_      							=			tokenizer_p.save_pretrained(UpperCAmelCase   )

							# Checks it save with the same files + the tokenizer.json file for the fast one
							self.assertTrue(any('''tokenizer.json''' in f for f in tokenizer_r_files   )   )
							lowerCAmelCase_      							=			tuple(f for f in tokenizer_r_files if '''tokenizer.json''' not in f   )
							self.assertSequenceEqual(UpperCAmelCase      ,	UpperCAmelCase   )

							# Checks everything loads correctly in the same way
							lowerCAmelCase_      							=			tokenizer_r.from_pretrained(UpperCAmelCase   )
							lowerCAmelCase_      							=			tokenizer_p.from_pretrained(UpperCAmelCase   )

							# Check special tokens are set accordingly on Rust and Python
							for key in tokenizer_pp.special_tokens_map:
								self.assertTrue(hasattr(UpperCAmelCase      ,	UpperCAmelCase   )   )

							shutil.rmtree(UpperCAmelCase   )

							# Save tokenizer rust, legacy_format=True
							lowerCAmelCase_      							=			tempfile.mkdtemp()

							lowerCAmelCase_      							=			tokenizer_r.save_pretrained(UpperCAmelCase      ,	legacy_format=UpperCAmelCase   )
							lowerCAmelCase_      							=			tokenizer_p.save_pretrained(UpperCAmelCase   )

							# Checks it save with the same files
							self.assertSequenceEqual(UpperCAmelCase      ,	UpperCAmelCase   )

							# Checks everything loads correctly in the same way
							lowerCAmelCase_      							=			tokenizer_r.from_pretrained(UpperCAmelCase   )
							lowerCAmelCase_      							=			tokenizer_p.from_pretrained(UpperCAmelCase   )

							# Check special tokens are set accordingly on Rust and Python
							for key in tokenizer_pp.special_tokens_map:
								self.assertTrue(hasattr(UpperCAmelCase      ,	UpperCAmelCase   )   )

							shutil.rmtree(UpperCAmelCase   )

							# Save tokenizer rust, legacy_format=False
							lowerCAmelCase_      							=			tempfile.mkdtemp()

							lowerCAmelCase_      							=			tokenizer_r.save_pretrained(UpperCAmelCase      ,	legacy_format=UpperCAmelCase   )
							lowerCAmelCase_      							=			tokenizer_p.save_pretrained(UpperCAmelCase   )

							# Checks it saved the tokenizer.json file
							self.assertTrue(any('''tokenizer.json''' in f for f in tokenizer_r_files   )   )

							# Checks everything loads correctly in the same way
							lowerCAmelCase_      							=			tokenizer_r.from_pretrained(UpperCAmelCase   )
							lowerCAmelCase_      							=			tokenizer_p.from_pretrained(UpperCAmelCase   )

							# Check special tokens are set accordingly on Rust and Python
							for key in tokenizer_pp.special_tokens_map:
								self.assertTrue(hasattr(UpperCAmelCase      ,	UpperCAmelCase   )   )

							shutil.rmtree(UpperCAmelCase   )







				@require_torch
				def      lowerCamelCase_     ( self   ):

					'''simple docstring'''






					if not self.test_seqaseq:
						return

					lowerCAmelCase_      							=			self.get_tokenizers()
					for tokenizer in tokenizers:
						with self.subTest(F"""{tokenizer.__class__.__name__}"""   ):
							# Longer text that will definitely require truncation.
							lowerCAmelCase_      							=			[
							    ''' UN Chief Says There Is No Military Solution in Syria''',
							    ''' Secretary-General Ban Ki-moon says his response to Russia\'s stepped up military support for'''
							    ''' Syria is that \'there is no military solution\' to the nearly five-year conflict and more weapons'''
							    ''' will only worsen the violence and misery for millions of people.''',
							]
							lowerCAmelCase_      							=			[
							    '''Şeful ONU declară că nu există o soluţie militară în Siria''',
							    '''Secretarul General Ban Ki-moon declară că răspunsul său la intensificarea sprijinului militar al'''
							    ''' Rusiei pentru Siria este că "nu există o soluţie militară" la conflictul de aproape cinci ani şi'''
							    ''' că noi arme nu vor face decât să înrăutăţească violenţele şi mizeria pentru milioane de oameni.''',
							]
							try:
								lowerCAmelCase_      							=			tokenizer.prepare_seqaseq_batch(
								    src_texts=UpperCAmelCase      ,	tgt_texts=UpperCAmelCase      ,	max_length=3      ,	max_target_length=10      ,	return_tensors='''pt'''      ,	src_lang='''eng_Latn'''      ,	tgt_lang='''ron_Latn'''      ,	)
							except NotImplementedError:
								return
							self.assertEqual(batch.input_ids.shape[1]      ,	3   )
							self.assertEqual(batch.labels.shape[1]      ,	10   )
							# max_target_length will default to max_length if not specified
							lowerCAmelCase_      							=			tokenizer.prepare_seqaseq_batch(
							    UpperCAmelCase      ,	tgt_texts=UpperCAmelCase      ,	max_length=3      ,	return_tensors='''pt'''   )
							self.assertEqual(batch.input_ids.shape[1]      ,	3   )
							self.assertEqual(batch.labels.shape[1]      ,	3   )

							lowerCAmelCase_      							=			tokenizer.prepare_seqaseq_batch(
							    src_texts=UpperCAmelCase      ,	max_length=3      ,	max_target_length=10      ,	return_tensors='''pt'''   )
							self.assertEqual(batch_encoder_only.input_ids.shape[1]      ,	3   )
							self.assertEqual(batch_encoder_only.attention_mask.shape[1]      ,	3   )
							self.assertNotIn('''decoder_input_ids'''      ,	UpperCAmelCase   )







				@unittest.skip('''Unfortunately way too slow to build a BPE with SentencePiece.'''   )
				def      lowerCamelCase_     ( self   ):

					'''simple docstring'''






					pass







				def      lowerCamelCase_     ( self   ):

					'''simple docstring'''






					for tokenizer, pretrained_name, kwargs in self.tokenizers_list:
						with self.subTest(F"""{tokenizer.__class__.__name__} ({pretrained_name})"""   ):
							lowerCAmelCase_      							=			[AddedToken('''<special>'''      ,	lstrip=UpperCAmelCase   )]

							lowerCAmelCase_      							=			self.rust_tokenizer_class.from_pretrained(
							    UpperCAmelCase      ,	additional_special_tokens=UpperCAmelCase      ,	**UpperCAmelCase   )
							lowerCAmelCase_      							=			tokenizer_r.encode('''Hey this is a <special> token'''   )

							lowerCAmelCase_      							=			tokenizer_r.encode('''<special>'''      ,	add_special_tokens=UpperCAmelCase   )[0]

							self.assertTrue(special_token_id in r_output   )

							if self.test_slow_tokenizer:
								lowerCAmelCase_      							=			self.rust_tokenizer_class.from_pretrained(
								    UpperCAmelCase      ,	additional_special_tokens=UpperCAmelCase      ,	**UpperCAmelCase      ,	)
								lowerCAmelCase_      							=			self.tokenizer_class.from_pretrained(
								    UpperCAmelCase      ,	additional_special_tokens=UpperCAmelCase      ,	**UpperCAmelCase   )

								lowerCAmelCase_      							=			tokenizer_p.encode('''Hey this is a <special> token'''   )

								lowerCAmelCase_      							=			tokenizer_cr.encode('''Hey this is a <special> token'''   )

								self.assertEqual(UpperCAmelCase      ,	UpperCAmelCase   )
								self.assertEqual(UpperCAmelCase      ,	UpperCAmelCase   )
								self.assertTrue(special_token_id in p_output   )
								self.assertTrue(special_token_id in cr_output   )







@require_torch
@require_sentencepiece
@require_tokenizers
class    __a						( unittest.TestCase							):
				lowerCamelCase     :       List[str]			       ='facebook/nllb-200-distilled-600M'
				lowerCamelCase     :       Any			       =[
				    ' UN Chief Says There Is No Military Solution in Syria',
				    ' Secretary-General Ban Ki-moon says his response to Russia\'s stepped up military support for Syria is that "there is no military solution" to the nearly five-year conflict and more weapons will only worsen the violence and misery for millions of people.',
				]
				lowerCamelCase     :       List[str]			       =[
				    'Şeful ONU declară că nu există o soluţie militară în Siria',
				    'Secretarul General Ban Ki-moon declară că răspunsul său la intensificarea sprijinului militar al Rusiei'
				    ' pentru Siria este că "nu există o soluţie militară" la conflictul de aproape cinci ani şi că noi arme nu vor'
				    ' face decât să înrăutăţească violenţele şi mizeria pentru milioane de oameni.',
				]
				lowerCamelCase     :       Any			       =[
				    25_6047,
				    1_6297,
				    13_4408,
				    8165,
				    24_8066,
				    1_4734,
				    950,
				    1135,
				    10_5721,
				    3573,
				    83,
				    2_7352,
				    108,
				    4_9486,
				    2,
				]







				@classmethod
				def      lowerCamelCase_     ( cls   ):

					'''simple docstring'''






					lowerCAmelCase_      							=			NllbTokenizer.from_pretrained(
					    cls.checkpoint_name      ,	src_lang='''eng_Latn'''      ,	tgt_lang='''ron_Latn'''   )
					lowerCAmelCase_      							=			1
					return cls







				def      lowerCamelCase_     ( self   ):

					'''simple docstring'''






					self.assertEqual(self.tokenizer.fairseq_tokens_to_ids['''ace_Arab''']      ,	25_6001   )
					self.assertEqual(self.tokenizer.fairseq_tokens_to_ids['''ace_Latn''']      ,	25_6002   )
					self.assertEqual(self.tokenizer.fairseq_tokens_to_ids['''fra_Latn''']      ,	25_6057   )







				def      lowerCamelCase_     ( self   ):

					'''simple docstring'''






					lowerCAmelCase_      							=			self.tokenizer.batch_encode_plus(self.src_text   ).input_ids[0]
					self.assertListEqual(self.expected_src_tokens      ,	UpperCAmelCase   )







				def      lowerCamelCase_     ( self   ):

					'''simple docstring'''






					self.assertIn(UpperCAmelCase      ,	self.tokenizer.all_special_ids   )
					# fmt: off
					lowerCAmelCase_      							=			[RO_CODE, 4254, 9_8068, 11_2923, 3_9072, 3909, 713, 10_2767, 26, 1_7314, 3_5642, 1_4683, 3_3118, 2022, 6_6987, 2, 25_6047]
					# fmt: on

					lowerCAmelCase_      							=			self.tokenizer.decode(UpperCAmelCase      ,	skip_special_tokens=UpperCAmelCase   )
					lowerCAmelCase_      							=			self.tokenizer.decode(generated_ids[1:]      ,	skip_special_tokens=UpperCAmelCase   )
					self.assertEqual(UpperCAmelCase      ,	UpperCAmelCase   )
					self.assertNotIn(self.tokenizer.eos_token      ,	UpperCAmelCase   )







				def      lowerCamelCase_     ( self   ):

					'''simple docstring'''






					lowerCAmelCase_      							=			['''this is gunna be a long sentence ''' * 20]
					assert isinstance(src_text[0]      ,	UpperCAmelCase   )
					lowerCAmelCase_      							=			10
					lowerCAmelCase_      							=			self.tokenizer(UpperCAmelCase      ,	max_length=UpperCAmelCase      ,	truncation=UpperCAmelCase   ).input_ids[0]
					self.assertEqual(ids[-1]      ,	2   )
					self.assertEqual(ids[0]      ,	UpperCAmelCase   )
					self.assertEqual(len(UpperCAmelCase   )      ,	UpperCAmelCase   )







				def      lowerCamelCase_     ( self   ):

					'''simple docstring'''






					self.assertListEqual(self.tokenizer.convert_tokens_to_ids(['''<mask>''', '''ar_AR''']   )      ,	[25_6203, 3]   )







				def      lowerCamelCase_     ( self   ):

					'''simple docstring'''






					lowerCAmelCase_      							=			tempfile.mkdtemp()
					lowerCAmelCase_      							=			self.tokenizer.fairseq_tokens_to_ids
					self.tokenizer.save_pretrained(UpperCAmelCase   )
					lowerCAmelCase_      							=			NllbTokenizer.from_pretrained(UpperCAmelCase   )
					self.assertDictEqual(new_tok.fairseq_tokens_to_ids      ,	UpperCAmelCase   )







				@require_torch
				def      lowerCamelCase_     ( self   ):

					'''simple docstring'''






					lowerCAmelCase_      							=			self.tokenizer(
					    self.src_text      ,	text_target=self.tgt_text      ,	padding=UpperCAmelCase      ,	truncation=UpperCAmelCase      ,	max_length=len(self.expected_src_tokens   )      ,	return_tensors='''pt'''      ,	)
					lowerCAmelCase_      							=			shift_tokens_right(
					    batch['''labels''']      ,	self.tokenizer.pad_token_id      ,	self.tokenizer.lang_code_to_id['''ron_Latn''']   )

					self.assertIsInstance(UpperCAmelCase      ,	UpperCAmelCase   )

					self.assertEqual((2, 15)      ,	batch.input_ids.shape   )
					self.assertEqual((2, 15)      ,	batch.attention_mask.shape   )
					lowerCAmelCase_      							=			batch.input_ids.tolist()[0]
					self.assertListEqual(self.expected_src_tokens      ,	UpperCAmelCase   )
					self.assertEqual(UpperCAmelCase      ,	batch.decoder_input_ids[0, 0]   )  # EOS
					# Test that special tokens are reset
					self.assertEqual(self.tokenizer.prefix_tokens      ,	[EN_CODE]   )
					self.assertEqual(self.tokenizer.suffix_tokens      ,	[self.tokenizer.eos_token_id]   )







				def      lowerCamelCase_     ( self   ):

					'''simple docstring'''






					lowerCAmelCase_      							=			self.tokenizer(self.src_text      ,	padding=UpperCAmelCase      ,	truncation=UpperCAmelCase      ,	max_length=3      ,	return_tensors='''pt'''   )
					lowerCAmelCase_      							=			self.tokenizer(
					    text_target=self.tgt_text      ,	padding=UpperCAmelCase      ,	truncation=UpperCAmelCase      ,	max_length=10      ,	return_tensors='''pt'''   )
					lowerCAmelCase_      							=			targets['''input_ids''']
					lowerCAmelCase_      							=			shift_tokens_right(
					    UpperCAmelCase      ,	self.tokenizer.pad_token_id      ,	decoder_start_token_id=self.tokenizer.lang_code_to_id[self.tokenizer.tgt_lang]      ,	)

					self.assertEqual(batch.input_ids.shape[1]      ,	3   )
					self.assertEqual(batch.decoder_input_ids.shape[1]      ,	10   )







				@require_torch
				def      lowerCamelCase_     ( self   ):

					'''simple docstring'''






					lowerCAmelCase_      							=			self.tokenizer._build_translation_inputs(
					    '''A test'''      ,	return_tensors='''pt'''      ,	src_lang='''eng_Latn'''      ,	tgt_lang='''fra_Latn'''   )

					self.assertEqual(
					    nested_simplify(UpperCAmelCase   )      ,	{
					        # A, test, EOS, en_XX
					        '''input_ids''': [[25_6047, 70, 7356, 2]],
					        '''attention_mask''': [[1, 1, 1, 1]],
					        # ar_AR
					        '''forced_bos_token_id''': 25_6057,
					    }      ,	)







				@require_torch
				def      lowerCamelCase_     ( self   ):

					'''simple docstring'''






					lowerCAmelCase_      							=			True
					lowerCAmelCase_      							=			self.tokenizer(
					    '''UN Chief says there is no military solution in Syria'''      ,	src_lang='''eng_Latn'''      ,	tgt_lang='''fra_Latn'''   )
					self.assertEqual(
					    inputs.input_ids      ,	[1_6297, 13_4408, 2_5653, 6370, 248, 254, 10_3929, 9_4995, 108, 4_9486, 2, 25_6047]   )

					lowerCAmelCase_      							=			False
					lowerCAmelCase_      							=			self.tokenizer(
					    '''UN Chief says there is no military solution in Syria'''      ,	src_lang='''eng_Latn'''      ,	tgt_lang='''fra_Latn'''   )
					self.assertEqual(
					    inputs.input_ids      ,	[25_6047, 1_6297, 13_4408, 2_5653, 6370, 248, 254, 10_3929, 9_4995, 108, 4_9486, 2]   )



"""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    __a						( __snake_case					,			__snake_case					,			__snake_case							):
				lowerCamelCase     :       Optional[int]			       =[R'h\.\d+\.attn\.bias', R'h\.\d+\.attn\.masked_bias']







				@register_to_config
				def __init__( self      ,	UpperCAmelCase      ,	UpperCAmelCase      ,	UpperCAmelCase = None      ,	UpperCAmelCase = 5_0257      ,	UpperCAmelCase = 1024      ,	UpperCAmelCase = 768      ,	UpperCAmelCase = 12      ,	UpperCAmelCase = 12      ,	UpperCAmelCase = None      ,	UpperCAmelCase = "gelu_new"      ,	UpperCAmelCase = 0.1      ,	UpperCAmelCase = 0.1      ,	UpperCAmelCase = 0.1      ,	UpperCAmelCase = 1E-5      ,	UpperCAmelCase = 0.0_2      ,	UpperCAmelCase = True      ,	UpperCAmelCase = True      ,	UpperCAmelCase = False      ,	UpperCAmelCase = False      ,	):

					'''simple docstring'''






					super().__init__()

					lowerCAmelCase_      							=			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."""   )

					lowerCAmelCase_      							=			prefix_inner_dim
					lowerCAmelCase_      							=			prefix_hidden_dim

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

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







				def      lowerCamelCase_     ( self      ,	UpperCAmelCase      ,	UpperCAmelCase      ,	UpperCAmelCase = None      ,	UpperCAmelCase = None      ,	):

					'''simple docstring'''






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

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







				def      lowerCamelCase_     ( self      ,	UpperCAmelCase      ,	UpperCAmelCase   ):

					'''simple docstring'''






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







				def      lowerCamelCase_     ( self      ,	UpperCAmelCase   ):

					'''simple docstring'''






					return self.encode_prefix(UpperCAmelCase   )







				@torch.no_grad()
				def      lowerCamelCase_     ( self      ,	UpperCAmelCase      ,	UpperCAmelCase      ,	UpperCAmelCase   ):

					'''simple docstring'''






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







				@torch.no_grad()
				def      lowerCamelCase_     ( self      ,	UpperCAmelCase=None      ,	UpperCAmelCase=None      ,	UpperCAmelCase=None      ,	UpperCAmelCase = 5      ,	UpperCAmelCase = 67      ,	UpperCAmelCase = 1.0      ,	UpperCAmelCase = None      ,	):

					'''simple docstring'''






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

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

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

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

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

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