



'''simple docstring'''





import logging
import math
from functools import partial
from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple, Union

import torch

from .tensor_utils import tensor_tree_map, tree_map




def lowerCAmelCase  (__A):



     """simple docstring"""



     _a      =  []
     if isinstance(__A  ,	__A):
          for v in tree.values():
               shapes.extend(_fetch_dims(__A))
     elif isinstance(__A  ,	(list, tuple)):
          for t in tree:
               shapes.extend(_fetch_dims(__A))
     elif isinstance(__A  ,	torch.Tensor):
          shapes.append(tree.shape)
     else:
          raise ValueError('''Not supported''')

     return shapes




@torch.jit.ignore
def lowerCAmelCase  (__A  ,	__A):



     """simple docstring"""



     _a      =  []
     for d in reversed(__A):
          idx.append(flat_idx % d)
          _a      =  flat_idx // d

     return tuple(reversed(__A))




@torch.jit.ignore
def lowerCAmelCase  (__A  ,	__A  ,	__A  ,	__A = None  ,	__A = None  ,	):



     """simple docstring"""



     def reduce_edge_list(__A) -> None:
          _a      =  True
          for i in range(len(__A)):
               _a      =  -1 * (i + 1)
               l[reversed_idx] &= tally
               _a      =  l[reversed_idx]

     if start_edges is None:
          _a      =  [s == 0 for s in start]
          reduce_edge_list(__A)
     if end_edges is None:
          _a      =  [e == (d - 1) for e, d in zip(__A  ,	__A)]
          reduce_edge_list(__A)

     # Base cases. Either start/end are empty and we're done, or the final,
     # one-dimensional tensor can be simply sliced
     if len(__A) == 0:
          return [()]
     elif len(__A) == 1:
          return [(slice(start[0]  ,	end[0] + 1),)]

     _a      =  []
     _a      =  []

     # Dimensions common to start and end can be selected directly
     for s, e in zip(__A  ,	__A):
          if s == e:
               path_list.append(slice(__A  ,	s + 1))
          else:
               break

     _a      =  tuple(__A)
     _a      =  len(__A)

     # start == end, and we're done
     if divergence_idx == len(__A):
          return [path]

     def upper() -> Tuple[Tuple[slice, ...], ...]:
          assert start_edges is not None
          assert end_edges is not None

          _a      =  start[divergence_idx]
          return tuple(
              path + (slice(__A  ,	sdi + 1),) + s
              for s in _get_minimal_slice_set(
                  start[divergence_idx + 1 :]  ,	[d - 1 for d in dims[divergence_idx + 1 :]]  ,	dims[divergence_idx + 1 :]  ,	start_edges=start_edges[divergence_idx + 1 :]  ,	end_edges=[True for _ in end_edges[divergence_idx + 1 :]]  ,	))

     def lower() -> Tuple[Tuple[slice, ...], ...]:
          assert start_edges is not None
          assert end_edges is not None

          _a      =  end[divergence_idx]
          return tuple(
              path + (slice(__A  ,	edi + 1),) + s
              for s in _get_minimal_slice_set(
                  [0 for _ in start[divergence_idx + 1 :]]  ,	end[divergence_idx + 1 :]  ,	dims[divergence_idx + 1 :]  ,	start_edges=[True for _ in start_edges[divergence_idx + 1 :]]  ,	end_edges=end_edges[divergence_idx + 1 :]  ,	))

     # If both start and end are at the edges of the subtree rooted at
     # divergence_idx, we can just select the whole subtree at once
     if start_edges[divergence_idx] and end_edges[divergence_idx]:
          slices.append(path + (slice(start[divergence_idx]  ,	end[divergence_idx] + 1),))
     # If just start is at the edge, we can grab almost all of the subtree,
     # treating only the ragged bottom edge as an edge case
     elif start_edges[divergence_idx]:
          slices.append(path + (slice(start[divergence_idx]  ,	end[divergence_idx]),))
          slices.extend(lower())
     # Analogous to the previous case, but the top is ragged this time
     elif end_edges[divergence_idx]:
          slices.extend(upper())
          slices.append(path + (slice(start[divergence_idx] + 1  ,	end[divergence_idx] + 1),))
     # If both sides of the range are ragged, we need to handle both sides
     # separately. If there's contiguous meat in between them, we can index it
     # in one big chunk
     else:
          slices.extend(upper())
          _a      =  end[divergence_idx] - start[divergence_idx]
          if middle_ground > 1:
               slices.append(path + (slice(start[divergence_idx] + 1  ,	end[divergence_idx]),))
          slices.extend(lower())

     return slices




