



"""simple docstring"""
import gc
import unittest

from transformers import CTRLConfig, 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, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
       import torch

       from transformers import (
           CTRL_PRETRAINED_MODEL_ARCHIVE_LIST,
           CTRLForSequenceClassification,
           CTRLLMHeadModel,
           CTRLModel,
       )


class    SCREAMING_SNAKE_CASE__ :

    def __init__(				self			:				int ,  lowerCAmelCase_			:				Optional[Any] ,  lowerCAmelCase_			:				Union[str, Any]=1_4 ,  lowerCAmelCase_			:				Tuple=7 ,  lowerCAmelCase_			:				Optional[Any]=True ,  lowerCAmelCase_			:				List[Any]=True ,  lowerCAmelCase_			:				Any=True ,  lowerCAmelCase_			:				Any=True ,  lowerCAmelCase_			:				int=True ,  lowerCAmelCase_			:				Optional[int]=9_9 ,  lowerCAmelCase_			:				str=3_2 ,  lowerCAmelCase_			:				List[Any]=5 ,  lowerCAmelCase_			:				List[str]=4 ,  lowerCAmelCase_			:				Any=3_7 ,  lowerCAmelCase_			:				List[str]="gelu" ,  lowerCAmelCase_			:				Tuple=0.1 ,  lowerCAmelCase_			:				str=0.1 ,  lowerCAmelCase_			:				Optional[Any]=5_1_2 ,  lowerCAmelCase_			:				int=1_6 ,  lowerCAmelCase_			:				Dict=2 ,  lowerCAmelCase_			:				Optional[int]=0.02 ,  lowerCAmelCase_			:				Union[str, Any]=3 ,  lowerCAmelCase_			:				Any=4 ,  lowerCAmelCase_			:				int=None ,  ):







      """simple docstring"""







      lowercase_								= parent
      lowercase_								= batch_size
      lowercase_								= seq_length
      lowercase_								= is_training
      lowercase_								= use_token_type_ids
      lowercase_								= use_input_mask
      lowercase_								= use_labels
      lowercase_								= use_mc_token_ids
      lowercase_								= vocab_size
      lowercase_								= hidden_size
      lowercase_								= num_hidden_layers
      lowercase_								= num_attention_heads
      lowercase_								= intermediate_size
      lowercase_								= hidden_act
      lowercase_								= hidden_dropout_prob
      lowercase_								= attention_probs_dropout_prob
      lowercase_								= max_position_embeddings
      lowercase_								= type_vocab_size
      lowercase_								= type_sequence_label_size
      lowercase_								= initializer_range
      lowercase_								= num_labels
      lowercase_								= num_choices
      lowercase_								= scope
      lowercase_								= self.vocab_size - 1

    def 					_UpperCAmelCase	(				self			:				Optional[int]):







      """simple docstring"""







      lowercase_								= ids_tensor([self.batch_size, self.seq_length] ,  self.vocab_size)

      lowercase_								= None
      if self.use_input_mask:
        lowercase_								= random_attention_mask([self.batch_size, self.seq_length])

      lowercase_								= None
      if self.use_token_type_ids:
        lowercase_								= ids_tensor([self.batch_size, self.seq_length] ,  self.type_vocab_size)

      lowercase_								= None
      if self.use_mc_token_ids:
        lowercase_								= ids_tensor([self.batch_size, self.num_choices] ,  self.seq_length)

      lowercase_								= None
      lowercase_								= None
      lowercase_								= None
      if self.use_labels:
        lowercase_								= ids_tensor([self.batch_size] ,  self.type_sequence_label_size)
        lowercase_								= ids_tensor([self.batch_size, self.seq_length] ,  self.num_labels)
        lowercase_								= ids_tensor([self.batch_size] ,  self.num_choices)

      lowercase_								= self.get_config()

      lowercase_								= 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 					_UpperCAmelCase	(				self			:				Union[str, Any]):







      """simple docstring"""







      return CTRLConfig(
          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 ,  )

    def 					_UpperCAmelCase	(				self			:				Any ,  lowerCAmelCase_			:				Dict ,  lowerCAmelCase_			:				Dict ,  lowerCAmelCase_			:				int ,  lowerCAmelCase_			:				Union[str, Any] ,  lowerCAmelCase_			:				Any ,  *lowerCAmelCase_			:				Any):







      """simple docstring"""







      lowercase_								= CTRLModel(config=lowerCAmelCase_)
      model.to(lowerCAmelCase_)
      model.eval()

      model(lowerCAmelCase_ ,  token_type_ids=lowerCAmelCase_ ,  head_mask=lowerCAmelCase_)
      model(lowerCAmelCase_ ,  token_type_ids=lowerCAmelCase_)
      lowercase_								= model(lowerCAmelCase_)
      self.parent.assertEqual(result.last_hidden_state.shape ,  (self.batch_size, self.seq_length, self.hidden_size))
      self.parent.assertEqual(len(result.past_key_values) ,  config.n_layer)

    def 					_UpperCAmelCase	(				self			:				Union[str, Any] ,  lowerCAmelCase_			:				Dict ,  lowerCAmelCase_			:				Any ,  lowerCAmelCase_			:				Union[str, Any] ,  lowerCAmelCase_			:				Union[str, Any] ,  lowerCAmelCase_			:				Optional[Any] ,  *lowerCAmelCase_			:				Any):







      """simple docstring"""







      lowercase_								= CTRLLMHeadModel(lowerCAmelCase_)
      model.to(lowerCAmelCase_)
      model.eval()

      lowercase_								= model(lowerCAmelCase_ ,  token_type_ids=lowerCAmelCase_ ,  labels=lowerCAmelCase_)
      self.parent.assertEqual(result.loss.shape ,  ())
      self.parent.assertEqual(result.logits.shape ,  (self.batch_size, self.seq_length, self.vocab_size))

    def 					_UpperCAmelCase	(				self			:				int):







      """simple docstring"""







      lowercase_								= self.prepare_config_and_inputs()

      (
          (
          lowercase_
      )					,    (
          lowercase_
      )					,    (
          lowercase_
      )					,    (
          lowercase_
      )					,    (
          lowercase_
      )					,    (
          lowercase_
      )					,    (
          lowercase_
      )					,    (
          lowercase_
      )					,    (
          lowercase_
      )					,    
      )								= config_and_inputs

      lowercase_								= {"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """head_mask""": head_mask}

      return config, inputs_dict




    def 					_UpperCAmelCase	(				self			:				str ,  lowerCAmelCase_			:				Dict ,  lowerCAmelCase_			:				List[str] ,  lowerCAmelCase_			:				Tuple ,  lowerCAmelCase_			:				Tuple ,  *lowerCAmelCase_			:				List[Any]):







      """simple docstring"""







      lowercase_								= self.num_labels
      lowercase_								= CTRLForSequenceClassification(lowerCAmelCase_)
      model.to(lowerCAmelCase_)
      model.eval()
      lowercase_								= ids_tensor([self.batch_size] ,  self.type_sequence_label_size)
      lowercase_								= model(lowerCAmelCase_ ,  token_type_ids=lowerCAmelCase_ ,  labels=lowerCAmelCase_)
      self.parent.assertEqual(result.logits.shape ,  (self.batch_size, self.num_labels))







@require_torch
class    SCREAMING_SNAKE_CASE__ (							__UpperCAmelCase      ,  __UpperCAmelCase      ,  __UpperCAmelCase      ,  unittest.TestCase   ):
    lowercase__      =   (CTRLModel, CTRLLMHeadModel, CTRLForSequenceClassification) if is_torch_available() else ()
    lowercase__      =   (CTRLLMHeadModel,) if is_torch_available() else ()
    lowercase__      =   (
        {
            "feature-extraction": CTRLModel,
            "text-classification": CTRLForSequenceClassification,
            "text-generation": CTRLLMHeadModel,
            "zero-shot": CTRLForSequenceClassification,
        }
        if is_torch_available()
        else {}
    )
    lowercase__      =   True
    lowercase__      =   False
    lowercase__      =   False

    def 					_UpperCAmelCase	(				self			:				List[str] ,  lowerCAmelCase_			:				Union[str, Any] ,  lowerCAmelCase_			:				int ,  lowerCAmelCase_			:				Union[str, Any] ,  lowerCAmelCase_			:				Optional[Any] ,  lowerCAmelCase_			:				Union[str, Any]):







      """simple docstring"""







      if pipeline_test_casse_name == "ZeroShotClassificationPipelineTests":
        # Get `tokenizer does not have a padding token` error for both fast/slow tokenizers.
        # `CTRLConfig` 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 					_UpperCAmelCase	(				self			:				str):







      """simple docstring"""







      lowercase_								= CTRLModelTester(self)
      lowercase_								= ConfigTester(self ,  config_class=lowerCAmelCase_ ,  n_embd=3_7)

    def 					_UpperCAmelCase	(				self			:				List[str]):







      """simple docstring"""







      super().tearDown()
      # clean-up as much as possible GPU memory occupied by PyTorch
      gc.collect()
      torch.cuda.empty_cache()

    def 					_UpperCAmelCase	(				self			:				Any):







      """simple docstring"""







      self.config_tester.run_common_tests()

    def 					_UpperCAmelCase	(				self			:				Optional[Any]):







      """simple docstring"""







      lowercase_								= self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_ctrl_model(*lowerCAmelCase_)

    def 					_UpperCAmelCase	(				self			:				Dict):







      """simple docstring"""







      lowercase_								= self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_lm_head_model(*lowerCAmelCase_)

    @unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests.""")
    def 					_UpperCAmelCase	(				self			:				Optional[Any]):







      """simple docstring"""







      pass

    @slow
    def 					_UpperCAmelCase	(				self			:				Any):







      """simple docstring"""







      for model_name in CTRL_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        lowercase_								= CTRLModel.from_pretrained(lowerCAmelCase_)
        self.assertIsNotNone(lowerCAmelCase_)




    @unittest.skip("""The model doesn't support left padding""")  # and it's not used enough to be worth fixing :)
    def 					_UpperCAmelCase	(				self			:				int):







      """simple docstring"""







      pass







