

import unittest

import numpy as np
import timeout_decorator  # noqa

from transformers import OPTConfig, is_flax_available
from transformers.testing_utils import require_flax, require_sentencepiece, slow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor


if is_flax_available():
  import os

  # The slow tests are often failing with OOM error on GPU
  # This makes JAX allocate exactly what is needed on demand, and deallocate memory that is no longer needed
  # but will be slower as stated here https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
  _a 	=       '''platform'''

  import jax
  import jax.numpy as jnp

  from transformers import FlaxOPTForCausalLM, FlaxOPTModel, GPTaTokenizer







def _a			(	SCREAMING_SNAKE_CASE		:      List[str]   , SCREAMING_SNAKE_CASE		:      Tuple   , SCREAMING_SNAKE_CASE		:      Any=None   , SCREAMING_SNAKE_CASE		:      int=None		)   ->   Optional[Any]:

 """simple docstring"""






 if attention_mask is None:
  __lowerCAmelCase:       List[str]            =				np.where(input_ids != config.pad_token_id   , 1   , 0		)
 return {
     "input_ids": input_ids,
     "attention_mask": attention_mask,
 }






@require_flax
class 							A_			:


     def __init__(							self      :		Optional[Any]      ,  UpperCAmelCase      :		Optional[Any]      ,  UpperCAmelCase      :		int=1_3      ,  UpperCAmelCase      :		Optional[int]=7      ,  UpperCAmelCase      :		Tuple=True      ,  UpperCAmelCase      :		str=False      ,  UpperCAmelCase      :		Dict=9_9      ,  UpperCAmelCase      :		List[str]=1_6      ,  UpperCAmelCase      :		Optional[int]=2      ,  UpperCAmelCase      :		Optional[Any]=4      ,  UpperCAmelCase      :		str=4      ,  UpperCAmelCase      :		int="gelu"      ,  UpperCAmelCase      :		int=0.1      ,  UpperCAmelCase      :		List[Any]=0.1      ,  UpperCAmelCase      :		Tuple=2_0      ,  UpperCAmelCase      :		Optional[int]=2      ,  UpperCAmelCase      :		Union[str, Any]=1      ,  UpperCAmelCase      :		Optional[Any]=0      ,  UpperCAmelCase      :		List[Any]=1_6      ,  UpperCAmelCase      :		Union[str, Any]=1_6      ,  UpperCAmelCase      :		Tuple=0.02      ,  )							->     List[Any]:
      __lowerCAmelCase:       int            =				parent
      __lowerCAmelCase:       List[str]            =				batch_size
      __lowerCAmelCase:       Dict            =				seq_length
      __lowerCAmelCase:       Dict            =				is_training
      __lowerCAmelCase:       List[Any]            =				use_labels
      __lowerCAmelCase:       str            =				vocab_size
      __lowerCAmelCase:       Union[str, Any]            =				hidden_size
      __lowerCAmelCase:       Optional[Any]            =				num_hidden_layers
      __lowerCAmelCase:       Tuple            =				num_attention_heads
      __lowerCAmelCase:       Dict            =				intermediate_size
      __lowerCAmelCase:       List[str]            =				hidden_act
      __lowerCAmelCase:       str            =				hidden_dropout_prob
      __lowerCAmelCase:       Union[str, Any]            =				attention_probs_dropout_prob
      __lowerCAmelCase:       str            =				max_position_embeddings
      __lowerCAmelCase:       Tuple            =				eos_token_id
      __lowerCAmelCase:       Dict            =				pad_token_id
      __lowerCAmelCase:       Tuple            =				bos_token_id
      __lowerCAmelCase:       Any            =				embed_dim
      __lowerCAmelCase:       Any            =				word_embed_proj_dim
      __lowerCAmelCase:       Optional[Any]            =				initializer_range
      __lowerCAmelCase:       List[str]            =				False


     def   UpperCAmelCase  (							self      :		Dict )							->     Optional[Any]:
      __lowerCAmelCase:       int            =				np.clip(ids_tensor([self.batch_size, self.seq_length - 1]      ,  self.vocab_size )      ,  3      ,  self.vocab_size )
      __lowerCAmelCase:       Union[str, Any]            =				np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1)      ,  dtype=np.intaa ))      ,  -1 )

      __lowerCAmelCase:       Union[str, Any]            =				OPTConfig(
          vocab_size=self.vocab_size      ,  hidden_size=self.hidden_size      ,  num_hidden_layers=self.num_hidden_layers      ,  num_attention_heads=self.num_attention_heads      ,  ffn_dim=self.intermediate_size      ,  dropout=self.hidden_dropout_prob      ,  attention_dropout=self.attention_probs_dropout_prob      ,  max_position_embeddings=self.max_position_embeddings      ,  eos_token_id=self.eos_token_id      ,  bos_token_id=self.bos_token_id      ,  pad_token_id=self.pad_token_id      ,  embed_dim=self.embed_dim      ,  is_encoder_decoder=_A      ,  word_embed_proj_dim=self.word_embed_proj_dim      ,  initializer_range=self.initializer_range      ,  use_cache=_A      ,  )
      __lowerCAmelCase:       int            =				prepare_opt_inputs_dict(_A      ,  _A )
      return config, inputs_dict


     def   UpperCAmelCase  (							self      :		Dict )							->     Union[str, Any]:
      __lowerCAmelCase:       Dict            =				self.prepare_config_and_inputs()
      return config, inputs_dict


     def   UpperCAmelCase  (							self      :		List[str]      ,  UpperCAmelCase      :		Optional[Any]      ,  UpperCAmelCase      :		Any      ,  UpperCAmelCase      :		str )							->     List[str]:
      __lowerCAmelCase:       Union[str, Any]            =				2_0
      __lowerCAmelCase:       str            =				model_class_name(_A )

      __lowerCAmelCase:       Dict            =				inputs_dict['input_ids']
      __lowerCAmelCase:       Optional[int]            =				inputs_dict['attention_mask']

      __lowerCAmelCase:       int            =				model.init_cache(input_ids.shape[0]      ,  _A )
      __lowerCAmelCase:       Tuple            =				jnp.ones((input_ids.shape[0], max_length)      ,  dtype='i4' )

      __lowerCAmelCase:       Optional[Any]            =				jnp.broadcast_to(
          jnp.arange(input_ids.shape[-1] - 1 )[None, :]      ,  (input_ids.shape[0], input_ids.shape[-1] - 1)      ,  )
      __lowerCAmelCase:       Tuple            =				model(
          input_ids[:, :-1]      ,  attention_mask=_A      ,  past_key_values=_A      ,  position_ids=_A      ,  )

      __lowerCAmelCase:       Optional[int]            =				jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]      ,  dtype='i4' )
      __lowerCAmelCase:       Any            =				model(
          input_ids[:, -1:]      ,  attention_mask=_A      ,  past_key_values=outputs_cache.past_key_values      ,  position_ids=_A      ,  )

      __lowerCAmelCase:       List[str]            =				model(_A )

      __lowerCAmelCase:       Tuple            =				np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5]) ) )
      self.parent.assertTrue(diff < 1E-3      ,  msg=F'''Max diff is {diff}''' )


     def   UpperCAmelCase  (							self      :		Optional[Any]      ,  UpperCAmelCase      :		Any      ,  UpperCAmelCase      :		List[Any]      ,  UpperCAmelCase      :		Optional[int] )							->     Optional[Any]:
      __lowerCAmelCase:       Tuple            =				2_0
      __lowerCAmelCase:       Optional[Any]            =				model_class_name(_A )

      __lowerCAmelCase:       Tuple            =				(
          inputs_dict['input_ids'],
          inputs_dict['attention_mask'],
      )

      __lowerCAmelCase:       List[Any]            =				jnp.concatenate(
          [
              attention_mask,
              jnp.zeros((attention_mask.shape[0], max_length - attention_mask.shape[1]) ),
          ]      ,  axis=-1      ,  )

      __lowerCAmelCase:       Tuple            =				model.init_cache(input_ids.shape[0]      ,  _A )
      __lowerCAmelCase:       Any            =				jnp.broadcast_to(
          jnp.arange(input_ids.shape[-1] - 1 )[None, :]      ,  (input_ids.shape[0], input_ids.shape[-1] - 1)      ,  )

      __lowerCAmelCase:       Tuple            =				model(
          input_ids[:, :-1]      ,  attention_mask=_A      ,  past_key_values=_A      ,  position_ids=_A      ,  )
      __lowerCAmelCase:       str            =				jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]      ,  dtype='i4' )
      __lowerCAmelCase:       int            =				model(
          input_ids[:, -1:]      ,  past_key_values=outputs_cache.past_key_values      ,  attention_mask=_A      ,  position_ids=_A      ,  )

      __lowerCAmelCase:       Optional[Any]            =				model(_A      ,  attention_mask=_A )

      __lowerCAmelCase:       str            =				np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5]) ) )
      self.parent.assertTrue(diff < 1E-3      ,  msg=F'''Max diff is {diff}''' )







