




import json
import os
from pathlib import Path
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple, Union

import sentencepiece

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


_UpperCAmelCase	:						List[Any]			=     logging.get_logger(__name__)

_UpperCAmelCase	:						Dict			=     "▁"

_UpperCAmelCase	:						Any			=     {
    "vocab_file": "vocab.json",
    "spm_file": "sentencepiece.bpe.model",
}

_UpperCAmelCase	:						Union[str, Any]			=     {
    "vocab_file": {
        "facebook/s2t-small-librispeech-asr": (
            "https://huggingface.co/facebook/s2t-small-librispeech-asr/resolve/main/vocab.json"
        ),
    },
    "spm_file": {
        "facebook/s2t-small-librispeech-asr": (
            "https://huggingface.co/facebook/s2t-small-librispeech-asr/resolve/main/sentencepiece.bpe.model"
        )
    },
}

_UpperCAmelCase	:						int			=     {
    "facebook/s2t-small-librispeech-asr": 1024,
}

_UpperCAmelCase	:						Optional[Any]			=     ["pt", "fr", "ru", "nl", "ro", "it", "es", "de"]

_UpperCAmelCase	:						int			=     {"mustc": MUSTC_LANGS}




class    __lowerCAmelCase       (							lowerCAmelCase):
       _a 					=						VOCAB_FILES_NAMES
       _a 					=						PRETRAINED_VOCAB_FILES_MAP
       _a 					=						MAX_MODEL_INPUT_SIZES
       _a 					=						['''input_ids''', '''attention_mask''']

       _a 					=						[]






       def __init__(   self:					Optional[Any]       ,					_lowerCAmelCase:					Optional[int]       ,					_lowerCAmelCase:					List[Any]       ,					_lowerCAmelCase:					str="<s>"       ,					_lowerCAmelCase:					Optional[int]="</s>"       ,					_lowerCAmelCase:					Any="<pad>"       ,					_lowerCAmelCase:					Tuple="<unk>"       ,					_lowerCAmelCase:					List[str]=False       ,					_lowerCAmelCase:					List[str]=False       ,					_lowerCAmelCase:					Union[str, Any]=None       ,					_lowerCAmelCase:					Dict=None       ,					_lowerCAmelCase:					Optional[Dict[str, Any]] = None       ,					**_lowerCAmelCase:					List[str]       ,					):
         lowercase    :Union[str, Any]      	=						{} if sp_model_kwargs is None else sp_model_kwargs

         super().__init__(
             bos_token=_lowerCAmelCase       ,					eos_token=_lowerCAmelCase       ,					unk_token=_lowerCAmelCase       ,					pad_token=_lowerCAmelCase       ,					do_upper_case=_lowerCAmelCase       ,					do_lower_case=_lowerCAmelCase       ,					tgt_lang=_lowerCAmelCase       ,					lang_codes=_lowerCAmelCase       ,					sp_model_kwargs=self.sp_model_kwargs       ,					**_lowerCAmelCase       ,					)
         lowercase    :Tuple      	=						do_upper_case
         lowercase    :int      	=						do_lower_case

         lowercase    :str      	=						load_json(_lowerCAmelCase				)
         lowercase    :int      	=						{v: k for k, v in self.encoder.items()}
         lowercase    :Dict      	=						spm_file
         lowercase    :str      	=						load_spm(_lowerCAmelCase       ,					self.sp_model_kwargs				)

         if lang_codes is not None:
           lowercase    :List[str]      	=						lang_codes
           lowercase    :Union[str, Any]      	=						LANGUAGES[lang_codes]
           lowercase    :Any      	=						[F"<lang:{lang}>" for lang in self.langs]
           lowercase    :List[str]      	=						{lang: self.sp_model.PieceToId(F"<lang:{lang}>"				) for lang in self.langs}

           lowercase    :int      	=						self.lang_tokens
           lowercase    :Dict      	=						tgt_lang if tgt_lang is not None else self.langs[0]

           self.set_tgt_lang_special_tokens(self._tgt_lang				)
         else:
           lowercase    :int      	=						{}






       @property
       def 					SCREAMING_SNAKE_CASE    (   self:					List[Any]				):
         return len(self.encoder				)






       @property
       def 					SCREAMING_SNAKE_CASE    (   self:					Any				):
         return self._tgt_lang






       @tgt_lang.setter
       def 					SCREAMING_SNAKE_CASE    (   self:					List[Any]       ,					_lowerCAmelCase:					str				):
         lowercase    :Optional[int]      	=						new_tgt_lang
         self.set_tgt_lang_special_tokens(_lowerCAmelCase				)






       def 					SCREAMING_SNAKE_CASE    (   self:					List[str]       ,					_lowerCAmelCase:					str				):
         lowercase    :Optional[Any]      	=						self.lang_code_to_id[tgt_lang]
         lowercase    :Tuple      	=						[lang_code_id]






       def 					SCREAMING_SNAKE_CASE    (   self:					List[Any]       ,					_lowerCAmelCase:					str				):
         return self.sp_model.encode(_lowerCAmelCase       ,					out_type=_lowerCAmelCase				)






       def 					SCREAMING_SNAKE_CASE    (   self:					Optional[int]       ,					_lowerCAmelCase:					int				):
         return self.encoder.get(_lowerCAmelCase       ,					self.encoder[self.unk_token]				)






       def 					SCREAMING_SNAKE_CASE    (   self:					int       ,					_lowerCAmelCase:					int				):
         return self.decoder.get(_lowerCAmelCase       ,					self.unk_token				)






       def 					SCREAMING_SNAKE_CASE    (   self:					int       ,					_lowerCAmelCase:					List[str]				):
         lowercase    :int      	=						[]
         lowercase    :Any      	=						""
         for token in tokens:
           # make sure that special tokens are not decoded using sentencepiece model
           if token in self.all_special_tokens:
             lowercase    :Optional[Any]      	=						self.sp_model.decode(_lowerCAmelCase				)
             out_string += (decoded.upper() if self.do_upper_case else decoded) + token + " "
             lowercase    :List[Any]      	=						[]
           else:
             current_sub_tokens.append(_lowerCAmelCase				)
         lowercase    :Dict      	=						self.sp_model.decode(_lowerCAmelCase				)
         out_string += decoded.upper() if self.do_upper_case else decoded
         return out_string.strip()






       def 					SCREAMING_SNAKE_CASE    (   self:					Union[str, Any]       ,					_lowerCAmelCase:					List[str]       ,					_lowerCAmelCase:					Any=None				):
         if token_ids_a is None:
           return self.prefix_tokens + token_ids_a + [self.eos_token_id]
         # We don't expect to process pairs, but leave the pair logic for API consistency
         return self.prefix_tokens + token_ids_a + token_ids_a + [self.eos_token_id]






       def 					SCREAMING_SNAKE_CASE    (   self:					Union[str, Any]       ,					_lowerCAmelCase:					List[int]       ,					_lowerCAmelCase:					Optional[List[int]] = None       ,					_lowerCAmelCase:					bool = False				):

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

         lowercase    :Union[str, Any]      	=						[1] * len(self.prefix_tokens				)
         lowercase    :Optional[int]      	=						[1]
         if token_ids_a is None:
           return prefix_ones + ([0] * len(_lowerCAmelCase				)) + suffix_ones
         return prefix_ones + ([0] * len(_lowerCAmelCase				)) + ([0] * len(_lowerCAmelCase				)) + suffix_ones






       def 					SCREAMING_SNAKE_CASE    (   self:					Dict				):
         lowercase    :int      	=						self.encoder.copy()
         vocab.update(self.added_tokens_encoder				)
         return vocab






       def __getstate__(   self:					Tuple				):
         lowercase    :Tuple      	=						self.__dict__.copy()
         lowercase    :Any      	=						None
         return state






       def __setstate__(   self:					Optional[int]       ,					_lowerCAmelCase:					Dict				):
         lowercase    :List[Any]      	=						d

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

         lowercase    :List[str]      	=						load_spm(self.spm_file       ,					self.sp_model_kwargs				)






       def 					SCREAMING_SNAKE_CASE    (   self:					int       ,					_lowerCAmelCase:					str       ,					_lowerCAmelCase:					Optional[str] = None				):
         lowercase    :Optional[Any]      	=						Path(_lowerCAmelCase				)
         assert save_dir.is_dir(), F"{save_directory} should be a directory"
         lowercase    :Union[str, Any]      	=						save_dir / (
             (filename_prefix + "-" if filename_prefix else "") + self.vocab_files_names["vocab_file"]
         )
         lowercase    :List[Any]      	=						save_dir / (
             (filename_prefix + "-" if filename_prefix else "") + self.vocab_files_names["spm_file"]
         )

         save_json(self.encoder       ,					_lowerCAmelCase				)

         if os.path.abspath(self.spm_file				) != os.path.abspath(_lowerCAmelCase				) and os.path.isfile(self.spm_file				):
           copyfile(self.spm_file       ,					_lowerCAmelCase				)
         elif not os.path.isfile(self.spm_file				):
           with open(_lowerCAmelCase       ,					"wb"				) as fi:
             lowercase    :Optional[int]      	=						self.sp_model.serialized_model_proto()
             fi.write(_lowerCAmelCase				)

         return (str(_lowerCAmelCase				), str(_lowerCAmelCase				))







