






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

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


_UpperCamelCase			:	str								=logging.get_logger(__name__)


_UpperCamelCase			:	Tuple								={
    'vocab_file': 'vocab.json',
    'tokenizer_config_file': 'tokenizer_config.json',
    'merges_file': 'merges.txt',
}

_UpperCamelCase			:	Tuple								={
    '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			:	str								='</w>'
_UpperCamelCase			:	Any								='@@ '



def  a__ (__lowercase							:int   )		-> List[Any]:
				_A	:				Any   =      set()
				_A	:				Any   =      word[0]
				for char in word[1:]:
								pairs.add((prev_char, char)   )
								_A	:				Optional[int]   =      char
				return pairs


# Speech2Text2 has no max input length
_UpperCamelCase			:	Optional[int]								={'facebook/s2t-wav2vec2-large-en-de': 1024}






class 						UpperCAmelCase__ (							__snake_case							):
		__snake_case      :       int     =   VOCAB_FILES_NAMES
		__snake_case      :       Optional[Any]     =   PRETRAINED_VOCAB_FILES_MAP
		__snake_case      :       List[str]     =   PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
		__snake_case      :       Any     =   ["input_ids", "attention_mask"]







		def __init__(					self       ,A__       ,A__="<s>"       ,A__="<pad>"       ,A__="</s>"       ,A__="<unk>"       ,A__=False       ,A__=None       ,**A__       ,):
						super().__init__(
						    unk_token=A__       ,bos_token=A__       ,eos_token=A__       ,pad_token=A__       ,do_lower_case=A__       ,**A__       ,)

						_A	:				int   =      do_lower_case

						with open(A__       ,encoding='''utf-8'''		) as vocab_handle:
										_A	:				str   =      json.load(A__		)
						_A	:				Union[str, Any]   =      {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."""		)

										_A	:				int   =      None
										_A	:				str   =      None
						else:
										with open(A__       ,encoding='''utf-8'''		) as merges_handle:
														_A	:				List[Any]   =      merges_handle.read().split('''\n'''		)[:-1]

										_A	:				int   =      [tuple(merge.split()[:2]		) for merge in merges]
										_A	:				Tuple   =      dict(zip(A__       ,range(len(A__		)		)		)		)
										_A	:				int   =      {}







		@property
		def  A__      (					self		):
						return len(self.decoder		)







		def  A__      (					self		):
						return dict(self.encoder       ,**self.added_tokens_encoder		)







		def  A__      (					self       ,A__		):
						_A	:				Tuple   =      tuple(token[:-1]		) + (token[-1] + BPE_TOKEN_MERGES,)
						if token in self.cache:
										return self.cache[token]
						_A	:				List[Any]   =      get_pairs(A__		)

						if not pairs:
										return token

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

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

						if word.endswith(A__		):
										_A	:				Any   =      word.replace(A__       ,''''''		)

						_A	:				str   =      word.replace(''' '''       ,A__		)
						_A	:				str   =      word
						return word







		def  A__      (					self       ,A__		):

						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:
										_A	:				Optional[int]   =      text.lower()

						_A	:				Optional[int]   =      text.split()

						_A	:				int   =      []
						for token in text:
										if token:
														split_tokens.extend(list(self.bpe(A__		).split(''' '''		)		)		)

						return split_tokens







		def  A__      (					self       ,A__		):
						return self.encoder.get(A__       ,self.encoder.get(self.unk_token		)		)







		def  A__      (					self       ,A__		):
						_A	:				List[str]   =      self.decoder.get(A__       ,self.unk_token		)
						return result







		def  A__      (					self       ,A__		):
						_A	:				str   =      ''' '''.join(A__		)

						# make sure @@ tokens are concatenated
						_A	:				int   =      ''''''.join(string.split(A__		)		)

						return string







		def  A__      (					self       ,A__       ,A__ = None		):
						if not os.path.isdir(A__		):
										logger.error(f"""Vocabulary path ({save_directory}) should be a directory"""		)
										return
						_A	:				Optional[int]   =      os.path.join(
						    A__       ,(filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''vocab_file''']		)
						_A	:				List[str]   =      os.path.join(
						    A__       ,(filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''merges_file''']		)

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

						_A	:				List[Any]   =      0
						if self.bpe_ranks is None:
										return (vocab_file,)

						with open(A__       ,'''w'''       ,encoding='''utf-8'''		) as writer:
										for bpe_tokens, token_index in sorted(self.bpe_ranks.items()       ,key=lambda A__		: 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!'''		)
																		_A	:				Tuple   =      token_index
														writer.write(''' '''.join(A__		) + '''\n'''		)
														index += 1

						return (vocab_file, merges_file)











import flax.linen as nn
import jax.numpy as jnp

from .attention_flax import FlaxTransformeraDModel
from .resnet_flax import FlaxDownsampleaD, FlaxResnetBlockaD, FlaxUpsampleaD






class 						UpperCAmelCase__ (							nn.Module							):
		__snake_case      :       int
		__snake_case      :       int
		__snake_case      :       float     =   0.0
		__snake_case      :       int     =   1
		__snake_case      :       int     =   1
		__snake_case      :       bool     =   True
		__snake_case      :       bool     =   False
		__snake_case      :       bool     =   False
		__snake_case      :       bool     =   False
		__snake_case      :       jnp.dtype     =   jnp.floataa







		def  A__      (					self		):
						_A	:				Optional[Any]   =      []
						_A	:				str   =      []

						for i in range(self.num_layers		):
										_A	:				Union[str, Any]   =      self.in_channels if i == 0 else self.out_channels

										_A	:				str   =      FlaxResnetBlockaD(
										    in_channels=A__       ,out_channels=self.out_channels       ,dropout_prob=self.dropout       ,dtype=self.dtype       ,)
										resnets.append(A__		)

										_A	:				List[str]   =      FlaxTransformeraDModel(
										    in_channels=self.out_channels       ,n_heads=self.num_attention_heads       ,d_head=self.out_channels // self.num_attention_heads       ,depth=1       ,use_linear_projection=self.use_linear_projection       ,only_cross_attention=self.only_cross_attention       ,use_memory_efficient_attention=self.use_memory_efficient_attention       ,dtype=self.dtype       ,)
										attentions.append(A__		)

						_A	:				List[str]   =      resnets
						_A	:				Any   =      attentions

						if self.add_downsample:
										_A	:				Any   =      FlaxDownsampleaD(self.out_channels       ,dtype=self.dtype		)







		def __call__(					self       ,A__       ,A__       ,A__       ,A__=True		):
						_A	:				List[str]   =      ()

						for resnet, attn in zip(self.resnets       ,self.attentions		):
										_A	:				Optional[Any]   =      resnet(A__       ,A__       ,deterministic=A__		)
										_A	:				Tuple   =      attn(A__       ,A__       ,deterministic=A__		)
										output_states += (hidden_states,)

						if self.add_downsample:
										_A	:				Optional[int]   =      self.downsamplers_a(A__		)
										output_states += (hidden_states,)

						return hidden_states, output_states






class 						UpperCAmelCase__ (							nn.Module							):
		__snake_case      :       int
		__snake_case      :       int
		__snake_case      :       float     =   0.0
		__snake_case      :       int     =   1
		__snake_case      :       bool     =   True
		__snake_case      :       jnp.dtype     =   jnp.floataa







		def  A__      (					self		):
						_A	:				List[Any]   =      []

						for i in range(self.num_layers		):
										_A	:				int   =      self.in_channels if i == 0 else self.out_channels

										_A	:				List[Any]   =      FlaxResnetBlockaD(
										    in_channels=A__       ,out_channels=self.out_channels       ,dropout_prob=self.dropout       ,dtype=self.dtype       ,)
										resnets.append(A__		)
						_A	:				Tuple   =      resnets

						if self.add_downsample:
										_A	:				Optional[Any]   =      FlaxDownsampleaD(self.out_channels       ,dtype=self.dtype		)







		def __call__(					self       ,A__       ,A__       ,A__=True		):
						_A	:				List[Any]   =      ()

						for resnet in self.resnets:
										_A	:				Optional[int]   =      resnet(A__       ,A__       ,deterministic=A__		)
										output_states += (hidden_states,)

						if self.add_downsample:
										_A	:				List[Any]   =      self.downsamplers_a(A__		)
										output_states += (hidden_states,)

						return hidden_states, output_states






class 						UpperCAmelCase__ (							nn.Module							):
		__snake_case      :       int
		__snake_case      :       int
		__snake_case      :       int
		__snake_case      :       float     =   0.0
		__snake_case      :       int     =   1
		__snake_case      :       int     =   1
		__snake_case      :       bool     =   True
		__snake_case      :       bool     =   False
		__snake_case      :       bool     =   False
		__snake_case      :       bool     =   False
		__snake_case      :       jnp.dtype     =   jnp.floataa







		def  A__      (					self		):
						_A	:				List[str]   =      []
						_A	:				Optional[Any]   =      []

						for i in range(self.num_layers		):
										_A	:				Dict   =      self.in_channels if (i == self.num_layers - 1) else self.out_channels
										_A	:				List[str]   =      self.prev_output_channel if i == 0 else self.out_channels

										_A	:				Optional[int]   =      FlaxResnetBlockaD(
										    in_channels=resnet_in_channels + res_skip_channels       ,out_channels=self.out_channels       ,dropout_prob=self.dropout       ,dtype=self.dtype       ,)
										resnets.append(A__		)

										_A	:				Union[str, Any]   =      FlaxTransformeraDModel(
										    in_channels=self.out_channels       ,n_heads=self.num_attention_heads       ,d_head=self.out_channels // self.num_attention_heads       ,depth=1       ,use_linear_projection=self.use_linear_projection       ,only_cross_attention=self.only_cross_attention       ,use_memory_efficient_attention=self.use_memory_efficient_attention       ,dtype=self.dtype       ,)
										attentions.append(A__		)

						_A	:				Dict   =      resnets
						_A	:				int   =      attentions

						if self.add_upsample:
										_A	:				str   =      FlaxUpsampleaD(self.out_channels       ,dtype=self.dtype		)







		def __call__(					self       ,A__       ,A__       ,A__       ,A__       ,A__=True		):
						for resnet, attn in zip(self.resnets       ,self.attentions		):
										# pop res hidden states
										_A	:				List[Any]   =      res_hidden_states_tuple[-1]
										_A	:				int   =      res_hidden_states_tuple[:-1]
										_A	:				List[Any]   =      jnp.concatenate((hidden_states, res_hidden_states)       ,axis=-1		)

										_A	:				str   =      resnet(A__       ,A__       ,deterministic=A__		)
										_A	:				Optional[int]   =      attn(A__       ,A__       ,deterministic=A__		)

						if self.add_upsample:
										_A	:				Union[str, Any]   =      self.upsamplers_a(A__		)

						return hidden_states






class 						UpperCAmelCase__ (							nn.Module							):
		__snake_case      :       int
		__snake_case      :       int
		__snake_case      :       int
		__snake_case      :       float     =   0.0
		__snake_case      :       int     =   1
		__snake_case      :       bool     =   True
		__snake_case      :       jnp.dtype     =   jnp.floataa







		def  A__      (					self		):
						_A	:				Optional[Any]   =      []

						for i in range(self.num_layers		):
										_A	:				Any   =      self.in_channels if (i == self.num_layers - 1) else self.out_channels
										_A	:				Optional[Any]   =      self.prev_output_channel if i == 0 else self.out_channels

										_A	:				List[Any]   =      FlaxResnetBlockaD(
										    in_channels=resnet_in_channels + res_skip_channels       ,out_channels=self.out_channels       ,dropout_prob=self.dropout       ,dtype=self.dtype       ,)
										resnets.append(A__		)

						_A	:				int   =      resnets

						if self.add_upsample:
										_A	:				List[str]   =      FlaxUpsampleaD(self.out_channels       ,dtype=self.dtype		)







		def __call__(					self       ,A__       ,A__       ,A__       ,A__=True		):
						for resnet in self.resnets:
										# pop res hidden states
										_A	:				Tuple   =      res_hidden_states_tuple[-1]
										_A	:				Optional[int]   =      res_hidden_states_tuple[:-1]
										_A	:				Any   =      jnp.concatenate((hidden_states, res_hidden_states)       ,axis=-1		)

										_A	:				Optional[Any]   =      resnet(A__       ,A__       ,deterministic=A__		)

						if self.add_upsample:
										_A	:				Tuple   =      self.upsamplers_a(A__		)

						return hidden_states






class 						UpperCAmelCase__ (							nn.Module							):
		__snake_case      :       int
		__snake_case      :       float     =   0.0
		__snake_case      :       int     =   1
		__snake_case      :       int     =   1
		__snake_case      :       bool     =   False
		__snake_case      :       bool     =   False
		__snake_case      :       jnp.dtype     =   jnp.floataa







		def  A__      (					self		):
						# there is always at least one resnet
						_A	:				List[str]   =      [
						    FlaxResnetBlockaD(
						        in_channels=self.in_channels       ,out_channels=self.in_channels       ,dropout_prob=self.dropout       ,dtype=self.dtype       ,)
						]

						_A	:				Union[str, Any]   =      []

						for _ in range(self.num_layers		):
										_A	:				Any   =      FlaxTransformeraDModel(
										    in_channels=self.in_channels       ,n_heads=self.num_attention_heads       ,d_head=self.in_channels // self.num_attention_heads       ,depth=1       ,use_linear_projection=self.use_linear_projection       ,use_memory_efficient_attention=self.use_memory_efficient_attention       ,dtype=self.dtype       ,)
										attentions.append(A__		)

										_A	:				Any   =      FlaxResnetBlockaD(
										    in_channels=self.in_channels       ,out_channels=self.in_channels       ,dropout_prob=self.dropout       ,dtype=self.dtype       ,)
										resnets.append(A__		)

						_A	:				List[str]   =      resnets
						_A	:				str   =      attentions







		def __call__(					self       ,A__       ,A__       ,A__       ,A__=True		):
						_A	:				Optional[int]   =      self.resnets[0](A__       ,A__		)
						for attn, resnet in zip(self.attentions       ,self.resnets[1:]		):
										_A	:				Dict   =      attn(A__       ,A__       ,deterministic=A__		)
										_A	:				Any   =      resnet(A__       ,A__       ,deterministic=A__		)

						return hidden_states



