





import json
import os
from typing import Dict, List, Optional, Tuple

from ...tokenization_utils import PreTrainedTokenizer
from ...utils import logging


UpperCAmelCase_							=				logging.get_logger(__name__)


UpperCAmelCase_							=				{
    """vocab_file""": """vocab.json""",
    """tokenizer_config_file""": """tokenizer_config.json""",
    """merges_file""": """merges.txt""",
}

UpperCAmelCase_							=				{
    """vocab_file""": {
        """facebook/s2t-wav2vec2-large-en-de""": (
            """https://huggingface.co/facebook/s2t-wav2vec2-large-en-de/resolve/main/vocab.json"""
        ),
    },
    """tokenizer_config_file""": {
        """facebook/s2t-wav2vec2-large-en-de""": (
            """https://huggingface.co/facebook/s2t-wav2vec2-large-en-de/resolve/main/tokenizer_config.json"""
        ),
    },
    """merges_file""": {
        """facebook/s2t-wav2vec2-large-en-de""": (
            """https://huggingface.co/facebook/s2t-wav2vec2-large-en-de/resolve/main/merges.txt"""
        ),
    },
}

UpperCAmelCase_							=				"""</w>"""
UpperCAmelCase_							=				"""@@ """





def  __magic_name__  (				lowercase				)				-> str:

							"""simple docstring"""



							lowercase_	:   Any  				=     set()
							lowercase_	:   Dict  				=     word[0]
							for char in word[1:]:
														pairs.add((prev_char, char)				)
														lowercase_	:   Optional[Any]  				=     char
							return pairs


# Speech2Text2 has no max input length
UpperCAmelCase_							=				{"""facebook/s2t-wav2vec2-large-en-de""": 1_024}




