



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




	'''simple docstring'''





	def __init__(  self:	Any ,    _SCREAMING_SNAKE_CASE:	Optional[Any] ,    _SCREAMING_SNAKE_CASE:	int=13 ,    _SCREAMING_SNAKE_CASE:	List[str]=7 ,    _SCREAMING_SNAKE_CASE:	Optional[int]=True ,    _SCREAMING_SNAKE_CASE:	Union[str, Any]=True ,    _SCREAMING_SNAKE_CASE:	Optional[int]=True ,    _SCREAMING_SNAKE_CASE:	List[Any]=99 ,    _SCREAMING_SNAKE_CASE:	List[str]=32 ,    _SCREAMING_SNAKE_CASE:	Optional[Any]=5 ,    _SCREAMING_SNAKE_CASE:	Tuple=4 ,    _SCREAMING_SNAKE_CASE:	Tuple=37 ,    _SCREAMING_SNAKE_CASE:	List[str]="gelu" ,    _SCREAMING_SNAKE_CASE:	Union[str, Any]=0.1 ,    _SCREAMING_SNAKE_CASE:	int=0.1 ,    _SCREAMING_SNAKE_CASE:	Tuple=512 ,    _SCREAMING_SNAKE_CASE:	Dict=16 ,    _SCREAMING_SNAKE_CASE:	Optional[int]=2 ,    _SCREAMING_SNAKE_CASE:	Union[str, Any]=0.02 ,    _SCREAMING_SNAKE_CASE:	Tuple=3 ,    _SCREAMING_SNAKE_CASE:	Optional[Any]=4 ,    _SCREAMING_SNAKE_CASE:	int=None ,    )       -> List[Any]:

				"""simple docstring"""

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

	def 							_SCREAMING_SNAKE_CASE     (  self:	Any)       -> int:

				"""simple docstring"""

				__lowerCAmelCase   :						Dict				= ids_tensor([self.batch_size, self.seq_length] ,    self.vocab_size)

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

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

				__lowerCAmelCase   :						Optional[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 ,    )

				__lowerCAmelCase   :						Optional[int]				= 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 							_SCREAMING_SNAKE_CASE     (  self:	Any ,    _SCREAMING_SNAKE_CASE:	List[str] ,    _SCREAMING_SNAKE_CASE:	Union[str, Any] ,    _SCREAMING_SNAKE_CASE:	int ,    _SCREAMING_SNAKE_CASE:	List[str] ,    *_SCREAMING_SNAKE_CASE:	Optional[int])       -> Dict:

				"""simple docstring"""

				__lowerCAmelCase   :						List[Any]				= OpenAIGPTModel(config=_SCREAMING_SNAKE_CASE)
				model.to(_SCREAMING_SNAKE_CASE)
				model.eval()

				__lowerCAmelCase   :						Optional[Any]				= model(_SCREAMING_SNAKE_CASE ,    token_type_ids=_SCREAMING_SNAKE_CASE ,    head_mask=_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase   :						int				= model(_SCREAMING_SNAKE_CASE ,    token_type_ids=_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase   :						str				= model(_SCREAMING_SNAKE_CASE)

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

	def 							_SCREAMING_SNAKE_CASE     (  self:	List[Any] ,    _SCREAMING_SNAKE_CASE:	Tuple ,    _SCREAMING_SNAKE_CASE:	Optional[Any] ,    _SCREAMING_SNAKE_CASE:	Union[str, Any] ,    _SCREAMING_SNAKE_CASE:	Dict ,    *_SCREAMING_SNAKE_CASE:	Union[str, Any])       -> Tuple:

				"""simple docstring"""

				__lowerCAmelCase   :						Union[str, Any]				= OpenAIGPTLMHeadModel(_SCREAMING_SNAKE_CASE)
				model.to(_SCREAMING_SNAKE_CASE)
				model.eval()

				__lowerCAmelCase   :						Union[str, Any]				= 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 							_SCREAMING_SNAKE_CASE     (  self:	int ,    _SCREAMING_SNAKE_CASE:	Dict ,    _SCREAMING_SNAKE_CASE:	Optional[int] ,    _SCREAMING_SNAKE_CASE:	Optional[int] ,    _SCREAMING_SNAKE_CASE:	Optional[int] ,    *_SCREAMING_SNAKE_CASE:	Any)       -> Optional[Any]:

				"""simple docstring"""

				__lowerCAmelCase   :						Union[str, Any]				= OpenAIGPTDoubleHeadsModel(_SCREAMING_SNAKE_CASE)
				model.to(_SCREAMING_SNAKE_CASE)
				model.eval()

				__lowerCAmelCase   :						Union[str, Any]				= 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 							_SCREAMING_SNAKE_CASE     (  self:	str ,    _SCREAMING_SNAKE_CASE:	Union[str, Any] ,    _SCREAMING_SNAKE_CASE:	Optional[int] ,    _SCREAMING_SNAKE_CASE:	str ,    _SCREAMING_SNAKE_CASE:	int ,    *_SCREAMING_SNAKE_CASE:	Tuple)       -> Dict:

				"""simple docstring"""

				__lowerCAmelCase   :						int				= self.num_labels
				__lowerCAmelCase   :						int				= OpenAIGPTForSequenceClassification(_SCREAMING_SNAKE_CASE)
				model.to(_SCREAMING_SNAKE_CASE)
				model.eval()

				__lowerCAmelCase   :						Optional[Any]				= ids_tensor([self.batch_size] ,    self.type_sequence_label_size)
				__lowerCAmelCase   :						List[Any]				= 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 							_SCREAMING_SNAKE_CASE     (  self:	Union[str, Any])       -> List[Any]:

				"""simple docstring"""

				__lowerCAmelCase   :						List[Any]				= self.prepare_config_and_inputs()
				(
				    (
				    __lowerCAmelCase
				)		,			(
				    __lowerCAmelCase
				)		,			(
				    __lowerCAmelCase
				)		,			(
				    __lowerCAmelCase
				)		,			(
				    __lowerCAmelCase
				)		,			(
				    __lowerCAmelCase
				)		,			(
				    __lowerCAmelCase
				)		,			
				)   :						Any				= config_and_inputs
				__lowerCAmelCase   :						List[str]				= {
				    "input_ids": input_ids,
				    "token_type_ids": token_type_ids,
				    "head_mask": head_mask,
				}

				return config, inputs_dict






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




	'''simple docstring'''





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

	def 							_SCREAMING_SNAKE_CASE     (  self:	Any ,    _SCREAMING_SNAKE_CASE:	Optional[Any] ,    _SCREAMING_SNAKE_CASE:	Tuple ,    _SCREAMING_SNAKE_CASE:	Optional[int] ,    _SCREAMING_SNAKE_CASE:	Any ,    _SCREAMING_SNAKE_CASE:	Optional[Any])       -> Optional[Any]:

				"""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 							_SCREAMING_SNAKE_CASE     (  self:	int ,    _SCREAMING_SNAKE_CASE:	Tuple ,    _SCREAMING_SNAKE_CASE:	Optional[int] ,    _SCREAMING_SNAKE_CASE:	int=False)       -> Optional[Any]:

				"""simple docstring"""

				__lowerCAmelCase   :						Optional[int]				= super()._prepare_for_class(_SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE ,    return_labels=_SCREAMING_SNAKE_CASE)

				if return_labels:
							if model_class.__name__ == "OpenAIGPTDoubleHeadsModel":
										__lowerCAmelCase   :						Any				= torch.zeros(
										    (self.model_tester.batch_size, self.model_tester.num_choices, self.model_tester.seq_length) ,    dtype=torch.long ,    device=_SCREAMING_SNAKE_CASE ,    )
										__lowerCAmelCase   :						Dict				= inputs_dict["labels"]
										__lowerCAmelCase   :						int				= inputs_dict["labels"]
										__lowerCAmelCase   :						Dict				= torch.zeros(
										    (self.model_tester.batch_size, self.model_tester.num_choices) ,    dtype=torch.long ,    device=_SCREAMING_SNAKE_CASE ,    )
										__lowerCAmelCase   :						Tuple				= torch.zeros(
										    self.model_tester.batch_size ,    dtype=torch.long ,    device=_SCREAMING_SNAKE_CASE)
				return inputs_dict

	def 							_SCREAMING_SNAKE_CASE     (  self:	Dict)       -> Optional[Any]:

				"""simple docstring"""

				__lowerCAmelCase   :						Any				= OpenAIGPTModelTester(self)
				__lowerCAmelCase   :						str				= ConfigTester(self ,    config_class=_SCREAMING_SNAKE_CASE ,    n_embd=37)

	def 							_SCREAMING_SNAKE_CASE     (  self:	Optional[int])       -> List[str]:

				"""simple docstring"""

				self.config_tester.run_common_tests()

	def 							_SCREAMING_SNAKE_CASE     (  self:	Optional[int])       -> Tuple:

				"""simple docstring"""

				__lowerCAmelCase   :						Any				= self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_openai_gpt_model(*_SCREAMING_SNAKE_CASE)

	def 							_SCREAMING_SNAKE_CASE     (  self:	Any)       -> List[str]:

				"""simple docstring"""

				__lowerCAmelCase   :						Any				= self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_lm_head_model(*_SCREAMING_SNAKE_CASE)

	def 							_SCREAMING_SNAKE_CASE     (  self:	List[Any])       -> Tuple:

				"""simple docstring"""

				__lowerCAmelCase   :						Union[str, Any]				= self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_double_lm_head_model(*_SCREAMING_SNAKE_CASE)

	def 							_SCREAMING_SNAKE_CASE     (  self:	Optional[int])       -> str:

				"""simple docstring"""

				__lowerCAmelCase   :						Any				= self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_openai_gpt_for_sequence_classification(*_SCREAMING_SNAKE_CASE)







	@slow
	def 							_SCREAMING_SNAKE_CASE     (  self:	Any)       -> Optional[Any]:

				"""simple docstring"""

				for model_name in OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
							__lowerCAmelCase   :						Tuple				= OpenAIGPTModel.from_pretrained(_SCREAMING_SNAKE_CASE)
							self.assertIsNotNone(_SCREAMING_SNAKE_CASE)






@require_torch
class     A__     (						unittest.TestCase							):




	'''simple docstring'''





	@slow
	def 							_SCREAMING_SNAKE_CASE     (  self:	int)       -> Any:

				"""simple docstring"""

				__lowerCAmelCase   :						Union[str, Any]				= OpenAIGPTLMHeadModel.from_pretrained("openai-gpt")
				model.to(_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase   :						str				= torch.tensor([[481, 4735, 544]] ,    dtype=torch.long ,    device=_SCREAMING_SNAKE_CASE)  # the president is
				__lowerCAmelCase   :						Any				= [
				    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[str]				= model.generate(_SCREAMING_SNAKE_CASE ,    do_sample=_SCREAMING_SNAKE_CASE)
				self.assertListEqual(output_ids[0].tolist() ,    _SCREAMING_SNAKE_CASE)




"""simple docstring"""






import torch
from torch import nn



class     A__     (						nn.Module							):




	'''simple docstring'''





	def __init__(  self:	Optional[Any] ,    _SCREAMING_SNAKE_CASE:	List[Any] ,    _SCREAMING_SNAKE_CASE:	str ,    _SCREAMING_SNAKE_CASE:	List[str] ,    _SCREAMING_SNAKE_CASE:	str ,    _SCREAMING_SNAKE_CASE:	Optional[Any]=1 ,    _SCREAMING_SNAKE_CASE:	Optional[Any]=False)       -> Dict:

				"""simple docstring"""

				super().__init__()

				__lowerCAmelCase   :						Optional[int]				= n_token
				__lowerCAmelCase   :						int				= d_embed
				__lowerCAmelCase   :						str				= d_proj

				__lowerCAmelCase   :						Optional[Any]				= cutoffs + [n_token]
				__lowerCAmelCase   :						Any				= [0] + self.cutoffs
				__lowerCAmelCase   :						Tuple				= div_val

				__lowerCAmelCase   :						Optional[Any]				= self.cutoffs[0]
				__lowerCAmelCase   :						Union[str, Any]				= len(self.cutoffs) - 1
				__lowerCAmelCase   :						List[str]				= self.shortlist_size + self.n_clusters

				if self.n_clusters > 0:
							__lowerCAmelCase   :						Any				= nn.Parameter(torch.zeros(self.n_clusters ,    self.d_embed))
							__lowerCAmelCase   :						Dict				= nn.Parameter(torch.zeros(self.n_clusters))

				__lowerCAmelCase   :						List[Any]				= nn.ModuleList()
				__lowerCAmelCase   :						Any				= nn.ParameterList()

				if div_val == 1:
							for i in range(len(self.cutoffs)):
										if d_proj != d_embed:
													self.out_projs.append(nn.Parameter(torch.FloatTensor(_SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE)))
										else:
													self.out_projs.append(_SCREAMING_SNAKE_CASE)

							self.out_layers.append(nn.Linear(_SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE))
				else:
							for i in range(len(self.cutoffs)):
										__lowerCAmelCase		,			__lowerCAmelCase   :						Tuple				= self.cutoff_ends[i], self.cutoff_ends[i + 1]
										__lowerCAmelCase   :						str				= d_embed // (div_val**i)

										self.out_projs.append(nn.Parameter(torch.FloatTensor(_SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE)))

										self.out_layers.append(nn.Linear(_SCREAMING_SNAKE_CASE ,    r_idx - l_idx))

				__lowerCAmelCase   :						str				= keep_order

	def 							_SCREAMING_SNAKE_CASE     (  self:	Optional[Any] ,    _SCREAMING_SNAKE_CASE:	Any ,    _SCREAMING_SNAKE_CASE:	Tuple ,    _SCREAMING_SNAKE_CASE:	int ,    _SCREAMING_SNAKE_CASE:	Dict)       -> int:

				"""simple docstring"""

				if proj is None:
							__lowerCAmelCase   :						Dict				= nn.functional.linear(_SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE ,    bias=_SCREAMING_SNAKE_CASE)
				else:
							# if CUDA_MAJOR <= 9 and CUDA_MINOR <= 1:
							__lowerCAmelCase   :						Optional[Any]				= nn.functional.linear(_SCREAMING_SNAKE_CASE ,    proj.t().contiguous())
							__lowerCAmelCase   :						List[Any]				= nn.functional.linear(_SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE ,    bias=_SCREAMING_SNAKE_CASE)
							# else:
							#     logit = torch.einsum('bd,de,ev->bv', (hidden, proj, weight.t()))
							#     if bias is not None:
							#         logit = logit + bias

				return logit

	def 							_SCREAMING_SNAKE_CASE     (  self:	Dict ,    _SCREAMING_SNAKE_CASE:	Tuple ,    _SCREAMING_SNAKE_CASE:	Tuple=None ,    _SCREAMING_SNAKE_CASE:	int=False)       -> List[str]:

				"""simple docstring"""

				if labels is not None:
							# Shift so that tokens < n predict n
							__lowerCAmelCase   :						List[Any]				= hidden[..., :-1, :].contiguous()
							__lowerCAmelCase   :						Tuple				= labels[..., 1:].contiguous()
							__lowerCAmelCase   :						Union[str, Any]				= hidden.view(-1 ,    hidden.size(-1))
							__lowerCAmelCase   :						Optional[int]				= labels.view(-1)
							if hidden.size(0) != labels.size(0):
										raise RuntimeError("Input and labels should have the same size in the batch dimension.")
				else:
							__lowerCAmelCase   :						List[Any]				= hidden.view(-1 ,    hidden.size(-1))

				if self.n_clusters == 0:
							__lowerCAmelCase   :						Any				= self._compute_logit(_SCREAMING_SNAKE_CASE ,    self.out_layers[0].weight ,    self.out_layers[0].bias ,    self.out_projs[0])
							if labels is not None:
										__lowerCAmelCase   :						Any				= labels != -100
										__lowerCAmelCase   :						Tuple				= torch.zeros_like(_SCREAMING_SNAKE_CASE ,    dtype=hidden.dtype ,    device=hidden.device)
										__lowerCAmelCase   :						List[Any]				= (
										    -nn.functional.log_softmax(_SCREAMING_SNAKE_CASE ,    dim=-1)[mask].gather(1 ,    labels[mask].unsqueeze(1)).squeeze(1)
										)
							else:
										__lowerCAmelCase   :						List[Any]				= nn.functional.log_softmax(_SCREAMING_SNAKE_CASE ,    dim=-1)
				else:
							# construct weights and biases
							__lowerCAmelCase		,			__lowerCAmelCase   :						Optional[Any]				= [], []
							for i in range(len(self.cutoffs)):
										if self.div_val == 1:
													__lowerCAmelCase		,			__lowerCAmelCase   :						int				= self.cutoff_ends[i], self.cutoff_ends[i + 1]
													__lowerCAmelCase   :						Optional[Any]				= self.out_layers[0].weight[l_idx:r_idx]
													__lowerCAmelCase   :						Optional[int]				= self.out_layers[0].bias[l_idx:r_idx]
										else:
													__lowerCAmelCase   :						int				= self.out_layers[i].weight
													__lowerCAmelCase   :						List[Any]				= self.out_layers[i].bias

										if i == 0:
													__lowerCAmelCase   :						int				= torch.cat([weight_i, self.cluster_weight] ,    dim=0)
													__lowerCAmelCase   :						Union[str, Any]				= torch.cat([bias_i, self.cluster_bias] ,    dim=0)

										weights.append(_SCREAMING_SNAKE_CASE)
										biases.append(_SCREAMING_SNAKE_CASE)

							__lowerCAmelCase		,			__lowerCAmelCase		,			__lowerCAmelCase   :						Dict				= weights[0], biases[0], self.out_projs[0]

							__lowerCAmelCase   :						Any				= self._compute_logit(_SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE)
							__lowerCAmelCase   :						Optional[Any]				= nn.functional.log_softmax(_SCREAMING_SNAKE_CASE ,    dim=1)

							if labels is None:
										__lowerCAmelCase   :						Optional[Any]				= hidden.new_empty((head_logit.size(0), self.n_token))
							else:
										__lowerCAmelCase   :						int				= torch.zeros_like(_SCREAMING_SNAKE_CASE ,    dtype=hidden.dtype ,    device=hidden.device)

							__lowerCAmelCase   :						str				= 0
							__lowerCAmelCase   :						List[str]				= [0] + self.cutoffs
							for i in range(len(_SCREAMING_SNAKE_CASE) - 1):
										__lowerCAmelCase		,			__lowerCAmelCase   :						List[Any]				= cutoff_values[i], cutoff_values[i + 1]

										if labels is not None:
													__lowerCAmelCase   :						List[str]				= (labels >= l_idx) & (labels < r_idx)
													__lowerCAmelCase   :						Any				= mask_i.nonzero().squeeze()

													if indices_i.numel() == 0:
																continue

													__lowerCAmelCase   :						Optional[Any]				= labels.index_select(0 ,    _SCREAMING_SNAKE_CASE) - l_idx
													__lowerCAmelCase   :						List[Any]				= head_logprob.index_select(0 ,    _SCREAMING_SNAKE_CASE)
													__lowerCAmelCase   :						Tuple				= hidden.index_select(0 ,    _SCREAMING_SNAKE_CASE)
										else:
													__lowerCAmelCase   :						List[str]				= hidden

										if i == 0:
													if labels is not None:
																__lowerCAmelCase   :						Optional[int]				= head_logprob_i.gather(1 ,    target_i[:, None]).squeeze(1)
													else:
																__lowerCAmelCase   :						Optional[Any]				= head_logprob[:, : self.cutoffs[0]]
										else:
													__lowerCAmelCase		,			__lowerCAmelCase		,			__lowerCAmelCase   :						List[str]				= weights[i], biases[i], self.out_projs[i]

													__lowerCAmelCase   :						Any				= self._compute_logit(_SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE)
													__lowerCAmelCase   :						Tuple				= nn.functional.log_softmax(_SCREAMING_SNAKE_CASE ,    dim=1)
													__lowerCAmelCase   :						int				= self.cutoffs[0] + i - 1  # No probability for the head cluster
													if labels is not None:
																__lowerCAmelCase   :						List[str]				= head_logprob_i[:, cluster_prob_idx] + tail_logprob_i.gather(
																    1 ,    target_i[:, None]).squeeze(1)
													else:
																__lowerCAmelCase   :						Optional[int]				= head_logprob[:, cluster_prob_idx, None] + tail_logprob_i
																__lowerCAmelCase   :						Optional[int]				= logprob_i

										if labels is not None:
													if (hasattr(self ,    "keep_order") and self.keep_order) or keep_order:
																out.index_copy_(0 ,    _SCREAMING_SNAKE_CASE ,    -logprob_i)
													else:
																out[offset : offset + logprob_i.size(0)].copy_(-logprob_i)
													offset += logprob_i.size(0)

				return out







	def 							_SCREAMING_SNAKE_CASE     (  self:	str ,    _SCREAMING_SNAKE_CASE:	List[Any])       -> Optional[Any]:

				"""simple docstring"""

				if self.n_clusters == 0:
							__lowerCAmelCase   :						int				= self._compute_logit(_SCREAMING_SNAKE_CASE ,    self.out_layers[0].weight ,    self.out_layers[0].bias ,    self.out_projs[0])
							return nn.functional.log_softmax(_SCREAMING_SNAKE_CASE ,    dim=-1)
				else:
							# construct weights and biases
							__lowerCAmelCase		,			__lowerCAmelCase   :						Dict				= [], []
							for i in range(len(self.cutoffs)):
										if self.div_val == 1:
													__lowerCAmelCase		,			__lowerCAmelCase   :						int				= self.cutoff_ends[i], self.cutoff_ends[i + 1]
													__lowerCAmelCase   :						Optional[Any]				= self.out_layers[0].weight[l_idx:r_idx]
													__lowerCAmelCase   :						Union[str, Any]				= self.out_layers[0].bias[l_idx:r_idx]
										else:
													__lowerCAmelCase   :						Dict				= self.out_layers[i].weight
													__lowerCAmelCase   :						List[str]				= self.out_layers[i].bias

										if i == 0:
													__lowerCAmelCase   :						int				= torch.cat([weight_i, self.cluster_weight] ,    dim=0)
													__lowerCAmelCase   :						int				= torch.cat([bias_i, self.cluster_bias] ,    dim=0)

										weights.append(_SCREAMING_SNAKE_CASE)
										biases.append(_SCREAMING_SNAKE_CASE)

							__lowerCAmelCase		,			__lowerCAmelCase		,			__lowerCAmelCase   :						List[Any]				= weights[0], biases[0], self.out_projs[0]
							__lowerCAmelCase   :						str				= self._compute_logit(_SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE)

							__lowerCAmelCase   :						Optional[Any]				= hidden.new_empty((head_logit.size(0), self.n_token))
							__lowerCAmelCase   :						Union[str, Any]				= nn.functional.log_softmax(_SCREAMING_SNAKE_CASE ,    dim=1)

							__lowerCAmelCase   :						Dict				= [0] + self.cutoffs
							for i in range(len(_SCREAMING_SNAKE_CASE) - 1):
										__lowerCAmelCase		,			__lowerCAmelCase   :						Optional[Any]				= cutoff_values[i], cutoff_values[i + 1]

										if i == 0:
													__lowerCAmelCase   :						Union[str, Any]				= head_logprob[:, : self.cutoffs[0]]
										else:
													__lowerCAmelCase		,			__lowerCAmelCase		,			__lowerCAmelCase   :						Dict				= weights[i], biases[i], self.out_projs[i]

													__lowerCAmelCase   :						Any				= self._compute_logit(_SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE)
													__lowerCAmelCase   :						List[Any]				= nn.functional.log_softmax(_SCREAMING_SNAKE_CASE ,    dim=1)

													__lowerCAmelCase   :						int				= head_logprob[:, -i] + tail_logprob_i
													__lowerCAmelCase   :						Tuple				= logprob_i

							return out