def       UpperCAmelCase__  (     lowerCamelCase,			lowerCamelCase   ):
  lowercase    :Tuple      	=						sentencepiece.SentencePieceProcessor(**lowerCamelCase   )
  spm.Load(str(lowerCamelCase   )   )
  return spm



def       UpperCAmelCase__  (     lowerCamelCase   ):
  with open(lowerCamelCase,			"r"   ) as f:
    return json.load(lowerCamelCase   )







def       UpperCAmelCase__  (     lowerCamelCase,			lowerCamelCase   ):
  with open(lowerCamelCase,			"w"   ) as f:
    json.dump(lowerCamelCase,			lowerCamelCase,			indent=2   )









import re

import jax.numpy as jnp
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.random import PRNGKey

from ..utils import logging


_UpperCAmelCase	:						List[Any]			=     logging.get_logger(__name__)



def       UpperCAmelCase__  (     lowerCamelCase   ):
  lowercase    :int      	=						r"\w+[.]\d+"
  lowercase    :Tuple      	=						re.findall(lowerCamelCase,			lowerCamelCase   )
  for pat in pats:
    lowercase    :List[str]      	=						key.replace(lowerCamelCase,			"_".join(pat.split("."   )   )   )
  return key



def       UpperCAmelCase__  (     lowerCamelCase,			lowerCamelCase,			lowerCamelCase   ):
  lowercase    :List[str]      	=						pt_tuple_key[:-1] + ("scale",)
  if (
      any("norm" in str_ for str_ in pt_tuple_key   )
      and (pt_tuple_key[-1] == "bias")
      and (pt_tuple_key[:-1] + ("bias",) not in random_flax_state_dict)
      and (pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict)
  ):
    lowercase    :Optional[int]      	=						pt_tuple_key[:-1] + ("scale",)
    return renamed_pt_tuple_key, pt_tensor
  elif pt_tuple_key[-1] in ["weight", "gamma"] and pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict:
    lowercase    :int      	=						pt_tuple_key[:-1] + ("scale",)
    return renamed_pt_tuple_key, pt_tensor

  # embedding
  if pt_tuple_key[-1] == "weight" and pt_tuple_key[:-1] + ("embedding",) in random_flax_state_dict:
    lowercase    :List[str]      	=						pt_tuple_key[:-1] + ("embedding",)
    return renamed_pt_tuple_key, pt_tensor

  # conv layer
  lowercase    :Union[str, Any]      	=						pt_tuple_key[:-1] + ("kernel",)
  if pt_tuple_key[-1] == "weight" and pt_tensor.ndim == 4:
    lowercase    :List[Any]      	=						pt_tensor.transpose(2,			3,			1,			0   )
    return renamed_pt_tuple_key, pt_tensor

  # linear layer
  lowercase    :Optional[Any]      	=						pt_tuple_key[:-1] + ("kernel",)
  if pt_tuple_key[-1] == "weight":
    lowercase    :Tuple      	=						pt_tensor.T
    return renamed_pt_tuple_key, pt_tensor

  # old PyTorch layer norm weight
  lowercase    :Tuple      	=						pt_tuple_key[:-1] + ("weight",)
  if pt_tuple_key[-1] == "gamma":
    return renamed_pt_tuple_key, pt_tensor

  # old PyTorch layer norm bias
  lowercase    :int      	=						pt_tuple_key[:-1] + ("bias",)
  if pt_tuple_key[-1] == "beta":
    return renamed_pt_tuple_key, pt_tensor

  return pt_tuple_key, pt_tensor