@require_flax
class 							A_			(		snake_case__							,     unittest.TestCase							,     snake_case__ ):
     _lowercase      :					str       =     (FlaxOPTModel, FlaxOPTForCausalLM) if is_flax_available() else ()
     _lowercase      :					Union[str, Any]       =     () if is_flax_available() else ()


     def   UpperCAmelCase  (							self      :		Union[str, Any] )							->     Optional[Any]:
      __lowerCAmelCase:       Union[str, Any]            =				FlaxOPTModelTester(self )


     def   UpperCAmelCase  (							self      :		str )							->     Any:
      __lowerCAmelCase:       Union[str, Any]            =				self.model_tester.prepare_config_and_inputs()
      for model_class in self.all_model_classes:
       self.model_tester.check_use_cache_forward(_A      ,  _A      ,  _A )


     def   UpperCAmelCase  (							self      :		Dict )							->     Any:
      __lowerCAmelCase:       Tuple            =				self.model_tester.prepare_config_and_inputs()
      for model_class in self.all_model_classes:
       self.model_tester.check_use_cache_forward_with_attn_mask(_A      ,  _A      ,  _A )


     @slow
     def   UpperCAmelCase  (							self      :		Union[str, Any] )							->     Tuple:
      for model_class_name in self.all_model_classes:
       __lowerCAmelCase:       Tuple            =				model_class_name.from_pretrained('facebook/opt-125m' )
       __lowerCAmelCase:       List[str]            =				np.ones((1, 1) ) * model.config.eos_token_id
       __lowerCAmelCase:       Any            =				model(_A )
       self.assertIsNotNone(_A )







