
import random
import unittest

import numpy as np

import transformers
from transformers import is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax


if is_flax_available():
		import os

		import jax.numpy as jnp
		from jax import jit

		from transformers import AutoTokenizer, FlaxAutoModelForCausalLM
		from transformers.modeling_flax_pytorch_utils import load_flax_weights_in_pytorch_model

		A		: str								=     '''0.12'''  # assumed parallelism: 8


if is_torch_available():
		import torch




def 	UpperCamelCase__  (     SCREAMING_SNAKE_CASE_	:	Optional[Any]		,   SCREAMING_SNAKE_CASE_	:	str		,   SCREAMING_SNAKE_CASE_	:	Dict=None       )  ->     int:
						if rng is None:
												_lowercase					       =					random.Random()

						_lowercase					       =					1
						for dim in shape:
												total_dims *= dim

						_lowercase					       =					[]
						for _ in range(SCREAMING_SNAKE_CASE_       ):
												values.append(rng.randint(0		,   vocab_size - 1       )       )

						_lowercase					       =					np.array(SCREAMING_SNAKE_CASE_		,   dtype=jnp.intaa       ).reshape(SCREAMING_SNAKE_CASE_       )

						return output




def 	UpperCamelCase__  (     SCREAMING_SNAKE_CASE_	:	int		,   SCREAMING_SNAKE_CASE_	:	Tuple=None       )  ->     Tuple:
						_lowercase					       =					ids_tensor(SCREAMING_SNAKE_CASE_		,   vocab_size=2		,   rng=SCREAMING_SNAKE_CASE_       )
						# make sure that at least one token is attended to for each batch
						_lowercase					       =					1
						return attn_mask



