



"""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 	__SCREAMING_SNAKE_CASE     :


					'''simple docstring'''





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



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

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

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

								lowerCamelCase__		:		Tuple		     =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,  )

								lowerCamelCase__		:		List[str]		     =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 		snake_case	(	self							:  Tuple,  lowerCamelCase							:  Dict,  lowerCamelCase							:  Union[str, Any],  lowerCamelCase							:  int,  lowerCamelCase							:  Tuple,  *lowerCamelCase							:  int					)-> List[Any]:
								lowerCamelCase__		:		List[str]		     =OpenAIGPTModel(config=lowerCamelCase					)
								model.to(lowerCamelCase					)
								model.eval()

								lowerCamelCase__		:		int		     =model(lowerCamelCase,  token_type_ids=lowerCamelCase,  head_mask=lowerCamelCase					)
								lowerCamelCase__		:		Optional[Any]		     =model(lowerCamelCase,  token_type_ids=lowerCamelCase					)
								lowerCamelCase__		:		int		     =model(lowerCamelCase					)

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



					def 		snake_case	(	self							:  List[str],  lowerCamelCase							:  List[Any],  lowerCamelCase							:  int,  lowerCamelCase							:  List[str],  lowerCamelCase							:  Optional[int],  *lowerCamelCase							:  Tuple					)-> List[Any]:
								lowerCamelCase__		:		Union[str, Any]		     =OpenAIGPTLMHeadModel(lowerCamelCase					)
								model.to(lowerCamelCase					)
								model.eval()

								lowerCamelCase__		:		Union[str, Any]		     =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 		snake_case	(	self							:  Any,  lowerCamelCase							:  int,  lowerCamelCase							:  int,  lowerCamelCase							:  Tuple,  lowerCamelCase							:  Union[str, Any],  *lowerCamelCase							:  str					)-> Tuple:
								lowerCamelCase__		:		Optional[int]		     =OpenAIGPTDoubleHeadsModel(lowerCamelCase					)
								model.to(lowerCamelCase					)
								model.eval()

								lowerCamelCase__		:		List[Any]		     =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 		snake_case	(	self							:  List[Any],  lowerCamelCase							:  Union[str, Any],  lowerCamelCase							:  Dict,  lowerCamelCase							:  Any,  lowerCamelCase							:  Union[str, Any],  *lowerCamelCase							:  Any					)-> Dict:
								lowerCamelCase__		:		List[Any]		     =self.num_labels
								lowerCamelCase__		:		Union[str, Any]		     =OpenAIGPTForSequenceClassification(lowerCamelCase					)
								model.to(lowerCamelCase					)
								model.eval()

								lowerCamelCase__		:		Union[str, Any]		     =ids_tensor([self.batch_size],  self.type_sequence_label_size					)
								lowerCamelCase__		:		Any		     =model(lowerCamelCase,  token_type_ids=lowerCamelCase,  labels=lowerCamelCase					)
								self.parent.assertEqual(result.logits.shape,  (self.batch_size, self.num_labels)					)




					def 		snake_case	(	self							:  int					)-> List[Any]:
								lowerCamelCase__		:		str		     =self.prepare_config_and_inputs()
								(
								    (
								    lowerCamelCase__
								)     ,       (
								    lowerCamelCase__
								)     ,       (
								    lowerCamelCase__
								)     ,       (
								    lowerCamelCase__
								)     ,       (
								    lowerCamelCase__
								)     ,       (
								    lowerCamelCase__
								)     ,       (
								    lowerCamelCase__
								)     ,       
								)		:		Dict		     =config_and_inputs
								lowerCamelCase__		:		Optional[int]		     ={
								    '''input_ids''': input_ids,
								    '''token_type_ids''': token_type_ids,
								    '''head_mask''': head_mask,
								}

								return config, inputs_dict





@require_torch
class 	__SCREAMING_SNAKE_CASE     (				lowerCAmelCase_		,       lowerCAmelCase_		,       lowerCAmelCase_		,       unittest.TestCase       ):


					'''simple docstring'''





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



					def 		snake_case	(	self							:  int,  lowerCamelCase							:  Union[str, Any],  lowerCamelCase							:  str,  lowerCamelCase							:  int,  lowerCamelCase							:  Union[str, Any],  lowerCamelCase							:  Any					)-> List[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 		snake_case	(	self							:  str,  lowerCamelCase							:  int,  lowerCamelCase							:  Union[str, Any],  lowerCamelCase							:  Optional[Any]=False					)-> int:
								lowerCamelCase__		:		List[Any]		     =super()._prepare_for_class(lowerCamelCase,  lowerCamelCase,  return_labels=lowerCamelCase					)

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



					def 		snake_case	(	self							:  Optional[int]					)-> Any:
								lowerCamelCase__		:		Any		     =OpenAIGPTModelTester(self					)
								lowerCamelCase__		:		List[Any]		     =ConfigTester(self,  config_class=lowerCamelCase,  n_embd=37					)



					def 		snake_case	(	self							:  Tuple					)-> Tuple:
								self.config_tester.run_common_tests()



					def 		snake_case	(	self							:  Dict					)-> Any:
								lowerCamelCase__		:		List[Any]		     =self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_openai_gpt_model(*lowerCamelCase					)



					def 		snake_case	(	self							:  Union[str, Any]					)-> List[str]:
								lowerCamelCase__		:		Union[str, Any]		     =self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_lm_head_model(*lowerCamelCase					)



					def 		snake_case	(	self							:  Tuple					)-> Optional[int]:
								lowerCamelCase__		:		int		     =self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_double_lm_head_model(*lowerCamelCase					)



					def 		snake_case	(	self							:  Optional[int]					)-> Dict:
								lowerCamelCase__		:		List[str]		     =self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_openai_gpt_for_sequence_classification(*lowerCamelCase					)




					@slow
					def 		snake_case	(	self							:  int					)-> Dict:
								for model_name in OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
											lowerCamelCase__		:		Tuple		     =OpenAIGPTModel.from_pretrained(lowerCamelCase					)
											self.assertIsNotNone(lowerCamelCase					)





@require_torch
class 	__SCREAMING_SNAKE_CASE     (				unittest.TestCase       ):


					'''simple docstring'''





					@slow
					def 		snake_case	(	self							:  Optional[Any]					)-> str:
								lowerCamelCase__		:		str		     =OpenAIGPTLMHeadModel.from_pretrained('''openai-gpt'''					)
								model.to(lowerCamelCase					)
								lowerCamelCase__		:		Optional[Any]		     =torch.tensor([[481, 4735, 544]],  dtype=torch.long,  device=lowerCamelCase					)  # the president is
								lowerCamelCase__		:		List[str]		     =[
								    481,
								    4735,
								    544,
								    246,
								    963,
								    870,
								    762,
								    239,
								    244,
								    4_0477,
								    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

								lowerCamelCase__		:		List[Any]		     =model.generate(lowerCamelCase,  do_sample=lowerCamelCase					)
								self.assertListEqual(output_ids[0].tolist(),  lowerCamelCase					)









"""simple docstring"""






import math

import flax.linen as nn
import jax.numpy as jnp






def 					snake_case__       (		__lowerCamelCase    :   jnp.ndarray							,		__lowerCamelCase    :   int							,		__lowerCamelCase    :   float = 1							,		__lowerCamelCase    :   float = 1							,		__lowerCamelCase    :   float = 1.0e4							,		__lowerCamelCase    :   bool = False							,		__lowerCamelCase    :   float = 1.0							,		):
			"""simple docstring"""
			assert timesteps.ndim == 1, "Timesteps should be a 1d-array"
			assert embedding_dim % 2 == 0, f'''Embedding dimension {embedding_dim} should be even'''
			lowerCamelCase__		:		Any		     =float(embedding_dim // 2     )
			lowerCamelCase__		:		List[str]		     =math.log(max_timescale / min_timescale     ) / (num_timescales - freq_shift)
			lowerCamelCase__		:		int		     =min_timescale * jnp.exp(jnp.arange(__lowerCamelCase							,		dtype=jnp.floataa     ) * -log_timescale_increment     )
			lowerCamelCase__		:		Tuple		     =jnp.expand_dims(__lowerCamelCase							,		1     ) * jnp.expand_dims(__lowerCamelCase							,		0     )

			# scale embeddings
			lowerCamelCase__		:		List[str]		     =scale * emb

			if flip_sin_to_cos:
						lowerCamelCase__		:		int		     =jnp.concatenate([jnp.cos(__lowerCamelCase     ), jnp.sin(__lowerCamelCase     )]							,		axis=1     )
			else:
						lowerCamelCase__		:		List[str]		     =jnp.concatenate([jnp.sin(__lowerCamelCase     ), jnp.cos(__lowerCamelCase     )]							,		axis=1     )
			lowerCamelCase__		:		str		     =jnp.reshape(__lowerCamelCase							,		[jnp.shape(__lowerCamelCase     )[0], embedding_dim]     )
			return signal

class 	__SCREAMING_SNAKE_CASE     (				nn.Module       ):


					'''simple docstring'''





					_a             =							3_2
					_a             =							jnp.floataa



					@nn.compact
					def __call__(	self							:  Optional[Any],  lowerCamelCase							:  int					)-> Any:
								lowerCamelCase__		:		Optional[Any]		     =nn.Dense(self.time_embed_dim,  dtype=self.dtype,  name='''linear_1'''					)(lowerCamelCase					)
								lowerCamelCase__		:		List[str]		     =nn.silu(lowerCamelCase					)
								lowerCamelCase__		:		Any		     =nn.Dense(self.time_embed_dim,  dtype=self.dtype,  name='''linear_2'''					)(lowerCamelCase					)
								return temb





class 	__SCREAMING_SNAKE_CASE     (				nn.Module       ):


					'''simple docstring'''





					_a             =							3_2
					_a             =							False
					_a             =							1



					@nn.compact
					def __call__(	self							:  Any,  lowerCamelCase							:  int					)-> int:
								return get_sinusoidal_embeddings(
								    lowerCamelCase,  embedding_dim=self.dim,  flip_sin_to_cos=self.flip_sin_to_cos,  freq_shift=self.freq_shift					)




