

"""simple docstring"""

import tensorflow as tf

from ...tf_utils import shape_list






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





       """simple docstring"""


       def __init__(   self			,					snake_case			,					snake_case			,					snake_case			,					snake_case			,					snake_case=1			,					snake_case=False			,					**snake_case   ):





        """simple docstring"""

        super().__init__(**snake_case   )

        lowerCAmelCase__   :				Dict 			=				vocab_size
        lowerCAmelCase__   :				Tuple 			=				d_embed
        lowerCAmelCase__   :				Dict 			=				d_proj

        lowerCAmelCase__   :				Any 			=				cutoffs + [vocab_size]
        lowerCAmelCase__   :				Any 			=				[0] + self.cutoffs
        lowerCAmelCase__   :				Optional[Any] 			=				div_val

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

        lowerCAmelCase__   :				int 			=				[]
        lowerCAmelCase__   :				List[Any] 			=				[]






       def 			SCREAMING_SNAKE_CASE_	(   self			,					snake_case   ):





        """simple docstring"""

        if self.n_clusters > 0:
         lowerCAmelCase__   :				str 			=				self.add_weight(
             shape=(self.n_clusters, self.d_embed)			,					initializer="zeros"			,					trainable=snake_case			,					name="cluster_weight"   )
         lowerCAmelCase__   :				int 			=				self.add_weight(
             shape=(self.n_clusters,)			,					initializer="zeros"			,					trainable=snake_case			,					name="cluster_bias"   )

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

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






       @staticmethod
       def 			SCREAMING_SNAKE_CASE_	(   snake_case			,					snake_case			,					snake_case			,					snake_case=None   ):





        """simple docstring"""

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






       @staticmethod
       def 			SCREAMING_SNAKE_CASE_	(   snake_case			,					snake_case   ):





        """simple docstring"""

        lowerCAmelCase__   :				int 			=				shape_list(snake_case   )
        lowerCAmelCase__   :				str 			=				tf.range(lp_size[0]			,					dtype=target.dtype   )
        lowerCAmelCase__   :				str 			=				tf.stack([r, target]			,					1   )
        return tf.gather_nd(snake_case			,					snake_case   )






       def 			SCREAMING_SNAKE_CASE_	(   self			,					snake_case			,					snake_case			,					snake_case=True			,					snake_case=False   ):





        """simple docstring"""

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

          if self.div_val == 1:
           lowerCAmelCase__   :				str 			=				self.out_layers[0][0][l_idx:r_idx]
           lowerCAmelCase__   :				List[Any] 			=				self.out_layers[0][1][l_idx:r_idx]
          else:
           lowerCAmelCase__   :				Tuple 			=				self.out_layers[i][0]
           lowerCAmelCase__   :				Union[str, Any] 			=				self.out_layers[i][1]

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

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

        if target is not None:
         if return_mean:
          lowerCAmelCase__   :				Optional[Any] 			=				tf.reduce_mean(snake_case   )
         # Add the training-time loss value to the layer using `self.add_loss()`.
         self.add_loss(snake_case   )

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

        return out









"""simple docstring"""

import os
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple

import sentencepiece as spm

from ...tokenization_utils import AddedToken, PreTrainedTokenizer
from ...utils import logging


_UpperCamelCase		     =    logging.get_logger(__name__)

_UpperCamelCase		     =    """▁"""

_UpperCamelCase		     =    {"""vocab_file""": """sentencepiece.bpe.model""", """monolingual_vocab_file""": """dict.txt"""}

_UpperCamelCase		     =    {
    """vocab_file""": {
        """vinai/bartpho-syllable""": """https://huggingface.co/vinai/bartpho-syllable/resolve/main/sentencepiece.bpe.model""",
    },
    """monolingual_vocab_file""": {
        """vinai/bartpho-syllable""": """https://huggingface.co/vinai/bartpho-syllable/resolve/main/dict.txt""",
    },
}

_UpperCamelCase		     =    {"""vinai/bartpho-syllable""": 1_0_2_4}






