






import tempfile
import unittest
from pathlib import Path
from shutil import copyfile

from transformers import MaMaaaTokenizer, is_torch_available
from transformers.testing_utils import (
    get_tests_dir,
    nested_simplify,
    require_sentencepiece,
    require_tokenizers,
    require_torch,
    slow,
)
from transformers.utils import is_sentencepiece_available


if is_sentencepiece_available():
							from transformers.models.mam_aaa.tokenization_mam_aaa import VOCAB_FILES_NAMES, save_json

from ...test_tokenization_common import TokenizerTesterMixin


if is_sentencepiece_available():
							A_       :Tuple       =						get_tests_dir('''fixtures/test_sentencepiece.model''')


if is_torch_available():
							from transformers.models.mam_aaa.modeling_mam_aaa import shift_tokens_right

A_       :Any       =						128022
A_       :List[Any]       =						128028




@require_sentencepiece
class 				__A     (			__lowerCAmelCase      ,     unittest.TestCase ):





				"""simple docstring"""







				UpperCamelCase__ :      Tuple			 =MaMaaaTokenizer
				UpperCamelCase__ :      List[Any]			 =False
				UpperCamelCase__ :      Optional[int]			 =False
				UpperCamelCase__ :      int			 =True


				def       __lowercase							( self    ):




									"""simple docstring"""

									super().setUp()

									__UpperCamelCase      :					List[str]       		=['</s>', '<unk>', '▁This', '▁is', '▁a', '▁t', 'est', '\u0120', '<pad>']
									__UpperCamelCase      :					Optional[Any]       		=dict(zip(lowerCAmelCase_							,  range(len(lowerCAmelCase_    )    )    )    )
									__UpperCamelCase      :					Dict       		=Path(self.tmpdirname    )
									save_json(lowerCAmelCase_							,  save_dir / VOCAB_FILES_NAMES['vocab_file']    )
									if not (save_dir / VOCAB_FILES_NAMES["spm_file"]).exists():
														copyfile(lowerCAmelCase_							,  save_dir / VOCAB_FILES_NAMES['spm_file']    )

									__UpperCamelCase      :					str       		=MaMaaaTokenizer.from_pretrained(self.tmpdirname    )
									tokenizer.save_pretrained(self.tmpdirname    )


				def       __lowercase							( self							,  **lowerCamelCase__    ):




									"""simple docstring"""

									return MaMaaaTokenizer.from_pretrained(self.tmpdirname							,  **lowerCAmelCase_    )


				def       __lowercase							( self							,  lowerCamelCase__    ):




									"""simple docstring"""

									return (
									    "This is a test",
									    "This is a test",
									)


				def       __lowercase							( self    ):




									"""simple docstring"""

									__UpperCamelCase      :					List[Any]       		='</s>'
									__UpperCamelCase      :					Tuple       		=0

									self.assertEqual(self.get_tokenizer()._convert_token_to_id(lowerCAmelCase_    )							,  lowerCAmelCase_    )
									self.assertEqual(self.get_tokenizer()._convert_id_to_token(lowerCAmelCase_    )							,  lowerCAmelCase_    )


				def       __lowercase							( self    ):




									"""simple docstring"""

									__UpperCamelCase      :					int       		=self.get_tokenizer()
									__UpperCamelCase      :					str       		=list(tokenizer.get_vocab().keys()    )

									self.assertEqual(vocab_keys[0]							,  '</s>'    )
									self.assertEqual(vocab_keys[1]							,  '<unk>'    )
									self.assertEqual(vocab_keys[-1]							,  '<s>'    )
									self.assertEqual(len(lowerCAmelCase_    )							,  tokenizer.vocab_size + len(tokenizer.get_added_vocab()    )    )


				@unittest.skip('Skip this test while all models are still to be uploaded.'    )
				def       __lowercase							( self    ):




									"""simple docstring"""

									pass


				def       __lowercase							( self    ):




									"""simple docstring"""

									__UpperCamelCase      :					Optional[Any]       		=self.get_tokenizer()

									__UpperCamelCase      :					str       		=tokenizer.tokenize('This is a test'    )
									self.assertListEqual(lowerCAmelCase_							,  ['▁This', '▁is', '▁a', '▁t', 'est']    )

									self.assertListEqual(
									    tokenizer.convert_tokens_to_ids(lowerCAmelCase_    )							,  [2, 3, 4, 5, 6]							,  )

									__UpperCamelCase      :					Dict       		=tokenizer.convert_ids_to_tokens([2, 3, 4, 5, 6]    )
									self.assertListEqual(lowerCAmelCase_							,  ['▁This', '▁is', '▁a', '▁t', 'est']    )

									__UpperCamelCase      :					List[Any]       		=tokenizer.convert_tokens_to_string(lowerCAmelCase_    )
									self.assertEqual(lowerCAmelCase_							,  'This is a test'    )







				@slow
				def       __lowercase							( self    ):




									"""simple docstring"""

									__UpperCamelCase      :					str       		={'input_ids': [[128022, 110108, 397, 11, 38272, 2247, 124811, 285, 18105, 1586, 207, 7, 39534, 4428, 397, 1019, 18105, 1586, 207, 7, 41337, 16786, 241, 7, 20214, 17, 125690, 10398, 7, 44378, 58069, 68342, 7798, 7343, 11, 299, 33310, 4, 158, 37350, 94077, 4569, 299, 33310, 90, 4, 52840, 290, 4, 31270, 112, 299, 682, 4, 52840, 39953, 14079, 193, 52519, 90894, 17894, 120697, 11, 40445, 551, 17, 1019, 52519, 90894, 17756, 963, 11, 40445, 480, 17, 9792, 1120, 5173, 1393, 6240, 16786, 241, 120996, 28, 1245, 1393, 118240, 11123, 1019, 93612, 2691, 10618, 98058, 120409, 1928, 279, 4, 40683, 367, 178, 207, 1019, 103, 103121, 506, 65296, 5, 2], [128022, 21217, 367, 117, 125450, 128, 719, 7, 7308, 40, 93612, 12669, 1116, 16704, 71, 17785, 3699, 15592, 35, 144, 9584, 241, 11943, 713, 950, 799, 2247, 88427, 150, 149, 118813, 120706, 1019, 106906, 81518, 28, 1224, 22799, 397, 5, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [128022, 1658, 123311, 5155, 5578, 4722, 279, 14947, 2366, 1120, 1197, 14, 1348, 9232, 5, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]], 'attention_mask': [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]}  # noqa: E501
									# fmt: on

									self.tokenizer_integration_test_util(
									    expected_encoding=lowerCAmelCase_							,  model_name='facebook/m2m100_418M'							,  revision='c168bae485c864188cf9aa0e4108b0b6934dc91e'							,  )




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





				"""simple docstring"""







				UpperCamelCase__ :      Union[str, Any]			 ='''facebook/m2m100_418M'''
				UpperCamelCase__ :      Optional[int]			 =[
				    '''In my opinion, there are two levels of response from the French government.''',
				    '''NSA Affair Emphasizes Complete Lack of Debate on Intelligence''',
				]
				UpperCamelCase__ :      str			 =[
				    '''Selon moi, il y a deux niveaux de réponse de la part du gouvernement français.''',
				    '''L\'affaire NSA souligne l\'absence totale de débat sur le renseignement''',
				]

				# fmt: off
				UpperCamelCase__ :      Tuple			 =[EN_CODE, 5_9_3, 1_9_4_9, 1_1_5_7_8_1, 4, 7_1_5_8_6, 4_2_3_4, 6_0_6_3_3, 1_2_6_2_3_3, 4_3_2, 1_2_3_8_0_8, 1_5_5_9_2, 1_1_9_7, 1_1_7_1_3_2, 1_2_0_6_1_8, 5, 2]


				@classmethod
				def       __lowercase							( cls    ):




									"""simple docstring"""

									__UpperCamelCase      :					Any       		=MaMaaaTokenizer.from_pretrained(
									    cls.checkpoint_name							,  src_lang='en'							,  tgt_lang='fr'    )
									__UpperCamelCase      :					List[Any]       		=1
									return cls


				def       __lowercase							( self    ):




									"""simple docstring"""

									self.assertEqual(self.tokenizer.get_lang_id('ar'    )							,  128006    )
									self.assertEqual(self.tokenizer.get_lang_id('en'    )							,  128022    )
									self.assertEqual(self.tokenizer.get_lang_id('ro'    )							,  128076    )
									self.assertEqual(self.tokenizer.get_lang_id('mr'    )							,  128063    )


				def       __lowercase							( self    ):




									"""simple docstring"""

									__UpperCamelCase      :					Optional[Any]       		=self.tokenizer.get_vocab()
									self.assertEqual(len(lowerCAmelCase_    )							,  self.tokenizer.vocab_size    )
									self.assertEqual(vocab['<unk>']							,  3    )
									self.assertIn(self.tokenizer.get_lang_token('en'    )							,  lowerCAmelCase_    )


				def       __lowercase							( self    ):




									"""simple docstring"""

									__UpperCamelCase      :					int       		='en'
									__UpperCamelCase      :					str       		=self.tokenizer.batch_encode_plus(self.src_text    ).input_ids[0]
									self.assertListEqual(self.expected_src_tokens							,  lowerCAmelCase_    )


				def       __lowercase							( self    ):




									"""simple docstring"""

									self.assertIn(lowerCAmelCase_							,  self.tokenizer.all_special_ids    )
									# fmt: off
									__UpperCamelCase      :					Union[str, Any]       		=[FR_CODE, 5364, 82, 8642, 4, 294, 47, 8, 14028, 136, 3286, 9706, 6, 90797, 6, 144012, 162, 88128, 30061, 5, 2]
									# fmt: on
									__UpperCamelCase      :					Optional[Any]       		=self.tokenizer.decode(lowerCAmelCase_							,  skip_special_tokens=lowerCAmelCase_    )
									__UpperCamelCase      :					Union[str, Any]       		=self.tokenizer.decode(generated_ids[1:]							,  skip_special_tokens=lowerCAmelCase_    )
									self.assertEqual(lowerCAmelCase_							,  lowerCAmelCase_    )
									self.assertNotIn(self.tokenizer.eos_token							,  lowerCAmelCase_    )


				def       __lowercase							( self    ):




									"""simple docstring"""

									__UpperCamelCase      :					Dict       		=tempfile.mkdtemp()
									__UpperCamelCase      :					str       		=self.tokenizer.lang_token_to_id
									self.tokenizer.save_pretrained(lowerCAmelCase_    )
									__UpperCamelCase      :					Optional[Any]       		=MaMaaaTokenizer.from_pretrained(lowerCAmelCase_    )
									self.assertDictEqual(new_tok.lang_token_to_id							,  lowerCAmelCase_    )


				@require_torch
				def       __lowercase							( self    ):




									"""simple docstring"""

									__UpperCamelCase      :					Optional[int]       		='en'
									__UpperCamelCase      :					Optional[int]       		='fr'

									__UpperCamelCase      :					Any       		=self.tokenizer(self.src_text							,  text_target=self.tgt_text							,  padding=lowerCAmelCase_							,  return_tensors='pt'    )

									__UpperCamelCase      :					int       		=shift_tokens_right(
									    batch['labels']							,  self.tokenizer.pad_token_id							,  self.tokenizer.eos_token_id    )

									for k in batch:
														__UpperCamelCase      :					List[Any]       		=batch[k].tolist()
									# batch = {k: v.tolist() for k,v in batch.items()}
									# fairseq batch: https://gist.github.com/sshleifer/cba08bc2109361a74ac3760a7e30e4f4
									# batch.decoder_inputs_ids[0][0] ==
									assert batch.input_ids[1][0] == EN_CODE
									assert batch.input_ids[1][-1] == 2
									assert batch.labels[1][0] == FR_CODE
									assert batch.labels[1][-1] == 2
									assert batch.decoder_input_ids[1][:2] == [2, FR_CODE]


				@require_torch
				def       __lowercase							( self    ):




									"""simple docstring"""

									__UpperCamelCase      :					Optional[Any]       		='mr'
									self.assertListEqual(self.tokenizer.prefix_tokens							,  [self.tokenizer.get_lang_id('mr'    )]    )
									self.assertListEqual(self.tokenizer.suffix_tokens							,  [self.tokenizer.eos_token_id]    )

									__UpperCamelCase      :					List[str]       		='zh'
									self.assertListEqual(self.tokenizer.prefix_tokens							,  [self.tokenizer.get_lang_id('zh'    )]    )
									self.assertListEqual(self.tokenizer.suffix_tokens							,  [self.tokenizer.eos_token_id]    )


				@require_torch
				def       __lowercase							( self    ):




									"""simple docstring"""

									__UpperCamelCase      :					Optional[int]       		='mr'
									self.tokenizer._switch_to_target_mode()
									self.assertListEqual(self.tokenizer.prefix_tokens							,  [self.tokenizer.get_lang_id('mr'    )]    )
									self.assertListEqual(self.tokenizer.suffix_tokens							,  [self.tokenizer.eos_token_id]    )
									self.tokenizer._switch_to_input_mode()
									self.assertListEqual(self.tokenizer.prefix_tokens							,  [self.tokenizer.get_lang_id(self.tokenizer.src_lang    )]    )

									__UpperCamelCase      :					int       		='zh'
									self.tokenizer._switch_to_target_mode()
									self.assertListEqual(self.tokenizer.prefix_tokens							,  [self.tokenizer.get_lang_id('zh'    )]    )
									self.assertListEqual(self.tokenizer.suffix_tokens							,  [self.tokenizer.eos_token_id]    )
									self.tokenizer._switch_to_input_mode()
									self.assertListEqual(self.tokenizer.prefix_tokens							,  [self.tokenizer.get_lang_id(self.tokenizer.src_lang    )]    )







				@require_torch
				def       __lowercase							( self    ):




									"""simple docstring"""

									__UpperCamelCase      :					Any       		=self.tokenizer._build_translation_inputs('A test'							,  return_tensors='pt'							,  src_lang='en'							,  tgt_lang='ar'    )

									self.assertEqual(
									    nested_simplify(lowerCAmelCase_    )							,  {
									        # en_XX, A, test, EOS
									        'input_ids': [[128022, 58, 4183, 2]],
									        'attention_mask': [[1, 1, 1, 1]],
									        # ar_AR
									        'forced_bos_token_id': 128006,
									    }							,  )


import torch
from torch import nn



class 				a							(			nn.Module       ):







 """simple docstring"""






 def __init__(					self		,				lowerCAmelCase_		,				lowerCAmelCase_		,				lowerCAmelCase_		,				lowerCAmelCase_		,				lowerCAmelCase_=1		,				lowerCAmelCase_=False							)  ->					Any:
        super().__init__()

        _A        =  n_token
        _A        =  d_embed
        _A        =  d_proj

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

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

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

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

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

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

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

        _A        =  keep_order



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

        return logit



 def      UpperCAmelCase							(					self		,				lowerCAmelCase_		,				lowerCAmelCase_=None		,				lowerCAmelCase_=False							)  ->					List[Any]:

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

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

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

                      weights.append(lowerCAmelCase_							)
                      biases.append(lowerCAmelCase_							)

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

               _A        =  self._compute_logit(lowerCAmelCase_		,				lowerCAmelCase_		,				lowerCAmelCase_		,				lowerCAmelCase_							)
               _A        =  nn.functional.log_softmax(lowerCAmelCase_		,				dim=1							)

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

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

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

                             if indices_i.numel() == 0:
                                    continue

                             _A        =  labels.index_select(0		,				lowerCAmelCase_							) - l_idx
                             _A        =  head_logprob.index_select(0		,				lowerCAmelCase_							)
                             _A        =  hidden.index_select(0		,				lowerCAmelCase_							)
                      else:
                             _A        =  hidden

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

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

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

        return out



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

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

                      weights.append(lowerCAmelCase_							)
                      biases.append(lowerCAmelCase_							)

               _A   ,  _A   ,  _A        =  weights[0], biases[0], self.out_projs[0]
               _A        =  self._compute_logit(lowerCAmelCase_		,				lowerCAmelCase_		,				lowerCAmelCase_		,				lowerCAmelCase_							)

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

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

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

                             _A        =  self._compute_logit(lowerCAmelCase_		,				lowerCAmelCase_		,				lowerCAmelCase_		,				lowerCAmelCase_							)
                             _A        =  nn.functional.log_softmax(lowerCAmelCase_		,				dim=1							)

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

               return out




