

import unittest

from transformers import AutoTokenizer, FalconConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
     import torch

     from transformers import (
         FalconForCausalLM,
         FalconForQuestionAnswering,
         FalconForSequenceClassification,
         FalconForTokenClassification,
         FalconModel,
     )





class  __UpperCAmelCase					:



   def __init__(							self							:					Optional[int],		__A							:					List[Any],		__A							:					List[str]=3,		__A							:					List[Any]=7,		__A							:					str=True,		__A							:					Optional[int]=True,		__A							:					List[Any]=False,		__A							:					str=True,		__A							:					Any=9_9,		__A							:					int=3_2,		__A							:					List[Any]=5,		__A							:					str=4,		__A							:					Optional[int]=3_7,		__A							:					int="gelu",		__A							:					List[str]=0.1,		__A							:					str=0.1,		__A							:					int=5_1_2,		__A							:					int=1_6,		__A							:					Optional[int]=2,		__A							:					Optional[Any]=0.0_2,		__A							:					Any=3,		__A							:					Optional[int]=4,		__A							:					Union[str, Any]=None,		):
     UpperCAmelCase	:							List[str]						=			parent
     UpperCAmelCase	:							List[Any]						=			batch_size
     UpperCAmelCase	:							Tuple						=			seq_length
     UpperCAmelCase	:							Optional[int]						=			is_training
     UpperCAmelCase	:							List[Any]						=			use_input_mask
     UpperCAmelCase	:							int						=			use_token_type_ids
     UpperCAmelCase	:							Tuple						=			use_labels
     UpperCAmelCase	:							Dict						=			vocab_size
     UpperCAmelCase	:							List[str]						=			hidden_size
     UpperCAmelCase	:							Dict						=			num_hidden_layers
     UpperCAmelCase	:							Tuple						=			num_attention_heads
     UpperCAmelCase	:							List[str]						=			intermediate_size
     UpperCAmelCase	:							Dict						=			hidden_act
     UpperCAmelCase	:							Any						=			hidden_dropout_prob
     UpperCAmelCase	:							Optional[Any]						=			attention_probs_dropout_prob
     UpperCAmelCase	:							Union[str, Any]						=			max_position_embeddings
     UpperCAmelCase	:							List[Any]						=			type_vocab_size
     UpperCAmelCase	:							str						=			type_sequence_label_size
     UpperCAmelCase	:							List[Any]						=			initializer_range
     UpperCAmelCase	:							Tuple						=			num_labels
     UpperCAmelCase	:							List[Any]						=			num_choices
     UpperCAmelCase	:							List[Any]						=			scope



   def 	__magic_name__   (							self							:					Dict						):
     UpperCAmelCase	:							List[str]						=			ids_tensor([self.batch_size, self.seq_length],		self.vocab_size						)

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

     UpperCAmelCase	:							int						=			None

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

     UpperCAmelCase	:							Optional[int]						=			self.get_config()

     return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels



   def 	__magic_name__   (							self							:					Optional[int]						):
     return FalconConfig(
         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,		is_decoder=__A,		initializer_range=self.initializer_range,		pad_token_id=1,		new_decoder_architecture=__A,		)



   def 	__magic_name__   (							self							:					List[str],		__A							:					Union[str, Any],		__A							:					Optional[Any],		__A							:					List[Any],		__A							:					Any,		__A							:					Dict,		__A							:					List[str],		__A							:					Dict						):
     UpperCAmelCase	:							str						=			FalconModel(config=__A						)
     model.to(__A						)
     model.eval()
     UpperCAmelCase	:							str						=			model(__A,		attention_mask=__A						)
     UpperCAmelCase	:							Any						=			model(__A						)
     self.parent.assertEqual(result.last_hidden_state.shape,		(self.batch_size, self.seq_length, self.hidden_size)						)



   def 	__magic_name__   (							self							:					int,		__A							:					Dict,		__A							:					Dict,		__A							:					List[Any],		__A							:					int,		__A							:					str,		__A							:					str,		__A							:					Any,		__A							:					Optional[Any],		__A							:					Union[str, Any],		):
     UpperCAmelCase	:							Dict						=			True
     UpperCAmelCase	:							List[str]						=			FalconModel(__A						)
     model.to(__A						)
     model.eval()
     UpperCAmelCase	:							Union[str, Any]						=			model(
         __A,		attention_mask=__A,		encoder_hidden_states=__A,		encoder_attention_mask=__A,		)
     UpperCAmelCase	:							int						=			model(
         __A,		attention_mask=__A,		encoder_hidden_states=__A,		)
     UpperCAmelCase	:							List[Any]						=			model(__A,		attention_mask=__A						)
     self.parent.assertEqual(result.last_hidden_state.shape,		(self.batch_size, self.seq_length, self.hidden_size)						)



   def 	__magic_name__   (							self							:					List[Any],		__A							:					List[Any],		__A							:					List[Any],		__A							:					Optional[int],		__A							:					Union[str, Any],		__A							:					Optional[Any],		__A							:					Optional[Any],		__A							:					List[Any],		__A							:					Any,		__A							:					Union[str, Any],		):
     UpperCAmelCase	:							Any						=			FalconForCausalLM(config=__A						)
     model.to(__A						)
     model.eval()
     UpperCAmelCase	:							str						=			model(__A,		attention_mask=__A,		labels=__A						)
     self.parent.assertEqual(result.logits.shape,		(self.batch_size, self.seq_length, self.vocab_size)						)



   def 	__magic_name__   (							self							:					List[str],		__A							:					Optional[Any],		__A							:					List[Any],		__A							:					List[str],		__A							:					str,		__A							:					Any,		__A							:					List[Any],		__A							:					Optional[int],		__A							:					Tuple,		__A							:					Tuple,		):
     UpperCAmelCase	:							List[Any]						=			True
     UpperCAmelCase	:							Optional[Any]						=			True
     UpperCAmelCase	:							Dict						=			FalconForCausalLM(config=__A						)
     model.to(__A						)
     model.eval()

     # first forward pass
     UpperCAmelCase	:							Union[str, Any]						=			model(
         __A,		attention_mask=__A,		encoder_hidden_states=__A,		encoder_attention_mask=__A,		use_cache=__A,		)
     UpperCAmelCase	:							Union[str, Any]						=			outputs.past_key_values

     # create hypothetical multiple next token and extent to next_input_ids
     UpperCAmelCase	:							Tuple						=			ids_tensor((self.batch_size, 3),		config.vocab_size						)
     UpperCAmelCase	:							Dict						=			ids_tensor((self.batch_size, 3),		vocab_size=2						)

     # append to next input_ids and
     UpperCAmelCase	:							Tuple						=			torch.cat([input_ids, next_tokens],		dim=-1						)
     UpperCAmelCase	:							List[Any]						=			torch.cat([input_mask, next_mask],		dim=-1						)

     UpperCAmelCase	:							Tuple						=			model(
         __A,		attention_mask=__A,		encoder_hidden_states=__A,		encoder_attention_mask=__A,		output_hidden_states=__A,		)['''hidden_states'''][0]
     UpperCAmelCase	:							Any						=			model(
         __A,		attention_mask=__A,		encoder_hidden_states=__A,		encoder_attention_mask=__A,		past_key_values=__A,		output_hidden_states=__A,		)['''hidden_states'''][0]

     # select random slice
     UpperCAmelCase	:							List[str]						=			ids_tensor((1,),		output_from_past.shape[-1]						).item()
     UpperCAmelCase	:							Dict						=			output_from_no_past[:, -3:, random_slice_idx].detach()
     UpperCAmelCase	:							int						=			output_from_past[:, :, random_slice_idx].detach()

     self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]						)

     # test that outputs are equal for slice
     self.parent.assertTrue(torch.allclose(__A,		__A,		atol=1E-3						)						)





   def 	__magic_name__   (							self							:					Union[str, Any]						):
     UpperCAmelCase	:							Dict						=			self.prepare_config_and_inputs()
     (
         (
         UpperCAmelCase
     )			,						(
         UpperCAmelCase
     )			,						(
         UpperCAmelCase
     )			,						(
         UpperCAmelCase
     )			,						(
         UpperCAmelCase
     )			,						(
         UpperCAmelCase
     )			,						(
         UpperCAmelCase
     )			,						
     )	:							Dict						=			config_and_inputs
     UpperCAmelCase	:							Tuple						=			{'''input_ids''': input_ids, '''attention_mask''': input_mask}
     return config, inputs_dict






