




'''simple docstring'''







from __future__ import annotations

import unittest

from transformers import RoFormerConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
   import tensorflow as tf

   from transformers import (
       TFRoFormerForCausalLM,
       TFRoFormerForMaskedLM,
       TFRoFormerForMultipleChoice,
       TFRoFormerForQuestionAnswering,
       TFRoFormerForSequenceClassification,
       TFRoFormerForTokenClassification,
       TFRoFormerModel,
   )
   from transformers.models.roformer.modeling_tf_roformer import (
       TFRoFormerSelfAttention,
       TFRoFormerSinusoidalPositionalEmbedding,
   )


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



           '''simple docstring'''

           UpperCAmelCase			       =		parent
           UpperCAmelCase			       =		13
           UpperCAmelCase			       =		7
           UpperCAmelCase			       =		True
           UpperCAmelCase			       =		True
           UpperCAmelCase			       =		True
           UpperCAmelCase			       =		True
           UpperCAmelCase			       =		99
           UpperCAmelCase			       =		32
           UpperCAmelCase			       =		2
           UpperCAmelCase			       =		4
           UpperCAmelCase			       =		37
           UpperCAmelCase			       =		'''gelu'''
           UpperCAmelCase			       =		0.1
           UpperCAmelCase			       =		0.1
           UpperCAmelCase			       =		512
           UpperCAmelCase			       =		16
           UpperCAmelCase			       =		2
           UpperCAmelCase			       =		0.02
           UpperCAmelCase			       =		3
           UpperCAmelCase			       =		4
           UpperCAmelCase			       =		None
       def     A      (	self						:	Union[str, Any]				):



           '''simple docstring'''

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

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

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

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

           UpperCAmelCase			       =		RoFormerConfig(
               vocab_size=self.vocab_size   ,				hidden_size=self.hidden_size   ,				num_hidden_layers=self.num_hidden_layers   ,				num_attention_heads=self.num_attention_heads   ,				intermediate_size=self.intermediate_size   ,				hidden_act=self.hidden_act   ,				hidden_dropout_prob=self.hidden_dropout_prob   ,				attention_probs_dropout_prob=self.attention_probs_dropout_prob   ,				max_position_embeddings=self.max_position_embeddings   ,				type_vocab_size=self.type_vocab_size   ,				initializer_range=self.initializer_range   ,				return_dict=lowercase   ,				)

           return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
       def     A      (	self						:	Dict   ,				lowercase						:	Optional[Any]   ,				lowercase						:	List[Any]   ,				lowercase						:	int   ,				lowercase						:	Optional[Any]   ,				lowercase						:	List[str]   ,				lowercase						:	str   ,				lowercase						:	Optional[int]				):



           '''simple docstring'''

           UpperCAmelCase			       =		TFRoFormerModel(config=lowercase				)
           UpperCAmelCase			       =		{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}

           UpperCAmelCase			       =		[input_ids, input_mask]
           UpperCAmelCase			       =		model(lowercase				)

           UpperCAmelCase			       =		model(lowercase				)

           self.parent.assertEqual(result.last_hidden_state.shape   ,				(self.batch_size, self.seq_length, self.hidden_size)				)
       def     A      (	self						:	str   ,				lowercase						:	Any   ,				lowercase						:	List[str]   ,				lowercase						:	Union[str, Any]   ,				lowercase						:	int   ,				lowercase						:	List[Any]   ,				lowercase						:	Optional[int]   ,				lowercase						:	List[Any]				):



           '''simple docstring'''

           UpperCAmelCase			       =		True
           UpperCAmelCase			       =		TFRoFormerForCausalLM(config=lowercase				)
           UpperCAmelCase			       =		{
               '''input_ids''': input_ids,
               '''attention_mask''': input_mask,
               '''token_type_ids''': token_type_ids,
           }
           UpperCAmelCase			       =		model(lowercase				)['''logits''']
           self.parent.assertListEqual(
               list(prediction_scores.numpy().shape				)   ,				[self.batch_size, self.seq_length, self.vocab_size]				)
       def     A      (	self						:	List[Any]   ,				lowercase						:	Optional[int]   ,				lowercase						:	List[str]   ,				lowercase						:	List[Any]   ,				lowercase						:	Optional[int]   ,				lowercase						:	Any   ,				lowercase						:	str   ,				lowercase						:	List[Any]				):



           '''simple docstring'''

           UpperCAmelCase			       =		TFRoFormerForMaskedLM(config=lowercase				)
           UpperCAmelCase			       =		{
               '''input_ids''': input_ids,
               '''attention_mask''': input_mask,
               '''token_type_ids''': token_type_ids,
           }
           UpperCAmelCase			       =		model(lowercase				)
           self.parent.assertEqual(result.logits.shape   ,				(self.batch_size, self.seq_length, self.vocab_size)				)
       def     A      (	self						:	List[Any]   ,				lowercase						:	Dict   ,				lowercase						:	List[Any]   ,				lowercase						:	Optional[int]   ,				lowercase						:	Optional[int]   ,				lowercase						:	Any   ,				lowercase						:	Tuple   ,				lowercase						:	int				):



           '''simple docstring'''

           UpperCAmelCase			       =		self.num_labels
           UpperCAmelCase			       =		TFRoFormerForSequenceClassification(config=lowercase				)
           UpperCAmelCase			       =		{
               '''input_ids''': input_ids,
               '''attention_mask''': input_mask,
               '''token_type_ids''': token_type_ids,
           }

           UpperCAmelCase			       =		model(lowercase				)
           self.parent.assertEqual(result.logits.shape   ,				(self.batch_size, self.num_labels)				)
       def     A      (	self						:	List[Any]   ,				lowercase						:	Tuple   ,				lowercase						:	List[str]   ,				lowercase						:	Union[str, Any]   ,				lowercase						:	Optional[Any]   ,				lowercase						:	Tuple   ,				lowercase						:	Optional[Any]   ,				lowercase						:	str				):



           '''simple docstring'''

           UpperCAmelCase			       =		self.num_choices
           UpperCAmelCase			       =		TFRoFormerForMultipleChoice(config=lowercase				)
           UpperCAmelCase			       =		tf.tile(tf.expand_dims(lowercase   ,				1				)   ,				(1, self.num_choices, 1)				)
           UpperCAmelCase			       =		tf.tile(tf.expand_dims(lowercase   ,				1				)   ,				(1, self.num_choices, 1)				)
           UpperCAmelCase			       =		tf.tile(tf.expand_dims(lowercase   ,				1				)   ,				(1, self.num_choices, 1)				)
           UpperCAmelCase			       =		{
               '''input_ids''': multiple_choice_inputs_ids,
               '''attention_mask''': multiple_choice_input_mask,
               '''token_type_ids''': multiple_choice_token_type_ids,
           }
           UpperCAmelCase			       =		model(lowercase				)
           self.parent.assertEqual(result.logits.shape   ,				(self.batch_size, self.num_choices)				)
       def     A      (	self						:	str   ,				lowercase						:	List[str]   ,				lowercase						:	Dict   ,				lowercase						:	Dict   ,				lowercase						:	Optional[int]   ,				lowercase						:	Any   ,				lowercase						:	Optional[int]   ,				lowercase						:	Tuple				):



           '''simple docstring'''

           UpperCAmelCase			       =		self.num_labels
           UpperCAmelCase			       =		TFRoFormerForTokenClassification(config=lowercase				)
           UpperCAmelCase			       =		{
               '''input_ids''': input_ids,
               '''attention_mask''': input_mask,
               '''token_type_ids''': token_type_ids,
           }
           UpperCAmelCase			       =		model(lowercase				)
           self.parent.assertEqual(result.logits.shape   ,				(self.batch_size, self.seq_length, self.num_labels)				)
       def     A      (	self						:	Dict   ,				lowercase						:	int   ,				lowercase						:	Union[str, Any]   ,				lowercase						:	Any   ,				lowercase						:	List[str]   ,				lowercase						:	Optional[int]   ,				lowercase						:	Optional[int]   ,				lowercase						:	str				):



           '''simple docstring'''

           UpperCAmelCase			       =		TFRoFormerForQuestionAnswering(config=lowercase				)
           UpperCAmelCase			       =		{
               '''input_ids''': input_ids,
               '''attention_mask''': input_mask,
               '''token_type_ids''': token_type_ids,
           }

           UpperCAmelCase			       =		model(lowercase				)
           self.parent.assertEqual(result.start_logits.shape   ,				(self.batch_size, self.seq_length)				)
           self.parent.assertEqual(result.end_logits.shape   ,				(self.batch_size, self.seq_length)				)







       def     A      (	self						:	Dict				):



           '''simple docstring'''

           UpperCAmelCase			       =		self.prepare_config_and_inputs()
           (
               (
               UpperCAmelCase
           )       ,     (
               UpperCAmelCase
           )       ,     (
               UpperCAmelCase
           )       ,     (
               UpperCAmelCase
           )       ,     (
               UpperCAmelCase
           )       ,     (
               UpperCAmelCase
           )       ,     (
               UpperCAmelCase
           )       ,     
           )			       =		config_and_inputs
           UpperCAmelCase			       =		{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
           return config, inputs_dict






@require_tf
class 			_a  (	__a							,							__a							,							unittest.TestCase ):
       __a				:     Any   				=					(
           (
               TFRoFormerModel,
               TFRoFormerForCausalLM,
               TFRoFormerForMaskedLM,
               TFRoFormerForQuestionAnswering,
               TFRoFormerForSequenceClassification,
               TFRoFormerForTokenClassification,
               TFRoFormerForMultipleChoice,
           )
           if is_tf_available()
           else ()
       )
       __a				:     Tuple   				=					(
           {
               """feature-extraction""": TFRoFormerModel,
               """fill-mask""": TFRoFormerForMaskedLM,
               """question-answering""": TFRoFormerForQuestionAnswering,
               """text-classification""": TFRoFormerForSequenceClassification,
               """text-generation""": TFRoFormerForCausalLM,
               """token-classification""": TFRoFormerForTokenClassification,
               """zero-shot""": TFRoFormerForSequenceClassification,
           }
           if is_tf_available()
           else {}
       )

       __a				:     List[Any]   				=					False
       __a				:     Dict   				=					False
       def     A      (	self						:	Union[str, Any]   ,				lowercase						:	Any   ,				lowercase						:	Tuple   ,				lowercase						:	Tuple   ,				lowercase						:	Optional[int]   ,				lowercase						:	Dict				):



           '''simple docstring'''

           if pipeline_test_casse_name == "TextGenerationPipelineTests":
               return True

           return False
       def     A      (	self						:	List[Any]				):



           '''simple docstring'''

           UpperCAmelCase			       =		TFRoFormerModelTester(self				)
           UpperCAmelCase			       =		ConfigTester(self   ,				config_class=lowercase   ,				hidden_size=37				)
       def     A      (	self						:	Any				):



           '''simple docstring'''

           self.config_tester.run_common_tests()
       def     A      (	self						:	List[str]				):



           '''simple docstring'''

           UpperCAmelCase			       =		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*lowercase				)
       def     A      (	self						:	int				):



           '''simple docstring'''

           UpperCAmelCase			       =		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_masked_lm(*lowercase				)
       def     A      (	self						:	str				):



           '''simple docstring'''

           UpperCAmelCase			       =		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_lm_head(*lowercase				)
       def     A      (	self						:	Tuple				):



           '''simple docstring'''

           UpperCAmelCase			       =		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_multiple_choice(*lowercase				)
       def     A      (	self						:	Optional[Any]				):



           '''simple docstring'''

           UpperCAmelCase			       =		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_question_answering(*lowercase				)
       def     A      (	self						:	Dict				):



           '''simple docstring'''

           UpperCAmelCase			       =		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_sequence_classification(*lowercase				)
       def     A      (	self						:	Optional[Any]				):



           '''simple docstring'''

           UpperCAmelCase			       =		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_token_classification(*lowercase				)







       @slow
       def     A      (	self						:	Optional[Any]				):



           '''simple docstring'''

           UpperCAmelCase			       =		TFRoFormerModel.from_pretrained('''junnyu/roformer_chinese_base'''				)
           self.assertIsNotNone(lowercase				)