@torch.jit.ignore
def lowerCAmelCase  (__A  ,	__A  ,	__A  ,	__A):



     """simple docstring"""



     _a      =  t.shape[:no_batch_dims]
     _a      =  list(_flat_idx_to_idx(__A  ,	__A))
     # _get_minimal_slice_set is inclusive
     _a      =  list(_flat_idx_to_idx(flat_end - 1  ,	__A))

     # Get an ordered list of slices to perform
     _a      =  _get_minimal_slice_set(
         __A  ,	__A  ,	__A  ,	)

     _a      =  [t[s] for s in slices]

     return torch.cat([s.view((-1,) + t.shape[no_batch_dims:]) for s in sliced_tensors])




def lowerCAmelCase  (__A  ,	__A  ,	__A  ,	__A  ,	__A = False  ,	__A = None  ,	__A = False  ,	):



     """simple docstring"""



     if not (len(__A) > 0):
          raise ValueError('''Must provide at least one input''')

     _a      =  [shape[:no_batch_dims] for shape in _fetch_dims(__A)]
     _a      =  tuple([max(__A) for s in zip(*__A)])

     def _prep_inputs(__A) -> torch.Tensor:
          if not low_mem:
               if not sum(t.shape[:no_batch_dims]) == no_batch_dims:
                    _a      =  t.expand(orig_batch_dims + t.shape[no_batch_dims:])
               _a      =  t.reshape(-1  ,	*t.shape[no_batch_dims:])
          else:
               _a      =  t.expand(orig_batch_dims + t.shape[no_batch_dims:])
          return t

     _a      =  tensor_tree_map(_prep_inputs  ,	__A)
     _a      =  None
     if _out is not None:
          _a      =  tensor_tree_map(lambda __A: t.view([-1] + list(t.shape[no_batch_dims:]))  ,	_out)

     _a      =  1
     for d in orig_batch_dims:
          flat_batch_dim *= d

     _a      =  flat_batch_dim // chunk_size + (flat_batch_dim % chunk_size != 0)

     def _select_chunk(__A) -> torch.Tensor:
          return t[i : i + chunk_size] if t.shape[0] != 1 else t

     _a      =  0
     _a      =  prepped_outputs
     for _ in range(__A):
          # Chunk the input
          if not low_mem:
               _a      =  _select_chunk
          else:
               _a      =  partial(
                   _chunk_slice  ,	flat_start=__A  ,	flat_end=min(__A  ,	i + chunk_size)  ,	no_batch_dims=len(__A)  ,	)

          _a      =  tensor_tree_map(__A  ,	__A)

          # Run the layer on the chunk
          _a      =  layer(**__A)

          # Allocate space for the output
          if out is None:
               _a      =  tensor_tree_map(lambda __A: t.new_zeros((flat_batch_dim,) + t.shape[1:])  ,	__A)

          # Put the chunk in its pre-allocated space
          if isinstance(__A  ,	__A):

               def assign(__A  ,	__A) -> None:
                    for k, v in da.items():
                         if isinstance(__A  ,	__A):
                              assign(__A  ,	da[k])
                         else:
                              if _add_into_out:
                                   v[i : i + chunk_size] += da[k]
                              else:
                                   _a      =  da[k]

               assign(__A  ,	__A)
          elif isinstance(__A  ,	__A):
               for xa, xa in zip(__A  ,	__A):
                    if _add_into_out:
                         xa[i : i + chunk_size] += xa
                    else:
                         _a      =  xa
          elif isinstance(__A  ,	torch.Tensor):
               if _add_into_out:
                    out[i : i + chunk_size] += output_chunk
               else:
                    _a      =  output_chunk
          else:
               raise ValueError('''Not supported''')

          i += chunk_size

     _a      =  tensor_tree_map(lambda __A: t.view(orig_batch_dims + t.shape[1:])  ,	__A)

     return out






