






import inspect

import jax
import jax.lax as lax
import jax.numpy as jnp

from ..utils import add_start_docstrings
from ..utils.logging import get_logger


__magic_name__   :  Dict        =							get_logger(__name__)


__magic_name__   :  Dict        =							R'\n    Args:\n        input_ids (`jnp.ndarray` of shape `(batch_size, sequence_length)`):\n            Indices of input sequence tokens in the vocabulary.\n\n            Indices can be obtained using [`PreTrainedTokenizer`]. See [`PreTrainedTokenizer.encode`] and\n            [`PreTrainedTokenizer.__call__`] for details.\n\n            [What are input IDs?](../glossary#input-ids)\n        scores (`jnp.ndarray` of shape `(batch_size, config.vocab_size)`):\n            Prediction scores of a language modeling head. These can be logits for each vocabulary when not using beam\n            search or log softmax for each vocabulary token when using beam search\n        kwargs (`Dict[str, Any]`, *optional*):\n            Additional logits processor specific kwargs.\n\n    Return:\n        `jnp.ndarray` of shape `(batch_size, config.vocab_size)`: The processed prediction scores.\n\n'





class 					A__     :

     '''simple docstring'''






     @add_start_docstrings(_SCREAMING_SNAKE_CASE )
     def __call__( self      :    str   ,						_SCREAMING_SNAKE_CASE      :    Optional[int]   ,						_SCREAMING_SNAKE_CASE      :    Any ):






           """simple docstring"""




           raise NotImplementedError(
               f'{self.__class__} is an abstract class. Only classes inheriting this class can be called.' )





class 					A__     :

     '''simple docstring'''






     @add_start_docstrings(_SCREAMING_SNAKE_CASE )
     def __call__( self      :    Dict   ,						_SCREAMING_SNAKE_CASE      :    int   ,						_SCREAMING_SNAKE_CASE      :    List[str] ):






           """simple docstring"""




           raise NotImplementedError(
               f'{self.__class__} is an abstract class. Only classes inheriting this class can be called.' )





class 					A__     (     _lowercase ):

     '''simple docstring'''






     @add_start_docstrings(_SCREAMING_SNAKE_CASE )
     def __call__( self      :    Tuple   ,						_SCREAMING_SNAKE_CASE      :    str   ,						_SCREAMING_SNAKE_CASE      :    Dict   ,						_SCREAMING_SNAKE_CASE      :    str   ,						**_SCREAMING_SNAKE_CASE      :    Optional[int] ):






           """simple docstring"""




           for processor in self:
                 UpperCamelCase					       =  inspect.signature(processor.__call__ ).parameters
                 if len(_SCREAMING_SNAKE_CASE ) > 3:
                       if not all(arg in kwargs for arg in list(function_args.keys() )[2:] ):
                             raise ValueError(
                                 f'Make sure that all the required parameters: {list(function_args.keys() )} for '
                                 f'{processor.__class__} are passed to the logits processor.' )
                       UpperCamelCase					       =  processor(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						**_SCREAMING_SNAKE_CASE )
                 else:
                       UpperCamelCase					       =  processor(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE )
           return scores





class 					A__     (     _lowercase ):

     '''simple docstring'''






     def __init__( self      :    Dict   ,						_SCREAMING_SNAKE_CASE      :    Any ):






           """simple docstring"""




           if not isinstance(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE ) or not (temperature > 0):
                 raise ValueError(f'`temperature` has to be a strictly positive float, but is {temperature}' )

           UpperCamelCase					       =  temperature







     def __call__( self      :    Union[str, Any]   ,						_SCREAMING_SNAKE_CASE      :    Optional[Any]   ,						_SCREAMING_SNAKE_CASE      :    List[Any]   ,						_SCREAMING_SNAKE_CASE      :    Optional[int] ):






           """simple docstring"""




           UpperCamelCase					       =  scores / self.temperature
           return scores





