





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   _SCREAMING_SNAKE_CASE		:







 '''simple docstring'''

 def __init__(self							:  Any			,					UpperCAmelCase_							:  Union[str, Any]			,					UpperCAmelCase_							:  List[Any]=99			,					UpperCAmelCase_							:  Any=13			,					UpperCAmelCase_							:  Dict=7			,					UpperCAmelCase_							:  str=9			,					UpperCAmelCase_							:  int=True			,					UpperCAmelCase_							:  int=True			,					UpperCAmelCase_							:  Dict=False			,					UpperCAmelCase_							:  Optional[Any]=32			,					UpperCAmelCase_							:  Optional[Any]=5			,					UpperCAmelCase_							:  Union[str, Any]=4			,					UpperCAmelCase_							:  List[Any]=37			,					UpperCAmelCase_							:  Tuple=8			,					UpperCAmelCase_							:  Optional[Any]=0.1			,					UpperCAmelCase_							:  Optional[Any]=0.002			,					UpperCAmelCase_							:  Optional[int]=1			,					UpperCAmelCase_							:  Union[str, Any]=0			,					UpperCAmelCase_							:  str=0			,					UpperCAmelCase_							:  Any=None			,					UpperCAmelCase_							:  Union[str, Any]=None			,					)			->int:



   '''simple docstring'''



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

 def 				SCREAMING_SNAKE_CASE_   (self							:  str)			->List[str]:



   '''simple docstring'''



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

 def 				SCREAMING_SNAKE_CASE_   (self							:  Union[str, Any]			,					UpperCAmelCase_							:  Optional[int]			,					UpperCAmelCase_							:  Tuple			,					UpperCAmelCase_							:  Union[str, Any]			,					UpperCAmelCase_							:  List[Any]=None			,					UpperCAmelCase_							:  Any=None			,					UpperCAmelCase_							:  int=None			,					UpperCAmelCase_							:  Union[str, Any]=None			,					UpperCAmelCase_							:  List[str]=None			,					)			->Tuple:



   '''simple docstring'''



   if attention_mask is None:
     lowerCamelCase__: int         =input_ids.ne(config.pad_token_id)
   if decoder_attention_mask is None:
     lowerCamelCase__: Tuple         =decoder_input_ids.ne(config.pad_token_id)
   if head_mask is None:
     lowerCamelCase__: List[Any]         =torch.ones(config.num_hidden_layers			,					config.num_attention_heads			,					device=UpperCAmelCase_)
   if decoder_head_mask is None:
     lowerCamelCase__: Union[str, Any]         =torch.ones(config.num_decoder_layers			,					config.num_attention_heads			,					device=UpperCAmelCase_)
   if cross_attn_head_mask is None:
     lowerCamelCase__: List[str]         =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 				SCREAMING_SNAKE_CASE_   (self							:  Optional[Any])			->Any:



   '''simple docstring'''



   lowerCamelCase__: Tuple         =ids_tensor([self.batch_size, self.encoder_seq_length]			,					self.vocab_size)
   lowerCamelCase__: int         =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
   lowerCamelCase__: Optional[int]         =input_ids.clamp(self.pad_token_id + 1)
   lowerCamelCase__: Union[str, Any]         =decoder_input_ids.clamp(self.pad_token_id + 1)

   lowerCamelCase__: Tuple         =self.get_config()
   lowerCamelCase__: Dict         =config.num_attention_heads
   lowerCamelCase__: List[str]         =self.prepare_inputs_dict(UpperCAmelCase_			,					UpperCAmelCase_			,					UpperCAmelCase_)
   return config, input_dict

 def 				SCREAMING_SNAKE_CASE_   (self							:  Tuple)			->Tuple:



   '''simple docstring'''



   lowerCamelCase__    ,   lowerCamelCase__: List[Any]         =self.prepare_config_and_inputs()
   return config, inputs_dict

 def 				SCREAMING_SNAKE_CASE_   (self							:  Optional[Any])			->Any:



   '''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 				SCREAMING_SNAKE_CASE_   (self							:  Tuple)			->Dict:



   '''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 				SCREAMING_SNAKE_CASE_   (self							:  List[Any]			,					UpperCAmelCase_							:  Dict			,					UpperCAmelCase_							:  List[str]			,					UpperCAmelCase_							:  List[Any]			,					UpperCAmelCase_							:  str			,					UpperCAmelCase_							:  Union[str, Any]			,					UpperCAmelCase_							:  Optional[int]			,					)			->Dict:



   '''simple docstring'''



   lowerCamelCase__: Dict         =UMTaModel(config=UpperCAmelCase_)
   model.to(UpperCAmelCase_)
   model.eval()
   lowerCamelCase__: Any         =model(
       input_ids=UpperCAmelCase_			,					decoder_input_ids=UpperCAmelCase_			,					attention_mask=UpperCAmelCase_			,					decoder_attention_mask=UpperCAmelCase_			,					)
   lowerCamelCase__: Optional[int]         =model(input_ids=UpperCAmelCase_			,					decoder_input_ids=UpperCAmelCase_)
   lowerCamelCase__: List[str]         =result.last_hidden_state
   lowerCamelCase__: Union[str, Any]         =result.past_key_values
   lowerCamelCase__: 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 				SCREAMING_SNAKE_CASE_   (self							:  List[str]			,					UpperCAmelCase_							:  Any			,					UpperCAmelCase_							:  Tuple			,					UpperCAmelCase_							:  Optional[Any]			,					UpperCAmelCase_							:  List[Any]			,					UpperCAmelCase_							:  Tuple			,					UpperCAmelCase_							:  Optional[Any]			,					)			->Any:



   '''simple docstring'''



   lowerCamelCase__: Dict         =UMTaModel(config=UpperCAmelCase_).get_decoder().to(UpperCAmelCase_).eval()
   # first forward pass
   lowerCamelCase__: List[Any]         =model(UpperCAmelCase_			,					use_cache=UpperCAmelCase_)
   lowerCamelCase__: Union[str, Any]         =model(UpperCAmelCase_)
   lowerCamelCase__: Optional[int]         =model(UpperCAmelCase_			,					use_cache=UpperCAmelCase_)

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

   lowerCamelCase__    ,   lowerCamelCase__: Dict         =outputs.to_tuple()

   # create hypothetical next token and extent to next_input_ids
   lowerCamelCase__: List[Any]         =ids_tensor((self.batch_size, 1)			,					config.vocab_size)

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

   lowerCamelCase__: List[str]         =model(UpperCAmelCase_)["last_hidden_state"]
   lowerCamelCase__: List[str]         =model(UpperCAmelCase_			,					past_key_values=UpperCAmelCase_)["last_hidden_state"]

   # select random slice
   lowerCamelCase__: Optional[Any]         =ids_tensor((1,)			,					output_from_past.shape[-1]).item()
   lowerCamelCase__: Optional[Any]         =output_from_no_past[:, -1, random_slice_idx].detach()
   lowerCamelCase__: 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 				SCREAMING_SNAKE_CASE_   (self							:  Tuple			,					UpperCAmelCase_							:  Optional[Any]			,					UpperCAmelCase_							:  Dict			,					)			->Dict:



   '''simple docstring'''



   lowerCamelCase__: Tuple         =UMTaModel(config=UpperCAmelCase_).to(UpperCAmelCase_).half().eval()
   lowerCamelCase__: str         =model(**UpperCAmelCase_)["last_hidden_state"]
   self.parent.assertFalse(torch.isnan(UpperCAmelCase_).any().item())



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







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

 def 				SCREAMING_SNAKE_CASE_   (self							:  Optional[int])			->List[Any]:



   '''simple docstring'''



   lowerCamelCase__: Optional[Any]         =UMTaModelTester(self)

 @unittest.skip("Test has a segmentation fault on torch 1.8.0")
 def 				SCREAMING_SNAKE_CASE_   (self							:  Dict)			->Optional[int]:



   '''simple docstring'''



   lowerCamelCase__: Dict         =self.model_tester.prepare_config_and_inputs()
   lowerCamelCase__: int         =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 				SCREAMING_SNAKE_CASE_   (self							:  Union[str, Any])			->Any:



   '''simple docstring'''



   lowerCamelCase__: Dict         =self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_model_fpaa_forward(*UpperCAmelCase_)

 def 				SCREAMING_SNAKE_CASE_   (self							:  Any)			->int:



   '''simple docstring'''



   lowerCamelCase__: Optional[Any]         =["encoder_attentions", "decoder_attentions", "cross_attentions"]
   lowerCamelCase__: Any         =self.model_tester.prepare_config_and_inputs()
   lowerCamelCase__: List[Any]         =config_and_inputs[0]
   lowerCamelCase__: Tuple         =UMTaForConditionalGeneration(UpperCAmelCase_).eval()
   model.to(UpperCAmelCase_)

   lowerCamelCase__: int         ={
       "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()):
     lowerCamelCase__: Any         ={name: mask}
     # Explicitly pass decoder_head_mask as it is required from T5 model when head_mask specified
     if name == "head_mask":
       lowerCamelCase__: Union[str, Any]         =torch.ones(
           config.num_decoder_layers			,					config.num_heads			,					device=UpperCAmelCase_)

     lowerCamelCase__: str         =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
     lowerCamelCase__: Union[str, Any]         =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 				SCREAMING_SNAKE_CASE_   (self							:  Any)			->str:



   '''simple docstring'''



   pass



@require_torch
@require_sentencepiece
@require_tokenizers
class   _SCREAMING_SNAKE_CASE		(      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 				SCREAMING_SNAKE_CASE_   (self							:  Tuple)			->List[str]:



   '''simple docstring'''



   lowerCamelCase__: str         =UMTaForConditionalGeneration.from_pretrained("google/umt5-small"			,					return_dict=UpperCAmelCase_).to(UpperCAmelCase_)
   lowerCamelCase__: Optional[int]         =AutoTokenizer.from_pretrained("google/umt5-small"			,					use_fast=UpperCAmelCase_			,					legacy=UpperCAmelCase_)
   lowerCamelCase__: List[str]         =[
       "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>.",
   ]
   lowerCamelCase__: Optional[int]         =tokenizer(UpperCAmelCase_			,					return_tensors="pt"			,					padding=UpperCAmelCase_).input_ids
   # fmt: off
   lowerCamelCase__: Optional[Any]         =torch.tensor(
       [
           [ 38_530, 210_703, 256_299, 1_410, 256_298, 274, 1, 0,0, 0, 0, 0, 0, 0, 0, 0,0, 0],
           [   826, 321, 671, 25_922, 256_299, 274, 1, 0,0, 0, 0, 0, 0, 0, 0, 0,0, 0],
           [  1_460, 339, 312, 19_014, 10_620, 758, 256_299, 2_355,274, 1, 0, 0, 0, 0, 0, 0,0, 0],
           [   517, 256_299, 14_869, 281, 301, 256_298, 275, 119_983,1, 0, 0, 0, 0, 0, 0, 0,0, 0],
           [   320, 256_299, 14_869, 281, 2_234, 289, 2_275, 333,61_391, 289, 256_298, 543, 256_297, 168_714, 329, 256_296,274, 1],
       ])
   # fmt: on
   torch.testing.assert_allclose(UpperCAmelCase_			,					UpperCAmelCase_)

   lowerCamelCase__: Union[str, Any]         =model.generate(input_ids.to(UpperCAmelCase_))
   lowerCamelCase__: Optional[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>",
   ]
   lowerCamelCase__: int         =tokenizer.batch_decode(UpperCAmelCase_)
   self.assertEqual(UpperCAmelCase_			,					UpperCAmelCase_)









import torch
from torch import nn





class 					A__     (     nn.Module ):

     '''simple docstring'''






     def __init__( self      :    List[str]   ,						_SCREAMING_SNAKE_CASE      :    List[Any]   ,						_SCREAMING_SNAKE_CASE      :    Optional[int]   ,						_SCREAMING_SNAKE_CASE      :    Tuple   ,						_SCREAMING_SNAKE_CASE      :    str   ,						_SCREAMING_SNAKE_CASE      :    Optional[int]=1   ,						_SCREAMING_SNAKE_CASE      :    List[str]=False ):






           """simple docstring"""




           super().__init__()

           UpperCamelCase					       =  n_token
           UpperCamelCase					       =  d_embed
           UpperCamelCase					       =  d_proj

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

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

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

           UpperCamelCase					       =  nn.ModuleList()
           UpperCamelCase					       =  nn.ParameterList()

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

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

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

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

           UpperCamelCase					       =  keep_order







     def   _SCREAMING_SNAKE_CASE     ( self      :    Dict   ,						_SCREAMING_SNAKE_CASE      :    Optional[Any]   ,						_SCREAMING_SNAKE_CASE      :    Dict   ,						_SCREAMING_SNAKE_CASE      :    List[Any]   ,						_SCREAMING_SNAKE_CASE      :    Dict ):






           """simple docstring"""




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

           return logit







     def   _SCREAMING_SNAKE_CASE     ( self      :    List[Any]   ,						_SCREAMING_SNAKE_CASE      :    Union[str, Any]   ,						_SCREAMING_SNAKE_CASE      :    Optional[Any]=None   ,						_SCREAMING_SNAKE_CASE      :    Dict=False ):






           """simple docstring"""




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

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

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

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

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

                 UpperCamelCase					       =  self._compute_logit(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE )
                 UpperCamelCase					       =  nn.functional.log_softmax(_SCREAMING_SNAKE_CASE   ,						dim=1 )

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

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

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

                             if indices_i.numel() == 0:
                                   continue

                             UpperCamelCase					       =  labels.index_select(0   ,						_SCREAMING_SNAKE_CASE ) - l_idx
                             UpperCamelCase					       =  head_logprob.index_select(0   ,						_SCREAMING_SNAKE_CASE )
                             UpperCamelCase					       =  hidden.index_select(0   ,						_SCREAMING_SNAKE_CASE )
                       else:
                             UpperCamelCase					       =  hidden

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

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

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

           return out







     def   _SCREAMING_SNAKE_CASE     ( self      :    int   ,						_SCREAMING_SNAKE_CASE      :    str ):






           """simple docstring"""




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

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

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

                 UpperCamelCase       ,  UpperCamelCase       ,  UpperCamelCase					       =  weights[0], biases[0], self.out_projs[0]
                 UpperCamelCase					       =  self._compute_logit(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE )

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

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

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

                             UpperCamelCase					       =  self._compute_logit(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE )
                             UpperCamelCase					       =  nn.functional.log_softmax(_SCREAMING_SNAKE_CASE   ,						dim=1 )

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

                 return out