class __A			:



 '''simple docstring'''







 def __init__(self ,   A = 512 ,   )    ->     List[str]:


      """simple docstring"""



      _a      =  max_chunk_size
      _a      =  None
      _a      =  None



 def 					a__  (self ,   A ,   A ,   A       )    ->     int:


      """simple docstring"""



      logging.info('''Tuning chunk size...'''       )

      if min_chunk_size >= self.max_chunk_size:
           return min_chunk_size

      _a      =  [2**l for l in range(int(math.log(self.max_chunk_size ,   2       )       ) + 1       )]
      _a      =  [c for c in candidates if c > min_chunk_size]
      _a      =  [min_chunk_size] + candidates
      candidates[-1] += 4

      def test_chunk_size(A       ) -> bool:
           try:
                with torch.no_grad():
                     fn(*A ,   chunk_size=A       )
                return True
           except RuntimeError:
                return False

      _a      =  0
      _a      =  len(A       ) - 1
      while i > min_viable_chunk_size_index:
           _a      =  test_chunk_size(candidates[i]       )
           if not viable:
                _a      =  (min_viable_chunk_size_index + i) // 2
           else:
                _a      =  i
                _a      =  (i + len(A       ) - 1) // 2

      return candidates[min_viable_chunk_size_index]



 def 					a__  (self ,   A ,   A       )    ->     bool:


      """simple docstring"""



      _a      =  True
      for aa, aa in zip(A ,   A       ):
           assert type(A       ) == type(A       )
           if isinstance(A ,   (list, tuple)       ):
                consistent &= self._compare_arg_caches(A ,   A       )
           elif isinstance(A ,   A       ):
                _a      =  [v for _, v in sorted(aa.items() ,   key=lambda A       : x[0]       )]
                _a      =  [v for _, v in sorted(aa.items() ,   key=lambda A       : x[0]       )]
                consistent &= self._compare_arg_caches(A ,   A       )
           else:
                consistent &= aa == aa

      return consistent



 def 					a__  (self ,   A ,   A ,   A ,   )    ->     int:


      """simple docstring"""



      _a      =  True
      _a      =  tree_map(lambda A       : a.shape if isinstance(A ,   torch.Tensor       ) else a ,   A ,   A       )
      if self.cached_arg_data is not None:
           # If args have changed shape/value, we need to re-tune
           assert len(self.cached_arg_data       ) == len(A       )
           _a      =  self._compare_arg_caches(self.cached_arg_data ,   A       )
      else:
           # Otherwise, we can reuse the precomputed value
           _a      =  False

      if not consistent:
           _a      =  self._determine_favorable_chunk_size(
               A ,   A ,   A ,   )
           _a      =  arg_data

      assert self.cached_chunk_size is not None

      return self.cached_chunk_size








'''simple docstring'''





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 BatchEncoding, PreTrainedTokenizer
from ...utils import logging


lowercase_		       =  logging.get_logger(__name__)

lowercase_		       =  "▁"

lowercase_		       =  {
    "vocab_file": "vocab.json",
    "spm_file": "sentencepiece.bpe.model",
    "tokenizer_config_file": "tokenizer_config.json",
}

lowercase_		       =  {
    "vocab_file": {
        "facebook/m2m100_418M": "https://huggingface.co/facebook/m2m100_418M/resolve/main/vocab.json",
        "facebook/m2m100_1.2B": "https://huggingface.co/facebook/m2m100_1.2B/resolve/main/vocab.json",
    },
    "spm_file": {
        "facebook/m2m100_418M": "https://huggingface.co/facebook/m2m100_418M/resolve/main/sentencepiece.bpe.model",
        "facebook/m2m100_1.2B": "https://huggingface.co/facebook/m2m100_1.2B/resolve/main/sentencepiece.bpe.model",
    },
    "tokenizer_config_file": {
        "facebook/m2m100_418M": "https://huggingface.co/facebook/m2m100_418M/resolve/main/tokenizer_config.json",
        "facebook/m2m100_1.2B": "https://huggingface.co/facebook/m2m100_1.2B/resolve/main/tokenizer_config.json",
    },
}

lowercase_		       =  {
    "facebook/m2m100_418M": 1_024,
}

# fmt: off
lowercase_		       =  {
    "m2m100": ["af", "am", "ar", "ast", "az", "ba", "be", "bg", "bn", "br", "bs", "ca", "ceb", "cs", "cy", "da", "de", "el", "en", "es", "et", "fa", "ff", "fi", "fr", "fy", "ga", "gd", "gl", "gu", "ha", "he", "hi", "hr", "ht", "hu", "hy", "id", "ig", "ilo", "is", "it", "ja", "jv", "ka", "kk", "km", "kn", "ko", "lb", "lg", "ln", "lo", "lt", "lv", "mg", "mk", "ml", "mn", "mr", "ms", "my", "ne", "nl", "no", "ns", "oc", "or", "pa", "pl", "ps", "pt", "ro", "ru", "sd", "si", "sk", "sl", "so", "sq", "sr", "ss", "su", "sv", "sw", "ta", "th", "tl", "tn", "tr", "uk", "ur", "uz", "vi", "wo", "xh", "yi", "yo", "zh", "zu"],
    "wmt21": ["en", "ha", "is", "ja", "cs", "ru", "zh", "de"]
}