class 			__a					(						__magic_name__       ):





       """simple docstring"""


       __UpperCamelCase					:      int    =						VOCAB_FILES_NAMES
       __UpperCamelCase					:      List[str]    =						PRETRAINED_VOCAB_FILES_MAP
       __UpperCamelCase					:      List[Any]    =						PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
       __UpperCamelCase					:      int    =						['input_ids', 'attention_mask']






       def __init__(   self			,					snake_case			,					snake_case			,					snake_case="<s>"			,					snake_case="</s>"			,					snake_case="</s>"			,					snake_case="<s>"			,					snake_case="<unk>"			,					snake_case="<pad>"			,					snake_case="<mask>"			,					snake_case = None			,					**snake_case			,					):





        """simple docstring"""

        lowerCAmelCase__   :				Union[str, Any] 			=				AddedToken(snake_case			,					lstrip=snake_case			,					rstrip=snake_case   ) if isinstance(snake_case			,					snake_case   ) else mask_token

        lowerCAmelCase__   :				Union[str, Any] 			=				{} if sp_model_kwargs is None else sp_model_kwargs

        super().__init__(
            bos_token=snake_case			,					eos_token=snake_case			,					unk_token=snake_case			,					sep_token=snake_case			,					cls_token=snake_case			,					pad_token=snake_case			,					mask_token=snake_case			,					sp_model_kwargs=self.sp_model_kwargs			,					**snake_case			,					)

        lowerCAmelCase__   :				Union[str, Any] 			=				vocab_file
        lowerCAmelCase__   :				Optional[Any] 			=				monolingual_vocab_file
        lowerCAmelCase__   :				Optional[Any] 			=				spm.SentencePieceProcessor(**self.sp_model_kwargs   )
        self.sp_model.Load(str(snake_case   )   )

        # Load the reduced vocab

        # Keep order of special tokens for backward compatibility
        lowerCAmelCase__   :				Optional[int] 			=				{}
        lowerCAmelCase__   :				int 			=				0
        for token in [bos_token, pad_token, eos_token, unk_token, sep_token, cls_token]:
         if str(snake_case   ) not in self.fairseq_tokens_to_ids:
          lowerCAmelCase__   :				Optional[int] 			=				cnt
          cnt += 1
        with open(snake_case			,					"r"			,					encoding="utf-8"   ) as f:
         for line in f.readlines():
          lowerCAmelCase__   :				Union[str, Any] 			=				line.strip().split()[0]
          lowerCAmelCase__   :				List[str] 			=				len(self.fairseq_tokens_to_ids   )
        if str(snake_case   ) not in self.fairseq_tokens_to_ids:
         lowerCAmelCase__   :				int 			=				len(self.fairseq_tokens_to_ids   )

        lowerCAmelCase__   :				Union[str, Any] 			=				{v: k for k, v in self.fairseq_tokens_to_ids.items()}






       def __getstate__(   self   ):





        """simple docstring"""

        lowerCAmelCase__   :				List[Any] 			=				self.__dict__.copy()
        lowerCAmelCase__   :				str 			=				None
        lowerCAmelCase__   :				Union[str, Any] 			=				self.sp_model.serialized_model_proto()
        return state






       def __setstate__(   self			,					snake_case   ):





        """simple docstring"""

        lowerCAmelCase__   :				Any 			=				d

        # for backward compatibility
        if not hasattr(self			,					"sp_model_kwargs"   ):
         lowerCAmelCase__   :				Any 			=				{}

        lowerCAmelCase__   :				Dict 			=				spm.SentencePieceProcessor(**self.sp_model_kwargs   )
        self.sp_model.LoadFromSerializedProto(self.sp_model_proto   )






       def 			SCREAMING_SNAKE_CASE_	(   self			,					snake_case			,					snake_case = None   ):





        """simple docstring"""

        if token_ids_a is None:
         return [self.cls_token_id] + token_ids_a + [self.sep_token_id]
        lowerCAmelCase__   :				Any 			=				[self.cls_token_id]
        lowerCAmelCase__   :				Dict 			=				[self.sep_token_id]
        return cls + token_ids_a + sep + sep + token_ids_a + sep






       def 			SCREAMING_SNAKE_CASE_	(   self			,					snake_case			,					snake_case = None			,					snake_case = False   ):





        """simple docstring"""

        if already_has_special_tokens:
         return super().get_special_tokens_mask(
             token_ids_a=snake_case			,					token_ids_a=snake_case			,					already_has_special_tokens=snake_case   )

        if token_ids_a is None:
         return [1] + ([0] * len(snake_case   )) + [1]
        return [1] + ([0] * len(snake_case   )) + [1, 1] + ([0] * len(snake_case   )) + [1]






       def 			SCREAMING_SNAKE_CASE_	(   self			,					snake_case			,					snake_case = None   ):





        """simple docstring"""

        lowerCAmelCase__   :				Optional[Any] 			=				[self.sep_token_id]
        lowerCAmelCase__   :				Optional[int] 			=				[self.cls_token_id]

        if token_ids_a is None:
         return len(cls + token_ids_a + sep   ) * [0]
        return len(cls + token_ids_a + sep + sep + token_ids_a + sep   ) * [0]






       @property
       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        return len(self.fairseq_ids_to_tokens   )






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        lowerCAmelCase__   :				List[str] 			=				{self.convert_ids_to_tokens(snake_case   ): i for i in range(self.vocab_size   )}
        vocab.update(self.added_tokens_encoder   )
        return vocab






       def 			SCREAMING_SNAKE_CASE_	(   self			,					snake_case   ):





        """simple docstring"""

        return self.sp_model.encode(snake_case			,					out_type=snake_case   )






       def 			SCREAMING_SNAKE_CASE_	(   self			,					snake_case   ):





        """simple docstring"""

        if token in self.fairseq_tokens_to_ids:
         return self.fairseq_tokens_to_ids[token]
        else:
         return self.unk_token_id






       def 			SCREAMING_SNAKE_CASE_	(   self			,					snake_case   ):





        """simple docstring"""

        return self.fairseq_ids_to_tokens[index]






       def 			SCREAMING_SNAKE_CASE_	(   self			,					snake_case   ):





        """simple docstring"""

        lowerCAmelCase__   :				List[str] 			=				"".join(snake_case   ).replace(snake_case			,					" "   ).strip()
        return out_string






       def 			SCREAMING_SNAKE_CASE_	(   self			,					snake_case			,					snake_case = None   ):





        """simple docstring"""

        if not os.path.isdir(snake_case   ):
         logger.error(F"""Vocabulary path ({save_directory}) should be a directory"""   )
         return
        lowerCAmelCase__   :				List[str] 			=				os.path.join(
            snake_case			,					(filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]   )
        lowerCAmelCase__   :				Optional[Any] 			=				os.path.join(
            snake_case			,					(filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["monolingual_vocab_file"]			,					)

        if os.path.abspath(self.vocab_file   ) != os.path.abspath(snake_case   ) and os.path.isfile(self.vocab_file   ):
         copyfile(self.vocab_file			,					snake_case   )
        elif not os.path.isfile(self.vocab_file   ):
         with open(snake_case			,					"wb"   ) as fi:
          lowerCAmelCase__   :				str 			=				self.sp_model.serialized_model_proto()
          fi.write(snake_case   )

        if os.path.abspath(self.monolingual_vocab_file   ) != os.path.abspath(
            snake_case   ) and os.path.isfile(self.monolingual_vocab_file   ):
         copyfile(self.monolingual_vocab_file			,					snake_case   )
        elif not os.path.isfile(self.monolingual_vocab_file   ):
         with open(snake_case			,					"w"			,					encoding="utf-8"   ) as fp:
          for token in self.fairseq_tokens_to_ids:
           if token not in self.all_special_tokens:
            fp.write(F"""{str(snake_case   )} \n"""   )

        return out_vocab_file, out_monolingual_vocab_file