class 					A__     (     _lowercase ):

     '''simple docstring'''






     def __init__( self      :    List[str]   ,						_SCREAMING_SNAKE_CASE      :    Optional[int]   ,						_SCREAMING_SNAKE_CASE      :    List[Any] = -float('Inf' )   ,						_SCREAMING_SNAKE_CASE      :    List[str] = 1 ):






           """simple docstring"""




           if not isinstance(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE ) or (top_p < 0 or top_p > 1.0):
                 raise ValueError(f'`top_p` has to be a float > 0 and < 1, but is {top_p}' )
           if not isinstance(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE ) or (min_tokens_to_keep < 1):
                 raise ValueError(f'`min_tokens_to_keep` has to be a positive integer, but is {min_tokens_to_keep}' )

           UpperCamelCase					       =  top_p
           UpperCamelCase					       =  filter_value
           UpperCamelCase					       =  min_tokens_to_keep







     def __call__( self      :    int   ,						_SCREAMING_SNAKE_CASE      :    Union[str, Any]   ,						_SCREAMING_SNAKE_CASE      :    int   ,						_SCREAMING_SNAKE_CASE      :    str ):






           """simple docstring"""




           UpperCamelCase       ,  UpperCamelCase					       =  lax.top_k(_SCREAMING_SNAKE_CASE   ,						scores.shape[-1] )

           UpperCamelCase					       =  jnp.full_like(_SCREAMING_SNAKE_CASE   ,						self.filter_value )
           UpperCamelCase					       =  jax.nn.softmax(_SCREAMING_SNAKE_CASE   ,						axis=-1 ).cumsum(axis=-1 )
           UpperCamelCase					       =  cumulative_probs < self.top_p

           # include the token that is higher than top_p as well
           UpperCamelCase					       =  jnp.roll(_SCREAMING_SNAKE_CASE   ,						1 )
           score_mask |= score_mask.at[:, 0].set(_SCREAMING_SNAKE_CASE )

           # min tokens to keep
           UpperCamelCase					       =  score_mask.at[:, : self.min_tokens_to_keep].set(_SCREAMING_SNAKE_CASE )

           UpperCamelCase					       =  jnp.where(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE )
           UpperCamelCase					       =  jax.lax.sort_key_val(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE )[-1]

           return next_scores





class 					A__     (     _lowercase ):

     '''simple docstring'''






     def __init__( self      :    Optional[Any]   ,						_SCREAMING_SNAKE_CASE      :    str   ,						_SCREAMING_SNAKE_CASE      :    Tuple = -float('Inf' )   ,						_SCREAMING_SNAKE_CASE      :    Dict = 1 ):






           """simple docstring"""




           if not isinstance(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE ) or top_k <= 0:
                 raise ValueError(f'`top_k` has to be a strictly positive integer, but is {top_k}' )

           UpperCamelCase					       =  max(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE )
           UpperCamelCase					       =  filter_value







     def __call__( self      :    Dict   ,						_SCREAMING_SNAKE_CASE      :    int   ,						_SCREAMING_SNAKE_CASE      :    str   ,						_SCREAMING_SNAKE_CASE      :    Optional[Any] ):






           """simple docstring"""




           UpperCamelCase       ,  UpperCamelCase					       =  scores.shape
           UpperCamelCase					       =  jnp.full(batch_size * vocab_size   ,						self.filter_value )

           UpperCamelCase					       =  min(self.top_k   ,						scores.shape[-1] )  # Safety check
           UpperCamelCase       ,  UpperCamelCase					       =  lax.top_k(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE )
           UpperCamelCase					       =  jnp.broadcast_to((jnp.arange(_SCREAMING_SNAKE_CASE ) * vocab_size)[:, None]   ,						(batch_size, topk) ).flatten()
           UpperCamelCase					       =  topk_scores.flatten()
           UpperCamelCase					       =  topk_indices.flatten() + shift

           UpperCamelCase					       =  next_scores_flat.at[topk_indices_flat].set(_SCREAMING_SNAKE_CASE )
           UpperCamelCase					       =  next_scores_flat.reshape(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE )
           return next_scores