class     UpperCamelCase__					(     lowerCamelCase__				):






		'''simple docstring'''






		__a       :       int       =	VOCAB_FILES_NAMES
		__a       :       Tuple       =	PRETRAINED_VOCAB_FILES_MAP
		__a       :       Tuple       =	PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
		__a       :       str       =	["""input_ids""", """attention_mask"""]






		def __init__(							self,   snake_case__,   snake_case__="<s>",   snake_case__="<pad>",   snake_case__="</s>",   snake_case__="<unk>",   snake_case__=False,   snake_case__=None,   **snake_case__,   )       -> Dict:






									"""simple docstring"""
									super().__init__(
									    unk_token=snake_case__,   bos_token=snake_case__,   eos_token=snake_case__,   pad_token=snake_case__,   do_lower_case=snake_case__,   **snake_case__,   )

									lowercase_	:   List[str]  				=     do_lower_case

									with open(snake_case__,   encoding="""utf-8"""  ) as vocab_handle:
																lowercase_	:   Dict  				=     json.load(snake_case__  )
									lowercase_	:   List[str]  				=     {v: k for k, v in self.encoder.items()}

									if merges_file is None:
																logger.info(f"""No merges files provided. {self.__class__.__name__} can only be used for decoding."""  )

																lowercase_	:   List[Any]  				=     None
																lowercase_	:   Optional[int]  				=     None
									else:
																with open(snake_case__,   encoding="""utf-8"""  ) as merges_handle:
																							lowercase_	:   List[Any]  				=     merges_handle.read().split("""\n"""  )[:-1]

																lowercase_	:   Dict  				=     [tuple(merge.split()[:2]  ) for merge in merges]
																lowercase_	:   Tuple  				=     dict(zip(snake_case__,   range(len(snake_case__  )  )  )  )
																lowercase_	:   List[Any]  				=     {}






		@property
		def 						snake_case__				(							self  )       -> int:






									"""simple docstring"""
									return len(self.decoder  )






		def 						snake_case__				(							self  )       -> Dict:






									"""simple docstring"""
									return dict(self.encoder,   **self.added_tokens_encoder  )






		def 						snake_case__				(							self,   snake_case__  )       -> Any:






									"""simple docstring"""
									lowercase_	:   Optional[int]  				=     tuple(token[:-1]  ) + (token[-1] + BPE_TOKEN_MERGES,)
									if token in self.cache:
																return self.cache[token]
									lowercase_	:   int  				=     get_pairs(snake_case__  )

									if not pairs:
																return token

									while True:
																lowercase_	:   List[Any]  				=     min(snake_case__,   key=lambda snake_case__  : self.bpe_ranks.get(snake_case__,   float("""inf"""  )  )  )
																if bigram not in self.bpe_ranks:
																							break
																lowercase_				,		lowercase_	:   str  				=     bigram
																lowercase_	:   Any  				=     []
																lowercase_	:   str  				=     0
																while i < len(snake_case__  ):
																							try:
																														lowercase_	:   Union[str, Any]  				=     word.index(snake_case__,   snake_case__  )
																							except ValueError:
																														new_word.extend(word[i:]  )
																														break
																							else:
																														new_word.extend(word[i:j]  )
																														lowercase_	:   Union[str, Any]  				=     j

																							if word[i] == first and i < len(snake_case__  ) - 1 and word[i + 1] == second:
																														new_word.append(first + second  )
																														i += 2
																							else:
																														new_word.append(word[i]  )
																														i += 1
																lowercase_	:   Optional[int]  				=     tuple(snake_case__  )
																lowercase_	:   Any  				=     new_word
																if len(snake_case__  ) == 1:
																							break
																else:
																							lowercase_	:   Union[str, Any]  				=     get_pairs(snake_case__  )
									lowercase_	:   Any  				=     """ """.join(snake_case__  )
									if word == "\n  " + BPE_TOKEN_MERGES:
																lowercase_	:   Optional[int]  				=     """\n""" + BPE_TOKEN_MERGES

									if word.endswith(snake_case__  ):
																lowercase_	:   Optional[Any]  				=     word.replace(snake_case__,   """"""  )

									lowercase_	:   Any  				=     word.replace(""" """,   snake_case__  )
									lowercase_	:   Optional[Any]  				=     word
									return word






		def 						snake_case__				(							self,   snake_case__  )       -> Optional[int]:






									"""simple docstring"""

									if self.bpe_ranks is None:
																raise ValueError(
																    """This tokenizer was instantiated without a `merges.txt` file, so"""
																    """ that it can only be used for decoding, not for encoding."""
																    """Make sure to provide `merges.txt` file at instantiation to enable """
																    """encoding."""  )

									if self.do_lower_case:
																lowercase_	:   Optional[Any]  				=     text.lower()

									lowercase_	:   str  				=     text.split()

									lowercase_	:   str  				=     []
									for token in text:
																if token:
																							split_tokens.extend(list(self.bpe(snake_case__  ).split(""" """  )  )  )

									return split_tokens






		def 						snake_case__				(							self,   snake_case__  )       -> int:






									"""simple docstring"""
									return self.encoder.get(snake_case__,   self.encoder.get(self.unk_token  )  )






		def 						snake_case__				(							self,   snake_case__  )       -> str:






									"""simple docstring"""
									lowercase_	:   Any  				=     self.decoder.get(snake_case__,   self.unk_token  )
									return result






		def 						snake_case__				(							self,   snake_case__  )       -> str:






									"""simple docstring"""
									lowercase_	:   Union[str, Any]  				=     """ """.join(snake_case__  )

									# make sure @@ tokens are concatenated
									lowercase_	:   Dict  				=     """""".join(string.split(snake_case__  )  )

									return string






		def 						snake_case__				(							self,   snake_case__,   snake_case__ = None  )       -> Tuple[str]:






									"""simple docstring"""
									if not os.path.isdir(snake_case__  ):
																logger.error(f"""Vocabulary path ({save_directory}) should be a directory"""  )
																return
									lowercase_	:   Union[str, Any]  				=     os.path.join(
									    snake_case__,   (filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""vocab_file"""]  )
									lowercase_	:   Any  				=     os.path.join(
									    snake_case__,   (filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""merges_file"""]  )

									with open(snake_case__,   """w""",   encoding="""utf-8"""  ) as f:
																f.write(json.dumps(self.encoder,   indent=2,   sort_keys=snake_case__,   ensure_ascii=snake_case__  ) + """\n"""  )

									lowercase_	:   Any  				=     0
									if self.bpe_ranks is None:
																return (vocab_file,)

									with open(snake_case__,   """w""",   encoding="""utf-8"""  ) as writer:
																for bpe_tokens, token_index in sorted(self.bpe_ranks.items(),   key=lambda snake_case__  : kv[1]  ):
																							if index != token_index:
																														logger.warning(
																														    f"""Saving vocabulary to {merges_file}: BPE merge indices are not consecutive."""
																														    """ Please check that the tokenizer is not corrupted!"""  )
																														lowercase_	:   Any  				=     token_index
																							writer.write(""" """.join(snake_case__  ) + """\n"""  )
																							index += 1

									return (vocab_file, merges_file)






import tensorflow as tf

