


"""simple docstring"""



import unittest

from transformers import BigBirdTokenizer, BigBirdTokenizerFast
from transformers.testing_utils import get_tests_dir, require_sentencepiece, require_tokenizers, require_torch, slow
from transformers.utils import cached_property

from ...test_tokenization_common import TokenizerTesterMixin


lowercase_							=							'▁'

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







@require_sentencepiece
@require_tokenizers
class    __a						( __snake_case					,			unittest.TestCase							):
				lowerCamelCase     :       List[Any]			       =BigBirdTokenizer
				lowerCamelCase     :       Optional[int]			       =BigBirdTokenizerFast
				lowerCamelCase     :       str			       =True
				lowerCamelCase     :       int			       =True







				def      lowerCamelCase_     ( self   ):

					'''simple docstring'''






					super().setUp()

					lowerCAmelCase_      							=			self.tokenizer_class(UpperCAmelCase      ,	keep_accents=UpperCAmelCase   )
					tokenizer.save_pretrained(self.tmpdirname   )







				def      lowerCamelCase_     ( self   ):

					'''simple docstring'''






					lowerCAmelCase_      							=			'''<s>'''
					lowerCAmelCase_      							=			1

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







				def      lowerCamelCase_     ( self   ):

					'''simple docstring'''






					lowerCAmelCase_      							=			list(self.get_tokenizer().get_vocab().keys()   )

					self.assertEqual(vocab_keys[0]      ,	'''<unk>'''   )
					self.assertEqual(vocab_keys[1]      ,	'''<s>'''   )
					self.assertEqual(vocab_keys[-1]      ,	'''[MASK]'''   )
					self.assertEqual(len(UpperCAmelCase   )      ,	1004   )







				def      lowerCamelCase_     ( self   ):

					'''simple docstring'''






					self.assertEqual(self.get_tokenizer().vocab_size      ,	1000   )







				def      lowerCamelCase_     ( self   ):

					'''simple docstring'''






					if not self.test_rust_tokenizer:
						return

					lowerCAmelCase_      							=			self.get_tokenizer()
					lowerCAmelCase_      							=			self.get_rust_tokenizer()

					lowerCAmelCase_      							=			'''I was born in 92000, and this is falsé.'''

					lowerCAmelCase_      							=			tokenizer.tokenize(UpperCAmelCase   )
					lowerCAmelCase_      							=			rust_tokenizer.tokenize(UpperCAmelCase   )
					self.assertListEqual(UpperCAmelCase      ,	UpperCAmelCase   )

					lowerCAmelCase_      							=			tokenizer.encode(UpperCAmelCase      ,	add_special_tokens=UpperCAmelCase   )
					lowerCAmelCase_      							=			rust_tokenizer.encode(UpperCAmelCase      ,	add_special_tokens=UpperCAmelCase   )
					self.assertListEqual(UpperCAmelCase      ,	UpperCAmelCase   )

					lowerCAmelCase_      							=			self.get_rust_tokenizer()
					lowerCAmelCase_      							=			tokenizer.encode(UpperCAmelCase   )
					lowerCAmelCase_      							=			rust_tokenizer.encode(UpperCAmelCase   )
					self.assertListEqual(UpperCAmelCase      ,	UpperCAmelCase   )







				def      lowerCamelCase_     ( self   ):

					'''simple docstring'''






					lowerCAmelCase_      							=			BigBirdTokenizer(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   )      ,	[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      ,	[8, 21, 84, 55, 24, 19, 7, 0, 602, 347, 347, 347, 3, 12, 66, 46, 72, 80, 6, 0, 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>''',
					        '''.''',
					    ]      ,	)







				@cached_property
				def      lowerCamelCase_     ( self   ):

					'''simple docstring'''






					return BigBirdTokenizer.from_pretrained('''google/bigbird-roberta-base'''   )







				@slow
				def      lowerCamelCase_     ( self   ):

					'''simple docstring'''






					lowerCAmelCase_      							=			'''Hello World!'''
					lowerCAmelCase_      							=			[65, 1_8536, 2260, 101, 66]

					self.assertListEqual(UpperCAmelCase      ,	self.big_tokenizer.encode(UpperCAmelCase   )   )







				@slow
				def      lowerCamelCase_     ( self   ):

					'''simple docstring'''






					lowerCAmelCase_      							=			(
					    '''This is a very long text with a lot of weird characters, such as: . , ~ ? ( ) " [ ] ! : - . Also we will'''
					    ''' add words that should not exsist and be tokenized to <unk>, such as saoneuhaoesuth'''
					)
					# fmt: off
					lowerCAmelCase_      							=			[65, 871, 419, 358, 946, 991, 2521, 452, 358, 1357, 387, 7751, 3536, 112, 985, 456, 126, 865, 938, 5400, 5734, 458, 1368, 467, 786, 2462, 5246, 1159, 633, 865, 4519, 457, 582, 852, 2557, 427, 916, 508, 405, 3_4324, 497, 391, 408, 1_1342, 1244, 385, 100, 938, 985, 456, 574, 362, 1_2597, 3200, 3129, 1172, 66]  # noqa: E231
					# fmt: on
					self.assertListEqual(UpperCAmelCase      ,	self.big_tokenizer.encode(UpperCAmelCase   )   )







				@require_torch
				@slow
				def      lowerCamelCase_     ( self   ):

					'''simple docstring'''






					import torch

					from transformers import BigBirdConfig, BigBirdModel

					# Build sequence
					lowerCAmelCase_      							=			list(self.big_tokenizer.get_vocab().keys()   )[:10]
					lowerCAmelCase_      							=			''' '''.join(UpperCAmelCase   )
					lowerCAmelCase_      							=			self.big_tokenizer.encode_plus(UpperCAmelCase      ,	return_tensors='''pt'''      ,	return_token_type_ids=UpperCAmelCase   )
					lowerCAmelCase_      							=			self.big_tokenizer.batch_encode_plus(
					    [sequence + ''' ''' + sequence]      ,	return_tensors='''pt'''      ,	return_token_type_ids=UpperCAmelCase   )

					lowerCAmelCase_      							=			BigBirdConfig(attention_type='''original_full'''   )
					lowerCAmelCase_      							=			BigBirdModel(UpperCAmelCase   )

					assert model.get_input_embeddings().weight.shape[0] >= self.big_tokenizer.vocab_size

					with torch.no_grad():
						model(**UpperCAmelCase   )
						model(**UpperCAmelCase   )







				@slow
				def      lowerCamelCase_     ( self   ):

					'''simple docstring'''






					lowerCAmelCase_      							=			BigBirdTokenizer.from_pretrained('''google/bigbird-roberta-base'''   )
					lowerCAmelCase_      							=			tokenizer.decode(tokenizer('''Paris is the [MASK].'''   ).input_ids   )

					self.assertTrue(decoded_text == '''[CLS] Paris is the[MASK].[SEP]'''   )







				@slow
				def      lowerCamelCase_     ( self   ):

					'''simple docstring'''






					lowerCAmelCase_      							=			{'''input_ids''': [[65, 3_9286, 458, 3_6335, 2001, 456, 1_3073, 1_3266, 455, 113, 7746, 1741, 1_1157, 391, 1_3073, 1_3266, 455, 113, 3967, 3_5412, 113, 4936, 109, 3870, 2377, 113, 3_0084, 4_5720, 458, 134, 1_7496, 112, 503, 1_1672, 113, 118, 112, 5665, 1_3347, 3_8687, 112, 1496, 3_1389, 112, 3268, 4_7264, 134, 962, 112, 1_6377, 8035, 2_3130, 430, 1_2169, 1_5518, 2_8592, 458, 146, 4_1697, 109, 391, 1_2169, 1_5518, 1_6689, 458, 146, 4_1358, 109, 452, 726, 4034, 111, 763, 3_5412, 5082, 388, 1903, 111, 9051, 391, 2870, 4_8918, 1900, 1123, 550, 998, 112, 9586, 1_5985, 455, 391, 410, 2_2955, 3_7636, 114, 66], [65, 448, 1_7496, 419, 3663, 385, 763, 113, 2_7533, 2870, 3283, 1_3043, 1639, 2_4713, 523, 656, 2_4013, 1_8550, 2521, 517, 2_7014, 2_1244, 420, 1212, 1465, 391, 927, 4833, 388, 578, 1_1786, 114, 66, 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], [65, 484, 2169, 7687, 2_1932, 1_8146, 726, 363, 1_7032, 3391, 114, 66, 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]], '''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, 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, 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=UpperCAmelCase      ,	model_name='''google/bigbird-roberta-base'''      ,	revision='''215c99f1600e06f83acce68422f2035b2b5c3510'''      ,	)



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