





"""simple docstring"""




import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
   import torch

   from transformers import (
       OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST,
       OpenAIGPTConfig,
       OpenAIGPTDoubleHeadsModel,
       OpenAIGPTForSequenceClassification,
       OpenAIGPTLMHeadModel,
       OpenAIGPTModel,
   )

class  _UpperCAmelCase  :



      '''simple docstring'''

      def __init__(     self					,      snake_case_					,      snake_case_=1_3					,      snake_case_=7					,      snake_case_=True					,      snake_case_=True					,      snake_case_=True					,      snake_case_=9_9					,      snake_case_=3_2					,      snake_case_=5					,      snake_case_=4					,      snake_case_=3_7					,      snake_case_="gelu"					,      snake_case_=0.1					,      snake_case_=0.1					,      snake_case_=5_1_2					,      snake_case_=1_6					,      snake_case_=2					,      snake_case_=0.02					,      snake_case_=3					,      snake_case_=4					,      snake_case_=None					,      ):

        """simple docstring"""

        A_  :       Union[str, Any]										=  parent
        A_  :       Union[str, Any]										=  batch_size
        A_  :       Optional[Any]										=  seq_length
        A_  :       int										=  is_training
        A_  :       Union[str, Any]										=  use_token_type_ids
        A_  :       Optional[int]										=  use_labels
        A_  :       List[str]										=  vocab_size
        A_  :       Optional[int]										=  hidden_size
        A_  :       Tuple										=  num_hidden_layers
        A_  :       Optional[Any]										=  num_attention_heads
        A_  :       Optional[int]										=  intermediate_size
        A_  :       List[str]										=  hidden_act
        A_  :       Dict										=  hidden_dropout_prob
        A_  :       List[Any]										=  attention_probs_dropout_prob
        A_  :       str										=  max_position_embeddings
        A_  :       List[str]										=  type_vocab_size
        A_  :       Dict										=  type_sequence_label_size
        A_  :       Tuple										=  initializer_range
        A_  :       Tuple										=  num_labels
        A_  :       List[Any]										=  num_choices
        A_  :       Tuple										=  scope
        A_  :       List[Any]										=  self.vocab_size - 1

      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        A_  :       List[Any]										=  ids_tensor([self.batch_size, self.seq_length]					,      self.vocab_size  )

        A_  :       Union[str, Any]										=  None
        if self.use_token_type_ids:
          A_  :       Any										=  ids_tensor([self.batch_size, self.seq_length]					,      self.type_vocab_size  )

        A_  :       Any										=  None
        A_  :       Any										=  None
        A_  :       Dict										=  None
        if self.use_labels:
          A_  :       str										=  ids_tensor([self.batch_size]					,      self.type_sequence_label_size  )
          A_  :       int										=  ids_tensor([self.batch_size, self.seq_length]					,      self.num_labels  )
          A_  :       Union[str, Any]										=  ids_tensor([self.batch_size]					,      self.num_choices  )

        A_  :       List[Any]										=  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					,      )

        A_  :       Optional[Any]										=  ids_tensor([self.num_hidden_layers, self.num_attention_heads]					,      2  )

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

      def 					lowerCamelCase_			(     self					,      snake_case_					,      snake_case_					,      snake_case_					,      snake_case_					,      *snake_case_  ):

        """simple docstring"""

        A_  :       List[str]										=  OpenAIGPTModel(config=snake_case_  )
        model.to(snake_case_  )
        model.eval()

        A_  :       List[str]										=  model(snake_case_					,      token_type_ids=snake_case_					,      head_mask=snake_case_  )
        A_  :       List[str]										=  model(snake_case_					,      token_type_ids=snake_case_  )
        A_  :       Dict										=  model(snake_case_  )

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

      def 					lowerCamelCase_			(     self					,      snake_case_					,      snake_case_					,      snake_case_					,      snake_case_					,      *snake_case_  ):

        """simple docstring"""

        A_  :       Optional[int]										=  OpenAIGPTLMHeadModel(snake_case_  )
        model.to(snake_case_  )
        model.eval()

        A_  :       str										=  model(snake_case_					,      token_type_ids=snake_case_					,      labels=snake_case_  )
        self.parent.assertEqual(result.loss.shape					,      ()  )
        self.parent.assertEqual(result.logits.shape					,      (self.batch_size, self.seq_length, self.vocab_size)  )

      def 					lowerCamelCase_			(     self					,      snake_case_					,      snake_case_					,      snake_case_					,      snake_case_					,      *snake_case_  ):

        """simple docstring"""

        A_  :       Optional[int]										=  OpenAIGPTDoubleHeadsModel(snake_case_  )
        model.to(snake_case_  )
        model.eval()

        A_  :       Optional[int]										=  model(snake_case_					,      token_type_ids=snake_case_					,      labels=snake_case_  )
        self.parent.assertEqual(result.loss.shape					,      ()  )
        self.parent.assertEqual(result.logits.shape					,      (self.batch_size, self.seq_length, self.vocab_size)  )

      def 					lowerCamelCase_			(     self					,      snake_case_					,      snake_case_					,      snake_case_					,      snake_case_					,      *snake_case_  ):

        """simple docstring"""

        A_  :       int										=  self.num_labels
        A_  :       List[Any]										=  OpenAIGPTForSequenceClassification(snake_case_  )
        model.to(snake_case_  )
        model.eval()

        A_  :       str										=  ids_tensor([self.batch_size]					,      self.type_sequence_label_size  )
        A_  :       List[Any]										=  model(snake_case_					,      token_type_ids=snake_case_					,      labels=snake_case_  )
        self.parent.assertEqual(result.logits.shape					,      (self.batch_size, self.num_labels)  )

      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        A_  :       Optional[Any]										=  self.prepare_config_and_inputs()
        (
            (
            A_
        )						,       (
            A_
        )						,       (
            A_
        )						,       (
            A_
        )						,       (
            A_
        )						,       (
            A_
        )						,       (
            A_
        )						,       
        )  :       List[str]										=  config_and_inputs
        A_  :       Tuple										=  {
            'input_ids': input_ids,
            'token_type_ids': token_type_ids,
            'head_mask': head_mask,
        }

        return config, inputs_dict


