

'''simple docstring'''






import gc
import unittest

from transformers import MODEL_FOR_MASKED_LM_MAPPING, TF_MODEL_FOR_MASKED_LM_MAPPING, FillMaskPipeline, pipeline
from transformers.pipelines import PipelineException
from transformers.testing_utils import (
    is_pipeline_test,
    is_torch_available,
    nested_simplify,
    require_tf,
    require_torch,
    require_torch_gpu,
    slow,
)

from .test_pipelines_common import ANY






@is_pipeline_test
class 			lowercase_	(unittest.TestCase      ):




  """simple docstring"""




  SCREAMING_SNAKE_CASE    :  Optional[int]							=			MODEL_FOR_MASKED_LM_MAPPING
  SCREAMING_SNAKE_CASE    :  int							=			TF_MODEL_FOR_MASKED_LM_MAPPING






  def 					SCREAMING_SNAKE_CASE      (      self   :     Any			):
   super().tearDown()
   # clean-up as much as possible GPU memory occupied by PyTorch
   gc.collect()
   if is_torch_available():
    import torch

    torch.cuda.empty_cache()






  @require_tf
  def 					SCREAMING_SNAKE_CASE      (      self   :     List[str]			):
   __lowercase							  = pipeline(task='''fill-mask'''					,model='''sshleifer/tiny-distilroberta-base'''					,top_k=2					,framework='''tf'''			)
   __lowercase							  = unmasker('''My name is <mask>'''			)
   self.assertEqual(
       nested_simplify(lowercase__					,decimals=6			)					,[
           {'''sequence''': '''My name is grouped''', '''score''': 2.1e-0_5, '''token''': 3_8_0_1_5, '''token_str''': ''' grouped'''},
           {'''sequence''': '''My name is accuser''', '''score''': 2.1e-0_5, '''token''': 2_5_5_0_6, '''token_str''': ''' accuser'''},
       ]					,)

   __lowercase							  = unmasker('''The largest city in France is <mask>'''			)
   self.assertEqual(
       nested_simplify(lowercase__					,decimals=6			)					,[
           {
               '''sequence''': '''The largest city in France is grouped''',
               '''score''': 2.1e-0_5,
               '''token''': 3_8_0_1_5,
               '''token_str''': ''' grouped''',
           },
           {
               '''sequence''': '''The largest city in France is accuser''',
               '''score''': 2.1e-0_5,
               '''token''': 2_5_5_0_6,
               '''token_str''': ''' accuser''',
           },
       ]					,)

   __lowercase							  = unmasker('''My name is <mask>'''					,targets=[''' Patrick''', ''' Clara''', ''' Teven''']					,top_k=3			)
   self.assertEqual(
       nested_simplify(lowercase__					,decimals=6			)					,[
           {'''sequence''': '''My name is Clara''', '''score''': 2e-0_5, '''token''': 1_3_6_0_6, '''token_str''': ''' Clara'''},
           {'''sequence''': '''My name is Patrick''', '''score''': 2e-0_5, '''token''': 3_4_9_9, '''token_str''': ''' Patrick'''},
           {'''sequence''': '''My name is Te''', '''score''': 1.9e-0_5, '''token''': 2_9_4_1, '''token_str''': ''' Te'''},
       ]					,)






  @require_torch
  def 					SCREAMING_SNAKE_CASE      (      self   :     Tuple			):
   __lowercase							  = pipeline(task='''fill-mask'''					,model='''sshleifer/tiny-distilroberta-base'''					,top_k=2					,framework='''pt'''			)

   __lowercase							  = unmasker('''My name is <mask>'''			)
   self.assertEqual(
       nested_simplify(lowercase__					,decimals=6			)					,[
           {'''sequence''': '''My name is Maul''', '''score''': 2.2e-0_5, '''token''': 3_5_6_7_6, '''token_str''': ''' Maul'''},
           {'''sequence''': '''My name isELS''', '''score''': 2.2e-0_5, '''token''': 1_6_4_1_6, '''token_str''': '''ELS'''},
       ]					,)

   __lowercase							  = unmasker('''The largest city in France is <mask>'''			)
   self.assertEqual(
       nested_simplify(lowercase__					,decimals=6			)					,[
           {
               '''sequence''': '''The largest city in France is Maul''',
               '''score''': 2.2e-0_5,
               '''token''': 3_5_6_7_6,
               '''token_str''': ''' Maul''',
           },
           {'''sequence''': '''The largest city in France isELS''', '''score''': 2.2e-0_5, '''token''': 1_6_4_1_6, '''token_str''': '''ELS'''},
       ]					,)

   __lowercase							  = unmasker('''My name is <mask>'''					,targets=[''' Patrick''', ''' Clara''', ''' Teven''']					,top_k=3			)
   self.assertEqual(
       nested_simplify(lowercase__					,decimals=6			)					,[
           {'''sequence''': '''My name is Patrick''', '''score''': 2.1e-0_5, '''token''': 3_4_9_9, '''token_str''': ''' Patrick'''},
           {'''sequence''': '''My name is Te''', '''score''': 2e-0_5, '''token''': 2_9_4_1, '''token_str''': ''' Te'''},
           {'''sequence''': '''My name is Clara''', '''score''': 2e-0_5, '''token''': 1_3_6_0_6, '''token_str''': ''' Clara'''},
       ]					,)

   __lowercase							  = unmasker('''My name is <mask> <mask>'''					,top_k=2			)

   self.assertEqual(
       nested_simplify(lowercase__					,decimals=6			)					,[
           [
               {
                   '''score''': 2.2e-0_5,
                   '''token''': 3_5_6_7_6,
                   '''token_str''': ''' Maul''',
                   '''sequence''': '''<s>My name is Maul<mask></s>''',
               },
               {'''score''': 2.2e-0_5, '''token''': 1_6_4_1_6, '''token_str''': '''ELS''', '''sequence''': '''<s>My name isELS<mask></s>'''},
           ],
           [
               {
                   '''score''': 2.2e-0_5,
                   '''token''': 3_5_6_7_6,
                   '''token_str''': ''' Maul''',
                   '''sequence''': '''<s>My name is<mask> Maul</s>''',
               },
               {'''score''': 2.2e-0_5, '''token''': 1_6_4_1_6, '''token_str''': '''ELS''', '''sequence''': '''<s>My name is<mask>ELS</s>'''},
           ],
       ]					,)






  @require_torch_gpu
  def 					SCREAMING_SNAKE_CASE      (      self   :     List[str]			):
   __lowercase							  = pipeline('''fill-mask'''					,model='''hf-internal-testing/tiny-random-distilbert'''					,device=0					,framework='''pt'''			)

   # convert model to fp16
   pipe.model.half()

   __lowercase							  = pipe('''Paris is the [MASK] of France.'''			)
   # We actually don't care about the result, we just want to make sure
   # it works, meaning the float16 tensor got casted back to float32
   # for postprocessing.
   self.assertIsInstance(lowercase__					,lowercase__			)






  @slow
  @require_torch
  def 					SCREAMING_SNAKE_CASE      (      self   :     str			):
   __lowercase							  = pipeline(task='''fill-mask'''					,model='''distilroberta-base'''					,top_k=2					,framework='''pt'''			)
   self.run_large_test(lowercase__			)






  @slow
  @require_tf
  def 					SCREAMING_SNAKE_CASE      (      self   :     List[str]			):
   __lowercase							  = pipeline(task='''fill-mask'''					,model='''distilroberta-base'''					,top_k=2					,framework='''tf'''			)
   self.run_large_test(lowercase__			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     List[str]					,lowercase__   :     str			):
   __lowercase							  = unmasker('''My name is <mask>'''			)
   self.assertEqual(
       nested_simplify(lowercase__			)					,[
           {'''sequence''': '''My name is John''', '''score''': 0.0_0_8, '''token''': 6_1_0, '''token_str''': ''' John'''},
           {'''sequence''': '''My name is Chris''', '''score''': 0.0_0_7, '''token''': 1_5_7_3, '''token_str''': ''' Chris'''},
       ]					,)
   __lowercase							  = unmasker('''The largest city in France is <mask>'''			)
   self.assertEqual(
       nested_simplify(lowercase__			)					,[
           {
               '''sequence''': '''The largest city in France is Paris''',
               '''score''': 0.2_5_1,
               '''token''': 2_2_0_1,
               '''token_str''': ''' Paris''',
           },
           {
               '''sequence''': '''The largest city in France is Lyon''',
               '''score''': 0.2_1_4,
               '''token''': 1_2_7_9_0,
               '''token_str''': ''' Lyon''',
           },
       ]					,)

   __lowercase							  = unmasker('''My name is <mask>'''					,targets=[''' Patrick''', ''' Clara''', ''' Teven''']					,top_k=3			)
   self.assertEqual(
       nested_simplify(lowercase__			)					,[
           {'''sequence''': '''My name is Patrick''', '''score''': 0.0_0_5, '''token''': 3_4_9_9, '''token_str''': ''' Patrick'''},
           {'''sequence''': '''My name is Clara''', '''score''': 0.0_0_0, '''token''': 1_3_6_0_6, '''token_str''': ''' Clara'''},
           {'''sequence''': '''My name is Te''', '''score''': 0.0_0_0, '''token''': 2_9_4_1, '''token_str''': ''' Te'''},
       ]					,)






  @require_torch
  def 					SCREAMING_SNAKE_CASE      (      self   :     Tuple			):
   __lowercase							  = pipeline(task='''fill-mask'''					,model='''sshleifer/tiny-distilroberta-base'''					,framework='''pt'''			)
   __lowercase							  = None
   __lowercase							  = None
   self.run_pipeline_test(lowercase__					,[]			)






  @require_tf
  def 					SCREAMING_SNAKE_CASE      (      self   :     int			):
   __lowercase							  = pipeline(task='''fill-mask'''					,model='''sshleifer/tiny-distilroberta-base'''					,framework='''tf'''			)
   __lowercase							  = None
   __lowercase							  = None
   self.run_pipeline_test(lowercase__					,[]			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     str					,lowercase__   :     int					,lowercase__   :     str					,lowercase__   :     Dict			):
   if tokenizer is None or tokenizer.mask_token_id is None:
    self.skipTest('''The provided tokenizer has no mask token, (probably reformer or wav2vec2)'''			)

   __lowercase							  = FillMaskPipeline(model=lowercase__					,tokenizer=lowercase__			)
   __lowercase							  = [
       F"This is another {tokenizer.mask_token} test",
   ]
   return fill_masker, examples






  def 					SCREAMING_SNAKE_CASE      (      self   :     Union[str, Any]					,lowercase__   :     str					,lowercase__   :     str			):
   __lowercase							  = fill_masker.tokenizer
   __lowercase							  = fill_masker.model

   __lowercase							  = fill_masker(
       F"This is a {tokenizer.mask_token}"					,)
   self.assertEqual(
       lowercase__					,[
           {'''sequence''': ANY(lowercase__			), '''score''': ANY(lowercase__			), '''token''': ANY(lowercase__			), '''token_str''': ANY(lowercase__			)},
           {'''sequence''': ANY(lowercase__			), '''score''': ANY(lowercase__			), '''token''': ANY(lowercase__			), '''token_str''': ANY(lowercase__			)},
           {'''sequence''': ANY(lowercase__			), '''score''': ANY(lowercase__			), '''token''': ANY(lowercase__			), '''token_str''': ANY(lowercase__			)},
           {'''sequence''': ANY(lowercase__			), '''score''': ANY(lowercase__			), '''token''': ANY(lowercase__			), '''token_str''': ANY(lowercase__			)},
           {'''sequence''': ANY(lowercase__			), '''score''': ANY(lowercase__			), '''token''': ANY(lowercase__			), '''token_str''': ANY(lowercase__			)},
       ]					,)

   __lowercase							  = fill_masker([F"This is a {tokenizer.mask_token}"]			)
   self.assertEqual(
       lowercase__					,[
           {'''sequence''': ANY(lowercase__			), '''score''': ANY(lowercase__			), '''token''': ANY(lowercase__			), '''token_str''': ANY(lowercase__			)},
           {'''sequence''': ANY(lowercase__			), '''score''': ANY(lowercase__			), '''token''': ANY(lowercase__			), '''token_str''': ANY(lowercase__			)},
           {'''sequence''': ANY(lowercase__			), '''score''': ANY(lowercase__			), '''token''': ANY(lowercase__			), '''token_str''': ANY(lowercase__			)},
           {'''sequence''': ANY(lowercase__			), '''score''': ANY(lowercase__			), '''token''': ANY(lowercase__			), '''token_str''': ANY(lowercase__			)},
           {'''sequence''': ANY(lowercase__			), '''score''': ANY(lowercase__			), '''token''': ANY(lowercase__			), '''token_str''': ANY(lowercase__			)},
       ]					,)

   __lowercase							  = fill_masker([F"This is a {tokenizer.mask_token}", F"Another {tokenizer.mask_token} great test."]			)
   self.assertEqual(
       lowercase__					,[
           [
               {'''sequence''': ANY(lowercase__			), '''score''': ANY(lowercase__			), '''token''': ANY(lowercase__			), '''token_str''': ANY(lowercase__			)},
               {'''sequence''': ANY(lowercase__			), '''score''': ANY(lowercase__			), '''token''': ANY(lowercase__			), '''token_str''': ANY(lowercase__			)},
               {'''sequence''': ANY(lowercase__			), '''score''': ANY(lowercase__			), '''token''': ANY(lowercase__			), '''token_str''': ANY(lowercase__			)},
               {'''sequence''': ANY(lowercase__			), '''score''': ANY(lowercase__			), '''token''': ANY(lowercase__			), '''token_str''': ANY(lowercase__			)},
               {'''sequence''': ANY(lowercase__			), '''score''': ANY(lowercase__			), '''token''': ANY(lowercase__			), '''token_str''': ANY(lowercase__			)},
           ],
           [
               {'''sequence''': ANY(lowercase__			), '''score''': ANY(lowercase__			), '''token''': ANY(lowercase__			), '''token_str''': ANY(lowercase__			)},
               {'''sequence''': ANY(lowercase__			), '''score''': ANY(lowercase__			), '''token''': ANY(lowercase__			), '''token_str''': ANY(lowercase__			)},
               {'''sequence''': ANY(lowercase__			), '''score''': ANY(lowercase__			), '''token''': ANY(lowercase__			), '''token_str''': ANY(lowercase__			)},
               {'''sequence''': ANY(lowercase__			), '''score''': ANY(lowercase__			), '''token''': ANY(lowercase__			), '''token_str''': ANY(lowercase__			)},
               {'''sequence''': ANY(lowercase__			), '''score''': ANY(lowercase__			), '''token''': ANY(lowercase__			), '''token_str''': ANY(lowercase__			)},
           ],
       ]					,)

   with self.assertRaises(lowercase__			):
    fill_masker([None]			)
   # No mask_token is not supported
   with self.assertRaises(lowercase__			):
    fill_masker('''This is'''			)

   self.run_test_top_k(lowercase__					,lowercase__			)
   self.run_test_targets(lowercase__					,lowercase__			)
   self.run_test_top_k_targets(lowercase__					,lowercase__			)
   self.fill_mask_with_duplicate_targets_and_top_k(lowercase__					,lowercase__			)
   self.fill_mask_with_multiple_masks(lowercase__					,lowercase__			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     Optional[int]					,lowercase__   :     Any					,lowercase__   :     str			):
   __lowercase							  = tokenizer.get_vocab()
   __lowercase							  = sorted(vocab.keys()			)[:2]
   # Pipeline argument
   __lowercase							  = FillMaskPipeline(model=lowercase__					,tokenizer=lowercase__					,targets=lowercase__			)
   __lowercase							  = fill_masker(F"This is a {tokenizer.mask_token}"			)
   self.assertEqual(
       lowercase__					,[
           {'''sequence''': ANY(lowercase__			), '''score''': ANY(lowercase__			), '''token''': ANY(lowercase__			), '''token_str''': ANY(lowercase__			)},
           {'''sequence''': ANY(lowercase__			), '''score''': ANY(lowercase__			), '''token''': ANY(lowercase__			), '''token_str''': ANY(lowercase__			)},
       ]					,)
   __lowercase							  = {vocab[el] for el in targets}
   self.assertEqual({el['''token'''] for el in outputs}					,lowercase__			)
   __lowercase							  = [tokenizer.decode([x]			) for x in target_ids]
   self.assertEqual({el['''token_str'''] for el in outputs}					,set(lowercase__			)			)

   # Call argument
   __lowercase							  = FillMaskPipeline(model=lowercase__					,tokenizer=lowercase__			)
   __lowercase							  = fill_masker(F"This is a {tokenizer.mask_token}"					,targets=lowercase__			)
   self.assertEqual(
       lowercase__					,[
           {'''sequence''': ANY(lowercase__			), '''score''': ANY(lowercase__			), '''token''': ANY(lowercase__			), '''token_str''': ANY(lowercase__			)},
           {'''sequence''': ANY(lowercase__			), '''score''': ANY(lowercase__			), '''token''': ANY(lowercase__			), '''token_str''': ANY(lowercase__			)},
       ]					,)
   __lowercase							  = {vocab[el] for el in targets}
   self.assertEqual({el['''token'''] for el in outputs}					,lowercase__			)
   __lowercase							  = [tokenizer.decode([x]			) for x in target_ids]
   self.assertEqual({el['''token_str'''] for el in outputs}					,set(lowercase__			)			)

   # Score equivalence
   __lowercase							  = fill_masker(F"This is a {tokenizer.mask_token}"					,targets=lowercase__			)
   __lowercase							  = [top_mask['''token_str'''] for top_mask in outputs]
   __lowercase							  = [top_mask['''score'''] for top_mask in outputs]

   # For some BPE tokenizers, `</w>` is removed during decoding, so `token_str` won't be the same as in `targets`.
   if set(lowercase__			) == set(lowercase__			):
    __lowercase							  = fill_masker(F"This is a {tokenizer.mask_token}"					,targets=lowercase__			)
    __lowercase							  = [top_mask['''score'''] for top_mask in unmasked_targets]
    self.assertEqual(nested_simplify(lowercase__			)					,nested_simplify(lowercase__			)			)

   # Raises with invalid
   with self.assertRaises(lowercase__			):
    __lowercase							  = fill_masker(F"This is a {tokenizer.mask_token}"					,targets=[]			)
   # For some tokenizers, `""` is actually in the vocabulary and the expected error won't raised
   if "" not in tokenizer.get_vocab():
    with self.assertRaises(lowercase__			):
     __lowercase							  = fill_masker(F"This is a {tokenizer.mask_token}"					,targets=['''''']			)
    with self.assertRaises(lowercase__			):
     __lowercase							  = fill_masker(F"This is a {tokenizer.mask_token}"					,targets=''''''			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     Union[str, Any]					,lowercase__   :     str					,lowercase__   :     List[Any]			):
   __lowercase							  = FillMaskPipeline(model=lowercase__					,tokenizer=lowercase__					,top_k=2			)
   __lowercase							  = fill_masker(F"This is a {tokenizer.mask_token}"			)
   self.assertEqual(
       lowercase__					,[
           {'''sequence''': ANY(lowercase__			), '''score''': ANY(lowercase__			), '''token''': ANY(lowercase__			), '''token_str''': ANY(lowercase__			)},
           {'''sequence''': ANY(lowercase__			), '''score''': ANY(lowercase__			), '''token''': ANY(lowercase__			), '''token_str''': ANY(lowercase__			)},
       ]					,)

   __lowercase							  = FillMaskPipeline(model=lowercase__					,tokenizer=lowercase__			)
   __lowercase							  = fill_masker(F"This is a {tokenizer.mask_token}"					,top_k=2			)
   self.assertEqual(
       lowercase__					,[
           {'''sequence''': ANY(lowercase__			), '''score''': ANY(lowercase__			), '''token''': ANY(lowercase__			), '''token_str''': ANY(lowercase__			)},
           {'''sequence''': ANY(lowercase__			), '''score''': ANY(lowercase__			), '''token''': ANY(lowercase__			), '''token_str''': ANY(lowercase__			)},
       ]					,)
   self.assertEqual(nested_simplify(lowercase__			)					,nested_simplify(lowercase__			)			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     Tuple					,lowercase__   :     Optional[Any]					,lowercase__   :     int			):
   __lowercase							  = tokenizer.get_vocab()
   __lowercase							  = FillMaskPipeline(model=lowercase__					,tokenizer=lowercase__			)

   # top_k=2, ntargets=3
   __lowercase							  = sorted(vocab.keys()			)[:3]
   __lowercase							  = fill_masker(F"This is a {tokenizer.mask_token}"					,top_k=2					,targets=lowercase__			)

   # If we use the most probably targets, and filter differently, we should still
   # have the same results
   __lowercase							  = [el['''token_str'''] for el in sorted(lowercase__					,key=lambda lowercase__			: x["score"]					,reverse=lowercase__			)]
   # For some BPE tokenizers, `</w>` is removed during decoding, so `token_str` won't be the same as in `targets`.
   if set(lowercase__			).issubset(lowercase__			):
    __lowercase							  = fill_masker(F"This is a {tokenizer.mask_token}"					,top_k=3					,targets=lowercase__			)
    # They should yield exactly the same result
    self.assertEqual(nested_simplify(lowercase__			)					,nested_simplify(lowercase__			)			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     str					,lowercase__   :     str					,lowercase__   :     Optional[int]			):
   __lowercase							  = FillMaskPipeline(model=lowercase__					,tokenizer=lowercase__			)
   __lowercase							  = tokenizer.get_vocab()
   # String duplicates + id duplicates
   __lowercase							  = sorted(vocab.keys()			)[:3]
   __lowercase							  = [targets[0], targets[1], targets[0], targets[2], targets[1]]
   __lowercase							  = fill_masker(F"My name is {tokenizer.mask_token}"					,targets=lowercase__					,top_k=1_0			)

   # The target list contains duplicates, so we can't output more
   # than them
   self.assertEqual(len(lowercase__			)					,3			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     str					,lowercase__   :     Union[str, Any]					,lowercase__   :     Union[str, Any]			):
   __lowercase							  = FillMaskPipeline(model=lowercase__					,tokenizer=lowercase__			)

   __lowercase							  = fill_masker(
       F"This is a {tokenizer.mask_token} {tokenizer.mask_token} {tokenizer.mask_token}"					,top_k=2			)
   self.assertEqual(
       lowercase__					,[
           [
               {'''sequence''': ANY(lowercase__			), '''score''': ANY(lowercase__			), '''token''': ANY(lowercase__			), '''token_str''': ANY(lowercase__			)},
               {'''sequence''': ANY(lowercase__			), '''score''': ANY(lowercase__			), '''token''': ANY(lowercase__			), '''token_str''': ANY(lowercase__			)},
           ],
           [
               {'''sequence''': ANY(lowercase__			), '''score''': ANY(lowercase__			), '''token''': ANY(lowercase__			), '''token_str''': ANY(lowercase__			)},
               {'''sequence''': ANY(lowercase__			), '''score''': ANY(lowercase__			), '''token''': ANY(lowercase__			), '''token_str''': ANY(lowercase__			)},
           ],
           [
               {'''sequence''': ANY(lowercase__			), '''score''': ANY(lowercase__			), '''token''': ANY(lowercase__			), '''token_str''': ANY(lowercase__			)},
               {'''sequence''': ANY(lowercase__			), '''score''': ANY(lowercase__			), '''token''': ANY(lowercase__			), '''token_str''': ANY(lowercase__			)},
           ],
       ]					,)