@require_torch
class  __UpperCAmelCase					(	lowerCamelCase__      ,  lowerCamelCase__      ,  lowerCamelCase__      ,  unittest.TestCase						):
   UpperCamelCase       					=					(
       (
           FalconModel,
           FalconForCausalLM,
           FalconForSequenceClassification,
           FalconForTokenClassification,
           FalconForQuestionAnswering,
       )
       if is_torch_available()
       else ()
   )
   UpperCamelCase       					=					(FalconForCausalLM,) if is_torch_available() else ()
   UpperCamelCase       					=					(
       {
           """feature-extraction""": FalconModel,
           """text-classification""": FalconForSequenceClassification,
           """text-generation""": FalconForCausalLM,
           """question-answering""": FalconForQuestionAnswering,
           """token-classification""": FalconForTokenClassification,
           """zero-shot""": FalconForSequenceClassification,
       }
       if is_torch_available()
       else {}
   )
   UpperCamelCase       					=					False
   UpperCamelCase       					=					False



   def 	__magic_name__   (							self							:					Dict						):
     UpperCAmelCase	:							Optional[int]						=			FalconModelTester(self						)
     UpperCAmelCase	:							Any						=			ConfigTester(self,		config_class=__A,		hidden_size=3_7						)



   def 	__magic_name__   (							self							:					Optional[Any]						):
     self.config_tester.run_common_tests()



   def 	__magic_name__   (							self							:					Union[str, Any]						):
     UpperCAmelCase	:							List[Any]						=			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*__A						)



   def 	__magic_name__   (							self							:					Optional[Any]						):
     UpperCAmelCase			,						*UpperCAmelCase	:							Dict						=			self.model_tester.prepare_config_and_inputs()
     for alibi in [True, False]:
       UpperCAmelCase	:							int						=			alibi
       self.model_tester.create_and_check_model(__A,		*__A						)



   def 	__magic_name__   (							self							:					List[str]						):
     UpperCAmelCase			,						UpperCAmelCase	:							Optional[Any]						=			self.model_tester.prepare_config_and_inputs_for_common()
     UpperCAmelCase	:							str						=			3
     UpperCAmelCase	:							Optional[int]						=			input_dict['''input_ids''']
     UpperCAmelCase	:							Dict						=			input_ids.ne(1						).to(__A						)
     UpperCAmelCase	:							List[Any]						=			ids_tensor([self.model_tester.batch_size],		self.model_tester.type_sequence_label_size						)
     UpperCAmelCase	:							Union[str, Any]						=			FalconForSequenceClassification(__A						)
     model.to(__A						)
     model.eval()
     UpperCAmelCase	:							str						=			model(__A,		attention_mask=__A,		labels=__A						)
     self.assertEqual(result.logits.shape,		(self.model_tester.batch_size, self.model_tester.num_labels)						)



   def 	__magic_name__   (							self							:					Optional[int]						):
     UpperCAmelCase			,						UpperCAmelCase	:							List[str]						=			self.model_tester.prepare_config_and_inputs_for_common()
     UpperCAmelCase	:							Any						=			3
     UpperCAmelCase	:							Dict						=			'''single_label_classification'''
     UpperCAmelCase	:							Tuple						=			input_dict['''input_ids''']
     UpperCAmelCase	:							Any						=			input_ids.ne(1						).to(__A						)
     UpperCAmelCase	:							str						=			ids_tensor([self.model_tester.batch_size],		self.model_tester.type_sequence_label_size						)
     UpperCAmelCase	:							str						=			FalconForSequenceClassification(__A						)
     model.to(__A						)
     model.eval()
     UpperCAmelCase	:							Dict						=			model(__A,		attention_mask=__A,		labels=__A						)
     self.assertEqual(result.logits.shape,		(self.model_tester.batch_size, self.model_tester.num_labels)						)



   def 	__magic_name__   (							self							:					Any						):
     UpperCAmelCase			,						UpperCAmelCase	:							Union[str, Any]						=			self.model_tester.prepare_config_and_inputs_for_common()
     UpperCAmelCase	:							Optional[int]						=			input_dict['''input_ids''']
     UpperCAmelCase	:							Union[str, Any]						=			FalconForCausalLM(__A						)
     model.to(__A						)
     model.eval()
     UpperCAmelCase	:							Optional[int]						=			model(__A,		use_cache=__A						)
     UpperCAmelCase	:							Union[str, Any]						=			input_ids.shape[0]
     UpperCAmelCase	:							str						=			model._convert_to_rw_cache(result.past_key_values						)
     UpperCAmelCase	:							Optional[Any]						=			model._convert_cache_to_standard_format(__A,		__A						)
     for layer in range(len(__A						)						):
       for tensor_idx in range(2						):
         self.assertTrue(rw_cache[layer][tensor_idx].ndim == 3						)
         self.assertTrue(result.past_key_values[layer][tensor_idx].ndim == 4						)
         self.assertTrue(
             torch.all(result.past_key_values[layer][tensor_idx] == standard_cache[layer][tensor_idx]						)						)



   def 	__magic_name__   (							self							:					Union[str, Any]						):
     UpperCAmelCase			,						UpperCAmelCase	:							List[str]						=			self.model_tester.prepare_config_and_inputs_for_common()
     UpperCAmelCase	:							List[str]						=			3
     UpperCAmelCase	:							Optional[Any]						=			'''multi_label_classification'''
     UpperCAmelCase	:							Optional[Any]						=			input_dict['''input_ids''']
     UpperCAmelCase	:							Optional[Any]						=			input_ids.ne(1						).to(__A						)
     UpperCAmelCase	:							Optional[Any]						=			ids_tensor(
         [self.model_tester.batch_size, config.num_labels],		self.model_tester.type_sequence_label_size						).to(torch.float						)
     UpperCAmelCase	:							str						=			FalconForSequenceClassification(__A						)
     model.to(__A						)
     model.eval()
     UpperCAmelCase	:							Tuple						=			model(__A,		attention_mask=__A,		labels=__A						)
     self.assertEqual(result.logits.shape,		(self.model_tester.batch_size, self.model_tester.num_labels)						)





   def 	__magic_name__   (							self							:					Optional[Any]						):
     # Falcon can have different numbers of KV-heads than the number of query heads, so we need
     # to override this test to use the right head counts.
     for model_class in self.all_generative_model_classes:
       UpperCAmelCase			,						UpperCAmelCase	:							Tuple						=			self.model_tester.prepare_config_and_inputs_for_common()

       # If it doesn't support cache, pass the test
       if not hasattr(__A,		'''use_cache'''						):
         return

       UpperCAmelCase	:							List[Any]						=			model_class(__A						).to(__A						)
       if "use_cache" not in inputs:
         UpperCAmelCase	:							Optional[Any]						=			True
       UpperCAmelCase	:							Optional[int]						=			model(**__A						)

       # If "past_key_values" is not returned, pass the test (e.g. RWKV uses a different cache name and format)
       if "past_key_values" not in outputs:
         return

       UpperCAmelCase	:							List[str]						=			(
           getattr(__A,		'''decoder_layers''',		__A						)
           or getattr(__A,		'''num_decoder_layers''',		__A						)
           or config.num_hidden_layers
       )
       UpperCAmelCase	:							Optional[int]						=			getattr(__A,		'''num_kv_heads''',		config.num_attention_heads						)
       UpperCAmelCase	:							List[Any]						=			getattr(__A,		'''d_model''',		config.hidden_size						)
       UpperCAmelCase	:							int						=			embed_dim // num_attention_heads

       UpperCAmelCase	:							int						=			outputs['''past_key_values''']
       self.assertEqual(len(__A						),		__A						)

       UpperCAmelCase			,						UpperCAmelCase	:							str						=			inputs['''input_ids'''].shape
       for i in range(__A						):
         if config.new_decoder_architecture:
           UpperCAmelCase	:							Dict						=			config.num_attention_heads
         elif config.multi_query:
           UpperCAmelCase	:							Tuple						=			1
         self.assertEqual(len(past_kv[0]						),		2						)  # K V for the decoder = 2
         self.assertEqual(
             past_kv[i][0].shape,		(batch_size, num_attention_heads, seq_length, per_head_embed_dim)						)
         self.assertEqual(
             past_kv[i][1].shape,		(batch_size, num_attention_heads, seq_length, per_head_embed_dim)						)






