

import unittest

from transformers import BigBirdTokenizer, BigBirdTokenizerFast
from transformers.testing_utils import get_tests_dir, require_sentencepiece, require_tokenizers, require_torch, slow
from transformers.utils import cached_property

from ...test_tokenization_common import TokenizerTesterMixin


_lowerCamelCase    :       List[str]        =							"▁"

_lowerCamelCase    :       Any        =							get_tests_dir("fixtures/test_sentencepiece.model")





@require_sentencepiece
@require_tokenizers
class  __UpperCAmelCase					(	lowerCamelCase__      ,  unittest.TestCase						):
   UpperCamelCase       					=					BigBirdTokenizer
   UpperCamelCase       					=					BigBirdTokenizerFast
   UpperCamelCase       					=					True
   UpperCamelCase       					=					True



   def 	__magic_name__   (							self							:					int						):
     super().setUp()

     UpperCAmelCase	:							str						=			self.tokenizer_class(__A,		keep_accents=__A						)
     tokenizer.save_pretrained(self.tmpdirname						)



   def 	__magic_name__   (							self							:					Optional[int]						):
     UpperCAmelCase	:							Tuple						=			'''<s>'''
     UpperCAmelCase	:							int						=			1

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



   def 	__magic_name__   (							self							:					Union[str, Any]						):
     UpperCAmelCase	:							int						=			list(self.get_tokenizer().get_vocab().keys()						)

     self.assertEqual(vocab_keys[0],		'''<unk>'''						)
     self.assertEqual(vocab_keys[1],		'''<s>'''						)
     self.assertEqual(vocab_keys[-1],		'''[MASK]'''						)
     self.assertEqual(len(__A						),		1_0_0_4						)



   def 	__magic_name__   (							self							:					Dict						):
     self.assertEqual(self.get_tokenizer().vocab_size,		1_0_0_0						)



   def 	__magic_name__   (							self							:					Optional[int]						):
     if not self.test_rust_tokenizer:
       return

     UpperCAmelCase	:							Tuple						=			self.get_tokenizer()
     UpperCAmelCase	:							Any						=			self.get_rust_tokenizer()

     UpperCAmelCase	:							str						=			'''I was born in 92000, and this is falsé.'''

     UpperCAmelCase	:							Union[str, Any]						=			tokenizer.tokenize(__A						)
     UpperCAmelCase	:							Union[str, Any]						=			rust_tokenizer.tokenize(__A						)
     self.assertListEqual(__A,		__A						)

     UpperCAmelCase	:							Dict						=			tokenizer.encode(__A,		add_special_tokens=__A						)
     UpperCAmelCase	:							str						=			rust_tokenizer.encode(__A,		add_special_tokens=__A						)
     self.assertListEqual(__A,		__A						)

     UpperCAmelCase	:							List[Any]						=			self.get_rust_tokenizer()
     UpperCAmelCase	:							Optional[int]						=			tokenizer.encode(__A						)
     UpperCAmelCase	:							str						=			rust_tokenizer.encode(__A						)
     self.assertListEqual(__A,		__A						)



   def 	__magic_name__   (							self							:					List[Any]						):
     UpperCAmelCase	:							Dict						=			BigBirdTokenizer(__A,		keep_accents=__A						)

     UpperCAmelCase	:							Optional[Any]						=			tokenizer.tokenize('''This is a test'''						)
     self.assertListEqual(__A,		['''▁This''', '''▁is''', '''▁a''', '''▁t''', '''est''']						)

     self.assertListEqual(
         tokenizer.convert_tokens_to_ids(__A						),		[2_8_5, 4_6, 1_0, 1_7_0, 3_8_2],		)

     UpperCAmelCase	:							Any						=			tokenizer.tokenize('''I was born in 92000, and this is falsé.'''						)
     self.assertListEqual(
         __A,		[
             SPIECE_UNDERLINE + '''I''',
             SPIECE_UNDERLINE + '''was''',
             SPIECE_UNDERLINE + '''b''',
             '''or''',
             '''n''',
             SPIECE_UNDERLINE + '''in''',
             SPIECE_UNDERLINE + '''''',
             '''9''',
             '''2''',
             '''0''',
             '''0''',
             '''0''',
             ''',''',
             SPIECE_UNDERLINE + '''and''',
             SPIECE_UNDERLINE + '''this''',
             SPIECE_UNDERLINE + '''is''',
             SPIECE_UNDERLINE + '''f''',
             '''al''',
             '''s''',
             '''é''',
             '''.''',
         ],		)
     UpperCAmelCase	:							List[str]						=			tokenizer.convert_tokens_to_ids(__A						)
     self.assertListEqual(
         __A,		[8, 2_1, 8_4, 5_5, 2_4, 1_9, 7, 0, 6_0_2, 3_4_7, 3_4_7, 3_4_7, 3, 1_2, 6_6, 4_6, 7_2, 8_0, 6, 0, 4],		)

     UpperCAmelCase	:							Union[str, Any]						=			tokenizer.convert_ids_to_tokens(__A						)
     self.assertListEqual(
         __A,		[
             SPIECE_UNDERLINE + '''I''',
             SPIECE_UNDERLINE + '''was''',
             SPIECE_UNDERLINE + '''b''',
             '''or''',
             '''n''',
             SPIECE_UNDERLINE + '''in''',
             SPIECE_UNDERLINE + '''''',
             '''<unk>''',
             '''2''',
             '''0''',
             '''0''',
             '''0''',
             ''',''',
             SPIECE_UNDERLINE + '''and''',
             SPIECE_UNDERLINE + '''this''',
             SPIECE_UNDERLINE + '''is''',
             SPIECE_UNDERLINE + '''f''',
             '''al''',
             '''s''',
             '''<unk>''',
             '''.''',
         ],		)



   @cached_property
   def 	__magic_name__   (							self							:					List[str]						):
     return BigBirdTokenizer.from_pretrained('''google/bigbird-roberta-base'''						)



   @slow
   def 	__magic_name__   (							self							:					int						):
     UpperCAmelCase	:							Dict						=			'''Hello World!'''
     UpperCAmelCase	:							Dict						=			[6_5, 1_8_5_3_6, 2_2_6_0, 1_0_1, 6_6]

     self.assertListEqual(__A,		self.big_tokenizer.encode(__A						)						)



   @slow
   def 	__magic_name__   (							self							:					List[Any]						):
     UpperCAmelCase	:							List[str]						=			(
         '''This is a very long text with a lot of weird characters, such as: . , ~ ? ( ) " [ ] ! : - . Also we will'''
         ''' add words that should not exsist and be tokenized to <unk>, such as saoneuhaoesuth'''
     )
     # fmt: off
     UpperCAmelCase	:							Tuple						=			[6_5, 8_7_1, 4_1_9, 3_5_8, 9_4_6, 9_9_1, 2_5_2_1, 4_5_2, 3_5_8, 1_3_5_7, 3_8_7, 7_7_5_1, 3_5_3_6, 1_1_2, 9_8_5, 4_5_6, 1_2_6, 8_6_5, 9_3_8, 5_4_0_0, 5_7_3_4, 4_5_8, 1_3_6_8, 4_6_7, 7_8_6, 2_4_6_2, 5_2_4_6, 1_1_5_9, 6_3_3, 8_6_5, 4_5_1_9, 4_5_7, 5_8_2, 8_5_2, 2_5_5_7, 4_2_7, 9_1_6, 5_0_8, 4_0_5, 3_4_3_2_4, 4_9_7, 3_9_1, 4_0_8, 1_1_3_4_2, 1_2_4_4, 3_8_5, 1_0_0, 9_3_8, 9_8_5, 4_5_6, 5_7_4, 3_6_2, 1_2_5_9_7, 3_2_0_0, 3_1_2_9, 1_1_7_2, 6_6]  # noqa: E231
     # fmt: on
     self.assertListEqual(__A,		self.big_tokenizer.encode(__A						)						)



   @require_torch
   @slow
   def 	__magic_name__   (							self							:					List[str]						):
     import torch

     from transformers import BigBirdConfig, BigBirdModel

     # Build sequence
     UpperCAmelCase	:							int						=			list(self.big_tokenizer.get_vocab().keys()						)[:1_0]
     UpperCAmelCase	:							str						=			''' '''.join(__A						)
     UpperCAmelCase	:							Optional[int]						=			self.big_tokenizer.encode_plus(__A,		return_tensors='''pt''',		return_token_type_ids=__A						)
     UpperCAmelCase	:							Tuple						=			self.big_tokenizer.batch_encode_plus(
         [sequence + ''' ''' + sequence],		return_tensors='''pt''',		return_token_type_ids=__A						)

     UpperCAmelCase	:							Any						=			BigBirdConfig(attention_type='''original_full'''						)
     UpperCAmelCase	:							str						=			BigBirdModel(__A						)

     assert model.get_input_embeddings().weight.shape[0] >= self.big_tokenizer.vocab_size

     with torch.no_grad():
       model(**__A						)
       model(**__A						)



   @slow
   def 	__magic_name__   (							self							:					Tuple						):
     UpperCAmelCase	:							Tuple						=			BigBirdTokenizer.from_pretrained('''google/bigbird-roberta-base'''						)
     UpperCAmelCase	:							Dict						=			tokenizer.decode(tokenizer('''Paris is the [MASK].'''						).input_ids						)

     self.assertTrue(decoded_text == '''[CLS] Paris is the[MASK].[SEP]'''						)





   @slow
   def 	__magic_name__   (							self							:					Dict						):
     # fmt: off
     UpperCAmelCase	:							List[Any]						=			{'''input_ids''': [[6_5, 3_9_2_8_6, 4_5_8, 3_6_3_3_5, 2_0_0_1, 4_5_6, 1_3_0_7_3, 1_3_2_6_6, 4_5_5, 1_1_3, 7_7_4_6, 1_7_4_1, 1_1_1_5_7, 3_9_1, 1_3_0_7_3, 1_3_2_6_6, 4_5_5, 1_1_3, 3_9_6_7, 3_5_4_1_2, 1_1_3, 4_9_3_6, 1_0_9, 3_8_7_0, 2_3_7_7, 1_1_3, 3_0_0_8_4, 4_5_7_2_0, 4_5_8, 1_3_4, 1_7_4_9_6, 1_1_2, 5_0_3, 1_1_6_7_2, 1_1_3, 1_1_8, 1_1_2, 5_6_6_5, 1_3_3_4_7, 3_8_6_8_7, 1_1_2, 1_4_9_6, 3_1_3_8_9, 1_1_2, 3_2_6_8, 4_7_2_6_4, 1_3_4, 9_6_2, 1_1_2, 1_6_3_7_7, 8_0_3_5, 2_3_1_3_0, 4_3_0, 1_2_1_6_9, 1_5_5_1_8, 2_8_5_9_2, 4_5_8, 1_4_6, 4_1_6_9_7, 1_0_9, 3_9_1, 1_2_1_6_9, 1_5_5_1_8, 1_6_6_8_9, 4_5_8, 1_4_6, 4_1_3_5_8, 1_0_9, 4_5_2, 7_2_6, 4_0_3_4, 1_1_1, 7_6_3, 3_5_4_1_2, 5_0_8_2, 3_8_8, 1_9_0_3, 1_1_1, 9_0_5_1, 3_9_1, 2_8_7_0, 4_8_9_1_8, 1_9_0_0, 1_1_2_3, 5_5_0, 9_9_8, 1_1_2, 9_5_8_6, 1_5_9_8_5, 4_5_5, 3_9_1, 4_1_0, 2_2_9_5_5, 3_7_6_3_6, 1_1_4, 6_6], [6_5, 4_4_8, 1_7_4_9_6, 4_1_9, 3_6_6_3, 3_8_5, 7_6_3, 1_1_3, 2_7_5_3_3, 2_8_7_0, 3_2_8_3, 1_3_0_4_3, 1_6_3_9, 2_4_7_1_3, 5_2_3, 6_5_6, 2_4_0_1_3, 1_8_5_5_0, 2_5_2_1, 5_1_7, 2_7_0_1_4, 2_1_2_4_4, 4_2_0, 1_2_1_2, 1_4_6_5, 3_9_1, 9_2_7, 4_8_3_3, 3_8_8, 5_7_8, 1_1_7_8_6, 1_1_4, 6_6, 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], [6_5, 4_8_4, 2_1_6_9, 7_6_8_7, 2_1_9_3_2, 1_8_1_4_6, 7_2_6, 3_6_3, 1_7_0_3_2, 3_3_9_1, 1_1_4, 6_6, 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]], '''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, 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, 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=__A,		model_name='''google/bigbird-roberta-base''',		revision='''215c99f1600e06f83acce68422f2035b2b5c3510''',		)








import torch
from torch import nn





class  __UpperCAmelCase					(	nn.Module						):



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

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

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

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

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

     UpperCAmelCase	:							Dict						=			nn.ModuleList()
     UpperCAmelCase	:							Optional[int]						=			nn.ParameterList()

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

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

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

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

     UpperCAmelCase	:							Optional[int]						=			keep_order



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

     return logit



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

     if labels is not None:
       # Shift so that tokens < n predict n
       UpperCAmelCase	:							List[Any]						=			hidden[..., :-1, :].contiguous()
       UpperCAmelCase	:							Any						=			labels[..., 1:].contiguous()
       UpperCAmelCase	:							Optional[Any]						=			hidden.view(-1,		hidden.size(-1						)						)
       UpperCAmelCase	:							Union[str, Any]						=			labels.view(-1						)
       if hidden.size(0						) != labels.size(0						):
         raise RuntimeError('''Input and labels should have the same size in the batch dimension.'''						)
     else:
       UpperCAmelCase	:							str						=			hidden.view(-1,		hidden.size(-1						)						)

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

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

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

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

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

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

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

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

           if indices_i.numel() == 0:
             continue

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

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

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

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

     return out





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

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

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

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

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

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

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

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

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

       return out





