


'''simple docstring'''
import tempfile
import unittest

from transformers import TaConfig, is_torch_available
from transformers.testing_utils import (
    require_sentencepiece,
    require_tokenizers,
    require_torch,
    slow,
    torch_device,
)

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


if is_torch_available():
		import torch

		from transformers import AutoTokenizer, UMTaForConditionalGeneration, UMTaForQuestionAnswering, UMTaModel
class      lowercase			:

		"""simple docstring"""

		def __init__(	self					,     UpperCamelCase_					,     UpperCamelCase_=99					,     UpperCamelCase_=13					,     UpperCamelCase_=7					,     UpperCamelCase_=9					,     UpperCamelCase_=True					,     UpperCamelCase_=True					,     UpperCamelCase_=False					,     UpperCamelCase_=32					,     UpperCamelCase_=5					,     UpperCamelCase_=4					,     UpperCamelCase_=37					,     UpperCamelCase_=8					,     UpperCamelCase_=0.1					,     UpperCamelCase_=0.002					,     UpperCamelCase_=1					,     UpperCamelCase_=0					,     UpperCamelCase_=0					,     UpperCamelCase_=None					,     UpperCamelCase_=None					,     ):

							'''simple docstring'''


							UpperCamelCase__     :Dict							      =				parent
							UpperCamelCase__     :Optional[Any]							      =				batch_size
							UpperCamelCase__     :int							      =				encoder_seq_length
							UpperCamelCase__     :Optional[Any]							      =				decoder_seq_length
							# For common tests
							UpperCamelCase__     :Optional[int]							      =				self.decoder_seq_length
							UpperCamelCase__     :int							      =				is_training
							UpperCamelCase__     :List[str]							      =				use_attention_mask
							UpperCamelCase__     :List[Any]							      =				use_labels
							UpperCamelCase__     :Tuple							      =				vocab_size
							UpperCamelCase__     :Union[str, Any]							      =				hidden_size
							UpperCamelCase__     :Dict							      =				num_hidden_layers
							UpperCamelCase__     :Optional[Any]							      =				num_attention_heads
							UpperCamelCase__     :Optional[int]							      =				d_ff
							UpperCamelCase__     :Optional[int]							      =				relative_attention_num_buckets
							UpperCamelCase__     :Union[str, Any]							      =				dropout_rate
							UpperCamelCase__     :Any							      =				initializer_factor
							UpperCamelCase__     :List[Any]							      =				eos_token_id
							UpperCamelCase__     :Union[str, Any]							      =				pad_token_id
							UpperCamelCase__     :Optional[int]							      =				decoder_start_token_id
							UpperCamelCase__     :int							      =				None
							UpperCamelCase__     :Optional[Any]							      =				decoder_layers




		def 					lowerCAmelCase__       (	self							):

							'''simple docstring'''


							return TaConfig.from_pretrained('''google/umt5-base'''							)




		def 					lowerCAmelCase__       (	self					,     UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_=None					,     UpperCamelCase_=None					,     UpperCamelCase_=None					,     UpperCamelCase_=None					,     UpperCamelCase_=None					,     ):

							'''simple docstring'''


							if attention_mask is None:
												UpperCamelCase__     :int							      =				input_ids.ne(config.pad_token_id							)
							if decoder_attention_mask is None:
												UpperCamelCase__     :int							      =				decoder_input_ids.ne(config.pad_token_id							)
							if head_mask is None:
												UpperCamelCase__     :Optional[Any]							      =				torch.ones(config.num_hidden_layers					,     config.num_attention_heads					,     device=UpperCamelCase_							)
							if decoder_head_mask is None:
												UpperCamelCase__     :List[str]							      =				torch.ones(config.num_decoder_layers					,     config.num_attention_heads					,     device=UpperCamelCase_							)
							if cross_attn_head_mask is None:
												UpperCamelCase__     :Optional[Any]							      =				torch.ones(
												    config.num_decoder_layers					,     config.num_attention_heads					,     device=UpperCamelCase_							)
							return {
							    "input_ids": input_ids,
							    "decoder_input_ids": decoder_input_ids,
							    "attention_mask": attention_mask,
							    "decoder_attention_mask": decoder_attention_mask,
							    "head_mask": head_mask,
							    "decoder_head_mask": decoder_head_mask,
							    "cross_attn_head_mask": cross_attn_head_mask,
							}




		def 					lowerCAmelCase__       (	self							):

							'''simple docstring'''


							UpperCamelCase__     :Optional[Any]							      =				ids_tensor([self.batch_size, self.encoder_seq_length]					,     self.vocab_size							)
							UpperCamelCase__     :List[Any]							      =				ids_tensor([self.batch_size, self.decoder_seq_length]					,     self.vocab_size							)

							# we need to clamp the input ids here to avoid having pad token in between
							# this is because for NllbMoe the position_ids are prepared such that
							# all pad tokens have pos id = 2 and rest are between 2..seq_length
							# and the seq_length here is seq_length - num_pad_tokens
							# but when using past, there is no way of knowing if the past input ids had
							# pad tokens in them, which results in incorrect seq_lenth and which in turn results in
							# position_ids being off by num_pad_tokens in past input
							UpperCamelCase__     :str							      =				input_ids.clamp(self.pad_token_id + 1							)
							UpperCamelCase__     :Optional[int]							      =				decoder_input_ids.clamp(self.pad_token_id + 1							)

							UpperCamelCase__     :Optional[int]							      =				self.get_config()
							UpperCamelCase__     :str							      =				config.num_attention_heads
							UpperCamelCase__     :Union[str, Any]							      =				self.prepare_inputs_dict(UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_							)
							return config, input_dict




		def 					lowerCAmelCase__       (	self							):

							'''simple docstring'''


							UpperCamelCase__      ,       UpperCamelCase__     :Tuple							      =				self.prepare_config_and_inputs()
							return config, inputs_dict




		def 					lowerCAmelCase__       (	self							):

							'''simple docstring'''


							return TaConfig(
							    vocab_size=166					,     d_model=self.hidden_size					,     d_ff=self.d_ff					,     d_kv=self.hidden_size // self.num_attention_heads					,     num_layers=self.num_hidden_layers					,     num_decoder_layers=self.decoder_layers					,     num_heads=self.num_attention_heads					,     relative_attention_num_buckets=self.relative_attention_num_buckets					,     dropout_rate=self.dropout_rate					,     initializer_factor=self.initializer_factor					,     eos_token_id=self.eos_token_id					,     bos_token_id=self.pad_token_id					,     pad_token_id=self.pad_token_id					,     decoder_start_token_id=self.decoder_start_token_id					,     )




		def 					lowerCAmelCase__       (	self							):

							'''simple docstring'''


							return TaConfig(
							    vocab_size=self.vocab_size					,     d_model=self.hidden_size					,     d_ff=self.d_ff					,     d_kv=self.hidden_size // self.num_attention_heads					,     num_layers=self.num_hidden_layers					,     num_decoder_layers=self.decoder_layers					,     num_heads=self.num_attention_heads					,     relative_attention_num_buckets=self.relative_attention_num_buckets					,     dropout_rate=self.dropout_rate					,     initializer_factor=self.initializer_factor					,     eos_token_id=self.eos_token_id					,     bos_token_id=self.pad_token_id					,     pad_token_id=self.pad_token_id					,     decoder_start_token_id=self.decoder_start_token_id					,     )




		def 					lowerCAmelCase__       (	self					,     UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_					,     ):

							'''simple docstring'''


							UpperCamelCase__     :Optional[Any]							      =				UMTaModel(config=UpperCamelCase_							)
							model.to(UpperCamelCase_							)
							model.eval()
							UpperCamelCase__     :str							      =				model(
							    input_ids=UpperCamelCase_					,     decoder_input_ids=UpperCamelCase_					,     attention_mask=UpperCamelCase_					,     decoder_attention_mask=UpperCamelCase_					,     )
							UpperCamelCase__     :Any							      =				model(input_ids=UpperCamelCase_					,     decoder_input_ids=UpperCamelCase_							)
							UpperCamelCase__     :int							      =				result.last_hidden_state
							UpperCamelCase__     :Optional[int]							      =				result.past_key_values
							UpperCamelCase__     :List[str]							      =				result.encoder_last_hidden_state

							self.parent.assertEqual(encoder_output.size()					,     (self.batch_size, self.encoder_seq_length, self.hidden_size)							)
							self.parent.assertEqual(decoder_output.size()					,     (self.batch_size, self.decoder_seq_length, self.hidden_size)							)
							# There should be `num_layers` key value embeddings stored in decoder_past
							self.parent.assertEqual(len(UpperCamelCase_							)					,     config.num_layers							)
							# There should be a self attn key, a self attn value, a cross attn key and a cross attn value stored in each decoder_past tuple
							self.parent.assertEqual(len(decoder_past[0]							)					,     4							)




		def 					lowerCAmelCase__       (	self					,     UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_					,     ):

							'''simple docstring'''


							UpperCamelCase__     :Union[str, Any]							      =				UMTaModel(config=UpperCamelCase_							).get_decoder().to(UpperCamelCase_							).eval()
							# first forward pass
							UpperCamelCase__     :Union[str, Any]							      =				model(UpperCamelCase_					,     use_cache=UpperCamelCase_							)
							UpperCamelCase__     :int							      =				model(UpperCamelCase_							)
							UpperCamelCase__     :str							      =				model(UpperCamelCase_					,     use_cache=UpperCamelCase_							)

							self.parent.assertTrue(len(UpperCamelCase_							) == len(UpperCamelCase_							)							)
							self.parent.assertTrue(len(UpperCamelCase_							) == len(UpperCamelCase_							) + 1							)

							UpperCamelCase__      ,       UpperCamelCase__     :Any							      =				outputs.to_tuple()

							# create hypothetical next token and extent to next_input_ids
							UpperCamelCase__     :int							      =				ids_tensor((self.batch_size, 1)					,     config.vocab_size							)

							# append to next input_ids and
							UpperCamelCase__     :List[str]							      =				torch.cat([input_ids, next_tokens]					,     dim=-1							)

							UpperCamelCase__     :Optional[Any]							      =				model(UpperCamelCase_							)['''last_hidden_state''']
							UpperCamelCase__     :Optional[int]							      =				model(UpperCamelCase_					,     past_key_values=UpperCamelCase_							)['''last_hidden_state''']

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

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




		def 					lowerCAmelCase__       (	self					,     UpperCamelCase_					,     UpperCamelCase_					,     ):

							'''simple docstring'''


							UpperCamelCase__     :int							      =				UMTaModel(config=UpperCamelCase_							).to(UpperCamelCase_							).half().eval()
							UpperCamelCase__     :Optional[int]							      =				model(**UpperCamelCase_							)['''last_hidden_state''']
							self.parent.assertFalse(torch.isnan(UpperCamelCase_							).any().item()							)



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

		"""simple docstring"""

		_a									=       (
		    (UMTaModel, UMTaForConditionalGeneration, UMTaForQuestionAnswering) if is_torch_available() else ()
		)
		_a									=       (UMTaForConditionalGeneration,) if is_torch_available() else ()
		_a									=       (
		    {
		        'conversational': UMTaForConditionalGeneration,
		        'feature-extraction': UMTaModel,
		        'summarization': UMTaForConditionalGeneration,
		        'text2text-generation': UMTaForConditionalGeneration,
		        'translation': UMTaForConditionalGeneration,
		        'question-answering': UMTaForQuestionAnswering,
		    }
		    if is_torch_available()
		    else {}
		)
		_a									=       True
		_a									=       False
		_a									=       False
		_a									=       True
		_a									=       True
		# The small UMT5 model needs higher percentages for CPU/MP tests
		_a									=       [0.8, 0.9]




		def 					lowerCAmelCase__       (	self							):

							'''simple docstring'''


							UpperCamelCase__     :Any							      =				UMTaModelTester(self							)




		@unittest.skip('''Test has a segmentation fault on torch 1.8.0'''							)
		def 					lowerCAmelCase__       (	self							):

							'''simple docstring'''


							UpperCamelCase__     :Union[str, Any]							      =				self.model_tester.prepare_config_and_inputs()
							UpperCamelCase__     :Tuple							      =				UMTaModel(config_and_inputs[0]							).to(UpperCamelCase_							)
							with tempfile.TemporaryDirectory() as tmpdirname:
												torch.onnx.export(
												    UpperCamelCase_					,     (config_and_inputs[1], config_and_inputs[3], config_and_inputs[2])					,     F'''{tmpdirname}/t5_test.onnx'''					,     export_params=UpperCamelCase_					,     opset_version=9					,     input_names=['''input_ids''', '''decoder_input_ids''']					,     )




		@unittest.skipIf(torch_device == '''cpu'''					,     '''Cant do half precision'''							)
		def 					lowerCAmelCase__       (	self							):

							'''simple docstring'''


							UpperCamelCase__     :Union[str, Any]							      =				self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model_fpaa_forward(*UpperCamelCase_							)




		def 					lowerCAmelCase__       (	self							):

							'''simple docstring'''


							UpperCamelCase__     :int							      =				['''encoder_attentions''', '''decoder_attentions''', '''cross_attentions''']
							UpperCamelCase__     :List[Any]							      =				self.model_tester.prepare_config_and_inputs()
							UpperCamelCase__     :Optional[int]							      =				config_and_inputs[0]
							UpperCamelCase__     :List[Any]							      =				UMTaForConditionalGeneration(UpperCamelCase_							).eval()
							model.to(UpperCamelCase_							)

							UpperCamelCase__     :Any							      =				{
							    '''head_mask''': torch.zeros(config.num_layers					,     config.num_heads					,     device=UpperCamelCase_							),
							    '''decoder_head_mask''': torch.zeros(config.num_decoder_layers					,     config.num_heads					,     device=UpperCamelCase_							),
							    '''cross_attn_head_mask''': torch.zeros(config.num_decoder_layers					,     config.num_heads					,     device=UpperCamelCase_							),
							}

							for attn_name, (name, mask) in zip(UpperCamelCase_					,     head_masking.items()							):
												UpperCamelCase__     :str							      =				{name: mask}
												# Explicitly pass decoder_head_mask as it is required from T5 model when head_mask specified
												if name == "head_mask":
																	UpperCamelCase__     :Union[str, Any]							      =				torch.ones(
																	    config.num_decoder_layers					,     config.num_heads					,     device=UpperCamelCase_							)

												UpperCamelCase__     :Optional[int]							      =				model.generate(
												    config_and_inputs[1]['''input_ids''']					,     num_beams=1					,     max_length=3					,     output_attentions=UpperCamelCase_					,     return_dict_in_generate=UpperCamelCase_					,     **UpperCamelCase_					,     )
												# We check the state of decoder_attentions and cross_attentions just from the last step
												UpperCamelCase__     :int							      =				out[attn_name] if attn_name == attention_names[0] else out[attn_name][-1]
												self.assertEqual(sum([w.sum().item() for w in attn_weights]							)					,     0.0							)




		@unittest.skip('''Does not work on the tiny model as we keep hitting edge cases.'''							)
		def 					lowerCAmelCase__       (	self							):

							'''simple docstring'''


							pass



@require_torch
@require_sentencepiece
@require_tokenizers
class      lowercase			(							unittest.TestCase      ):

		"""simple docstring"""

		@slow
		@unittest.skip(
		    '''Unless we stop stripping left and right by default for all special tokens, the expected ids obtained here will not match the original ones. Wait for https://github.com/huggingface/transformers/pull/23909 to be merged'''							)
		def 					lowerCAmelCase__       (	self							):

							'''simple docstring'''


							UpperCamelCase__     :Optional[int]							      =				UMTaForConditionalGeneration.from_pretrained('''google/umt5-small'''					,     return_dict=UpperCamelCase_							).to(UpperCamelCase_							)
							UpperCamelCase__     :List[Any]							      =				AutoTokenizer.from_pretrained('''google/umt5-small'''					,     use_fast=UpperCamelCase_					,     legacy=UpperCamelCase_							)
							UpperCamelCase__     :Tuple							      =				[
							    '''Bonjour monsieur <extra_id_0> bien <extra_id_1>.''',
							    '''No se como puedo <extra_id_0>.''',
							    '''This is the reason why we <extra_id_0> them.''',
							    '''The <extra_id_0> walks in <extra_id_1>, seats''',
							    '''A <extra_id_0> walks into a bar and orders a <extra_id_1> with <extra_id_2> pinch of <extra_id_3>.''',
							]
							UpperCamelCase__     :Union[str, Any]							      =				tokenizer(UpperCamelCase_					,     return_tensors='''pt'''					,     padding=UpperCamelCase_							).input_ids
							# fmt: off
							UpperCamelCase__     :Any							      =				torch.tensor(
							    [
							        [ 38530, 210703, 256299, 1410, 256298, 274, 1, 0,0, 0, 0, 0, 0, 0, 0, 0,0, 0],
							        [   826, 321, 671, 25922, 256299, 274, 1, 0,0, 0, 0, 0, 0, 0, 0, 0,0, 0],
							        [  1460, 339, 312, 19014, 10620, 758, 256299, 2355,274, 1, 0, 0, 0, 0, 0, 0,0, 0],
							        [   517, 256299, 14869, 281, 301, 256298, 275, 119983,1, 0, 0, 0, 0, 0, 0, 0,0, 0],
							        [   320, 256299, 14869, 281, 2234, 289, 2275, 333,61391, 289, 256298, 543, 256297, 168714, 329, 256296,274, 1],
							    ]							)
							# fmt: on
							torch.testing.assert_allclose(UpperCamelCase_					,     UpperCamelCase_							)

							UpperCamelCase__     :Dict							      =				model.generate(input_ids.to(UpperCamelCase_							)							)
							UpperCamelCase__     :Union[str, Any]							      =				[
							    '''<pad><extra_id_0> et<extra_id_1> [eod] <extra_id_2><extra_id_55>.. [eod] 💐 💐 💐 💐 💐 💐 💐 💐 💐 💐 💐 <extra_id_56>ajšietosto<extra_id_56>lleux<extra_id_19><extra_id_6>ajšie</s>''',
							    '''<pad><extra_id_0>.<extra_id_1>.,<0x0A>...spech <0x0A><extra_id_20> <extra_id_21></s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>''',
							    '''<pad><extra_id_0> are not going to be a part of the world. We are not going to be a part of<extra_id_1> and<extra_id_2><0x0A><extra_id_48>.<extra_id_48></s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>''',
							    '''<pad><extra_id_0> door<extra_id_1>, the door<extra_id_2> 피해[/</s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>''',
							    '''<pad><extra_id_0>nyone who<extra_id_1> drink<extra_id_2> a<extra_id_3> alcohol<extra_id_4> A<extra_id_5> A. This<extra_id_6> I<extra_id_7><extra_id_52><extra_id_53></s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>''',
							]
							UpperCamelCase__     :int							      =				tokenizer.batch_decode(UpperCamelCase_							)
							self.assertEqual(UpperCamelCase_					,     UpperCamelCase_							)



'''simple docstring'''
import torch
from torch import nn
class      lowercase			(							nn.Module      ):

		"""simple docstring"""

		def __init__(	self					,     UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_=1					,     UpperCamelCase_=False							):

							'''simple docstring'''


							super().__init__()

							UpperCamelCase__     :Dict							      =				n_token
							UpperCamelCase__     :List[Any]							      =				d_embed
							UpperCamelCase__     :Dict							      =				d_proj

							UpperCamelCase__     :Dict							      =				cutoffs + [n_token]
							UpperCamelCase__     :Union[str, Any]							      =				[0] + self.cutoffs
							UpperCamelCase__     :Any							      =				div_val

							UpperCamelCase__     :int							      =				self.cutoffs[0]
							UpperCamelCase__     :List[Any]							      =				len(self.cutoffs							) - 1
							UpperCamelCase__     :List[Any]							      =				self.shortlist_size + self.n_clusters

							if self.n_clusters > 0:
												UpperCamelCase__     :Any							      =				nn.Parameter(torch.zeros(self.n_clusters					,     self.d_embed							)							)
												UpperCamelCase__     :Optional[Any]							      =				nn.Parameter(torch.zeros(self.n_clusters							)							)

							UpperCamelCase__     :Union[str, Any]							      =				nn.ModuleList()
							UpperCamelCase__     :str							      =				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(UpperCamelCase_					,     UpperCamelCase_							)							)							)
																	else:
																						self.out_projs.append(UpperCamelCase_							)

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

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

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

							UpperCamelCase__     :Tuple							      =				keep_order




		def 					lowerCAmelCase__       (	self					,     UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_							):

							'''simple docstring'''


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

							return logit




		def 					lowerCAmelCase__       (	self					,     UpperCamelCase_					,     UpperCamelCase_=None					,     UpperCamelCase_=False							):

							'''simple docstring'''


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

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

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

																	weights.append(UpperCamelCase_							)
																	biases.append(UpperCamelCase_							)

												UpperCamelCase__      ,       UpperCamelCase__      ,       UpperCamelCase__     :Optional[Any]							      =				weights[0], biases[0], self.out_projs[0]

												UpperCamelCase__     :str							      =				self._compute_logit(UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_							)
												UpperCamelCase__     :List[Any]							      =				nn.functional.log_softmax(UpperCamelCase_					,     dim=1							)

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

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

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

																						if indices_i.numel() == 0:
																											continue

																						UpperCamelCase__     :Tuple							      =				labels.index_select(0					,     UpperCamelCase_							) - l_idx
																						UpperCamelCase__     :str							      =				head_logprob.index_select(0					,     UpperCamelCase_							)
																						UpperCamelCase__     :int							      =				hidden.index_select(0					,     UpperCamelCase_							)
																	else:
																						UpperCamelCase__     :Dict							      =				hidden

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

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

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

							return out




		def 					lowerCAmelCase__       (	self					,     UpperCamelCase_							):

							'''simple docstring'''


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

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

																	weights.append(UpperCamelCase_							)
																	biases.append(UpperCamelCase_							)

												UpperCamelCase__      ,       UpperCamelCase__      ,       UpperCamelCase__     :str							      =				weights[0], biases[0], self.out_projs[0]
												UpperCamelCase__     :Optional[int]							      =				self._compute_logit(UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_							)

												UpperCamelCase__     :int							      =				hidden.new_empty((head_logit.size(0							), self.n_token)							)
												UpperCamelCase__     :Union[str, Any]							      =				nn.functional.log_softmax(UpperCamelCase_					,     dim=1							)

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

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

																						UpperCamelCase__     :List[str]							      =				self._compute_logit(UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_							)
																						UpperCamelCase__     :Tuple							      =				nn.functional.log_softmax(UpperCamelCase_					,     dim=1							)

																						UpperCamelCase__     :Optional[Any]							      =				head_logprob[:, -i] + tail_logprob_i
																						UpperCamelCase__     :Union[str, Any]							      =				logprob_i

												return out