'''simple docstring'''




import unittest

from transformers import ElectraConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, slow, torch_device

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 (
      MODEL_FOR_PRETRAINING_MAPPING,
      ElectraForCausalLM,
      ElectraForMaskedLM,
      ElectraForMultipleChoice,
      ElectraForPreTraining,
      ElectraForQuestionAnswering,
      ElectraForSequenceClassification,
      ElectraForTokenClassification,
      ElectraModel,
  )
  from transformers.models.electra.modeling_electra import ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST

class      lowerCamelCase_     :
      '''simple docstring'''

      def __init__(	self						:    Any						, A						:    List[str]						, A						:    List[str]=13						, A						:    Optional[Any]=7						, A						:    str=True						, A						:    List[Any]=True						, A						:    Optional[Any]=True						, A						:    List[Any]=True						, A						:    int=99						, A						:    Any=32						, A						:    List[Any]=5						, A						:    int=4						, A						:    List[Any]=37						, A						:    Any="gelu"						, A						:    List[str]=0.1						, A						:    str=0.1						, A						:    str=512						, A						:    Optional[Any]=16						, A						:    Optional[Any]=2						, A						:    Dict=0.02						, A						:    Any=3						, A						:    Optional[int]=4						, A						:    Optional[int]=None						, ):
             _UpperCAmelCase					:  Optional[int]											=   parent
             _UpperCAmelCase					:  str											=   batch_size
             _UpperCAmelCase					:  Any											=   seq_length
             _UpperCAmelCase					:  Optional[int]											=   is_training
             _UpperCAmelCase					:  Optional[int]											=   use_input_mask
             _UpperCAmelCase					:  Any											=   use_token_type_ids
             _UpperCAmelCase					:  List[str]											=   use_labels
             _UpperCAmelCase					:  str											=   vocab_size
             _UpperCAmelCase					:  Dict											=   hidden_size
             _UpperCAmelCase					:  Optional[Any]											=   num_hidden_layers
             _UpperCAmelCase					:  Dict											=   num_attention_heads
             _UpperCAmelCase					:  str											=   intermediate_size
             _UpperCAmelCase					:  Union[str, Any]											=   hidden_act
             _UpperCAmelCase					:  Optional[Any]											=   hidden_dropout_prob
             _UpperCAmelCase					:  List[str]											=   attention_probs_dropout_prob
             _UpperCAmelCase					:  int											=   max_position_embeddings
             _UpperCAmelCase					:  Tuple											=   type_vocab_size
             _UpperCAmelCase					:  str											=   type_sequence_label_size
             _UpperCAmelCase					:  Union[str, Any]											=   initializer_range
             _UpperCAmelCase					:  List[Any]											=   num_labels
             _UpperCAmelCase					:  List[Any]											=   num_choices
             _UpperCAmelCase					:  str											=   scope




      def 					_A     (	self						:    List[str]	):
             _UpperCAmelCase					:  Any											=   ids_tensor([self.batch_size, self.seq_length]						, self.vocab_size	)

             _UpperCAmelCase					:  Tuple											=   None
             if self.use_input_mask:
                    _UpperCAmelCase					:  List[Any]											=   random_attention_mask([self.batch_size, self.seq_length]	)

             _UpperCAmelCase					:  Optional[Any]											=   None
             if self.use_token_type_ids:
                    _UpperCAmelCase					:  Optional[int]											=   ids_tensor([self.batch_size, self.seq_length]						, self.type_vocab_size	)

             _UpperCAmelCase					:  Union[str, Any]											=   None
             _UpperCAmelCase					:  Tuple											=   None
             _UpperCAmelCase					:  Optional[int]											=   None
             if self.use_labels:
                    _UpperCAmelCase					:  List[str]											=   ids_tensor([self.batch_size]						, self.type_sequence_label_size	)
                    _UpperCAmelCase					:  List[str]											=   ids_tensor([self.batch_size, self.seq_length]						, self.num_labels	)
                    _UpperCAmelCase					:  Any											=   ids_tensor([self.batch_size]						, self.num_choices	)
                    _UpperCAmelCase					:  str											=   ids_tensor([self.batch_size, self.seq_length]						, 1	)

             _UpperCAmelCase					:  List[Any]											=   self.get_config()

             return (
                 config,
                 input_ids,
                 token_type_ids,
                 input_mask,
                 sequence_labels,
                 token_labels,
                 choice_labels,
                 fake_token_labels,
             )




      def 					_A     (	self						:    Optional[Any]	):
             return ElectraConfig(
                 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						, type_vocab_size=self.type_vocab_size						, is_decoder=A						, initializer_range=self.initializer_range						, )




      def 					_A     (	self						:    Optional[Any]	):
             (
                 (
                 _UpperCAmelCase
             )				,  (
                 _UpperCAmelCase
             )				,  (
                 _UpperCAmelCase
             )				,  (
                 _UpperCAmelCase
             )				,  (
                 _UpperCAmelCase
             )				,  (
                 _UpperCAmelCase
             )				,  (
                 _UpperCAmelCase
             )				,  (
                 _UpperCAmelCase
             )				,  
             )					:  str											=   self.prepare_config_and_inputs()

             _UpperCAmelCase					:  List[Any]											=   True
             _UpperCAmelCase					:  int											=   floats_tensor([self.batch_size, self.seq_length, self.hidden_size]	)
             _UpperCAmelCase					:  List[Any]											=   ids_tensor([self.batch_size, self.seq_length]						, vocab_size=2	)

             return (
                 config,
                 input_ids,
                 token_type_ids,
                 input_mask,
                 sequence_labels,
                 token_labels,
                 choice_labels,
                 encoder_hidden_states,
                 encoder_attention_mask,
             )




      def 					_A     (	self						:    Any						, A						:    Union[str, Any]						, A						:    Dict						, A						:    str						, A						:    int						, A						:    Dict						, A						:    Union[str, Any]						, A						:    List[Any]						, A						:    Optional[Any]						, ):
             _UpperCAmelCase					:  List[str]											=   ElectraModel(config=A	)
             model.to(A	)
             model.eval()
             _UpperCAmelCase					:  str											=   model(A						, attention_mask=A						, token_type_ids=A	)
             _UpperCAmelCase					:  Dict											=   model(A						, token_type_ids=A	)
             _UpperCAmelCase					:  Any											=   model(A	)
             self.parent.assertEqual(result.last_hidden_state.shape						, (self.batch_size, self.seq_length, self.hidden_size)	)




      def 					_A     (	self						:    Any						, A						:    Optional[Any]						, A						:    Dict						, A						:    str						, A						:    Any						, A						:    List[Any]						, A						:    int						, A						:    Tuple						, A						:    str						, A						:    Union[str, Any]						, ):
             _UpperCAmelCase					:  List[Any]											=   True
             _UpperCAmelCase					:  Optional[int]											=   ElectraModel(A	)
             model.to(A	)
             model.eval()
             _UpperCAmelCase					:  Optional[int]											=   model(
                 A						, attention_mask=A						, token_type_ids=A						, encoder_hidden_states=A						, encoder_attention_mask=A						, )
             _UpperCAmelCase					:  Tuple											=   model(
                 A						, attention_mask=A						, token_type_ids=A						, encoder_hidden_states=A						, )
             _UpperCAmelCase					:  List[Any]											=   model(A						, attention_mask=A						, token_type_ids=A	)
             self.parent.assertEqual(result.last_hidden_state.shape						, (self.batch_size, self.seq_length, self.hidden_size)	)




      def 					_A     (	self						:    Optional[Any]						, A						:    int						, A						:    Union[str, Any]						, A						:    List[Any]						, A						:    Any						, A						:    Dict						, A						:    Tuple						, A						:    Union[str, Any]						, A						:    List[Any]						, ):
             _UpperCAmelCase					:  List[Any]											=   ElectraForMaskedLM(config=A	)
             model.to(A	)
             model.eval()
             _UpperCAmelCase					:  Dict											=   model(A						, attention_mask=A						, token_type_ids=A						, labels=A	)
             self.parent.assertEqual(result.logits.shape						, (self.batch_size, self.seq_length, self.vocab_size)	)




      def 					_A     (	self						:    Tuple						, A						:    int						, A						:    Tuple						, A						:    Optional[int]						, A						:    str						, A						:    List[str]						, A						:    Union[str, Any]						, A						:    List[Any]						, A						:    List[Any]						, A						:    Dict						, ):
             _UpperCAmelCase					:  Union[str, Any]											=   ElectraForCausalLM(config=A	)
             model.to(A	)
             model.eval()
             _UpperCAmelCase					:  Union[str, Any]											=   model(A						, attention_mask=A						, token_type_ids=A						, labels=A	)
             self.parent.assertEqual(result.logits.shape						, (self.batch_size, self.seq_length, self.vocab_size)	)




      def 					_A     (	self						:    Tuple						, A						:    Any						, A						:    Any						, A						:    List[str]						, A						:    List[str]						, A						:    int						, A						:    Union[str, Any]						, A						:    List[str]						, A						:    Dict						, ):
             _UpperCAmelCase					:  Any											=   self.num_labels
             _UpperCAmelCase					:  Optional[Any]											=   ElectraForTokenClassification(config=A	)
             model.to(A	)
             model.eval()
             _UpperCAmelCase					:  Any											=   model(A						, attention_mask=A						, token_type_ids=A						, labels=A	)
             self.parent.assertEqual(result.logits.shape						, (self.batch_size, self.seq_length, self.num_labels)	)




      def 					_A     (	self						:    List[str]						, A						:    Tuple						, A						:    Dict						, A						:    List[Any]						, A						:    int						, A						:    str						, A						:    Tuple						, A						:    int						, A						:    List[str]						, ):
             _UpperCAmelCase					:  Optional[int]											=   self.num_labels
             _UpperCAmelCase					:  str											=   ElectraForPreTraining(config=A	)
             model.to(A	)
             model.eval()
             _UpperCAmelCase					:  Dict											=   model(A						, attention_mask=A						, token_type_ids=A						, labels=A	)
             self.parent.assertEqual(result.logits.shape						, (self.batch_size, self.seq_length)	)




      def 					_A     (	self						:    Tuple						, A						:    Dict						, A						:    Any						, A						:    Union[str, Any]						, A						:    Optional[int]						, A						:    int						, A						:    int						, A						:    Union[str, Any]						, A						:    Optional[Any]						, ):
             _UpperCAmelCase					:  Union[str, Any]											=   self.num_labels
             _UpperCAmelCase					:  Tuple											=   ElectraForSequenceClassification(A	)
             model.to(A	)
             model.eval()
             _UpperCAmelCase					:  int											=   model(A						, attention_mask=A						, token_type_ids=A						, labels=A	)
             self.parent.assertEqual(result.logits.shape						, (self.batch_size, self.num_labels)	)




      def 					_A     (	self						:    Optional[Any]						, A						:    Tuple						, A						:    Optional[Any]						, A						:    Optional[Any]						, A						:    Tuple						, A						:    Optional[Any]						, A						:    Any						, A						:    List[Any]						, A						:    Optional[Any]						, ):
             _UpperCAmelCase					:  str											=   ElectraForQuestionAnswering(config=A	)
             model.to(A	)
             model.eval()
             _UpperCAmelCase					:  List[Any]											=   model(
                 A						, attention_mask=A						, token_type_ids=A						, start_positions=A						, end_positions=A						, )
             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 					_A     (	self						:    int						, A						:    Union[str, Any]						, A						:    Optional[Any]						, A						:    Tuple						, A						:    List[str]						, A						:    Union[str, Any]						, A						:    Any						, A						:    Tuple						, A						:    Tuple						, ):
             _UpperCAmelCase					:  str											=   self.num_choices
             _UpperCAmelCase					:  int											=   ElectraForMultipleChoice(config=A	)
             model.to(A	)
             model.eval()
             _UpperCAmelCase					:  Union[str, Any]											=   input_ids.unsqueeze(1	).expand(-1						, self.num_choices						, -1	).contiguous()
             _UpperCAmelCase					:  Any											=   token_type_ids.unsqueeze(1	).expand(-1						, self.num_choices						, -1	).contiguous()
             _UpperCAmelCase					:  List[str]											=   input_mask.unsqueeze(1	).expand(-1						, self.num_choices						, -1	).contiguous()
             _UpperCAmelCase					:  int											=   model(
                 A						, attention_mask=A						, token_type_ids=A						, labels=A						, )
             self.parent.assertEqual(result.logits.shape						, (self.batch_size, self.num_choices)	)





      def 					_A     (	self						:    int	):
             _UpperCAmelCase					:  Dict											=   self.prepare_config_and_inputs()
             (
                 (
                 _UpperCAmelCase
             )				,  (
                 _UpperCAmelCase
             )				,  (
                 _UpperCAmelCase
             )				,  (
                 _UpperCAmelCase
             )				,  (
                 _UpperCAmelCase
             )				,  (
                 _UpperCAmelCase
             )				,  (
                 _UpperCAmelCase
             )				,  (
                 _UpperCAmelCase
             )				,  
             )					:  str											=   config_and_inputs
             _UpperCAmelCase					:  str											=   {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
             return config, inputs_dict

@require_torch
class      lowerCamelCase_     (snake_case__							,			snake_case__							,			unittest.TestCase   ):
      '''simple docstring'''

      __UpperCamelCase:  Dict      =						(
          (
              ElectraModel,
              ElectraForPreTraining,
              ElectraForMaskedLM,
              ElectraForCausalLM,
              ElectraForMultipleChoice,
              ElectraForTokenClassification,
              ElectraForSequenceClassification,
              ElectraForQuestionAnswering,
          )
          if is_torch_available()
          else ()
      )
      __UpperCamelCase:  Tuple      =						(
          {
              "feature-extraction": ElectraModel,
              "fill-mask": ElectraForMaskedLM,
              "question-answering": ElectraForQuestionAnswering,
              "text-classification": ElectraForSequenceClassification,
              "text-generation": ElectraForCausalLM,
              "token-classification": ElectraForTokenClassification,
              "zero-shot": ElectraForSequenceClassification,
          }
          if is_torch_available()
          else {}
      )
      __UpperCamelCase:  str      =						True




      def 					_A     (	self						:    Tuple						, A						:    int						, A						:    Optional[Any]						, A						:    str=False	):
             _UpperCAmelCase					:  Union[str, Any]											=   super()._prepare_for_class(A						, A						, return_labels=A	)

             if return_labels:
                    if model_class in get_values(A	):
                           _UpperCAmelCase					:  Any											=   torch.zeros(
                               (self.model_tester.batch_size, self.model_tester.seq_length)						, dtype=torch.long						, device=A	)
             return inputs_dict




      def 					_A     (	self						:    Optional[int]	):
             _UpperCAmelCase					:  Optional[int]											=   ElectraModelTester(self	)
             _UpperCAmelCase					:  Optional[Any]											=   ConfigTester(self						, config_class=A						, hidden_size=37	)




      def 					_A     (	self						:    Tuple	):
             self.config_tester.run_common_tests()




      def 					_A     (	self						:    Tuple	):
             _UpperCAmelCase					:  Any											=   self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_electra_model(*A	)




      def 					_A     (	self						:    List[Any]	):
             _UpperCAmelCase					:  Any											=   self.model_tester.prepare_config_and_inputs_for_decoder()
             self.model_tester.create_and_check_electra_model_as_decoder(*A	)




      def 					_A     (	self						:    Optional[int]	):
             _UpperCAmelCase					:  Union[str, Any]											=   self.model_tester.prepare_config_and_inputs()
             for type in ["absolute", "relative_key", "relative_key_query"]:
                    _UpperCAmelCase					:  Dict											=   type
                    self.model_tester.create_and_check_electra_model(*A	)




      def 					_A     (	self						:    str	):
             _UpperCAmelCase					:  Optional[Any]											=   self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_electra_for_masked_lm(*A	)




      def 					_A     (	self						:    Tuple	):
             _UpperCAmelCase					:  str											=   self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_electra_for_token_classification(*A	)




      def 					_A     (	self						:    Any	):
             _UpperCAmelCase					:  Optional[Any]											=   self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_electra_for_pretraining(*A	)




      def 					_A     (	self						:    List[str]	):
             _UpperCAmelCase					:  Any											=   self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_electra_for_sequence_classification(*A	)




      def 					_A     (	self						:    Tuple	):
             _UpperCAmelCase					:  Optional[Any]											=   self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_electra_for_question_answering(*A	)




      def 					_A     (	self						:    Optional[Any]	):
             _UpperCAmelCase					:  Tuple											=   self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_electra_for_multiple_choice(*A	)




      @slow
      def 					_A     (	self						:    Any	):
             for model_name in ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                    _UpperCAmelCase					:  str											=   ElectraModel.from_pretrained(A	)
                    self.assertIsNotNone(A	)





      def 					_A     (	self						:    Dict	):
             _UpperCAmelCase					:  List[Any]											=   self.model_tester.prepare_config_and_inputs_for_decoder()
             self.model_tester.create_and_check_electra_for_causal_lm(*A	)

@require_torch
class      lowerCamelCase_     (unittest.TestCase   ):
      '''simple docstring'''

      @slow
      def 					_A     (	self						:    str	):
             _UpperCAmelCase					:  List[Any]											=   ElectraModel.from_pretrained("google/electra-small-discriminator"	)
             _UpperCAmelCase					:  Optional[Any]											=   torch.tensor([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]	)
             _UpperCAmelCase					:  int											=   torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]	)
             _UpperCAmelCase					:  Any											=   model(A						, attention_mask=A	)[0]
             _UpperCAmelCase					:  str											=   torch.Size((1, 11, 256)	)
             self.assertEqual(output.shape						, A	)
             _UpperCAmelCase					:  Dict											=   torch.tensor(
                 [[[0.4_471, 0.6_821, -0.3_265], [0.4_627, 0.5_255, -0.3_668], [0.4_532, 0.3_313, -0.4_344]]]	)

             self.assertTrue(torch.allclose(output[:, 1:4, 1:4]						, A						, atol=1E-4	)	)


'''simple docstring'''




from itertools import compress, repeat
from math import ceil, sqrt

def      UpperCamelCase_							( _UpperCAmelCase  :				int   )   ->   list[int]:


       """simple docstring"""







       if num <= 2:
              return []
       if num == 3:
              return [2]

       # Odd sieve for numbers in range [3, num - 1]
       _UpperCAmelCase					:  Optional[int]											=   bytearray(b"\x01"   ) * ((num >> 1) - 1)

       for i in range(3      ,     int(sqrt(_UpperCAmelCase   )   ) + 1      ,     2   ):
              if sieve[(i >> 1) - 1]:
                     _UpperCAmelCase					:  List[Any]											=   i**2
                     _UpperCAmelCase					:  Optional[int]											=   repeat(
                         0      ,     ceil((num - i_squared) / (i << 1)   )   )

       return [2] + list(compress(range(3      ,     _UpperCAmelCase      ,     2   )      ,     _UpperCAmelCase   )   )


if __name__ == "__main__":
  import doctest

  doctest.testmod()