@require_torch
class    SCREAMING_SNAKE_CASE__ (							unittest.TestCase   ):

    def 					_UpperCAmelCase	(				self			:				Dict):







      """simple docstring"""







      super().tearDown()
      # clean-up as much as possible GPU memory occupied by PyTorch
      gc.collect()
      torch.cuda.empty_cache()




    @slow
    def 					_UpperCAmelCase	(				self			:				Dict):







      """simple docstring"""







      lowercase_								= CTRLLMHeadModel.from_pretrained("""ctrl""")
      model.to(lowerCAmelCase_)
      lowercase_								= torch.tensor(
          [[1_1_8_5_9, 0, 1_6_1_1, 8]] ,  dtype=torch.long ,  device=lowerCAmelCase_)  # Legal the president is
      lowercase_								= [
          1_1_8_5_9,
          0,
          1_6_1_1,
          8,
          5,
          1_5_0,
          2_6_4_4_9,
          2,
          1_9,
          3_4_8,
          4_6_9,
          3,
          2_5_9_5,
          4_8,
          2_0_7_4_0,
          2_4_6_5_3_3,
          2_4_6_5_3_3,
          1_9,
          3_0,
          5,
      ]  # Legal the president is a good guy and I don't want to lose my job. \n \n I have a

      lowercase_								= model.generate(lowerCAmelCase_ ,  do_sample=lowerCAmelCase_)
      self.assertListEqual(output_ids[0].tolist() ,  lowerCAmelCase_)