@require_flax
class    a_ :
			a	:			Dict					   =							None
			a	:			List[Any]					   =							()






			def    UpperCamelCase_				(      self  ):
									_lowercase       ,						_lowercase					       =					self.model_tester.prepare_config_and_inputs_for_common()

									# cut to half length & take max batch_size 3
									_lowercase					       =					2
									_lowercase					       =					inputs["""input_ids"""].shape[-1] // 2
									_lowercase					       =					inputs["""input_ids"""][:max_batch_size, :sequence_length]

									_lowercase					       =					jnp.ones_like(__UpperCamelCase  )
									_lowercase					       =					attention_mask[:max_batch_size, :sequence_length]

									# generate max 5 tokens
									_lowercase					       =					input_ids.shape[-1] + 5
									if config.eos_token_id is not None and config.pad_token_id is None:
															# hack to allow generate for models such as GPT2 as is done in `generate()`
															_lowercase					       =					config.eos_token_id
									return config, input_ids, attention_mask, max_length






			@is_pt_flax_cross_test
			def    UpperCamelCase_				(      self  ):
									_lowercase       ,						_lowercase       ,						_lowercase       ,						_lowercase					       =					self._get_input_ids_and_config()
									_lowercase					       =					False
									_lowercase					       =					max_length
									_lowercase					       =					0

									for model_class in self.all_generative_model_classes:
															_lowercase					       =					model_class(__UpperCamelCase  )

															_lowercase					       =					model_class.__name__[4:]  # Skip the "Flax" at the beginning
															_lowercase					       =					getattr(__UpperCamelCase				,			__UpperCamelCase  )
															_lowercase					       =					pt_model_class(__UpperCamelCase  ).eval()
															_lowercase					       =					load_flax_weights_in_pytorch_model(__UpperCamelCase				,			flax_model.params  )

															_lowercase					       =					flax_model.generate(__UpperCamelCase  ).sequences
															_lowercase					       =					pt_model.generate(torch.tensor(__UpperCamelCase				,			dtype=torch.long  )  )

															if flax_generation_outputs.shape[-1] > pt_generation_outputs.shape[-1]:
																					_lowercase					       =					flax_generation_outputs[:, : pt_generation_outputs.shape[-1]]

															self.assertListEqual(pt_generation_outputs.numpy().tolist()				,			flax_generation_outputs.tolist()  )






			def    UpperCamelCase_				(      self  ):
									_lowercase       ,						_lowercase       ,						_lowercase       ,						_lowercase					       =					self._get_input_ids_and_config()
									_lowercase					       =					False
									_lowercase					       =					max_length

									for model_class in self.all_generative_model_classes:
															_lowercase					       =					model_class(__UpperCamelCase  )

															_lowercase					       =					model.generate(__UpperCamelCase  ).sequences
															self.assertEqual(generation_outputs.shape[-1]				,			__UpperCamelCase  )

															_lowercase					       =					jit(model.generate  )
															_lowercase					       =					jit_generate(__UpperCamelCase  ).sequences

															self.assertListEqual(generation_outputs.tolist()				,			jit_generation_outputs.tolist()  )






			def    UpperCamelCase_				(      self  ):
									_lowercase       ,						_lowercase       ,						_lowercase       ,						_lowercase					       =					self._get_input_ids_and_config()
									_lowercase					       =					True
									_lowercase					       =					max_length

									for model_class in self.all_generative_model_classes:
															_lowercase					       =					model_class(__UpperCamelCase  )

															_lowercase					       =					model.generate(__UpperCamelCase  ).sequences
															self.assertEqual(generation_outputs.shape[-1]				,			__UpperCamelCase  )

															_lowercase					       =					jit(model.generate  )
															_lowercase					       =					jit_generate(__UpperCamelCase  ).sequences

															self.assertListEqual(generation_outputs.tolist()				,			jit_generation_outputs.tolist()  )






			def    UpperCamelCase_				(      self  ):
									_lowercase       ,						_lowercase       ,						_lowercase       ,						_lowercase					       =					self._get_input_ids_and_config()
									_lowercase					       =					False
									_lowercase					       =					max_length
									_lowercase					       =					2

									for model_class in self.all_generative_model_classes:
															_lowercase					       =					model_class(__UpperCamelCase  )

															_lowercase					       =					model.generate(__UpperCamelCase  ).sequences
															self.assertEqual(generation_outputs.shape[-1]				,			__UpperCamelCase  )

															_lowercase					       =					jit(model.generate  )
															_lowercase					       =					jit_generate(__UpperCamelCase  ).sequences

															self.assertListEqual(generation_outputs.tolist()				,			jit_generation_outputs.tolist()  )






			def    UpperCamelCase_				(      self  ):
									_lowercase       ,						_lowercase       ,						_lowercase       ,						_lowercase					       =					self._get_input_ids_and_config()
									_lowercase					       =					False
									_lowercase					       =					max_length
									_lowercase					       =					2
									_lowercase					       =					2

									for model_class in self.all_generative_model_classes:
															_lowercase					       =					model_class(__UpperCamelCase  )

															_lowercase					       =					model.generate(__UpperCamelCase  ).sequences
															self.assertEqual(generation_outputs.shape[0]				,			input_ids.shape[0] * config.num_return_sequences  )






			def    UpperCamelCase_				(      self  ):
									_lowercase       ,						_lowercase       ,						_lowercase       ,						_lowercase					       =					self._get_input_ids_and_config()
									_lowercase					       =					True
									_lowercase					       =					max_length
									_lowercase					       =					0.8
									_lowercase					       =					10
									_lowercase					       =					0.3
									_lowercase					       =					1
									_lowercase					       =					8
									_lowercase					       =					9

									for model_class in self.all_generative_model_classes:
															_lowercase					       =					model_class(__UpperCamelCase  )

															_lowercase					       =					model.generate(__UpperCamelCase  ).sequences
															self.assertEqual(generation_outputs.shape[-1]				,			__UpperCamelCase  )

															_lowercase					       =					jit(model.generate  )
															_lowercase					       =					jit_generate(__UpperCamelCase  ).sequences

															self.assertListEqual(generation_outputs.tolist()				,			jit_generation_outputs.tolist()  )






			def    UpperCamelCase_				(      self  ):
									_lowercase       ,						_lowercase       ,						_lowercase       ,						_lowercase					       =					self._get_input_ids_and_config()
									_lowercase					       =					max_length
									_lowercase					       =					1
									_lowercase					       =					8
									_lowercase					       =					9

									for model_class in self.all_generative_model_classes:
															_lowercase					       =					model_class(__UpperCamelCase  )

															_lowercase					       =					model.generate(__UpperCamelCase  ).sequences
															self.assertEqual(generation_outputs.shape[-1]				,			__UpperCamelCase  )

															_lowercase					       =					jit(model.generate  )
															_lowercase					       =					jit_generate(__UpperCamelCase  ).sequences

															self.assertListEqual(generation_outputs.tolist()				,			jit_generation_outputs.tolist()  )






			def    UpperCamelCase_				(      self  ):
									_lowercase       ,						_lowercase       ,						_lowercase       ,						_lowercase					       =					self._get_input_ids_and_config()
									_lowercase					       =					max_length
									_lowercase					       =					2
									_lowercase					       =					1
									_lowercase					       =					8
									_lowercase					       =					9

									for model_class in self.all_generative_model_classes:
															_lowercase					       =					model_class(__UpperCamelCase  )

															_lowercase					       =					model.generate(__UpperCamelCase  ).sequences
															self.assertEqual(generation_outputs.shape[-1]				,			__UpperCamelCase  )

															_lowercase					       =					jit(model.generate  )
															_lowercase					       =					jit_generate(__UpperCamelCase  ).sequences

															self.assertListEqual(generation_outputs.tolist()				,			jit_generation_outputs.tolist()  )






			def    UpperCamelCase_				(      self  ):
									_lowercase       ,						_lowercase       ,						_lowercase       ,						_lowercase					       =					self._get_input_ids_and_config()

									# pad attention mask on the left
									_lowercase					       =					attention_mask.at[(0, 0)].set(0  )

									_lowercase					       =					False
									_lowercase					       =					max_length

									for model_class in self.all_generative_model_classes:
															_lowercase					       =					model_class(__UpperCamelCase  )

															_lowercase					       =					model.generate(__UpperCamelCase				,			attention_mask=__UpperCamelCase  ).sequences
															self.assertEqual(generation_outputs.shape[-1]				,			__UpperCamelCase  )

															_lowercase					       =					jit(model.generate  )
															_lowercase					       =					jit_generate(__UpperCamelCase				,			attention_mask=__UpperCamelCase  ).sequences

															self.assertListEqual(generation_outputs.tolist()				,			jit_generation_outputs.tolist()  )






			def    UpperCamelCase_				(      self  ):
									_lowercase       ,						_lowercase       ,						_lowercase       ,						_lowercase					       =					self._get_input_ids_and_config()

									# pad attention mask on the left
									_lowercase					       =					attention_mask.at[(0, 0)].set(0  )

									_lowercase					       =					True
									_lowercase					       =					max_length

									for model_class in self.all_generative_model_classes:
															_lowercase					       =					model_class(__UpperCamelCase  )

															_lowercase					       =					model.generate(__UpperCamelCase				,			attention_mask=__UpperCamelCase  ).sequences
															self.assertEqual(generation_outputs.shape[-1]				,			__UpperCamelCase  )

															_lowercase					       =					jit(model.generate  )
															_lowercase					       =					jit_generate(__UpperCamelCase				,			attention_mask=__UpperCamelCase  ).sequences

															self.assertListEqual(generation_outputs.tolist()				,			jit_generation_outputs.tolist()  )






			def    UpperCamelCase_				(      self  ):
									_lowercase       ,						_lowercase       ,						_lowercase       ,						_lowercase					       =					self._get_input_ids_and_config()

									# pad attention mask on the left
									_lowercase					       =					attention_mask.at[(0, 0)].set(0  )

									_lowercase					       =					2
									_lowercase					       =					max_length

									for model_class in self.all_generative_model_classes:
															_lowercase					       =					model_class(__UpperCamelCase  )

															_lowercase					       =					model.generate(__UpperCamelCase				,			attention_mask=__UpperCamelCase  ).sequences
															self.assertEqual(generation_outputs.shape[-1]				,			__UpperCamelCase  )

															_lowercase					       =					jit(model.generate  )
															_lowercase					       =					jit_generate(__UpperCamelCase				,			attention_mask=__UpperCamelCase  ).sequences

															self.assertListEqual(generation_outputs.tolist()				,			jit_generation_outputs.tolist()  )







