


import tensorflow as tf

from ...tf_utils import shape_list



class 							_UpperCamelCase (						tf.keras.layers.Layer						):

   '''simple docstring'''







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






         '''simple docstring'''



         super().__init__(**_lowerCamelCase							)

         __lowerCamelCase     :      Union[str, Any]       						=						vocab_size
         __lowerCamelCase     :      Union[str, Any]       						=						d_embed
         __lowerCamelCase     :      str       						=						d_proj

         __lowerCamelCase     :      Dict       						=						cutoffs + [vocab_size]
         __lowerCamelCase     :      List[Any]       						=						[0] + self.cutoffs
         __lowerCamelCase     :      str       						=						div_val

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

         __lowerCamelCase     :      Dict       						=						[]
         __lowerCamelCase     :      List[str]       						=						[]

   def   _snake_case    (  self  : Optional[Any]       ,  _lowerCamelCase  : Union[str, Any]							):






         '''simple docstring'''



         if self.n_clusters > 0:
               __lowerCamelCase     :      List[str]       						=						self.add_weight(
                   shape=(self.n_clusters, self.d_embed)       ,  initializer="""zeros"""       ,  trainable=_lowerCamelCase       ,  name="""cluster_weight"""							)
               __lowerCamelCase     :      Any       						=						self.add_weight(
                   shape=(self.n_clusters,)       ,  initializer="""zeros"""       ,  trainable=_lowerCamelCase       ,  name="""cluster_bias"""							)

         if self.div_val == 1:
               for i in range(len(self.cutoffs							)							):
                     if self.d_proj != self.d_embed:
                           __lowerCamelCase     :      List[Any]       						=						self.add_weight(
                               shape=(self.d_embed, self.d_proj)       ,  initializer="""zeros"""       ,  trainable=_lowerCamelCase       ,  name=F"""out_projs_._{i}"""       ,  )
                           self.out_projs.append(_lowerCamelCase							)
                     else:
                           self.out_projs.append(_lowerCamelCase							)
                     __lowerCamelCase     :      int       						=						self.add_weight(
                         shape=(self.vocab_size, self.d_embed)       ,  initializer="""zeros"""       ,  trainable=_lowerCamelCase       ,  name=F"""out_layers_._{i}_._weight"""       ,  )
                     __lowerCamelCase     :      Tuple       						=						self.add_weight(
                         shape=(self.vocab_size,)       ,  initializer="""zeros"""       ,  trainable=_lowerCamelCase       ,  name=F"""out_layers_._{i}_._bias"""       ,  )
                     self.out_layers.append((weight, bias)							)
         else:
               for i in range(len(self.cutoffs							)							):
                     __lowerCamelCase						, __lowerCamelCase     :      int       						=						self.cutoff_ends[i], self.cutoff_ends[i + 1]
                     __lowerCamelCase     :      Optional[Any]       						=						self.d_embed // (self.div_val**i)

                     __lowerCamelCase     :      str       						=						self.add_weight(
                         shape=(d_emb_i, self.d_proj)       ,  initializer="""zeros"""       ,  trainable=_lowerCamelCase       ,  name=F"""out_projs_._{i}"""							)
                     self.out_projs.append(_lowerCamelCase							)
                     __lowerCamelCase     :      Union[str, Any]       						=						self.add_weight(
                         shape=(r_idx - l_idx, d_emb_i)       ,  initializer="""zeros"""       ,  trainable=_lowerCamelCase       ,  name=F"""out_layers_._{i}_._weight"""       ,  )
                     __lowerCamelCase     :      Optional[Any]       						=						self.add_weight(
                         shape=(r_idx - l_idx,)       ,  initializer="""zeros"""       ,  trainable=_lowerCamelCase       ,  name=F"""out_layers_._{i}_._bias"""       ,  )
                     self.out_layers.append((weight, bias)							)
         super().build(_lowerCamelCase							)

   @staticmethod
   def   _snake_case    (  _lowerCamelCase  : Dict       ,  _lowerCamelCase  : int       ,  _lowerCamelCase  : List[Any]       ,  _lowerCamelCase  : List[Any]=None							):






         '''simple docstring'''



         __lowerCamelCase     :      Dict       						=						x
         if proj is not None:
               __lowerCamelCase     :      List[Any]       						=						tf.einsum("""ibd,ed->ibe"""       ,  _lowerCamelCase       ,  _lowerCamelCase							)
         return tf.einsum("""ibd,nd->ibn"""       ,  _lowerCamelCase       ,  _lowerCamelCase							) + b

   @staticmethod
   def   _snake_case    (  _lowerCamelCase  : Any       ,  _lowerCamelCase  : Optional[int]							):






         '''simple docstring'''



         __lowerCamelCase     :      Optional[int]       						=						shape_list(_lowerCamelCase							)
         __lowerCamelCase     :      Optional[int]       						=						tf.range(lp_size[0]       ,  dtype=target.dtype							)
         __lowerCamelCase     :      Optional[int]       						=						tf.stack([r, target]       ,  1							)
         return tf.gather_nd(_lowerCamelCase       ,  _lowerCamelCase							)


   def   _snake_case    (  self  : Optional[Any]       ,  _lowerCamelCase  : int       ,  _lowerCamelCase  : Any       ,  _lowerCamelCase  : Tuple=True       ,  _lowerCamelCase  : Optional[int]=False							):






         '''simple docstring'''



         __lowerCamelCase     :      List[str]       						=						0
         if self.n_clusters == 0:
               __lowerCamelCase     :      int       						=						self._logit(_lowerCamelCase       ,  self.out_layers[0][0]       ,  self.out_layers[0][1]       ,  self.out_projs[0]							)
               if target is not None:
                     __lowerCamelCase     :      Union[str, Any]       						=						tf.nn.sparse_softmax_cross_entropy_with_logits(labels=_lowerCamelCase       ,  logits=_lowerCamelCase							)
               __lowerCamelCase     :      Optional[Any]       						=						tf.nn.log_softmax(_lowerCamelCase       ,  axis=-1							)
         else:
               __lowerCamelCase     :      List[str]       						=						shape_list(_lowerCamelCase							)
               __lowerCamelCase     :      str       						=						[]
               __lowerCamelCase     :      Tuple       						=						tf.zeros(hidden_sizes[:2]							)
               for i in range(len(self.cutoffs							)							):
                     __lowerCamelCase						, __lowerCamelCase     :      Optional[Any]       						=						self.cutoff_ends[i], self.cutoff_ends[i + 1]
                     if target is not None:
                           __lowerCamelCase     :      Optional[int]       						=						(target >= l_idx) & (target < r_idx)
                           __lowerCamelCase     :      List[str]       						=						tf.where(_lowerCamelCase							)
                           __lowerCamelCase     :      List[Any]       						=						tf.boolean_mask(_lowerCamelCase       ,  _lowerCamelCase							) - l_idx

                     if self.div_val == 1:
                           __lowerCamelCase     :      int       						=						self.out_layers[0][0][l_idx:r_idx]
                           __lowerCamelCase     :      int       						=						self.out_layers[0][1][l_idx:r_idx]
                     else:
                           __lowerCamelCase     :      Tuple       						=						self.out_layers[i][0]
                           __lowerCamelCase     :      Any       						=						self.out_layers[i][1]

                     if i == 0:
                           __lowerCamelCase     :      List[str]       						=						tf.concat([cur_W, self.cluster_weight]       ,  0							)
                           __lowerCamelCase     :      Dict       						=						tf.concat([cur_b, self.cluster_bias]       ,  0							)

                           __lowerCamelCase     :      Optional[int]       						=						self._logit(_lowerCamelCase       ,  _lowerCamelCase       ,  _lowerCamelCase       ,  self.out_projs[0]							)
                           __lowerCamelCase     :      List[str]       						=						tf.nn.log_softmax(_lowerCamelCase							)
                           out.append(head_logprob[..., : self.cutoffs[0]]							)
                           if target is not None:
                                 __lowerCamelCase     :      Any       						=						tf.boolean_mask(_lowerCamelCase       ,  _lowerCamelCase							)
                                 __lowerCamelCase     :      str       						=						self._gather_logprob(_lowerCamelCase       ,  _lowerCamelCase							)
                     else:
                           __lowerCamelCase     :      Tuple       						=						self._logit(_lowerCamelCase       ,  _lowerCamelCase       ,  _lowerCamelCase       ,  self.out_projs[i]							)
                           __lowerCamelCase     :      List[Any]       						=						tf.nn.log_softmax(_lowerCamelCase							)
                           __lowerCamelCase     :      str       						=						self.cutoffs[0] + i - 1  # No probability for the head cluster
                           __lowerCamelCase     :      Union[str, Any]       						=						head_logprob[..., cluster_prob_idx, None] + tail_logprob
                           out.append(_lowerCamelCase							)
                           if target is not None:
                                 __lowerCamelCase     :      int       						=						tf.boolean_mask(_lowerCamelCase       ,  _lowerCamelCase							)
                                 __lowerCamelCase     :      int       						=						tf.boolean_mask(_lowerCamelCase       ,  _lowerCamelCase							)
                                 __lowerCamelCase     :      List[Any]       						=						self._gather_logprob(_lowerCamelCase       ,  _lowerCamelCase							)
                                 cur_logprob += cur_head_logprob[:, self.cutoff_ends[1] + i - 1]
                     if target is not None:
                           loss += tf.scatter_nd(_lowerCamelCase       ,  -cur_logprob       ,  shape_list(_lowerCamelCase							)							)
               __lowerCamelCase     :      Union[str, Any]       						=						tf.concat(_lowerCamelCase       ,  axis=-1							)

         if target is not None:
               if return_mean:
                     __lowerCamelCase     :      str       						=						tf.reduce_mean(_lowerCamelCase							)
               # Add the training-time loss value to the layer using `self.add_loss()`.
               self.add_loss(_lowerCamelCase							)

               # Log the loss as a metric (we could log arbitrary metrics,
               # including different metrics for training and inference.
               self.add_metric(_lowerCamelCase       ,  name=self.name       ,  aggregation="""mean""" if return_mean else """"""							)

         return out