class __A			(  A						):



 '''simple docstring'''







 __lowerCamelCase    :       Union[str, Any]          =       VOCAB_FILES_NAMES
 __lowerCamelCase    :       Tuple          =       PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
 __lowerCamelCase    :       Tuple          =       PRETRAINED_VOCAB_FILES_MAP
 __lowerCamelCase    :       Dict          =       ['input_ids', 'attention_mask']

 __lowerCamelCase    :       List[int]          =       []
 __lowerCamelCase    :       List[int]          =       []



 def __init__(self ,   A ,   A ,   A=None ,   A=None ,   A="<s>" ,   A="</s>" ,   A="</s>" ,   A="<pad>" ,   A="<unk>" ,   A="m2m100" ,   A = None ,   A=8 ,   **A ,   )    ->     None:


      """simple docstring"""



      _a      =  {} if sp_model_kwargs is None else sp_model_kwargs

      _a      =  language_codes
      _a      =  FAIRSEQ_LANGUAGE_CODES[language_codes]
      _a      =  {lang_code: f'''__{lang_code}__''' for lang_code in fairseq_language_code}

      _a      =  kwargs.get('''additional_special_tokens''' ,   []       )
      kwargs["additional_special_tokens"] += [
          self.get_lang_token(A       )
          for lang_code in fairseq_language_code
          if self.get_lang_token(A       ) not in kwargs["additional_special_tokens"]
      ]

      super().__init__(
          src_lang=A ,   tgt_lang=A ,   bos_token=A ,   eos_token=A ,   sep_token=A ,   unk_token=A ,   pad_token=A ,   language_codes=A ,   sp_model_kwargs=self.sp_model_kwargs ,   num_madeup_words=A ,   **A ,   )

      _a      =  vocab_file
      _a      =  load_json(A       )
      _a      =  {v: k for k, v in self.encoder.items()}
      _a      =  spm_file
      _a      =  load_spm(A ,   self.sp_model_kwargs       )

      _a      =  len(self.encoder       )

      _a      =  {
          self.get_lang_token(A       ): self.encoder_size + i for i, lang_code in enumerate(A       )
      }
      _a      =  {lang_code: self.encoder_size + i for i, lang_code in enumerate(A       )}
      _a      =  {v: k for k, v in self.lang_token_to_id.items()}

      _a      =  src_lang if src_lang is not None else '''en'''
      _a      =  tgt_lang
      _a      =  self.get_lang_id(self._src_lang       )
      self.set_src_lang_special_tokens(self._src_lang       )

      _a      =  num_madeup_words



 @property
 def 					a__  (self       )    ->     int:


      """simple docstring"""



      return len(self.encoder       ) + len(self.lang_token_to_id       )



 @property
 def 					a__  (self       )    ->     str:


      """simple docstring"""



      return self._src_lang



 @src_lang.setter
 def 					a__  (self ,   A       )    ->     None:


      """simple docstring"""



      _a      =  new_src_lang
      self.set_src_lang_special_tokens(self._src_lang       )



 def 					a__  (self ,   A       )    ->     List[str]:


      """simple docstring"""



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



 def 					a__  (self ,   A       )    ->     Union[str, Any]:


      """simple docstring"""



      if token in self.lang_token_to_id:
           return self.lang_token_to_id[token]
      return self.encoder.get(A ,   self.encoder[self.unk_token]       )



 def 					a__  (self ,   A       )    ->     str:


      """simple docstring"""



      if index in self.id_to_lang_token:
           return self.id_to_lang_token[index]
      return self.decoder.get(A ,   self.unk_token       )



 def 					a__  (self ,   A       )    ->     Dict:


      """simple docstring"""



      _a      =  []
      _a      =  ''''''
      for token in tokens:
           # make sure that special tokens are not decoded using sentencepiece model
           if token in self.all_special_tokens:
                out_string += self.sp_model.decode(A       ) + token
                _a      =  []
           else:
                current_sub_tokens.append(A       )
      out_string += self.sp_model.decode(A       )
      return out_string.strip()



 def 					a__  (self ,   A ,   A = None ,   A = False       )    ->     List[int]:


      """simple docstring"""



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

      _a      =  [1] * len(self.prefix_tokens       )
      _a      =  [1] * len(self.suffix_tokens       )
      if token_ids_a is None:
           return prefix_ones + ([0] * len(A       )) + suffix_ones
      return prefix_ones + ([0] * len(A       )) + ([0] * len(A       )) + suffix_ones



 def 					a__  (self ,   A ,   A = None       )    ->     List[int]:


      """simple docstring"""



      if token_ids_a is None:
           return self.prefix_tokens + token_ids_a + self.suffix_tokens
      # 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.suffix_tokens



 def 					a__  (self       )    ->     Dict:


      """simple docstring"""



      _a      =  {self.convert_ids_to_tokens(A       ): i for i in range(self.vocab_size       )}
      vocab.update(self.added_tokens_encoder       )
      return vocab



 def __getstate__(self       )    ->     Dict:


      """simple docstring"""



      _a      =  self.__dict__.copy()
      _a      =  None
      return state



 def __setstate__(self ,   A       )    ->     None:


      """simple docstring"""



      _a      =  d

      # for backward compatibility
      if not hasattr(self ,   '''sp_model_kwargs'''       ):
           _a      =  {}

      _a      =  load_spm(self.spm_file ,   self.sp_model_kwargs       )



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


      """simple docstring"""



      _a      =  Path(A       )
      if not save_dir.is_dir():
           raise OSError(f'''{save_directory} should be a directory'''       )
      _a      =  save_dir / (
          (filename_prefix + '''-''' if filename_prefix else '''''') + self.vocab_files_names['''vocab_file''']
      )
      _a      =  save_dir / (
          (filename_prefix + '''-''' if filename_prefix else '''''') + self.vocab_files_names['''spm_file''']
      )

      save_json(self.encoder ,   A       )

      if os.path.abspath(self.spm_file       ) != os.path.abspath(A       ) and os.path.isfile(self.spm_file       ):
           copyfile(self.spm_file ,   A       )
      elif not os.path.isfile(self.spm_file       ):
           with open(A ,   '''wb'''       ) as fi:
                _a      =  self.sp_model.serialized_model_proto()
                fi.write(A       )

      return (str(A       ), str(A       ))



 def 					a__  (self ,   A ,   A = "en" ,   A = None ,   A = "ro" ,   **A ,   )    ->     BatchEncoding:


      """simple docstring"""



      _a      =  src_lang
      _a      =  tgt_lang
      self.set_src_lang_special_tokens(self.src_lang       )
      return super().prepare_seqaseq_batch(A ,   A ,   **A       )



 def 					a__  (self ,   A ,   A ,   A ,   **A       )    ->     Union[str, Any]:


      """simple docstring"""



      if src_lang is None or tgt_lang is None:
           raise ValueError('''Translation requires a `src_lang` and a `tgt_lang` for this model'''       )
      _a      =  src_lang
      _a      =  self(A ,   add_special_tokens=A ,   **A       )
      _a      =  self.get_lang_id(A       )
      _a      =  tgt_lang_id
      return inputs



 def 					a__  (self       )    ->     Optional[Any]:


      """simple docstring"""



      self.set_src_lang_special_tokens(self.src_lang       )



 def 					a__  (self       )    ->     Tuple:


      """simple docstring"""



      self.set_tgt_lang_special_tokens(self.tgt_lang       )



 def 					a__  (self ,   A       )    ->     None:


      """simple docstring"""



      _a      =  self.get_lang_token(A       )
      _a      =  self.lang_token_to_id[lang_token]
      _a      =  [self.cur_lang_id]
      _a      =  [self.eos_token_id]



 def 					a__  (self ,   A       )    ->     None:


      """simple docstring"""



      _a      =  self.get_lang_token(A       )
      _a      =  self.lang_token_to_id[lang_token]
      _a      =  [self.cur_lang_id]
      _a      =  [self.eos_token_id]



 def 					a__  (self ,   A       )    ->     str:


      """simple docstring"""



      return self.lang_code_to_token[lang]



 def 					a__  (self ,   A       )    ->     int:


      """simple docstring"""



      _a      =  self.get_lang_token(A       )
      return self.lang_token_to_id[lang_token]





def lowerCAmelCase  (__A  ,	__A):



     """simple docstring"""



     _a      =  sentencepiece.SentencePieceProcessor(**__A)
     spm.Load(str(__A))
     return spm




def lowerCAmelCase  (__A):



     """simple docstring"""



     with open(__A  ,	'''r''') as f:
          return json.load(__A)




def lowerCAmelCase  (__A  ,	__A):



     """simple docstring"""



     with open(__A  ,	'''w''') as f:
          json.dump(__A  ,	__A  ,	indent=2)