@require_sentencepiece
@require_flax
class 							A_			(		unittest.TestCase ):


     @slow
     def   UpperCAmelCase  (							self      :		int )							->     List[Any]:
      __lowerCAmelCase:       Dict            =				FlaxOPTModel.from_pretrained('facebook/opt-350m' )
      __lowerCAmelCase:       Any            =				jnp.array([[0, 3_1_4_1_4, 2_3_2, 3_2_8, 7_4_0, 1_1_4_0, 1_2_6_9_5, 6_9, 4_6_0_7_8, 1_5_8_8, 2]] )
      __lowerCAmelCase:       int            =				model(input_ids=_A ).last_hidden_state
      __lowerCAmelCase:       str            =				(1, 1_1, 5_1_2)
      self.assertEqual(output.shape      ,  _A )
      __lowerCAmelCase:       Optional[int]            =				jnp.array(
          [[-0.2867, -1.9256, -0.3062], [-1.2711, -0.1337, -0.1897], [0.4109, 0.1187, -1.3142]] )
      self.assertTrue(jnp.allclose(output[:, :3, :3]      ,  _A      ,  atol=4E-2 ) )







@require_flax
@slow
class 							A_			(		unittest.TestCase ):


     def   UpperCAmelCase  (							self      :		Any )							->     Union[str, Any]:
      super().setUp()
      __lowerCAmelCase:       int            =				'facebook/opt-350m'


     def   UpperCAmelCase  (							self      :		Tuple )							->     int:
      __lowerCAmelCase:       int            =				FlaxOPTForCausalLM.from_pretrained(self.path_model )
      __lowerCAmelCase:       Dict            =				GPTaTokenizer.from_pretrained(self.path_model )

      __lowerCAmelCase:       List[str]            =				[
          'Today is a beautiful day and I want to',
          'In the city of',
          'Paris is the capital of France and',
          'Computers and mobile phones have taken',
      ]
      # verify that prompt without BOS token is identical to Metaseq -> add_special_tokens=False
      __lowerCAmelCase:       Tuple            =				tokenizer(_A      ,  return_tensors='jax'      ,  padding=_A      ,  add_special_tokens=_A )
      __lowerCAmelCase:       List[Any]            =				model(inputs.input_ids      ,  attention_mask=inputs.attention_mask )[0].mean(axis=-1 )
      __lowerCAmelCase:       Optional[int]            =				jnp.array(
          [
              [1.3851, -1_3.8_9_2_3, -1_0.5_2_2_9, -1_0.7_5_3_3, -0.2309, -1_0.2_3_8_4, -0.5365, -9.0947, -5.1670],
              [-4.7073, -1_0.6_2_7_6, -3.9415, -2_1.5_2_4_2, -0.2822, -0.2822, -0.2822, -0.2822, -0.2822],
              [0.6247, -3.4229, -8.9179, -1.4297, -1_4.1_6_5_0, 1.4146, -9.0218, -0.2703, -0.2703],
              [6.4783, -1.9913, -1_0.7_9_2_6, -2.3336, 1.5092, -0.9974, -6.8213, 1.3477, 1.3477],
          ] )
      self.assertTrue(jnp.allclose(_A      ,  _A      ,  atol=4E-2 ) )

      __lowerCAmelCase:       List[Any]            =				jax.jit(_A )
      __lowerCAmelCase:       List[str]            =				model(inputs.input_ids      ,  attention_mask=inputs.attention_mask )[0].mean(axis=-1 )
      self.assertTrue(jnp.allclose(_A      ,  _A      ,  atol=4E-2 ) )