"""simple docstring"""
from typing import Optional

import numpy as np
import torch
from torch import nn
from transformers import GPTaConfig, GPTaLMHeadModel
from transformers.modeling_utils import ModuleUtilsMixin

from ...configuration_utils import ConfigMixin, register_to_config
from ...models import ModelMixin


class    SCREAMING_SNAKE_CASE__ (							__UpperCAmelCase      ,  __UpperCAmelCase      ,  __UpperCAmelCase   ):
    lowercase__      =   [r"h\.\d+\.attn\.bias", r"h\.\d+\.attn\.masked_bias"]

    @register_to_config
    def __init__(				self			:				int ,  lowerCAmelCase_			:				int ,  lowerCAmelCase_			:				int ,  lowerCAmelCase_			:				Optional[int] = None ,  lowerCAmelCase_			:				int = 5_0_2_5_7 ,  lowerCAmelCase_			:				int = 1_0_2_4 ,  lowerCAmelCase_			:				int = 7_6_8 ,  lowerCAmelCase_			:				int = 1_2 ,  lowerCAmelCase_			:				int = 1_2 ,  lowerCAmelCase_			:				Optional[int] = None ,  lowerCAmelCase_			:				str = "gelu_new" ,  lowerCAmelCase_			:				float = 0.1 ,  lowerCAmelCase_			:				float = 0.1 ,  lowerCAmelCase_			:				float = 0.1 ,  lowerCAmelCase_			:				float = 1E-5 ,  lowerCAmelCase_			:				float = 0.02 ,  lowerCAmelCase_			:				bool = True ,  lowerCAmelCase_			:				bool = True ,  lowerCAmelCase_			:				bool = False ,  lowerCAmelCase_			:				bool = False ,  ):







      """simple docstring"""







      super().__init__()

      lowercase_								= prefix_length

      if prefix_inner_dim != n_embd and prefix_hidden_dim is None:
        raise ValueError(
            F'''`prefix_hidden_dim` cannot be `None` when `prefix_inner_dim`: {prefix_hidden_dim} and'''
            F''' `n_embd`: {n_embd} are not equal.''')

      lowercase_								= prefix_inner_dim
      lowercase_								= prefix_hidden_dim

      lowercase_								= (
          nn.Linear(self.prefix_inner_dim ,  self.prefix_hidden_dim)
          if self.prefix_hidden_dim is not None
          else nn.Identity()
      )
      lowercase_								= (
          nn.Linear(self.prefix_hidden_dim ,  lowerCAmelCase_) if self.prefix_hidden_dim is not None else nn.Identity()
      )

      lowercase_								= GPTaConfig(
          vocab_size=lowerCAmelCase_ ,  n_positions=lowerCAmelCase_ ,  n_embd=lowerCAmelCase_ ,  n_layer=lowerCAmelCase_ ,  n_head=lowerCAmelCase_ ,  n_inner=lowerCAmelCase_ ,  activation_function=lowerCAmelCase_ ,  resid_pdrop=lowerCAmelCase_ ,  embd_pdrop=lowerCAmelCase_ ,  attn_pdrop=lowerCAmelCase_ ,  layer_norm_epsilon=lowerCAmelCase_ ,  initializer_range=lowerCAmelCase_ ,  scale_attn_weights=lowerCAmelCase_ ,  use_cache=lowerCAmelCase_ ,  scale_attn_by_inverse_layer_idx=lowerCAmelCase_ ,  reorder_and_upcast_attn=lowerCAmelCase_ ,  )
      lowercase_								= GPTaLMHeadModel(lowerCAmelCase_)

    def 					_UpperCAmelCase	(				self			:				int ,  lowerCAmelCase_			:				torch.Tensor ,  lowerCAmelCase_			:				torch.Tensor ,  lowerCAmelCase_			:				Optional[torch.Tensor] = None ,  lowerCAmelCase_			:				Optional[torch.Tensor] = None ,  ):







      """simple docstring"""







      lowercase_								= self.transformer.transformer.wte(lowerCAmelCase_)
      lowercase_								= self.encode_prefix(lowerCAmelCase_)
      lowercase_								= self.decode_prefix(lowerCAmelCase_)
      lowercase_								= torch.cat((prefix_embeds, embedding_text) ,  dim=1)

      if labels is not None:
        lowercase_								= self.get_dummy_token(input_ids.shape[0] ,  input_ids.device)
        lowercase_								= torch.cat((dummy_token, input_ids) ,  dim=1)
      lowercase_								= self.transformer(inputs_embeds=lowerCAmelCase_ ,  labels=lowerCAmelCase_ ,  attention_mask=lowerCAmelCase_)
      if self.prefix_hidden_dim is not None:
        return out, hidden
      else:
        return out

    def 					_UpperCAmelCase	(				self			:				Optional[int] ,  lowerCAmelCase_			:				int ,  lowerCAmelCase_			:				torch.device):







      """simple docstring"""







      return torch.zeros(lowerCAmelCase_ ,  self.prefix_length ,  dtype=torch.intaa ,  device=lowerCAmelCase_)

    def 					_UpperCAmelCase	(				self			:				List[str] ,  lowerCAmelCase_			:				List[Any]):







      """simple docstring"""







      return self.encode_prefix(lowerCAmelCase_)

    @torch.no_grad()
    def 					_UpperCAmelCase	(				self			:				Tuple ,  lowerCAmelCase_			:				str ,  lowerCAmelCase_			:				Dict ,  lowerCAmelCase_			:				Union[str, Any]):







      """simple docstring"""







      lowercase_								= torch.split(lowerCAmelCase_ ,  1 ,  dim=0)
      lowercase_								= []
      lowercase_								= []
      for feature in features:
        lowercase_								= self.decode_prefix(feature.to(lowerCAmelCase_))  # back to the clip feature
        # Only support beam search for now
        lowercase_					,    lowercase_								= self.generate_beam(
            input_embeds=lowerCAmelCase_ ,  device=lowerCAmelCase_ ,  eos_token_id=lowerCAmelCase_)
        generated_tokens.append(output_tokens[0])
        generated_seq_lengths.append(seq_lengths[0])
      lowercase_								= torch.stack(lowerCAmelCase_)
      lowercase_								= torch.stack(lowerCAmelCase_)
      return generated_tokens, generated_seq_lengths




    @torch.no_grad()
    def 					_UpperCAmelCase	(				self			:				int ,  lowerCAmelCase_			:				List[str]=None ,  lowerCAmelCase_			:				Tuple=None ,  lowerCAmelCase_			:				str=None ,  lowerCAmelCase_			:				int = 5 ,  lowerCAmelCase_			:				int = 6_7 ,  lowerCAmelCase_			:				float = 1.0 ,  lowerCAmelCase_			:				Optional[int] = None ,  ):







      """simple docstring"""







      lowercase_								= eos_token_id
      lowercase_								= None
      lowercase_								= None
      lowercase_								= torch.ones(lowerCAmelCase_ ,  device=lowerCAmelCase_ ,  dtype=torch.int)
      lowercase_								= torch.zeros(lowerCAmelCase_ ,  device=lowerCAmelCase_ ,  dtype=torch.bool)

      if input_embeds is not None:
        lowercase_								= input_embeds
      else:
        lowercase_								= self.transformer.transformer.wte(lowerCAmelCase_)

      for i in range(lowerCAmelCase_):
        lowercase_								= self.transformer(inputs_embeds=lowerCAmelCase_)
        lowercase_								= outputs.logits
        lowercase_								= logits[:, -1, :] / (temperature if temperature > 0 else 1.0)
        lowercase_								= logits.softmax(-1).log()

        if scores is None:
          lowercase_					,    lowercase_								= logits.topk(lowerCAmelCase_ ,  -1)
          lowercase_								= generated.expand(lowerCAmelCase_ ,  *generated.shape[1:])
          lowercase_					,    lowercase_								= next_tokens.permute(1 ,  0), scores.squeeze(0)
          if tokens is None:
            lowercase_								= next_tokens
          else:
            lowercase_								= tokens.expand(lowerCAmelCase_ ,  *tokens.shape[1:])
            lowercase_								= torch.cat((tokens, next_tokens) ,  dim=1)
        else:
          lowercase_								= -float(np.inf)
          lowercase_								= 0
          lowercase_								= scores[:, None] + logits
          seq_lengths[~is_stopped] += 1
          lowercase_								= scores_sum / seq_lengths[:, None]
          lowercase_					,    lowercase_								= scores_sum_average.view(-1).topk(lowerCAmelCase_ ,  -1)
          lowercase_								= next_tokens // scores_sum.shape[1]
          lowercase_								= seq_lengths[next_tokens_source]
          lowercase_								= next_tokens % scores_sum.shape[1]
          lowercase_								= next_tokens.unsqueeze(1)
          lowercase_								= tokens[next_tokens_source]
          lowercase_								= torch.cat((tokens, next_tokens) ,  dim=1)
          lowercase_								= generated[next_tokens_source]
          lowercase_								= scores_sum_average * seq_lengths
          lowercase_								= is_stopped[next_tokens_source]

        lowercase_								= self.transformer.transformer.wte(next_tokens.squeeze()).view(generated.shape[0] ,  1 ,  -1)
        lowercase_								= torch.cat((generated, next_token_embed) ,  dim=1)
        lowercase_								= is_stopped + next_tokens.eq(lowerCAmelCase_).squeeze()
        if is_stopped.all():
          break

      lowercase_								= scores / seq_lengths
      lowercase_								= scores.argsort(descending=lowerCAmelCase_)
      # tokens tensors are already padded to max_seq_length
      lowercase_								= [tokens[i] for i in order]
      lowercase_								= torch.stack(lowerCAmelCase_ ,  dim=0)
      lowercase_								= torch.tensor([seq_lengths[i] for i in order] ,  dtype=seq_lengths.dtype)
      return output_texts, seq_lengths