@require_flax
class    a_ (			unittest.TestCase    ):






			def    UpperCamelCase_				(      self  ):
									_lowercase					       =					AutoTokenizer.from_pretrained("""hf-internal-testing/tiny-bert"""  )
									_lowercase					       =					FlaxAutoModelForCausalLM.from_pretrained("""hf-internal-testing/tiny-bert-flax-only"""  )

									_lowercase					       =					"""Hello world"""
									_lowercase					       =					tokenizer(__UpperCamelCase				,			return_tensors="""np"""  ).input_ids

									# typos are quickly detected (the correct argument is `do_sample`)
									with self.assertRaisesRegex(__UpperCamelCase				,			"""do_samples"""  ):
															model.generate(__UpperCamelCase				,			do_samples=__UpperCamelCase  )

									# arbitrary arguments that will not be used anywhere are also not accepted
									with self.assertRaisesRegex(__UpperCamelCase				,			"""foo"""  ):
															_lowercase					       =					{"""foo""": """bar"""}
															model.generate(__UpperCamelCase				,			**__UpperCamelCase  )

import os

# Precomputes a list of the 100 first triangular numbers
A		: List[Any]								=     [int(0.5 * n * (n + 1)) for n in range(1, 101)]




def 	UpperCamelCase__  (     )  ->     str:
						_lowercase					       =					os.path.dirname(os.path.realpath(SCREAMING_SNAKE_CASE_       )       )
						_lowercase					       =					os.path.join(SCREAMING_SNAKE_CASE_		,   """words.txt"""       )

						_lowercase					       =					""""""
						with open(SCREAMING_SNAKE_CASE_       ) as f:
												_lowercase					       =					f.readline()

						_lowercase					       =					[word.strip("""\""""       ) for word in words.strip("""\r\n"""       ).split(""","""       )]
						_lowercase					       =					[
						    word
						    for word in [sum(ord(SCREAMING_SNAKE_CASE_       ) - 64 for x in word       ) for word in words]
						    if word in TRIANGULAR_NUMBERS
						]
						return len(SCREAMING_SNAKE_CASE_       )


if __name__ == "__main__":
		print(solution())