
import unittest

from transformers import BertGenerationConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
				import torch

				from transformers import BertGenerationDecoder, BertGenerationEncoder




class 				_lowerCamelCase :





					"""simple docstring"""







					def __init__(							self						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE=1_3						,						__SCREAMING_SNAKE_CASE=7						,						__SCREAMING_SNAKE_CASE=True						,						__SCREAMING_SNAKE_CASE=True						,						__SCREAMING_SNAKE_CASE=9_9						,						__SCREAMING_SNAKE_CASE=3_2						,						__SCREAMING_SNAKE_CASE=5						,						__SCREAMING_SNAKE_CASE=4						,						__SCREAMING_SNAKE_CASE=3_7						,						__SCREAMING_SNAKE_CASE="gelu"						,						__SCREAMING_SNAKE_CASE=0.1						,						__SCREAMING_SNAKE_CASE=0.1						,						__SCREAMING_SNAKE_CASE=5_0						,						__SCREAMING_SNAKE_CASE=0.02						,						__SCREAMING_SNAKE_CASE=True						,						__SCREAMING_SNAKE_CASE=None						,						)  ->   Optional[Any]:




											"""simple docstring"""







											UpperCamelCase__    :   Optional[int]					=				parent
											UpperCamelCase__    :   str					=				batch_size
											UpperCamelCase__    :   Optional[int]					=				seq_length
											UpperCamelCase__    :   Any					=				is_training
											UpperCamelCase__    :   Optional[int]					=				use_input_mask
											UpperCamelCase__    :   int					=				vocab_size
											UpperCamelCase__    :   Union[str, Any]					=				hidden_size
											UpperCamelCase__    :   str					=				num_hidden_layers
											UpperCamelCase__    :   int					=				num_attention_heads
											UpperCamelCase__    :   List[Any]					=				intermediate_size
											UpperCamelCase__    :   Any					=				hidden_act
											UpperCamelCase__    :   Union[str, Any]					=				hidden_dropout_prob
											UpperCamelCase__    :   str					=				attention_probs_dropout_prob
											UpperCamelCase__    :   Optional[Any]					=				max_position_embeddings
											UpperCamelCase__    :   Dict					=				initializer_range
											UpperCamelCase__    :   Any					=				use_labels
											UpperCamelCase__    :   Optional[Any]					=				scope






					def    __SCREAMING_SNAKE_CASE			(							self  )  ->   List[str]:




											"""simple docstring"""







											UpperCamelCase__    :   str					=				ids_tensor([self.batch_size, self.seq_length]						,						self.vocab_size  )

											UpperCamelCase__    :   Tuple					=				None
											if self.use_input_mask:
																	UpperCamelCase__    :   Optional[int]					=				random_attention_mask([self.batch_size, self.seq_length]  )

											if self.use_labels:
																	UpperCamelCase__    :   List[str]					=				ids_tensor([self.batch_size, self.seq_length]						,						self.vocab_size  )

											UpperCamelCase__    :   int					=				self.get_config()

											return config, input_ids, input_mask, token_labels






					def    __SCREAMING_SNAKE_CASE			(							self  )  ->   List[str]:




											"""simple docstring"""







											return BertGenerationConfig(
											    vocab_size=self.vocab_size						,						hidden_size=self.hidden_size						,						num_hidden_layers=self.num_hidden_layers						,						num_attention_heads=self.num_attention_heads						,						intermediate_size=self.intermediate_size						,						hidden_act=self.hidden_act						,						hidden_dropout_prob=self.hidden_dropout_prob						,						attention_probs_dropout_prob=self.attention_probs_dropout_prob						,						max_position_embeddings=self.max_position_embeddings						,						is_decoder=__SCREAMING_SNAKE_CASE						,						initializer_range=self.initializer_range						,						)






					def    __SCREAMING_SNAKE_CASE			(							self  )  ->   int:




											"""simple docstring"""







											(
											    UpperCamelCase__
											)    :   Union[str, Any]					=				self.prepare_config_and_inputs()

											UpperCamelCase__    :   Dict					=				True
											UpperCamelCase__    :   Dict					=				floats_tensor([self.batch_size, self.seq_length, self.hidden_size]  )
											UpperCamelCase__    :   Dict					=				ids_tensor([self.batch_size, self.seq_length]						,						vocab_size=2  )

											return (
											    config,
											    input_ids,
											    input_mask,
											    token_labels,
											    encoder_hidden_states,
											    encoder_attention_mask,
											)






					def    __SCREAMING_SNAKE_CASE			(							self						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE						,						**__SCREAMING_SNAKE_CASE						,						)  ->   int:




											"""simple docstring"""







											UpperCamelCase__    :   Optional[int]					=				BertGenerationEncoder(config=__SCREAMING_SNAKE_CASE  )
											model.to(__SCREAMING_SNAKE_CASE  )
											model.eval()
											UpperCamelCase__    :   List[Any]					=				model(__SCREAMING_SNAKE_CASE						,						attention_mask=__SCREAMING_SNAKE_CASE  )
											UpperCamelCase__    :   Optional[Any]					=				model(__SCREAMING_SNAKE_CASE  )
											self.parent.assertEqual(result.last_hidden_state.shape						,						(self.batch_size, self.seq_length, self.hidden_size)  )






					def    __SCREAMING_SNAKE_CASE			(							self						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE						,						**__SCREAMING_SNAKE_CASE						,						)  ->   Any:




											"""simple docstring"""







											UpperCamelCase__    :   Dict					=				True
											UpperCamelCase__    :   List[Any]					=				BertGenerationEncoder(config=__SCREAMING_SNAKE_CASE  )
											model.to(__SCREAMING_SNAKE_CASE  )
											model.eval()
											UpperCamelCase__    :   int					=				model(
											    __SCREAMING_SNAKE_CASE						,						attention_mask=__SCREAMING_SNAKE_CASE						,						encoder_hidden_states=__SCREAMING_SNAKE_CASE						,						encoder_attention_mask=__SCREAMING_SNAKE_CASE						,						)
											UpperCamelCase__    :   Dict					=				model(
											    __SCREAMING_SNAKE_CASE						,						attention_mask=__SCREAMING_SNAKE_CASE						,						encoder_hidden_states=__SCREAMING_SNAKE_CASE						,						)
											self.parent.assertEqual(result.last_hidden_state.shape						,						(self.batch_size, self.seq_length, self.hidden_size)  )






					def    __SCREAMING_SNAKE_CASE			(							self						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE						,						**__SCREAMING_SNAKE_CASE						,						)  ->   str:




											"""simple docstring"""







											UpperCamelCase__    :   Dict					=				True
											UpperCamelCase__    :   List[str]					=				True
											UpperCamelCase__    :   List[Any]					=				BertGenerationDecoder(config=__SCREAMING_SNAKE_CASE  ).to(__SCREAMING_SNAKE_CASE  ).eval()

											# first forward pass
											UpperCamelCase__    :   int					=				model(
											    __SCREAMING_SNAKE_CASE						,						attention_mask=__SCREAMING_SNAKE_CASE						,						encoder_hidden_states=__SCREAMING_SNAKE_CASE						,						encoder_attention_mask=__SCREAMING_SNAKE_CASE						,						use_cache=__SCREAMING_SNAKE_CASE						,						)
											UpperCamelCase__    :   Tuple					=				outputs.past_key_values

											# create hypothetical multiple next token and extent to next_input_ids
											UpperCamelCase__    :   Dict					=				ids_tensor((self.batch_size, 3)						,						config.vocab_size  )
											UpperCamelCase__    :   List[Any]					=				ids_tensor((self.batch_size, 3)						,						vocab_size=2  )

											# append to next input_ids and
											UpperCamelCase__    :   int					=				torch.cat([input_ids, next_tokens]						,						dim=-1  )
											UpperCamelCase__    :   Dict					=				torch.cat([input_mask, next_mask]						,						dim=-1  )

											UpperCamelCase__    :   Tuple					=				model(
											    __SCREAMING_SNAKE_CASE						,						attention_mask=__SCREAMING_SNAKE_CASE						,						encoder_hidden_states=__SCREAMING_SNAKE_CASE						,						encoder_attention_mask=__SCREAMING_SNAKE_CASE						,						output_hidden_states=__SCREAMING_SNAKE_CASE						,						)['''hidden_states'''][0]
											UpperCamelCase__    :   List[Any]					=				model(
											    __SCREAMING_SNAKE_CASE						,						attention_mask=__SCREAMING_SNAKE_CASE						,						encoder_hidden_states=__SCREAMING_SNAKE_CASE						,						encoder_attention_mask=__SCREAMING_SNAKE_CASE						,						past_key_values=__SCREAMING_SNAKE_CASE						,						output_hidden_states=__SCREAMING_SNAKE_CASE						,						)['''hidden_states'''][0]

											# select random slice
											UpperCamelCase__    :   List[Any]					=				ids_tensor((1,)						,						output_from_past.shape[-1]  ).item()
											UpperCamelCase__    :   Union[str, Any]					=				output_from_no_past[:, -3:, random_slice_idx].detach()
											UpperCamelCase__    :   Optional[int]					=				output_from_past[:, :, random_slice_idx].detach()

											self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]  )

											# test that outputs are equal for slice
											self.parent.assertTrue(torch.allclose(__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE						,						atol=1e-3  )  )






					def    __SCREAMING_SNAKE_CASE			(							self						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE						,						*__SCREAMING_SNAKE_CASE						,						)  ->   Optional[int]:




											"""simple docstring"""







											UpperCamelCase__    :   List[str]					=				BertGenerationDecoder(__SCREAMING_SNAKE_CASE  )
											model.to(__SCREAMING_SNAKE_CASE  )
											model.eval()
											UpperCamelCase__    :   Any					=				model(__SCREAMING_SNAKE_CASE						,						attention_mask=__SCREAMING_SNAKE_CASE						,						labels=__SCREAMING_SNAKE_CASE  )
											self.parent.assertEqual(result.logits.shape						,						(self.batch_size, self.seq_length, self.vocab_size)  )






					def    __SCREAMING_SNAKE_CASE			(							self  )  ->   Optional[int]:




											"""simple docstring"""







											UpperCamelCase__    :   str					=				self.prepare_config_and_inputs()
											UpperCamelCase__    :   Optional[int]					=				{'''input_ids''': input_ids, '''attention_mask''': input_mask}
											return config, inputs_dict