class 					A__     (     _lowercase ):

     '''simple docstring'''






     def __init__( self      :    Any   ,						_SCREAMING_SNAKE_CASE      :    Any ):






           """simple docstring"""




           UpperCamelCase					       =  bos_token_id







     def __call__( self      :    Optional[Any]   ,						_SCREAMING_SNAKE_CASE      :    List[str]   ,						_SCREAMING_SNAKE_CASE      :    Optional[Any]   ,						_SCREAMING_SNAKE_CASE      :    int ):






           """simple docstring"""




           UpperCamelCase					       =  jnp.full(scores.shape   ,						-float('inf' ) )

           UpperCamelCase					       =  1 - jnp.bool_(cur_len - 1 )

           UpperCamelCase					       =  jnp.where(_SCREAMING_SNAKE_CASE   ,						new_scores.at[:, self.bos_token_id].set(0 )   ,						_SCREAMING_SNAKE_CASE )

           return scores





class 					A__     (     _lowercase ):

     '''simple docstring'''






     def __init__( self      :    Tuple   ,						_SCREAMING_SNAKE_CASE      :    Tuple   ,						_SCREAMING_SNAKE_CASE      :    Any ):






           """simple docstring"""




           UpperCamelCase					       =  max_length
           UpperCamelCase					       =  eos_token_id







     def __call__( self      :    Dict   ,						_SCREAMING_SNAKE_CASE      :    Dict   ,						_SCREAMING_SNAKE_CASE      :    List[Any]   ,						_SCREAMING_SNAKE_CASE      :    Dict ):






           """simple docstring"""




           UpperCamelCase					       =  jnp.full(scores.shape   ,						-float('inf' ) )

           UpperCamelCase					       =  1 - jnp.bool_(cur_len - self.max_length + 1 )

           UpperCamelCase					       =  jnp.where(_SCREAMING_SNAKE_CASE   ,						new_scores.at[:, self.eos_token_id].set(0 )   ,						_SCREAMING_SNAKE_CASE )

           return scores





class 					A__     (     _lowercase ):

     '''simple docstring'''






     def __init__( self      :    List[str]   ,						_SCREAMING_SNAKE_CASE      :    str   ,						_SCREAMING_SNAKE_CASE      :    Optional[int] ):






           """simple docstring"""




           if not isinstance(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE ) or min_length < 0:
                 raise ValueError(f'`min_length` has to be a positive integer, but is {min_length}' )

           if not isinstance(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE ) or eos_token_id < 0:
                 raise ValueError(f'`eos_token_id` has to be a positive integer, but is {eos_token_id}' )

           UpperCamelCase					       =  min_length
           UpperCamelCase					       =  eos_token_id







     def __call__( self      :    Optional[int]   ,						_SCREAMING_SNAKE_CASE      :    Dict   ,						_SCREAMING_SNAKE_CASE      :    List[Any]   ,						_SCREAMING_SNAKE_CASE      :    Union[str, Any] ):






           """simple docstring"""




           UpperCamelCase					       =  1 - jnp.clip(cur_len - self.min_length   ,						0   ,						1 )

           UpperCamelCase					       =  jnp.where(_SCREAMING_SNAKE_CASE   ,						scores.at[:, self.eos_token_id].set(-float('inf' ) )   ,						_SCREAMING_SNAKE_CASE )

           return scores