@require_torch
class  __UpperCAmelCase					(	unittest.TestCase						):



   @slow
   def 	__magic_name__   (							self							:					Union[str, Any]						):
     UpperCAmelCase	:							Tuple						=			AutoTokenizer.from_pretrained('''Rocketknight1/falcon-rw-1b'''						)
     UpperCAmelCase	:							Tuple						=			FalconForCausalLM.from_pretrained('''Rocketknight1/falcon-rw-1b'''						)
     model.eval()
     model.to(__A						)
     UpperCAmelCase	:							Tuple						=			tokenizer('''My favorite food is''',		return_tensors='''pt'''						).to(__A						)

     UpperCAmelCase	:							int						=			(
         '''My favorite food is pizza. I love it so much that I have a pizza party every year for my birthday.'''
     )

     UpperCAmelCase	:							int						=			model.generate(**__A,		do_sample=__A,		max_new_tokens=1_9						)
     UpperCAmelCase	:							int						=			tokenizer.batch_decode(__A						)[0]

     self.assertEqual(__A,		__A						)



   @slow
   def 	__magic_name__   (							self							:					List[Any]						):
     # The big models are way too big for the CI, so we use tiny random models that resemble their
     # architectures but with much smaller and fewer layers
     for repo in ["Rocketknight1/tiny-random-falcon-7b", "Rocketknight1/tiny-random-falcon-40b"]:
       UpperCAmelCase	:							int						=			AutoTokenizer.from_pretrained(__A						)
       UpperCAmelCase	:							List[str]						=			FalconForCausalLM.from_pretrained(__A						)
       model.eval()
       model.to(__A						)
       UpperCAmelCase	:							Optional[int]						=			tokenizer('''My favorite food is''',		return_tensors='''pt'''						).to(__A						)

       # We just test that these run without errors - the models are randomly initialized
       # and so the actual text outputs will be garbage
       model.generate(**__A,		do_sample=__A,		max_new_tokens=4						)
       model.generate(**__A,		do_sample=__A,		max_new_tokens=4						)
       model.generate(**__A,		num_beams=2,		max_new_tokens=4						)





   @slow
   def 	__magic_name__   (							self							:					Union[str, Any]						):
     # The big models are way too big for the CI, so we use tiny random models that resemble their
     # architectures but with much smaller and fewer layers
     with torch.no_grad():
       for repo in [
           "Rocketknight1/falcon-rw-1b",
           "Rocketknight1/tiny-random-falcon-7b",
           "Rocketknight1/tiny-random-falcon-40b",
       ]:
         UpperCAmelCase	:							Optional[int]						=			AutoTokenizer.from_pretrained(__A						)
         UpperCAmelCase	:							Union[str, Any]						=			FalconForCausalLM.from_pretrained(__A						)
         model.eval()
         model.to(device=__A						)
         UpperCAmelCase	:							Dict						=			tokenizer('''My favorite food is''',		return_tensors='''pt'''						).to(__A						)

         # Test results are the same with and without cache
         UpperCAmelCase	:							int						=			model.generate(**__A,		do_sample=__A,		max_new_tokens=2_0,		use_cache=__A						)
         UpperCAmelCase	:							Dict						=			model.generate(**__A,		do_sample=__A,		max_new_tokens=2_0,		use_cache=__A						)
         self.assertTrue((outputs_cache - outputs_no_cache).sum().item() == 0						)