@require_torch
class  _UpperCAmelCase  (   UpperCAmelCase__			,   UpperCAmelCase__			,   UpperCAmelCase__			,   unittest.TestCase		):



      '''simple docstring'''
      lowercase_      : Optional[Any]	    =   (
          (OpenAIGPTModel, OpenAIGPTLMHeadModel, OpenAIGPTDoubleHeadsModel, OpenAIGPTForSequenceClassification)
          if is_torch_available()
          else ()
      )
      lowercase_      : str	    =   (
          (OpenAIGPTLMHeadModel,) if is_torch_available() else ()
      )  # TODO (PVP): Add Double HeadsModel when generate() function is changed accordingly
      lowercase_      : int	    =   (
          {
              """feature-extraction""": OpenAIGPTModel,
              """text-classification""": OpenAIGPTForSequenceClassification,
              """text-generation""": OpenAIGPTLMHeadModel,
              """zero-shot""": OpenAIGPTForSequenceClassification,
          }
          if is_torch_available()
          else {}
      )

      def 					lowerCamelCase_			(     self					,      snake_case_					,      snake_case_					,      snake_case_					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

        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 					lowerCamelCase_			(     self					,      snake_case_					,      snake_case_					,      snake_case_=False  ):

        """simple docstring"""

        A_  :       str										=  super()._prepare_for_class(snake_case_					,      snake_case_					,      return_labels=snake_case_  )

        if return_labels:
          if model_class.__name__ == "OpenAIGPTDoubleHeadsModel":
            A_  :       Dict										=  torch.zeros(
                (self.model_tester.batch_size, self.model_tester.num_choices, self.model_tester.seq_length)					,      dtype=torch.long					,      device=snake_case_					,      )
            A_  :       int										=  inputs_dict['labels']
            A_  :       Optional[int]										=  inputs_dict['labels']
            A_  :       List[Any]										=  torch.zeros(
                (self.model_tester.batch_size, self.model_tester.num_choices)					,      dtype=torch.long					,      device=snake_case_					,      )
            A_  :       Any										=  torch.zeros(
                self.model_tester.batch_size					,      dtype=torch.long					,      device=snake_case_  )
        return inputs_dict

      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        A_  :       Tuple										=  OpenAIGPTModelTester(self  )
        A_  :       Tuple										=  ConfigTester(self					,      config_class=snake_case_					,      n_embd=3_7  )

      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        self.config_tester.run_common_tests()

      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        A_  :       List[str]										=  self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_openai_gpt_model(*snake_case_  )

      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        A_  :       Union[str, Any]										=  self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_lm_head_model(*snake_case_  )

      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        A_  :       Any										=  self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_double_lm_head_model(*snake_case_  )

      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        A_  :       int										=  self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_openai_gpt_for_sequence_classification(*snake_case_  )

      @slow
      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        for model_name in OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
          A_  :       Any										=  OpenAIGPTModel.from_pretrained(snake_case_  )
          self.assertIsNotNone(snake_case_  )


@require_torch
class  _UpperCAmelCase  (   unittest.TestCase		):



      '''simple docstring'''

      @slow
      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        A_  :       Optional[int]										=  OpenAIGPTLMHeadModel.from_pretrained('openai-gpt'  )
        model.to(snake_case_  )
        A_  :       Optional[Any]										=  torch.tensor([[4_8_1, 4_7_3_5, 5_4_4]]					,      dtype=torch.long					,      device=snake_case_  )  # the president is
        A_  :       List[str]										=  [
            4_8_1,
            4_7_3_5,
            5_4_4,
            2_4_6,
            9_6_3,
            8_7_0,
            7_6_2,
            2_3_9,
            2_4_4,
            4_0_4_7_7,
            2_4_4,
            2_4_9,
            7_1_9,
            8_8_1,
            4_8_7,
            5_4_4,
            2_4_0,
            2_4_4,
            6_0_3,
            4_8_1,
        ]  # the president is a very good man. " \n " i\'m sure he is, " said the

        A_  :       List[str]										=  model.generate(snake_case_					,      do_sample=snake_case_  )
        self.assertListEqual(output_ids[0].tolist()					,      snake_case_  )






"""simple docstring"""




from copy import deepcopy

class  _UpperCAmelCase  :



      '''simple docstring'''

      def __init__(     self					,      snake_case_ = None					,      snake_case_ = None  ):

        """simple docstring"""

        if arr is None and size is not None:
          A_  :       Union[str, Any]										=  size
          A_  :       List[str]										=  [0] * size
        elif arr is not None:
          self.init(snake_case_  )
        else:
          raise ValueError('Either arr or size must be specified'  )

      def 					lowerCamelCase_			(     self					,      snake_case_  ):

        """simple docstring"""

        A_  :       Union[str, Any]										=  len(snake_case_  )
        A_  :       Optional[int]										=  deepcopy(snake_case_  )
        for i in range(1					,      self.size  ):
          A_  :       Optional[Any]										=  self.next_(snake_case_  )
          if j < self.size:
            self.tree[j] += self.tree[i]

      def 					lowerCamelCase_			(     self  ):

        """simple docstring"""

        A_  :       int										=  self.tree[:]
        for i in range(self.size - 1					,      0					,      -1  ):
          A_  :       Optional[int]										=  self.next_(snake_case_  )
          if j < self.size:
            arr[j] -= arr[i]
        return arr

      @staticmethod
      def 					lowerCamelCase_			(     snake_case_  ):

        """simple docstring"""

        return index + (index & (-index))

      @staticmethod
      def 					lowerCamelCase_			(     snake_case_  ):

        """simple docstring"""

        return index - (index & (-index))

      def 					lowerCamelCase_			(     self					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

        if index == 0:
          self.tree[0] += value
          return
        while index < self.size:
          self.tree[index] += value
          A_  :       List[str]										=  self.next_(snake_case_  )

      def 					lowerCamelCase_			(     self					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

        self.add(snake_case_					,      value - self.get(snake_case_  )  )

      def 					lowerCamelCase_			(     self					,      snake_case_  ):

        """simple docstring"""

        if right == 0:
          return 0
        A_  :       Any										=  self.tree[0]
        right -= 1  # make right inclusive
        while right > 0:
          result += self.tree[right]
          A_  :       Tuple										=  self.prev(snake_case_  )
        return result

      def 					lowerCamelCase_			(     self					,      snake_case_					,      snake_case_  ):

        """simple docstring"""

        return self.prefix(snake_case_  ) - self.prefix(snake_case_  )

      def 					lowerCamelCase_			(     self					,      snake_case_  ):

        """simple docstring"""

        return self.query(snake_case_					,      index + 1  )

      def 					lowerCamelCase_			(     self					,      snake_case_  ):

        """simple docstring"""

        value -= self.tree[0]
        if value < 0:
          return -1

        A_  :       List[Any]										=  1  # Largest power of 2 <= size
        while j * 2 < self.size:
          j *= 2

        A_  :       Tuple										=  0

        while j > 0:
          if i + j < self.size and self.tree[i + j] <= value:
            value -= self.tree[i + j]
            i += j
          j //= 2
        return i


if __name__ == "__main__":
   import doctest

   doctest.testmod()