class 					A__     (     _lowercase ):

     '''simple docstring'''






     def __init__( self      :    List[Any]   ,						_SCREAMING_SNAKE_CASE      :    Optional[int]   ,						_SCREAMING_SNAKE_CASE      :    Union[str, Any] ):






           """simple docstring"""




           UpperCamelCase					       =  list(_SCREAMING_SNAKE_CASE )
           UpperCamelCase					       =  begin_index







     def __call__( self      :    Optional[Any]   ,						_SCREAMING_SNAKE_CASE      :    int   ,						_SCREAMING_SNAKE_CASE      :    Dict   ,						_SCREAMING_SNAKE_CASE      :    Optional[int] ):






           """simple docstring"""




           UpperCamelCase					       =  1 - jnp.bool_(cur_len - self.begin_index )

           UpperCamelCase					       =  jnp.where(_SCREAMING_SNAKE_CASE   ,						scores.at[:, self.begin_suppress_tokens].set(-float('inf' ) )   ,						_SCREAMING_SNAKE_CASE )

           return scores





class 					A__     (     _lowercase ):

     '''simple docstring'''






     def __init__( self      :    Optional[Any]   ,						_SCREAMING_SNAKE_CASE      :    Optional[Any] ):






           """simple docstring"""




           UpperCamelCase					       =  list(_SCREAMING_SNAKE_CASE )







     def __call__( self      :    Any   ,						_SCREAMING_SNAKE_CASE      :    List[Any]   ,						_SCREAMING_SNAKE_CASE      :    Optional[Any]   ,						_SCREAMING_SNAKE_CASE      :    Dict ):






           """simple docstring"""




           UpperCamelCase					       =  scores.at[..., self.suppress_tokens].set(-float('inf' ) )

           return scores





class 					A__     (     _lowercase ):

     '''simple docstring'''






     def __init__( self      :    Tuple   ,						_SCREAMING_SNAKE_CASE      :    str ):






           """simple docstring"""




           UpperCamelCase					       =  dict(_SCREAMING_SNAKE_CASE )
           # Converts the dictionary of format {index: token} containing the tokens to be forced to an array, where the
           # index of the array corresponds to the index of the token to be forced, for XLA compatibility.
           # Indexes without forced tokens will have a negative value.
           UpperCamelCase					       =  jnp.ones((max(force_token_map.keys() ) + 1)   ,						dtype=jnp.intaa ) * -1
           for index, token in force_token_map.items():
                 if token is not None:
                       UpperCamelCase					       =  force_token_array.at[index].set(_SCREAMING_SNAKE_CASE )
           UpperCamelCase					       =  jnp.intaa(_SCREAMING_SNAKE_CASE )







     def __call__( self      :    Any   ,						_SCREAMING_SNAKE_CASE      :    Dict   ,						_SCREAMING_SNAKE_CASE      :    List[Any]   ,						_SCREAMING_SNAKE_CASE      :    str ):






           """simple docstring"""




           def _force_token(_SCREAMING_SNAKE_CASE      :    List[Any] ):
                 UpperCamelCase					       =  scores.shape[0]
                 UpperCamelCase					       =  self.force_token_array[generation_idx]

                 UpperCamelCase					       =  jnp.ones_like(_SCREAMING_SNAKE_CASE   ,						dtype=scores.dtype ) * -float('inf' )
                 UpperCamelCase					       =  jnp.zeros((batch_size, 1)   ,						dtype=scores.dtype )
                 UpperCamelCase					       =  lax.dynamic_update_slice(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						(0, current_token) )
                 return new_scores

           UpperCamelCase					       =  lax.cond(
               cur_len >= self.force_token_array.shape[0]   ,						lambda: scores   ,						lambda: lax.cond(
                   self.force_token_array[cur_len] >= 0   ,						lambda: _force_token(_SCREAMING_SNAKE_CASE )   ,						lambda: scores   ,						)   ,						)
           return scores





