

from __future__ import annotations

import unittest

from transformers import is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
       import tensorflow as tf

       from transformers import (
           TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST,
           TFXLMForMultipleChoice,
           TFXLMForQuestionAnsweringSimple,
           TFXLMForSequenceClassification,
           TFXLMForTokenClassification,
           TFXLMModel,
           TFXLMWithLMHeadModel,
           XLMConfig,
       )







class 		__magic_name__     :

   '''simple docstring'''




   def __init__(							self,       lowercase_,       )				->       Union[str, Any]:






     """simple docstring"""






     a__    =parent
     a__    =13
     a__    =7
     a__    =True
     a__    =True
     a__    =True
     a__    =True
     a__    =True
     a__    =False
     a__    =False
     a__    =False
     a__    =2
     a__    =99
     a__    =0
     a__    =32
     a__    =2
     a__    =4
     a__    =0.1
     a__    =0.1
     a__    =512
     a__    =16
     a__    =2
     a__    =0.02
     a__    =3
     a__    =4
     a__    ='''last'''
     a__    =True
     a__    =None
     a__    =0



   def 					_UpperCAmelCase					(							self			)				->       Tuple:






     """simple docstring"""






     a__    =ids_tensor([self.batch_size, self.seq_length],       self.vocab_size			)
     a__    =random_attention_mask([self.batch_size, self.seq_length],       dtype=tf.floataa			)

     a__    =None
     if self.use_input_lengths:
       a__    =(
           ids_tensor([self.batch_size],       vocab_size=2			) + self.seq_length - 2
       )  # small variation of seq_length

     a__    =None
     if self.use_token_type_ids:
       a__    =ids_tensor([self.batch_size, self.seq_length],       self.n_langs			)

     a__    =None
     a__    =None
     a__    =None
     if self.use_labels:
       a__    =ids_tensor([self.batch_size],       self.type_sequence_label_size			)
       a__    =ids_tensor([self.batch_size, self.seq_length],       self.num_labels			)
       a__    =ids_tensor([self.batch_size],       2,       dtype=tf.floataa			)
       a__    =ids_tensor([self.batch_size],       self.num_choices			)

     a__    =XLMConfig(
         vocab_size=self.vocab_size,       n_special=self.n_special,       emb_dim=self.hidden_size,       n_layers=self.num_hidden_layers,       n_heads=self.num_attention_heads,       dropout=self.hidden_dropout_prob,       attention_dropout=self.attention_probs_dropout_prob,       gelu_activation=self.gelu_activation,       sinusoidal_embeddings=self.sinusoidal_embeddings,       asm=self.asm,       causal=self.causal,       n_langs=self.n_langs,       max_position_embeddings=self.max_position_embeddings,       initializer_range=self.initializer_range,       summary_type=self.summary_type,       use_proj=self.use_proj,       bos_token_id=self.bos_token_id,       )

     return (
         config,
         input_ids,
         token_type_ids,
         input_lengths,
         sequence_labels,
         token_labels,
         is_impossible_labels,
         choice_labels,
         input_mask,
     )



   def 					_UpperCAmelCase					(							self,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       )				->       int:






     """simple docstring"""






     a__    =TFXLMModel(config=_SCREAMING_SNAKE_CASE			)
     a__    ={'''input_ids''': input_ids, '''lengths''': input_lengths, '''langs''': token_type_ids}
     a__    =model(_SCREAMING_SNAKE_CASE			)

     a__    =[input_ids, input_mask]
     a__    =model(_SCREAMING_SNAKE_CASE			)
     self.parent.assertEqual(result.last_hidden_state.shape,       (self.batch_size, self.seq_length, self.hidden_size)			)



   def 					_UpperCAmelCase					(							self,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       )				->       Tuple:






     """simple docstring"""






     a__    =TFXLMWithLMHeadModel(_SCREAMING_SNAKE_CASE			)

     a__    ={'''input_ids''': input_ids, '''lengths''': input_lengths, '''langs''': token_type_ids}
     a__    =model(_SCREAMING_SNAKE_CASE			)

     a__    =outputs

     self.parent.assertEqual(result.logits.shape,       (self.batch_size, self.seq_length, self.vocab_size)			)



   def 					_UpperCAmelCase					(							self,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       )				->       Any:






     """simple docstring"""






     a__    =TFXLMForQuestionAnsweringSimple(_SCREAMING_SNAKE_CASE			)

     a__    ={'''input_ids''': input_ids, '''lengths''': input_lengths}

     a__    =model(_SCREAMING_SNAKE_CASE			)

     self.parent.assertEqual(result.start_logits.shape,       (self.batch_size, self.seq_length)			)
     self.parent.assertEqual(result.end_logits.shape,       (self.batch_size, self.seq_length)			)



   def 					_UpperCAmelCase					(							self,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       )				->       str:






     """simple docstring"""






     a__    =TFXLMForSequenceClassification(_SCREAMING_SNAKE_CASE			)

     a__    ={'''input_ids''': input_ids, '''lengths''': input_lengths}

     a__    =model(_SCREAMING_SNAKE_CASE			)

     self.parent.assertEqual(result.logits.shape,       (self.batch_size, self.type_sequence_label_size)			)



   def 					_UpperCAmelCase					(							self,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       )				->       Tuple:






     """simple docstring"""






     a__    =self.num_labels
     a__    =TFXLMForTokenClassification(config=_SCREAMING_SNAKE_CASE			)
     a__    ={'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
     a__    =model(_SCREAMING_SNAKE_CASE			)
     self.parent.assertEqual(result.logits.shape,       (self.batch_size, self.seq_length, self.num_labels)			)



   def 					_UpperCAmelCase					(							self,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       )				->       List[Any]:






     """simple docstring"""






     a__    =self.num_choices
     a__    =TFXLMForMultipleChoice(config=_SCREAMING_SNAKE_CASE			)
     a__    =tf.tile(tf.expand_dims(_SCREAMING_SNAKE_CASE,       1			),       (1, self.num_choices, 1)			)
     a__    =tf.tile(tf.expand_dims(_SCREAMING_SNAKE_CASE,       1			),       (1, self.num_choices, 1)			)
     a__    =tf.tile(tf.expand_dims(_SCREAMING_SNAKE_CASE,       1			),       (1, self.num_choices, 1)			)
     a__    ={
         '''input_ids''': multiple_choice_inputs_ids,
         '''attention_mask''': multiple_choice_input_mask,
         '''token_type_ids''': multiple_choice_token_type_ids,
     }
     a__    =model(_SCREAMING_SNAKE_CASE			)
     self.parent.assertEqual(result.logits.shape,       (self.batch_size, self.num_choices)			)



   def 					_UpperCAmelCase					(							self			)				->       int:






     """simple docstring"""






     a__    =self.prepare_config_and_inputs()
     (
         (
         a__
     ), (
         a__
     ), (
         a__
     ), (
         a__
     ), (
         a__
     ), (
         a__
     ), (
         a__
     ), (
         a__
     ), (
         a__
     ), 
     )    =config_and_inputs
     a__    ={
         '''input_ids''': input_ids,
         '''token_type_ids''': token_type_ids,
         '''langs''': token_type_ids,
         '''lengths''': input_lengths,
     }
     return config, inputs_dict







@require_tf
class 		__magic_name__     (				_UpperCamelCase						,			_UpperCamelCase						,			unittest.TestCase				):

   '''simple docstring'''




   lowerCamelCase__					:      str        =  (
       (
           TFXLMModel,
           TFXLMWithLMHeadModel,
           TFXLMForSequenceClassification,
           TFXLMForQuestionAnsweringSimple,
           TFXLMForTokenClassification,
           TFXLMForMultipleChoice,
       )
       if is_tf_available()
       else ()
   )
   lowerCamelCase__					:      Tuple        =  (
       (TFXLMWithLMHeadModel,) if is_tf_available() else ()
   )  # TODO (PVP): Check other models whether language generation is also applicable
   lowerCamelCase__					:      Any        =  (
       {
           'feature-extraction': TFXLMModel,
           'fill-mask': TFXLMWithLMHeadModel,
           'question-answering': TFXLMForQuestionAnsweringSimple,
           'text-classification': TFXLMForSequenceClassification,
           'text-generation': TFXLMWithLMHeadModel,
           'token-classification': TFXLMForTokenClassification,
           'zero-shot': TFXLMForSequenceClassification,
       }
       if is_tf_available()
       else {}
   )
   lowerCamelCase__					:      List[str]        =  False
   lowerCamelCase__					:      List[str]        =  False



   def 					_UpperCAmelCase					(							self,       lowercase_,       lowercase_,       lowercase_,       lowercase_,       lowercase_			)				->       List[Any]:






     """simple docstring"""






     if (
         pipeline_test_casse_name == "QAPipelineTests"
         and tokenizer_name is not None
         and not tokenizer_name.endswith('''Fast'''			)
     ):
       # `QAPipelineTests` fails for a few models when the slower tokenizer are used.
       # (The slower tokenizers were never used for pipeline tests before the pipeline testing rework)
       # TODO: check (and possibly fix) the `QAPipelineTests` with slower tokenizer
       return True

     return False



   def 					_UpperCAmelCase					(							self			)				->       Union[str, Any]:






     """simple docstring"""






     a__    =TFXLMModelTester(self			)
     a__    =ConfigTester(self,       config_class=_SCREAMING_SNAKE_CASE,       emb_dim=37			)



   def 					_UpperCAmelCase					(							self			)				->       int:






     """simple docstring"""






     self.config_tester.run_common_tests()



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






     """simple docstring"""






     a__    =self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_xlm_model(*_SCREAMING_SNAKE_CASE			)



   def 					_UpperCAmelCase					(							self			)				->       int:






     """simple docstring"""






     a__    =self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_xlm_lm_head(*_SCREAMING_SNAKE_CASE			)



   def 					_UpperCAmelCase					(							self			)				->       Union[str, Any]:






     """simple docstring"""






     a__    =self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_xlm_qa(*_SCREAMING_SNAKE_CASE			)



   def 					_UpperCAmelCase					(							self			)				->       Dict:






     """simple docstring"""






     a__    =self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_xlm_sequence_classif(*_SCREAMING_SNAKE_CASE			)



   def 					_UpperCAmelCase					(							self			)				->       List[Any]:






     """simple docstring"""






     a__    =self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_xlm_for_token_classification(*_SCREAMING_SNAKE_CASE			)



   def 					_UpperCAmelCase					(							self			)				->       Tuple:






     """simple docstring"""






     a__    =self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_xlm_for_multiple_choice(*_SCREAMING_SNAKE_CASE			)



   @slow
   def 					_UpperCAmelCase					(							self			)				->       Optional[int]:






     """simple docstring"""






     for model_name in TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
       a__    =TFXLMModel.from_pretrained(_SCREAMING_SNAKE_CASE			)
       self.assertIsNotNone(_SCREAMING_SNAKE_CASE			)







@require_tf
class 		__magic_name__     (				unittest.TestCase				):

   '''simple docstring'''




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






     """simple docstring"""






     a__    =TFXLMWithLMHeadModel.from_pretrained('''xlm-mlm-en-2048'''			)
     a__    =tf.convert_to_tensor([[14, 447]],       dtype=tf.intaa			)  # the president
     a__    =[
         14,
         447,
         14,
         447,
         14,
         447,
         14,
         447,
         14,
         447,
         14,
         447,
         14,
         447,
         14,
         447,
         14,
         447,
         14,
         447,
     ]  # the president the president the president the president the president the president the president the president the president the president
     # TODO(PVP): this and other input_ids I tried for generation give pretty bad results. Not sure why. Model might just not be made for auto-regressive inference
     a__    =model.generate(_SCREAMING_SNAKE_CASE,       do_sample=_SCREAMING_SNAKE_CASE			)
     self.assertListEqual(output_ids[0].numpy().tolist(),       _SCREAMING_SNAKE_CASE			)

'''simple docstring'''




import itertools





def 	_UpperCAmelCase       (				_UpperCamelCase       :      int = 1_00_00_00  )      -> int:
    A_      =  [1]

    for i in itertools.count(len(_UpperCamelCase  )  ):
        A_      =  0
        for j in itertools.count(1  ):
            A_      =  -1 if j % 2 == 0 else +1
            A_      =  (j * j * 3 - j) // 2
            if index > i:
                break
            item += partitions[i - index] * sign
            item %= number
            index += j
            if index > i:
                break
            item += partitions[i - index] * sign
            item %= number

        if item == 0:
            return i
        partitions.append(_UpperCamelCase  )

    return 0


if __name__ == "__main__":
  import doctest

  doctest.testmod()

  print(F"""{solution() = }""")