import torch
from torch import nn





class  __UpperCAmelCase					(	nn.Module						):



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

     UpperCAmelCase	:							Union[str, Any]						=			n_token
     UpperCAmelCase	:							List[str]						=			d_embed
     UpperCAmelCase	:							Dict						=			d_proj

     UpperCAmelCase	:							List[Any]						=			cutoffs + [n_token]
     UpperCAmelCase	:							Dict						=			[0] + self.cutoffs
     UpperCAmelCase	:							int						=			div_val

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

     if self.n_clusters > 0:
       UpperCAmelCase	:							str						=			nn.Parameter(torch.zeros(self.n_clusters,		self.d_embed						)						)
       UpperCAmelCase	:							List[str]						=			nn.Parameter(torch.zeros(self.n_clusters						)						)

     UpperCAmelCase	:							Dict						=			nn.ModuleList()
     UpperCAmelCase	:							Optional[int]						=			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(__A,		__A						)						)						)
         else:
           self.out_projs.append(__A						)

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

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

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

     UpperCAmelCase	:							Optional[int]						=			keep_order



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

     return logit



   def 	__magic_name__   (							self							:					int,		__A							:					int,		__A							:					List[Any]=None,		__A							:					Dict=False						):

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

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

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

         weights.append(__A						)
         biases.append(__A						)

       UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase	:							str						=			weights[0], biases[0], self.out_projs[0]

       UpperCAmelCase	:							Dict						=			self._compute_logit(__A,		__A,		__A,		__A						)
       UpperCAmelCase	:							int						=			nn.functional.log_softmax(__A,		dim=1						)

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

       UpperCAmelCase	:							Union[str, Any]						=			0
       UpperCAmelCase	:							Any						=			[0] + self.cutoffs
       for i in range(len(__A						) - 1						):
         UpperCAmelCase			,						UpperCAmelCase	:							Optional[int]						=			cutoff_values[i], cutoff_values[i + 1]

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

           if indices_i.numel() == 0:
             continue

           UpperCAmelCase	:							Any						=			labels.index_select(0,		__A						) - l_idx
           UpperCAmelCase	:							Dict						=			head_logprob.index_select(0,		__A						)
           UpperCAmelCase	:							List[str]						=			hidden.index_select(0,		__A						)
         else:
           UpperCAmelCase	:							Tuple						=			hidden

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

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

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

     return out





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

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

         weights.append(__A						)
         biases.append(__A						)

       UpperCAmelCase			,						UpperCAmelCase			,						UpperCAmelCase	:							Tuple						=			weights[0], biases[0], self.out_projs[0]
       UpperCAmelCase	:							int						=			self._compute_logit(__A,		__A,		__A,		__A						)

       UpperCAmelCase	:							Optional[int]						=			hidden.new_empty((head_logit.size(0						), self.n_token)						)
       UpperCAmelCase	:							Dict						=			nn.functional.log_softmax(__A,		dim=1						)

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

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

           UpperCAmelCase	:							Tuple						=			self._compute_logit(__A,		__A,		__A,		__A						)
           UpperCAmelCase	:							List[Any]						=			nn.functional.log_softmax(__A,		dim=1						)

           UpperCAmelCase	:							Optional[int]						=			head_logprob[:, -i] + tail_logprob_i
           UpperCAmelCase	:							Optional[Any]						=			logprob_i

       return out