class 					A__     (     _lowercase ):

     '''simple docstring'''






     def __init__( self      :    Dict   ,						_SCREAMING_SNAKE_CASE      :    Union[str, Any]   ,						_SCREAMING_SNAKE_CASE      :    Optional[Any]   ,						_SCREAMING_SNAKE_CASE      :    Optional[int] ):






           """simple docstring"""




           UpperCamelCase					       =  generate_config.eos_token_id
           UpperCamelCase					       =  generate_config.no_timestamps_token_id
           UpperCamelCase					       =  generate_config.no_timestamps_token_id + 1

           UpperCamelCase					       =  decoder_input_length + 1

           if generate_config.is_multilingual:
                 # room for language token and task token
                 self.begin_index += 2
           if hasattr(_SCREAMING_SNAKE_CASE   ,						'max_initial_timestamp_index' ):
                 UpperCamelCase					       =  generate_config.max_initial_timestamp_index
           else:
                 UpperCamelCase					       =  model_config.vocab_size
           if self.max_initial_timestamp_index is None:
                 UpperCamelCase					       =  model_config.vocab_size







     def __call__( self      :    List[str]   ,						_SCREAMING_SNAKE_CASE      :    Tuple   ,						_SCREAMING_SNAKE_CASE      :    Union[str, Any]   ,						_SCREAMING_SNAKE_CASE      :    List[str] ):






           """simple docstring"""




           UpperCamelCase					       =  scores.at[:, self.no_timestamps_token_id].set(-float('inf' ) )

           def handle_pairs(_SCREAMING_SNAKE_CASE      :    Optional[int]   ,						_SCREAMING_SNAKE_CASE      :    Optional[int] ):
                 UpperCamelCase					       =  jnp.where((cur_len - self.begin_index) >= 1   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE )
                 UpperCamelCase					       =  jnp.where(
                     input_ids_k[cur_len - 1] >= self.timestamp_begin   ,						True and last_was_timestamp   ,						_SCREAMING_SNAKE_CASE   ,						)

                 UpperCamelCase					       =  jnp.where((cur_len - self.begin_index) < 2   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE )
                 UpperCamelCase					       =  jnp.where(
                     input_ids_k[cur_len - 2] >= self.timestamp_begin   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						)

                 return jnp.where(
                     _SCREAMING_SNAKE_CASE   ,						jnp.where(
                         penultimate_was_timestamp > 0   ,						scores_k.at[self.timestamp_begin :].set(-float('inf' ) )   ,						scores_k.at[: self.eos_token_id].set(-float('inf' ) )   ,						)   ,						_SCREAMING_SNAKE_CASE   ,						)

           UpperCamelCase					       =  jax.vmap(_SCREAMING_SNAKE_CASE )(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE )

           UpperCamelCase					       =  jnp.where(cur_len == self.begin_index   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE )
           UpperCamelCase					       =  jnp.where(
               self.max_initial_timestamp_index is not None   ,						True and apply_max_initial_timestamp   ,						_SCREAMING_SNAKE_CASE   ,						)

           UpperCamelCase					       =  self.timestamp_begin + self.max_initial_timestamp_index

           UpperCamelCase					       =  jnp.where(
               _SCREAMING_SNAKE_CASE   ,						scores.at[:, last_allowed + 1 :].set(-float('inf' ) )   ,						_SCREAMING_SNAKE_CASE   ,						)

           # if sum of probability over timestamps is above any other token, sample timestamp
           UpperCamelCase					       =  jax.nn.log_softmax(_SCREAMING_SNAKE_CASE   ,						axis=-1 )

           def handle_cumulative_probs(_SCREAMING_SNAKE_CASE      :    Tuple   ,						_SCREAMING_SNAKE_CASE      :    int ):
                 UpperCamelCase					       =  jax.nn.logsumexp(logprobs_k[self.timestamp_begin :]   ,						axis=-1 )
                 UpperCamelCase					       =  jnp.max(logprobs_k[: self.timestamp_begin] )
                 return jnp.where(
                     timestamp_logprob > max_text_token_logprob   ,						scores_k.at[: self.timestamp_begin].set(-float('inf' ) )   ,						_SCREAMING_SNAKE_CASE   ,						)

           UpperCamelCase					       =  jax.vmap(_SCREAMING_SNAKE_CASE )(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE )

           return scores








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		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE=1_3		,				SCREAMING_SNAKE_CASE=7		,				SCREAMING_SNAKE_CASE=True		,				SCREAMING_SNAKE_CASE=True		,				SCREAMING_SNAKE_CASE=True		,				SCREAMING_SNAKE_CASE=9_9		,				SCREAMING_SNAKE_CASE=3_2		,				SCREAMING_SNAKE_CASE=5		,				SCREAMING_SNAKE_CASE=4		,				SCREAMING_SNAKE_CASE=3_7		,				SCREAMING_SNAKE_CASE="gelu"		,				SCREAMING_SNAKE_CASE=0.1		,				SCREAMING_SNAKE_CASE=0.1		,				SCREAMING_SNAKE_CASE=5_1_2		,				SCREAMING_SNAKE_CASE=1_6		,				SCREAMING_SNAKE_CASE=2		,				SCREAMING_SNAKE_CASE=0.02		,				SCREAMING_SNAKE_CASE=3		,				SCREAMING_SNAKE_CASE=4		,				SCREAMING_SNAKE_CASE=None		,				) ->			Union[str, Any]:
  a__												= parent
  a__												= batch_size
  a__												= seq_length
  a__												= is_training
  a__												= use_token_type_ids
  a__												= use_labels
  a__												= vocab_size
  a__												= hidden_size
  a__												= num_hidden_layers
  a__												= num_attention_heads
  a__												= intermediate_size
  a__												= hidden_act
  a__												= hidden_dropout_prob
  a__												= attention_probs_dropout_prob
  a__												= max_position_embeddings
  a__												= type_vocab_size
  a__												= type_sequence_label_size
  a__												= initializer_range
  a__												= num_labels
  a__												= num_choices
  a__												= scope
  a__												= self.vocab_size - 1

 def 						_UpperCAmelCase       (	self      ) ->			int:
  a__												= ids_tensor([self.batch_size, self.seq_length]		,				self.vocab_size      )

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

  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]		,				self.num_choices      )

  a__												= 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__												= 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 						_UpperCAmelCase       (	self		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				*SCREAMING_SNAKE_CASE      ) ->			int:
  a__												= OpenAIGPTModel(config=SCREAMING_SNAKE_CASE      )
  model.to(SCREAMING_SNAKE_CASE      )
  model.eval()

  a__												= model(SCREAMING_SNAKE_CASE		,				token_type_ids=SCREAMING_SNAKE_CASE		,				head_mask=SCREAMING_SNAKE_CASE      )
  a__												= model(SCREAMING_SNAKE_CASE		,				token_type_ids=SCREAMING_SNAKE_CASE      )
  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		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				*SCREAMING_SNAKE_CASE      ) ->			Tuple:
  a__												= OpenAIGPTLMHeadModel(SCREAMING_SNAKE_CASE      )
  model.to(SCREAMING_SNAKE_CASE      )
  model.eval()

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

 def 						_UpperCAmelCase       (	self		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				*SCREAMING_SNAKE_CASE      ) ->			Dict:
  a__												= OpenAIGPTDoubleHeadsModel(SCREAMING_SNAKE_CASE      )
  model.to(SCREAMING_SNAKE_CASE      )
  model.eval()

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

 def 						_UpperCAmelCase       (	self		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				*SCREAMING_SNAKE_CASE      ) ->			Optional[Any]:
  a__												= self.num_labels
  a__												= OpenAIGPTForSequenceClassification(SCREAMING_SNAKE_CASE      )
  model.to(SCREAMING_SNAKE_CASE      )
  model.eval()

  a__												= ids_tensor([self.batch_size]		,				self.type_sequence_label_size      )
  a__												= model(SCREAMING_SNAKE_CASE		,				token_type_ids=SCREAMING_SNAKE_CASE		,				labels=SCREAMING_SNAKE_CASE      )
  self.parent.assertEqual(result.logits.shape		,				(self.batch_size, self.num_labels)      )



 def 						_UpperCAmelCase       (	self      ) ->			str:
  a__												= self.prepare_config_and_inputs()
  (
      (
      a__
  )							,			(
      a__
  )							,			(
      a__
  )							,			(
      a__
  )							,			(
      a__
  )							,			(
      a__
  )							,			(
      a__
  )							,			
  )												= config_and_inputs
  a__												= {
      '''input_ids''': input_ids,
      '''token_type_ids''': token_type_ids,
      '''head_mask''': head_mask,
  }

  return config, inputs_dict





