'''simple docstring'''




from __future__ import annotations

import unittest

from transformers import OpenAIGPTConfig, 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.models.openai.modeling_tf_openai import (
      TF_OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST,
      TFOpenAIGPTDoubleHeadsModel,
      TFOpenAIGPTForSequenceClassification,
      TFOpenAIGPTLMHeadModel,
      TFOpenAIGPTModel,
  )

class      lowerCamelCase_     :
      '''simple docstring'''

      def __init__(	self						:    Dict						, A						:    Tuple						, ):
             _UpperCAmelCase					:  Optional[Any]											=   parent
             _UpperCAmelCase					:  Optional[Any]											=   13
             _UpperCAmelCase					:  str											=   7
             _UpperCAmelCase					:  Any											=   True
             _UpperCAmelCase					:  Dict											=   True
             _UpperCAmelCase					:  Union[str, Any]											=   True
             _UpperCAmelCase					:  Tuple											=   True
             _UpperCAmelCase					:  List[Any]											=   True
             _UpperCAmelCase					:  Tuple											=   99
             _UpperCAmelCase					:  str											=   32
             _UpperCAmelCase					:  List[str]											=   2
             _UpperCAmelCase					:  Dict											=   4
             _UpperCAmelCase					:  int											=   37
             _UpperCAmelCase					:  Union[str, Any]											=   "gelu"
             _UpperCAmelCase					:  Any											=   0.1
             _UpperCAmelCase					:  List[str]											=   0.1
             _UpperCAmelCase					:  List[Any]											=   512
             _UpperCAmelCase					:  int											=   16
             _UpperCAmelCase					:  Dict											=   2
             _UpperCAmelCase					:  Any											=   0.02
             _UpperCAmelCase					:  List[str]											=   3
             _UpperCAmelCase					:  Optional[Any]											=   4
             _UpperCAmelCase					:  str											=   None
             _UpperCAmelCase					:  Union[str, Any]											=   self.vocab_size - 1




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

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

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

             _UpperCAmelCase					:  Tuple											=   None
             if self.use_mc_token_ids:
                    _UpperCAmelCase					:  List[Any]											=   ids_tensor([self.batch_size, self.num_choices]						, self.seq_length	)

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

             _UpperCAmelCase					:  int											=   OpenAIGPTConfig(
                 vocab_size=self.vocab_size						, n_embd=self.hidden_size						, n_layer=self.num_hidden_layers						, n_head=self.num_attention_heads						, n_positions=self.max_position_embeddings						, pad_token_id=self.pad_token_id						, )

             _UpperCAmelCase					:  Union[str, Any]											=   ids_tensor([self.num_hidden_layers, self.num_attention_heads]						, 2	)

             return (
                 config,
                 input_ids,
                 input_mask,
                 head_mask,
                 token_type_ids,
                 mc_token_ids,
                 sequence_labels,
                 token_labels,
                 choice_labels,
             )




      def 					_A     (	self						:    List[str]						, A						:    Optional[int]						, A						:    int						, A						:    Dict						, A						:    Tuple						, A						:    str						, *A						:    int	):
             _UpperCAmelCase					:  Any											=   TFOpenAIGPTModel(config=A	)
             _UpperCAmelCase					:  int											=   {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
             _UpperCAmelCase					:  Dict											=   model(A	)

             _UpperCAmelCase					:  Optional[Any]											=   [input_ids, input_mask]
             _UpperCAmelCase					:  Optional[int]											=   model(A	)

             _UpperCAmelCase					:  List[str]											=   model(A	)

             self.parent.assertEqual(result.last_hidden_state.shape						, (self.batch_size, self.seq_length, self.hidden_size)	)




      def 					_A     (	self						:    Optional[int]						, A						:    Tuple						, A						:    Optional[Any]						, A						:    List[Any]						, A						:    str						, A						:    Tuple						, *A						:    int	):
             _UpperCAmelCase					:  Any											=   TFOpenAIGPTLMHeadModel(config=A	)
             _UpperCAmelCase					:  Optional[int]											=   {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
             _UpperCAmelCase					:  Union[str, Any]											=   model(A	)
             self.parent.assertEqual(result.logits.shape						, (self.batch_size, self.seq_length, self.vocab_size)	)




      def 					_A     (	self						:    Optional[Any]						, A						:    List[str]						, A						:    Tuple						, A						:    str						, A						:    Optional[Any]						, A						:    Dict						, A						:    Union[str, Any]						, *A						:    Dict	):
             _UpperCAmelCase					:  List[str]											=   TFOpenAIGPTDoubleHeadsModel(config=A	)

             _UpperCAmelCase					:  List[Any]											=   tf.tile(tf.expand_dims(A						, 1	)						, (1, self.num_choices, 1)	)
             _UpperCAmelCase					:  List[Any]											=   tf.tile(tf.expand_dims(A						, 1	)						, (1, self.num_choices, 1)	)
             _UpperCAmelCase					:  Any											=   tf.tile(tf.expand_dims(A						, 1	)						, (1, self.num_choices, 1)	)

             _UpperCAmelCase					:  List[str]											=   {
                 "input_ids": multiple_choice_inputs_ids,
                 "mc_token_ids": mc_token_ids,
                 "attention_mask": multiple_choice_input_mask,
                 "token_type_ids": multiple_choice_token_type_ids,
             }
             _UpperCAmelCase					:  Tuple											=   model(A	)
             self.parent.assertEqual(
                 result.logits.shape						, (self.batch_size, self.num_choices, self.seq_length, self.vocab_size)	)
             self.parent.assertEqual(result.mc_logits.shape						, (self.batch_size, self.num_choices)	)




      def 					_A     (	self						:    Dict						, A						:    Any						, A						:    Union[str, Any]						, A						:    Optional[int]						, A						:    Union[str, Any]						, A						:    Optional[int]						, *A						:    Optional[int]	):
             _UpperCAmelCase					:  Union[str, Any]											=   self.num_labels
             _UpperCAmelCase					:  List[str]											=   ids_tensor([self.batch_size]						, self.type_sequence_label_size	)
             _UpperCAmelCase					:  Tuple											=   {
                 "input_ids": input_ids,
                 "attention_mask": input_mask,
                 "token_type_ids": token_type_ids,
                 "labels": sequence_labels,
             }
             _UpperCAmelCase					:  List[Any]											=   TFOpenAIGPTForSequenceClassification(A	)
             _UpperCAmelCase					:  Optional[Any]											=   model(A	)
             self.parent.assertEqual(result.logits.shape						, (self.batch_size, self.num_labels)	)





      def 					_A     (	self						:    Tuple	):
             _UpperCAmelCase					:  List[Any]											=   self.prepare_config_and_inputs()

             (
                 (
                 _UpperCAmelCase
             )				,  (
                 _UpperCAmelCase
             )				,  (
                 _UpperCAmelCase
             )				,  (
                 _UpperCAmelCase
             )				,  (
                 _UpperCAmelCase
             )				,  (
                 _UpperCAmelCase
             )				,  (
                 _UpperCAmelCase
             )				,  (
                 _UpperCAmelCase
             )				,  (
                 _UpperCAmelCase
             )				,  
             )					:  Union[str, Any]											=   config_and_inputs

             _UpperCAmelCase					:  List[str]											=   {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
             return config, inputs_dict

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

      __UpperCamelCase:  int      =						(
          (TFOpenAIGPTModel, TFOpenAIGPTLMHeadModel, TFOpenAIGPTDoubleHeadsModel, TFOpenAIGPTForSequenceClassification)
          if is_tf_available()
          else ()
      )
      __UpperCamelCase:  Optional[Any]      =						(
          (TFOpenAIGPTLMHeadModel,) if is_tf_available() else ()
      )  # TODO (PVP): Add Double HeadsModel when generate() function is changed accordingly
      __UpperCamelCase:  int      =						(
          {
              "feature-extraction": TFOpenAIGPTModel,
              "text-classification": TFOpenAIGPTForSequenceClassification,
              "text-generation": TFOpenAIGPTLMHeadModel,
              "zero-shot": TFOpenAIGPTForSequenceClassification,
          }
          if is_tf_available()
          else {}
      )
      __UpperCamelCase:  str      =						False
      __UpperCamelCase:  Dict      =						False




      def 					_A     (	self						:    Any						, A						:    str						, A						:    str						, A						:    Tuple						, A						:    Dict						, A						:    Any	):
             if pipeline_test_casse_name == "ZeroShotClassificationPipelineTests":
                    # Get `tokenizer does not have a padding token` error for both fast/slow tokenizers.
                    # `OpenAIGPTConfig` was never used in pipeline tests, either because of a missing checkpoint or because a
                    # tiny config could not be created.
                    return True

             return False




      def 					_A     (	self						:    Any	):
             _UpperCAmelCase					:  List[str]											=   TFOpenAIGPTModelTester(self	)
             _UpperCAmelCase					:  List[str]											=   ConfigTester(self						, config_class=A						, n_embd=37	)




      def 					_A     (	self						:    Union[str, Any]	):
             self.config_tester.run_common_tests()




      def 					_A     (	self						:    str	):
             _UpperCAmelCase					:  Any											=   self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_openai_gpt_model(*A	)




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




      def 					_A     (	self						:    Tuple	):
             _UpperCAmelCase					:  int											=   self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_openai_gpt_double_head(*A	)




      def 					_A     (	self						:    Dict	):
             _UpperCAmelCase					:  int											=   self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_openai_gpt_for_sequence_classification(*A	)





      @slow
      def 					_A     (	self						:    Dict	):
             for model_name in TF_OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                    _UpperCAmelCase					:  int											=   TFOpenAIGPTModel.from_pretrained(A	)
                    self.assertIsNotNone(A	)

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

      @slow
      def 					_A     (	self						:    Optional[Any]	):
             _UpperCAmelCase					:  Optional[int]											=   TFOpenAIGPTLMHeadModel.from_pretrained("openai-gpt"	)
             _UpperCAmelCase					:  List[str]											=   tf.convert_to_tensor([[481, 4735, 544]]						, dtype=tf.intaa	)  # the president is
             _UpperCAmelCase					:  int											=   [
                 481,
                 4735,
                 544,
                 246,
                 963,
                 870,
                 762,
                 239,
                 244,
                 40477,
                 244,
                 249,
                 719,
                 881,
                 487,
                 544,
                 240,
                 244,
                 603,
                 481,
             ]  # the president is a very good man. " \n " i\'m sure he is, " said the

             _UpperCAmelCase					:  Optional[int]											=   model.generate(A						, do_sample=A	)
             self.assertListEqual(output_ids[0].numpy().tolist()						, A	)


'''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()