@require_flax
@slow
class 							A_			(		unittest.TestCase ):


     @property
     def   UpperCAmelCase  (							self      :		Any )							->     List[str]:
      return [
          "Today is a beautiful day and I want",
          "In the city of",
          "Paris is the capital of France and",
          "Computers and mobile phones have taken",
      ]


     def   UpperCAmelCase  (							self      :		Union[str, Any] )							->     Optional[int]:
      __lowerCAmelCase:       Optional[int]            =				'facebook/opt-125m'

      __lowerCAmelCase:       Optional[Any]            =				[
          'Today is a beautiful day and I want to',
          'In the city of New York, the city',
          'Paris is the capital of France and the capital',
          'Computers and mobile phones have taken over the',
      ]

      __lowerCAmelCase:       List[str]            =				[]

      __lowerCAmelCase:       List[Any]            =				FlaxOPTForCausalLM.from_pretrained(_A )
      __lowerCAmelCase:       Optional[int]            =				GPTaTokenizer.from_pretrained(_A )

      for prompt in self.prompts:
       __lowerCAmelCase:       Dict            =				tokenizer(_A      ,  return_tensors='jax' ).input_ids

       __lowerCAmelCase:       Any            =				model.generate(_A      ,  max_length=1_0 )
       __lowerCAmelCase:       Any            =				generated_ids[0]

       __lowerCAmelCase:       Optional[Any]            =				tokenizer.batch_decode(_A      ,  skip_special_tokens=_A )
       predicted_outputs += generated_string

      self.assertListEqual(_A      ,  _A )


     def   UpperCAmelCase  (							self      :		Dict )							->     Tuple:
      __lowerCAmelCase:       Dict            =				'facebook/opt-350m'

      __lowerCAmelCase:       List[str]            =				[
          'Today is a beautiful day and I want to',
          'In the city of San Francisco, the city',
          'Paris is the capital of France and the capital',
          'Computers and mobile phones have taken over the',
      ]

      __lowerCAmelCase:       Dict            =				[]
      __lowerCAmelCase:       int            =				FlaxOPTForCausalLM.from_pretrained(_A )
      __lowerCAmelCase:       str            =				GPTaTokenizer.from_pretrained(_A )

      for prompt in self.prompts:
       __lowerCAmelCase:       str            =				tokenizer(_A      ,  return_tensors='jax' ).input_ids

       __lowerCAmelCase:       int            =				model.generate(_A      ,  max_length=1_0 )
       __lowerCAmelCase:       List[Any]            =				generated_ids[0]

       __lowerCAmelCase:       List[Any]            =				tokenizer.batch_decode(_A      ,  skip_special_tokens=_A )
       predicted_outputs += generated_string

      self.assertListEqual(_A      ,  _A )


     def   UpperCAmelCase  (							self      :		Optional[int] )							->     Optional[int]:
      __lowerCAmelCase:       Union[str, Any]            =				'facebook/opt-125m'
      __lowerCAmelCase:       int            =				[
          'Today is a beautiful day and I want to thank',
          'In the city of Rome Canaver Canaver Canaver Canaver',
      ]
      __lowerCAmelCase:       Dict            =				FlaxOPTForCausalLM.from_pretrained(_A )
      __lowerCAmelCase:       Dict            =				GPTaTokenizer.from_pretrained(_A )
      __lowerCAmelCase:       Dict            =				tokenizer(
          [
              'Today is a beautiful day and I want to',
              'In the city of',
          ]      ,  return_tensors='jax'      ,  padding=_A      ,  )

      __lowerCAmelCase:       str            =				jax.jit(model.generate )

      __lowerCAmelCase:       Union[str, Any]            =				jit_generate(inputs['input_ids']      ,  attention_mask=inputs['attention_mask'] ).sequences

      __lowerCAmelCase:       Optional[Any]            =				tokenizer.batch_decode(_A      ,  skip_special_tokens=_A )

      self.assertIsNotNone(_A      ,  _A )


     def   UpperCAmelCase  (							self      :		int )							->     int:
      __lowerCAmelCase:       Tuple            =				'facebook/opt-350m'

      __lowerCAmelCase:       Optional[Any]            =				GPTaTokenizer.from_pretrained(_A )
      __lowerCAmelCase:       Dict            =				FlaxOPTForCausalLM.from_pretrained(_A )

      __lowerCAmelCase:       int            =				'left'

      # use different length sentences to test batching
      __lowerCAmelCase:       Optional[Any]            =				[
          'Hello, my dog is a little',
          'Today, I',
      ]

      __lowerCAmelCase:       List[str]            =				tokenizer(_A      ,  return_tensors='jax'      ,  padding=_A )
      __lowerCAmelCase:       Tuple            =				inputs['input_ids']

      __lowerCAmelCase:       Tuple            =				model.generate(input_ids=_A      ,  attention_mask=inputs['attention_mask']      ,  trace=_A )

      __lowerCAmelCase:       Any            =				tokenizer(sentences[0]      ,  return_tensors='jax' ).input_ids
      __lowerCAmelCase:       List[str]            =				model.generate(input_ids=_A )

      __lowerCAmelCase:       str            =				inputs_non_padded.shape[-1] - inputs['attention_mask'][-1].sum()
      __lowerCAmelCase:       Union[str, Any]            =				tokenizer(sentences[1]      ,  return_tensors='jax' ).input_ids
      __lowerCAmelCase:       Optional[Any]            =				model.generate(input_ids=_A      ,  max_length=model.config.max_length - num_paddings )

      __lowerCAmelCase:       Tuple            =				tokenizer.batch_decode(outputs[0]      ,  skip_special_tokens=_A )
      __lowerCAmelCase:       List[Any]            =				tokenizer.decode(output_non_padded[0][0]      ,  skip_special_tokens=_A )
      __lowerCAmelCase:       Tuple            =				tokenizer.decode(output_padded[0][0]      ,  skip_special_tokens=_A )

      __lowerCAmelCase:       int            =				[
          'Hello, my dog is a little bit of a dork.\nI\'m a little bit',
          'Today, I was in the middle of a conversation with a friend about the',
      ]
      self.assertListEqual(_A      ,  _A )
      self.assertListEqual(_A      ,  [non_padded_sentence, padded_sentence] )