@require_torch
class 				_lowerCamelCase (      UpperCamelCase_	,							UpperCamelCase_	,							UpperCamelCase_	,							unittest.TestCase ):





					"""simple docstring"""







					SCREAMING_SNAKE_CASE_		   =  (BertGenerationEncoder, BertGenerationDecoder) if is_torch_available() else ()
					SCREAMING_SNAKE_CASE_		   =  (BertGenerationDecoder,) if is_torch_available() else ()
					SCREAMING_SNAKE_CASE_		   =  (
					    {'''feature-extraction''': BertGenerationEncoder, '''text-generation''': BertGenerationDecoder}
					    if is_torch_available()
					    else {}
					)






					def    __SCREAMING_SNAKE_CASE			(							self  )  ->   Dict:




											"""simple docstring"""







											UpperCamelCase__    :   List[Any]					=				BertGenerationEncoderTester(self  )
											UpperCamelCase__    :   Tuple					=				ConfigTester(self						,						config_class=__SCREAMING_SNAKE_CASE						,						hidden_size=3_7  )






					def    __SCREAMING_SNAKE_CASE			(							self  )  ->   Optional[Any]:




											"""simple docstring"""







											self.config_tester.run_common_tests()






					def    __SCREAMING_SNAKE_CASE			(							self  )  ->   Any:




											"""simple docstring"""







											UpperCamelCase__    :   Dict					=				self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*__SCREAMING_SNAKE_CASE  )






					def    __SCREAMING_SNAKE_CASE			(							self  )  ->   Dict:




											"""simple docstring"""







											UpperCamelCase__    :   Optional[Any]					=				self.model_tester.prepare_config_and_inputs()
											UpperCamelCase__    :   List[str]					=				'''bert'''
											self.model_tester.create_and_check_model(__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE  )






					def    __SCREAMING_SNAKE_CASE			(							self  )  ->   List[str]:




											"""simple docstring"""







											UpperCamelCase__    :   Any					=				self.model_tester.prepare_config_and_inputs_for_decoder()
											self.model_tester.create_and_check_model_as_decoder(*__SCREAMING_SNAKE_CASE  )






					def    __SCREAMING_SNAKE_CASE			(							self  )  ->   str:




											"""simple docstring"""







											UpperCamelCase__    :   List[Any]					=				self.model_tester.prepare_config_and_inputs_for_decoder()
											self.model_tester.create_and_check_decoder_model_past_large_inputs(*__SCREAMING_SNAKE_CASE  )






					def    __SCREAMING_SNAKE_CASE			(							self  )  ->   Tuple:




											"""simple docstring"""







											(
											    UpperCamelCase__
											)    :   Optional[Any]					=				self.model_tester.prepare_config_and_inputs_for_decoder()

											UpperCamelCase__    :   Tuple					=				None

											self.model_tester.create_and_check_model_as_decoder(
											    __SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE						,						__SCREAMING_SNAKE_CASE						,						)






					def    __SCREAMING_SNAKE_CASE			(							self  )  ->   str:




											"""simple docstring"""







											UpperCamelCase__    :   int					=				self.model_tester.prepare_config_and_inputs_for_decoder()
											self.model_tester.create_and_check_for_causal_lm(*__SCREAMING_SNAKE_CASE  )






					@slow
					def    __SCREAMING_SNAKE_CASE			(							self  )  ->   Any:




											"""simple docstring"""







											UpperCamelCase__    :   List[str]					=				BertGenerationEncoder.from_pretrained('''google/bert_for_seq_generation_L-24_bbc_encoder'''  )
											self.assertIsNotNone(__SCREAMING_SNAKE_CASE  )