@require_tf
class 			_a  (	unittest.TestCase ):
       @slow
       def     A      (	self						:	Union[str, Any]				):



           '''simple docstring'''

           UpperCAmelCase			       =		TFRoFormerForMaskedLM.from_pretrained('''junnyu/roformer_chinese_base'''				)
           UpperCAmelCase			       =		tf.constant([[0, 1, 2, 3, 4, 5]]				)
           UpperCAmelCase			       =		model(lowercase				)[0]

           # TODO Replace vocab size
           UpperCAmelCase			       =		50_000

           UpperCAmelCase			       =		[1, 6, vocab_size]
           self.assertEqual(output.shape   ,				lowercase				)

           print(output[:, :3, :3]				)

           # TODO Replace values below with what was printed above.
           UpperCAmelCase			       =		tf.constant(
               [
                   [
                       [-0.1205_3341, -1.026_4901, 0.2922_1946],
                       [-1.513_3783, 0.19_7433, 0.1519_0607],
                       [-5.013_5403, -3.90_0256, -0.8403_8764],
                   ]
               ]				)
           tf.debugging.assert_near(output[:, :3, :3]   ,				lowercase   ,				atol=1E-4				)






@require_tf
class 			_a  (	unittest.TestCase ):
       __a				:     List[Any]   				=					1e-4
       def     A      (	self						:	str				):



           '''simple docstring'''

           UpperCAmelCase			       =		tf.constant([[4, 10]]				)
           UpperCAmelCase			       =		TFRoFormerSinusoidalPositionalEmbedding(num_positions=6   ,				embedding_dim=6				)

           UpperCAmelCase			       =		emba(input_ids.shape				)
           UpperCAmelCase			       =		tf.constant(
               [[0.0000, 0.0000, 0.0000, 1.0000, 1.0000, 1.0000], [0.8415, 0.0464, 0.0022, 0.5403, 0.9989, 1.0000]]				)

           tf.debugging.assert_near(lowercase   ,				lowercase   ,				atol=self.tolerance				)







       def     A      (	self						:	Union[str, Any]				):



           '''simple docstring'''

           UpperCAmelCase			       =		tf.constant(
               [
                   [0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
                   [0.8415, 0.8219, 0.8020, 0.7819, 0.7617],
                   [0.9093, 0.9364, 0.9581, 0.9749, 0.9870],
               ]				)
           UpperCAmelCase			       =		TFRoFormerSinusoidalPositionalEmbedding(num_positions=512   ,				embedding_dim=512				)
           emba([2, 16, 512]				)
           UpperCAmelCase			       =		emba.weight[:3, :5]

           tf.debugging.assert_near(lowercase   ,				lowercase   ,				atol=self.tolerance				)






@require_tf
class 			_a  (	unittest.TestCase ):
       __a				:     Dict   				=					1e-4
       def     A      (	self						:	Optional[int]				):



           '''simple docstring'''

           UpperCAmelCase			       =		tf.reshape(tf.range(2 * 12 * 16 * 64   ,				dtype=tf.floataa				)   ,				shape=(2, 12, 16, 64)				) / 100

           UpperCAmelCase			       =		-tf.reshape(tf.range(2 * 12 * 16 * 64   ,				dtype=tf.floataa				)   ,				shape=(2, 12, 16, 64)				) / 100

           UpperCAmelCase			       =		TFRoFormerSinusoidalPositionalEmbedding(num_positions=32   ,				embedding_dim=64				)
           UpperCAmelCase			       =		embed_positions([2, 16, 768]				)[None, None, :, :]

           UpperCAmelCase       ,     UpperCAmelCase			       =		TFRoFormerSelfAttention.apply_rotary_position_embeddings(
               lowercase   ,				lowercase   ,				lowercase				)

           UpperCAmelCase			       =		tf.constant(
               [
                   [0.0000, 0.0100, 0.0200, 0.0300, 0.0400, 0.0500, 0.0600, 0.0700],
                   [-0.2012, 0.8897, 0.0263, 0.9401, 0.2074, 0.9463, 0.3481, 0.9343],
                   [-1.7057, 0.6271, -1.2145, 1.3897, -0.6303, 1.7647, -0.1173, 1.8985],
                   [-2.1731, -1.6397, -2.7358, 0.2854, -2.1840, 1.7183, -1.3018, 2.4871],
                   [0.2717, -3.6173, -2.9206, -2.1988, -3.6638, 0.3858, -2.9155, 2.2980],
                   [3.9859, -2.1580, -0.7984, -4.4904, -4.1181, -2.0252, -4.4782, 1.1253],
               ]				)
           UpperCAmelCase			       =		tf.constant(
               [
                   [0.0000, -0.0100, -0.0200, -0.0300, -0.0400, -0.0500, -0.0600, -0.0700],
                   [0.2012, -0.8897, -0.0263, -0.9401, -0.2074, -0.9463, -0.3481, -0.9343],
                   [1.7057, -0.6271, 1.2145, -1.3897, 0.6303, -1.7647, 0.1173, -1.8985],
                   [2.1731, 1.6397, 2.7358, -0.2854, 2.1840, -1.7183, 1.3018, -2.4871],
                   [-0.2717, 3.6173, 2.9206, 2.1988, 3.6638, -0.3858, 2.9155, -2.2980],
                   [-3.9859, 2.1580, 0.7984, 4.4904, 4.1181, 2.0252, 4.4782, -1.1253],
               ]				)

           tf.debugging.assert_near(query_layer[0, 0, :6, :8]   ,				lowercase   ,				atol=self.tolerance				)
           tf.debugging.assert_near(key_layer[0, 0, :6, :8]   ,				lowercase   ,				atol=self.tolerance				)







'''simple docstring'''






import torch
from torch import nn






class 			lowercase_	(nn.Module      ):




  """simple docstring"""




  def __init__(      self   :     Optional[int]					,lowercase__   :     List[str]					,lowercase__   :     Any					,lowercase__   :     Union[str, Any]					,lowercase__   :     Optional[int]					,lowercase__   :     Optional[int]=1					,lowercase__   :     Optional[Any]=False			):
   super().__init__()

   __lowercase							  = n_token
   __lowercase							  = d_embed
   __lowercase							  = d_proj

   __lowercase							  = cutoffs + [n_token]
   __lowercase							  = [0] + self.cutoffs
   __lowercase							  = div_val

   __lowercase							  = self.cutoffs[0]
   __lowercase							  = len(self.cutoffs			) - 1
   __lowercase							  = self.shortlist_size + self.n_clusters

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

   __lowercase							  = nn.ModuleList()
   __lowercase							  = 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(lowercase__					,lowercase__			)			)			)
     else:
      self.out_projs.append(lowercase__			)

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

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

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

   __lowercase							  = keep_order






  def 					SCREAMING_SNAKE_CASE      (      self   :     Optional[int]					,lowercase__   :     List[Any]					,lowercase__   :     Dict					,lowercase__   :     List[Any]					,lowercase__   :     Any			):
   if proj is None:
    __lowercase							  = nn.functional.linear(lowercase__					,lowercase__					,bias=lowercase__			)
   else:
    # if CUDA_MAJOR <= 9 and CUDA_MINOR <= 1:
    __lowercase							  = nn.functional.linear(lowercase__					,proj.t().contiguous()			)
    __lowercase							  = nn.functional.linear(lowercase__					,lowercase__					,bias=lowercase__			)
    # 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   :     Union[str, Any]					,lowercase__   :     Optional[Any]					,lowercase__   :     Any=None					,lowercase__   :     List[str]=False			):

   if labels is not None:
    # Shift so that tokens < n predict n
    __lowercase							  = hidden[..., :-1, :].contiguous()
    __lowercase							  = labels[..., 1:].contiguous()
    __lowercase							  = hidden.view(-1					,hidden.size(-1			)			)
    __lowercase							  = 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:
    __lowercase							  = hidden.view(-1					,hidden.size(-1			)			)

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

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

     weights.append(lowercase__			)
     biases.append(lowercase__			)

    __lowercase  ,	__lowercase  ,	__lowercase							  = weights[0], biases[0], self.out_projs[0]

    __lowercase							  = self._compute_logit(lowercase__					,lowercase__					,lowercase__					,lowercase__			)
    __lowercase							  = nn.functional.log_softmax(lowercase__					,dim=1			)

    if labels is None:
     __lowercase							  = hidden.new_empty((head_logit.size(0			), self.n_token)			)
    else:
     __lowercase							  = torch.zeros_like(lowercase__					,dtype=hidden.dtype					,device=hidden.device			)

    __lowercase							  = 0
    __lowercase							  = [0] + self.cutoffs
    for i in range(len(lowercase__			) - 1			):
     __lowercase  ,	__lowercase							  = cutoff_values[i], cutoff_values[i + 1]

     if labels is not None:
      __lowercase							  = (labels >= l_idx) & (labels < r_idx)
      __lowercase							  = mask_i.nonzero().squeeze()

      if indices_i.numel() == 0:
       continue

      __lowercase							  = labels.index_select(0					,lowercase__			) - l_idx
      __lowercase							  = head_logprob.index_select(0					,lowercase__			)
      __lowercase							  = hidden.index_select(0					,lowercase__			)
     else:
      __lowercase							  = hidden

     if i == 0:
      if labels is not None:
       __lowercase							  = head_logprob_i.gather(1					,target_i[:, None]			).squeeze(1			)
      else:
       __lowercase							  = head_logprob[:, : self.cutoffs[0]]
     else:
      __lowercase  ,	__lowercase  ,	__lowercase							  = weights[i], biases[i], self.out_projs[i]

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

     if labels is not None:
      if (hasattr(self					,'''keep_order'''			) and self.keep_order) or keep_order:
       out.index_copy_(0					,lowercase__					,-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   :     Any					,lowercase__   :     Union[str, Any]			):
   if self.n_clusters == 0:
    __lowercase							  = self._compute_logit(lowercase__					,self.out_layers[0].weight					,self.out_layers[0].bias					,self.out_projs[0]			)
    return nn.functional.log_softmax(lowercase__					,dim=-1			)
   else:
    # construct weights and biases
    __lowercase  ,	__lowercase							  = [], []
    for i in range(len(self.cutoffs			)			):
     if self.div_val == 1:
      __lowercase  ,	__lowercase							  = self.cutoff_ends[i], self.cutoff_ends[i + 1]
      __lowercase							  = self.out_layers[0].weight[l_idx:r_idx]
      __lowercase							  = self.out_layers[0].bias[l_idx:r_idx]
     else:
      __lowercase							  = self.out_layers[i].weight
      __lowercase							  = self.out_layers[i].bias

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

     weights.append(lowercase__			)
     biases.append(lowercase__			)

    __lowercase  ,	__lowercase  ,	__lowercase							  = weights[0], biases[0], self.out_projs[0]
    __lowercase							  = self._compute_logit(lowercase__					,lowercase__					,lowercase__					,lowercase__			)

    __lowercase							  = hidden.new_empty((head_logit.size(0			), self.n_token)			)
    __lowercase							  = nn.functional.log_softmax(lowercase__					,dim=1			)

    __lowercase							  = [0] + self.cutoffs
    for i in range(len(lowercase__			) - 1			):
     __lowercase  ,	__lowercase							  = cutoff_values[i], cutoff_values[i + 1]

     if i == 0:
      __lowercase							  = head_logprob[:, : self.cutoffs[0]]
     else:
      __lowercase  ,	__lowercase  ,	__lowercase							  = weights[i], biases[i], self.out_projs[i]

      __lowercase							  = self._compute_logit(lowercase__					,lowercase__					,lowercase__					,lowercase__			)
      __lowercase							  = nn.functional.log_softmax(lowercase__					,dim=1			)

      __lowercase							  = head_logprob[:, -i] + tail_logprob_i
      __lowercase							  = logprob_i

    return out