import json
from typing import Iterator, List, Union

from tokenizers import AddedToken, Regex, Tokenizer, decoders, normalizers, pre_tokenizers, trainers
from tokenizers.implementations.base_tokenizer import BaseTokenizer
from tokenizers.models import Unigram
from tokenizers.processors import TemplateProcessing



class 							_UpperCamelCase (						A						):

   '''simple docstring'''







   def __init__(  self  : str       ,  _lowerCamelCase  : str = "▁"       ,  _lowerCamelCase  : bool = True       ,  _lowerCamelCase  : Union[str, AddedToken] = "<unk>"       ,  _lowerCamelCase  : Union[str, AddedToken] = "</s>"       ,  _lowerCamelCase  : Union[str, AddedToken] = "<pad>"       ,  ):






         '''simple docstring'''



         __lowerCamelCase     :      Tuple       						=						{
             """pad""": {"""id""": 0, """token""": pad_token},
             """eos""": {"""id""": 1, """token""": eos_token},
             """unk""": {"""id""": 2, """token""": unk_token},
         }

         __lowerCamelCase     :      Optional[Any]       						=						[None] * len(self.special_tokens							)
         for token_dict in self.special_tokens.values():
               __lowerCamelCase     :      Tuple       						=						token_dict["""token"""]

         __lowerCamelCase     :      Optional[int]       						=						Tokenizer(Unigram()							)

         __lowerCamelCase     :      str       						=						normalizers.Sequence(
             [
                 normalizers.Nmt(),
                 normalizers.NFKC(),
                 normalizers.Replace(Regex(""" {2,}"""							)       ,  """ """							),
                 normalizers.Lowercase(),
             ]							)
         __lowerCamelCase     :      int       						=						pre_tokenizers.Sequence(
             [
                 pre_tokenizers.Metaspace(replacement=_lowerCamelCase       ,  add_prefix_space=_lowerCamelCase							),
                 pre_tokenizers.Digits(individual_digits=_lowerCamelCase							),
                 pre_tokenizers.Punctuation(),
             ]							)
         __lowerCamelCase     :      Tuple       						=						decoders.Metaspace(replacement=_lowerCamelCase       ,  add_prefix_space=_lowerCamelCase							)

         __lowerCamelCase     :      Tuple       						=						TemplateProcessing(
             single=F"""$A {self.special_tokens["eos"]["token"]}"""       ,  special_tokens=[(self.special_tokens["""eos"""]["""token"""], self.special_tokens["""eos"""]["""id"""])]       ,  )

         __lowerCamelCase     :      Tuple       						=						{
             """model""": """SentencePieceUnigram""",
             """replacement""": replacement,
             """add_prefix_space""": add_prefix_space,
         }

         super().__init__(_lowerCamelCase       ,  _lowerCamelCase							)

   def   _snake_case    (  self  : Dict       ,  _lowerCamelCase  : Union[str, List[str]]       ,  _lowerCamelCase  : int = 8_0_0_0       ,  _lowerCamelCase  : bool = True       ,  ):






         '''simple docstring'''



         __lowerCamelCase     :      Optional[Any]       						=						trainers.UnigramTrainer(
             vocab_size=_lowerCamelCase       ,  special_tokens=self.special_tokens_list       ,  show_progress=_lowerCamelCase       ,  )

         if isinstance(_lowerCamelCase       ,  _lowerCamelCase							):
               __lowerCamelCase     :      Optional[int]       						=						[files]
         self._tokenizer.train(_lowerCamelCase       ,  trainer=_lowerCamelCase							)

         self.add_unk_id()

   def   _snake_case    (  self  : List[Any]       ,  _lowerCamelCase  : Union[Iterator[str], Iterator[Iterator[str]]]       ,  _lowerCamelCase  : int = 8_0_0_0       ,  _lowerCamelCase  : bool = True       ,  ):






         '''simple docstring'''



         __lowerCamelCase     :      Any       						=						trainers.UnigramTrainer(
             vocab_size=_lowerCamelCase       ,  special_tokens=self.special_tokens_list       ,  show_progress=_lowerCamelCase       ,  )

         self._tokenizer.train_from_iterator(_lowerCamelCase       ,  trainer=_lowerCamelCase							)

         self.add_unk_id()


   def   _snake_case    (  self  : List[str]							):






         '''simple docstring'''



         __lowerCamelCase     :      List[str]       						=						json.loads(self._tokenizer.to_str()							)

         __lowerCamelCase     :      Optional[int]       						=						self.special_tokens["""unk"""]["""id"""]

         __lowerCamelCase     :      List[str]       						=						Tokenizer.from_str(json.dumps(_lowerCamelCase							)							)