def _a			(	SCREAMING_SNAKE_CASE		:      int = 10_00		)   ->   int:

 """simple docstring"""






 __lowerCAmelCase:       Union[str, Any]            =				[0, 3, 3, 5, 4, 4, 3, 5, 5, 4, 3, 6, 6, 8, 8, 7, 7, 9, 8, 8]
 # number of letters in twenty, thirty, ..., ninety (0 for numbers less than
 # 20 due to inconsistency in teens)
 __lowerCAmelCase:       Union[str, Any]            =				[0, 0, 6, 6, 5, 5, 5, 7, 6, 6]

 __lowerCAmelCase:       Optional[int]            =				0

 for i in range(1   , n + 1		):
  if i < 10_00:
   if i >= 1_00:
    # add number of letters for "n hundred"
    count += ones_counts[i // 1_00] + 7

    if i % 1_00 != 0:
     # add number of letters for "and" if number is not multiple
     # of 100
     count += 3

   if 0 < i % 1_00 < 20:
    # add number of letters for one, two, three, ..., nineteen
    # (could be combined with below if not for inconsistency in
    # teens)
    count += ones_counts[i % 1_00]
   else:
    # add number of letters for twenty, twenty one, ..., ninety
    # nine
    count += ones_counts[i % 10]
    count += tens_counts[(i % 1_00 - i % 10) // 10]
  else:
   count += ones_counts[i // 10_00] + 8
 return count


if __name__ == "__main__":
  print(solution(int(input().strip())))