@require_torch
class    __UpperCamelCase			(       _lowercase	,	_lowercase	,	_lowercase	,	unittest.TestCase				):

 """simple docstring"""



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

 def 						_UpperCAmelCase       (	self		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE      ) ->			Optional[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 						_UpperCAmelCase       (	self		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE=False      ) ->			List[str]:
  a__												= super()._prepare_for_class(SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				return_labels=SCREAMING_SNAKE_CASE      )

  if return_labels:
   if model_class.__name__ == "OpenAIGPTDoubleHeadsModel":
    a__												= torch.zeros(
        (self.model_tester.batch_size, self.model_tester.num_choices, self.model_tester.seq_length)		,				dtype=torch.long		,				device=SCREAMING_SNAKE_CASE		,				)
    a__												= inputs_dict['''labels''']
    a__												= inputs_dict['''labels''']
    a__												= torch.zeros(
        (self.model_tester.batch_size, self.model_tester.num_choices)		,				dtype=torch.long		,				device=SCREAMING_SNAKE_CASE		,				)
    a__												= torch.zeros(
        self.model_tester.batch_size		,				dtype=torch.long		,				device=SCREAMING_SNAKE_CASE      )
  return inputs_dict

 def 						_UpperCAmelCase       (	self      ) ->			List[str]:
  a__												= OpenAIGPTModelTester(self      )
  a__												= ConfigTester(self		,				config_class=SCREAMING_SNAKE_CASE		,				n_embd=3_7      )

 def 						_UpperCAmelCase       (	self      ) ->			Optional[Any]:
  self.config_tester.run_common_tests()

 def 						_UpperCAmelCase       (	self      ) ->			Optional[int]:
  a__												= self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_openai_gpt_model(*SCREAMING_SNAKE_CASE      )

 def 						_UpperCAmelCase       (	self      ) ->			int:
  a__												= self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_lm_head_model(*SCREAMING_SNAKE_CASE      )

 def 						_UpperCAmelCase       (	self      ) ->			Optional[int]:
  a__												= self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_double_lm_head_model(*SCREAMING_SNAKE_CASE      )

 def 						_UpperCAmelCase       (	self      ) ->			Optional[int]:
  a__												= self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_openai_gpt_for_sequence_classification(*SCREAMING_SNAKE_CASE      )



 @slow
 def 						_UpperCAmelCase       (	self      ) ->			Union[str, Any]:
  for model_name in OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
   a__												= OpenAIGPTModel.from_pretrained(SCREAMING_SNAKE_CASE      )
   self.assertIsNotNone(SCREAMING_SNAKE_CASE      )





@require_torch
class    __UpperCamelCase			(       unittest.TestCase				):

 """simple docstring"""



 @slow
 def 						_UpperCAmelCase       (	self      ) ->			Optional[int]:
  a__												= OpenAIGPTLMHeadModel.from_pretrained('''openai-gpt'''      )
  model.to(SCREAMING_SNAKE_CASE      )
  a__												= torch.tensor([[4_8_1, 4_7_3_5, 5_4_4]]		,				dtype=torch.long		,				device=SCREAMING_SNAKE_CASE      )  # the president is
  a__												= [
      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__												= model.generate(SCREAMING_SNAKE_CASE		,				do_sample=SCREAMING_SNAKE_CASE      )
  self.assertListEqual(output_ids[0].tolist()		,				SCREAMING_SNAKE_CASE      )