from ...tf_utils import shape_list
class     UpperCamelCase__					(     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__  )       -> Optional[int]:






									"""simple docstring"""
									super().__init__(**snake_case__  )

									lowercase_	:   Tuple  				=     vocab_size
									lowercase_	:   Union[str, Any]  				=     d_embed
									lowercase_	:   Optional[Any]  				=     d_proj

									lowercase_	:   Optional[Any]  				=     cutoffs + [vocab_size]
									lowercase_	:   Optional[int]  				=     [0] + self.cutoffs
									lowercase_	:   Union[str, Any]  				=     div_val

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

									lowercase_	:   Dict  				=     []
									lowercase_	:   int  				=     []






		def 						snake_case__				(							self,   snake_case__  )       -> Optional[int]:






									"""simple docstring"""
									if self.n_clusters > 0:
																lowercase_	:   List[str]  				=     self.add_weight(
																    shape=(self.n_clusters, self.d_embed),   initializer="""zeros""",   trainable=snake_case__,   name="""cluster_weight"""  )
																lowercase_	:   List[Any]  				=     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:
																														lowercase_	:   Any  				=     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__  )
																							lowercase_	:   List[str]  				=     self.add_weight(
																							    shape=(self.vocab_size, self.d_embed),   initializer="""zeros""",   trainable=snake_case__,   name=f"""out_layers_._{i}_._weight""",   )
																							lowercase_	:   List[str]  				=     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  )  ):
																							lowercase_				,		lowercase_	:   Optional[Any]  				=     self.cutoff_ends[i], self.cutoff_ends[i + 1]
																							lowercase_	:   Union[str, Any]  				=     self.d_embed // (self.div_val**i)

																							lowercase_	:   int  				=     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__  )
																							lowercase_	:   Any  				=     self.add_weight(
																							    shape=(r_idx - l_idx, d_emb_i),   initializer="""zeros""",   trainable=snake_case__,   name=f"""out_layers_._{i}_._weight""",   )
																							lowercase_	:   Tuple  				=     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 						snake_case__				(							snake_case__,   snake_case__,   snake_case__,   snake_case__=None  )       -> Optional[Any]:






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






		@staticmethod
		def 						snake_case__				(							snake_case__,   snake_case__  )       -> List[str]:






									"""simple docstring"""
									lowercase_	:   Optional[int]  				=     shape_list(snake_case__  )
									lowercase_	:   Optional[Any]  				=     tf.range(lp_size[0],   dtype=target.dtype  )
									lowercase_	:   str  				=     tf.stack([r, target],   1  )
									return tf.gather_nd(snake_case__,   snake_case__  )






		def 						snake_case__				(							self,   snake_case__,   snake_case__,   snake_case__=True,   snake_case__=False  )       -> Dict:






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

																							if self.div_val == 1:
																														lowercase_	:   Optional[Any]  				=     self.out_layers[0][0][l_idx:r_idx]
																														lowercase_	:   Union[str, Any]  				=     self.out_layers[0][1][l_idx:r_idx]
																							else:
																														lowercase_	:   Dict  				=     self.out_layers[i][0]
																														lowercase_	:   int  				=     self.out_layers[i][1]

																							if i == 0:
																														lowercase_	:   Optional[int]  				=     tf.concat([cur_W, self.cluster_weight],   0  )
																														lowercase_	:   Optional[Any]  				=     tf.concat([cur_b, self.cluster_bias],   0  )

																														lowercase_	:   List[str]  				=     self._logit(snake_case__,   snake_case__,   snake_case__,   self.out_projs[0]  )
																														lowercase_	:   List[str]  				=     tf.nn.log_softmax(snake_case__  )
																														out.append(head_logprob[..., : self.cutoffs[0]]  )
																														if target is not None:
																																					lowercase_	:   int  				=     tf.boolean_mask(snake_case__,   snake_case__  )
																																					lowercase_	:   Optional[int]  				=     self._gather_logprob(snake_case__,   snake_case__  )
																							else:
																														lowercase_	:   List[str]  				=     self._logit(snake_case__,   snake_case__,   snake_case__,   self.out_projs[i]  )
																														lowercase_	:   Dict  				=     tf.nn.log_softmax(snake_case__  )
																														lowercase_	:   Optional[int]  				=     self.cutoffs[0] + i - 1  # No probability for the head cluster
																														lowercase_	:   int  				=     head_logprob[..., cluster_prob_idx, None] + tail_logprob
																														out.append(snake_case__  )
																														if target is not None:
																																					lowercase_	:   Optional[Any]  				=     tf.boolean_mask(snake_case__,   snake_case__  )
																																					lowercase_	:   Optional[int]  				=     tf.boolean_mask(snake_case__,   snake_case__  )
																																					lowercase_	:   Union[str, Any]  				=     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__  )  )
																lowercase_	:   List[str]  				=     tf.concat(snake_case__,   axis=-1  )

									if target is not None:
																if return_mean:
																							lowercase_	:   Tuple  				=     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