def       UpperCAmelCase__  (     lowerCamelCase,			lowerCamelCase,			lowerCamelCase=42   ):
  # Step 1: Convert pytorch tensor to numpy
  lowercase    :Optional[int]      	=						{k: v.numpy() for k, v in pt_state_dict.items()}

  # Step 2: Since the model is stateless, get random Flax params
  lowercase    :str      	=						flax_model.init_weights(PRNGKey(lowerCamelCase   )   )

  lowercase    :Tuple      	=						flatten_dict(lowerCamelCase   )
  lowercase    :Optional[Any]      	=						{}

  # Need to change some parameters name to match Flax names
  for pt_key, pt_tensor in pt_state_dict.items():
    lowercase    :List[Any]      	=						rename_key(lowerCamelCase   )
    lowercase    :List[Any]      	=						tuple(renamed_pt_key.split("."   )   )

    # Correctly rename weight parameters
    lowercase	,    lowercase    :List[str]      	=						rename_key_and_reshape_tensor(lowerCamelCase,			lowerCamelCase,			lowerCamelCase   )

    if flax_key in random_flax_state_dict:
      if flax_tensor.shape != random_flax_state_dict[flax_key].shape:
        raise ValueError(
            F"PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape "
            F"{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}."   )

        # also add unexpected weight so that warning is thrown
    lowercase    :List[str]      	=						jnp.asarray(lowerCamelCase   )

  return unflatten_dict(lowerCamelCase   )