'''simple docstring'''






import torch
from torch import nn






class 			lowercase_	(nn.Module      ):




  """simple docstring"""




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

   __lowercase							  = n_token
   __lowercase							  = d_embed
   __lowercase							  = d_proj

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

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

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

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

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

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

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

   __lowercase							  = keep_order






  def 					SCREAMING_SNAKE_CASE      (      self   :     Optional[int]					,lowercase__   :     List[Any]					,lowercase__   :     Dict					,lowercase__   :     List[Any]					,lowercase__   :     Any			):
   if proj is None:
    __lowercase							  = nn.functional.linear(lowercase__					,lowercase__					,bias=lowercase__			)
   else:
    # if CUDA_MAJOR <= 9 and CUDA_MINOR <= 1:
    __lowercase							  = nn.functional.linear(lowercase__					,proj.t().contiguous()			)
    __lowercase							  = nn.functional.linear(lowercase__					,lowercase__					,bias=lowercase__			)
    # 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   :     Union[str, Any]					,lowercase__   :     Optional[Any]					,lowercase__   :     Any=None					,lowercase__   :     List[str]=False			):

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

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

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

     weights.append(lowercase__			)
     biases.append(lowercase__			)

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

    __lowercase							  = self._compute_logit(lowercase__					,lowercase__					,lowercase__					,lowercase__			)
    __lowercase							  = nn.functional.log_softmax(lowercase__					,dim=1			)

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

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

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

      if indices_i.numel() == 0:
       continue

      __lowercase							  = labels.index_select(0					,lowercase__			) - l_idx
      __lowercase							  = head_logprob.index_select(0					,lowercase__			)
      __lowercase							  = hidden.index_select(0					,lowercase__			)
     else:
      __lowercase							  = hidden

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

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

     if labels is not None:
      if (hasattr(self					,'''keep_order'''			) and self.keep_order) or keep_order:
       out.index_copy_(0					,lowercase__					,-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   :     Any					,lowercase__   :     Union[str, Any]			):
   if self.n_clusters == 0:
    __lowercase							  = self._compute_logit(lowercase__					,self.out_layers[0].weight					,self.out_layers[0].bias					,self.out_projs[0]			)
    return nn.functional.log_softmax(lowercase__					,dim=-1			)
   else:
    # construct weights and biases
    __lowercase  ,	__lowercase							  = [], []
    for i in range(len(self.cutoffs			)			):
     if self.div_val == 1:
      __lowercase  ,	__lowercase							  = self.cutoff_ends[i], self.cutoff_ends[i + 1]
      __lowercase							  = self.out_layers[0].weight[l_idx:r_idx]
      __lowercase							  = self.out_layers[0].bias[l_idx:r_idx]
     else:
      __lowercase							  = self.out_layers[i].weight
      __lowercase							  = self.out_layers[i].bias

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

     weights.append(lowercase__			)
     biases.append(lowercase__			)

    __lowercase  ,	__lowercase  ,	__lowercase							  = weights[0], biases[0], self.out_projs[0]
    __lowercase							  = self._compute_logit(lowercase__					,lowercase__					,lowercase__					,lowercase__			)

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

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

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

      __lowercase							  = self._compute_logit(lowercase__					,lowercase__					,lowercase__					,lowercase__			)
      __lowercase							  = nn.functional.log_softmax(lowercase__					,dim=1			)

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

    return out