@require_torch
class 				_lowerCamelCase (      unittest.TestCase ):





					"""simple docstring"""







					@slow
					def    __SCREAMING_SNAKE_CASE			(							self  )  ->   Optional[Any]:




											"""simple docstring"""







											UpperCamelCase__    :   List[str]					=				BertGenerationEncoder.from_pretrained('''google/bert_for_seq_generation_L-24_bbc_encoder'''  )
											UpperCamelCase__    :   List[Any]					=				torch.tensor([[1_0_1, 7_5_9_2, 1_0_1_0, 2_0_2_6, 3_8_9_9, 2_0_0_3, 1_0_1_4_0, 1_0_2]]  )
											with torch.no_grad():
																	UpperCamelCase__    :   int					=				model(__SCREAMING_SNAKE_CASE  )[0]
											UpperCamelCase__    :   Union[str, Any]					=				torch.Size([1, 8, 1_0_2_4]  )
											self.assertEqual(output.shape						,						__SCREAMING_SNAKE_CASE  )
											UpperCamelCase__    :   Dict					=				torch.tensor(
											    [[[0.1775, 0.0083, -0.0321], [1.6002, 0.1287, 0.3912], [2.1473, 0.5791, 0.6066]]]  )
											self.assertTrue(torch.allclose(output[:, :3, :3]						,						__SCREAMING_SNAKE_CASE						,						atol=1e-4  )  )




@require_torch
class 				_lowerCamelCase (      unittest.TestCase ):





					"""simple docstring"""







					@slow
					def    __SCREAMING_SNAKE_CASE			(							self  )  ->   List[Any]:




											"""simple docstring"""







											UpperCamelCase__    :   List[str]					=				BertGenerationDecoder.from_pretrained('''google/bert_for_seq_generation_L-24_bbc_encoder'''  )
											UpperCamelCase__    :   Optional[int]					=				torch.tensor([[1_0_1, 7_5_9_2, 1_0_1_0, 2_0_2_6, 3_8_9_9, 2_0_0_3, 1_0_1_4_0, 1_0_2]]  )
											with torch.no_grad():
																	UpperCamelCase__    :   str					=				model(__SCREAMING_SNAKE_CASE  )[0]
											UpperCamelCase__    :   Tuple					=				torch.Size([1, 8, 5_0_3_5_8]  )
											self.assertEqual(output.shape						,						__SCREAMING_SNAKE_CASE  )
											UpperCamelCase__    :   Dict					=				torch.tensor(
											    [[[-0.5788, -2.5994, -3.7054], [0.0438, 4.7997, 1.8795], [1.5862, 6.6409, 4.4638]]]  )
											self.assertTrue(torch.allclose(output[:, :3, :3]						,						__SCREAMING_SNAKE_CASE						,						atol=1e-4  )  )




def 							SCREAMING_SNAKE_CASE_      (			UpperCamelCase__ = 5_0_0_0_0_0_0_0		):
      UpperCamelCase__    :   Any					=				set()
      UpperCamelCase__    :   Any					=				int((limit - 2_4) ** (1 / 2)		)

      UpperCamelCase__    :   Dict					=				set(range(3 ,		prime_square_limit + 1 ,		2		)		)
      primes.add(2		)
      for p in range(3 ,		prime_square_limit + 1 ,		2		):
            if p not in primes:
                  continue
            primes.difference_update(set(range(p * p ,		prime_square_limit + 1 ,		UpperCamelCase__		)		)		)

      for primea in primes:
            UpperCamelCase__    :   Dict					=				primea * primea
            for primea in primes:
                  UpperCamelCase__    :   str					=				primea * primea * primea
                  if square + cube >= limit - 1_6:
                        break
                  for primea in primes:
                        UpperCamelCase__    :   Dict					=				primea * primea * primea * primea
                        UpperCamelCase__    :   Tuple					=				square + cube + tetr
                        if total >= limit:
                              break
                        ret.add(UpperCamelCase__		)

      return len(UpperCamelCase__		)


if __name__ == "__main__":
    print(F'''{solution() = }''')


