






import json
import os
import re
import warnings
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											=      logging.get_logger(__name__)

UpperCamelCase											=      {
    '''source_spm''': '''source.spm''',
    '''target_spm''': '''target.spm''',
    '''vocab''': '''vocab.json''',
    '''target_vocab_file''': '''target_vocab.json''',
    '''tokenizer_config_file''': '''tokenizer_config.json''',
}

UpperCamelCase											=      {
    '''source_spm''': {
        '''Helsinki-NLP/opus-mt-en-de''': '''https://huggingface.co/Helsinki-NLP/opus-mt-en-de/resolve/main/source.spm'''
    },
    '''target_spm''': {
        '''Helsinki-NLP/opus-mt-en-de''': '''https://huggingface.co/Helsinki-NLP/opus-mt-en-de/resolve/main/target.spm'''
    },
    '''vocab''': {
        '''Helsinki-NLP/opus-mt-en-de''': '''https://huggingface.co/Helsinki-NLP/opus-mt-en-de/resolve/main/vocab.json'''
    },
    '''tokenizer_config_file''': {
        '''Helsinki-NLP/opus-mt-en-de''': (
            '''https://huggingface.co/Helsinki-NLP/opus-mt-en-de/resolve/main/tokenizer_config.json'''
        )
    },
}

UpperCamelCase											=      {'''Helsinki-NLP/opus-mt-en-de''': 5_1_2}
UpperCamelCase											=      {}




class 	lowerCamelCase__				(							UpperCAmelCase ):
      lowerCamelCase_							:   Optional[Any]											=	VOCAB_FILES_NAMES
      lowerCamelCase_							:   List[Any]											=	PRETRAINED_VOCAB_FILES_MAP
      lowerCamelCase_							:   Optional[int]											=	PRETRAINED_INIT_CONFIGURATION
      lowerCamelCase_							:   str											=	PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
      lowerCamelCase_							:   Union[str, Any]											=	['input_ids', 'attention_mask']
      lowerCamelCase_							:   str											=	re.compile('>>.+<<' )  # type: re.Pattern






      def __init__(self					:    int		,		_snake_case					:    List[Any]		,		_snake_case					:    List[Any]		,		_snake_case					:    Any		,		_snake_case					:    Union[str, Any]=None		,		_snake_case					:    List[str]=None		,		_snake_case					:    Dict=None		,		_snake_case					:    Optional[Any]="<unk>"		,		_snake_case					:    Dict="</s>"		,		_snake_case					:    Dict="<pad>"		,		_snake_case					:    List[str]=512		,		_snake_case					:    Optional[Dict[str, Any]] = None		,		_snake_case					:    Tuple=False		,		**_snake_case					:    str		,		)       ->   None:


       """simple docstring"""






       lowerCamelCase_							:    int							=     {} if sp_model_kwargs is None else sp_model_kwargs

       super().__init__(
           # bos_token=bos_token,  unused. Start decoding with config.decoder_start_token_id
           source_lang=_snake_case		,		target_lang=_snake_case		,		unk_token=_snake_case		,		eos_token=_snake_case		,		pad_token=_snake_case		,		model_max_length=_snake_case		,		sp_model_kwargs=self.sp_model_kwargs		,		target_vocab_file=_snake_case		,		separate_vocabs=_snake_case		,		**_snake_case		,		)
       assert Path(_snake_case		).exists(), f'cannot find spm source {source_spm}'

       lowerCamelCase_							:    Optional[Any]							=     separate_vocabs
       lowerCamelCase_							:    str							=     load_json(_snake_case		)
       if self.unk_token not in self.encoder:
        raise KeyError('<unk> token must be in vocab'		)
       assert self.pad_token in self.encoder

       if separate_vocabs:
        lowerCamelCase_							:    Dict							=     load_json(_snake_case		)
        lowerCamelCase_							:    Dict							=     {v: k for k, v in self.target_encoder.items()}
        lowerCamelCase_							:    Union[str, Any]							=     []
       else:
        lowerCamelCase_							:    int							=     {v: k for k, v in self.encoder.items()}
        lowerCamelCase_							:    list							=     [k for k in self.encoder if k.startswith('>>'		) and k.endswith('<<'		)]

       lowerCamelCase_							:    Dict							=     source_lang
       lowerCamelCase_							:    Tuple							=     target_lang
       lowerCamelCase_							:    Tuple							=     [source_spm, target_spm]

       # load SentencePiece model for pre-processing
       lowerCamelCase_							:    Optional[Any]							=     load_spm(_snake_case		,		self.sp_model_kwargs		)
       lowerCamelCase_							:    str							=     load_spm(_snake_case		,		self.sp_model_kwargs		)
       lowerCamelCase_							:    Dict							=     self.spm_source
       lowerCamelCase_							:    Optional[int]							=     self.encoder

       # Multilingual target side: default to using first supported language code.

       self._setup_normalizer()






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


       """simple docstring"""






       try:
        from sacremoses import MosesPunctNormalizer

        lowerCamelCase_							:    Tuple							=     MosesPunctNormalizer(self.source_lang		).normalize
       except (ImportError, FileNotFoundError):
        warnings.warn('Recommended: pip install sacremoses.'		)
        lowerCamelCase_							:    Tuple							=     lambda _snake_case		: x






      def      UpperCAmelCase_     (self					:    str		,		_snake_case					:    str		)       ->   str:


       """simple docstring"""






       return self.punc_normalizer(_snake_case		) if x else ""






      def      UpperCAmelCase_     (self					:    Optional[int]		,		_snake_case					:    List[str]		)       ->   Any:


       """simple docstring"""






       return self.current_encoder.get(_snake_case		,		self.current_encoder[self.unk_token]		)






      def      UpperCAmelCase_     (self					:    Optional[Any]		,		_snake_case					:    str		)       ->   Dict:


       """simple docstring"""






       lowerCamelCase_							:    Union[str, Any]							=     self.language_code_re.match(_snake_case		)
       lowerCamelCase_							:    list							=     [match.group(0		)] if match else []
       return code, self.language_code_re.sub(''		,		_snake_case		)






      def      UpperCAmelCase_     (self					:    int		,		_snake_case					:    str		)       ->   List[str]:


       """simple docstring"""






       lowerCamelCase_		,      lowerCamelCase_							:    Union[str, Any]							=     self.remove_language_code(_snake_case		)
       lowerCamelCase_							:    Optional[int]							=     self.current_spm.encode(_snake_case		,		out_type=_snake_case		)
       return code + pieces






      def      UpperCAmelCase_     (self					:    Union[str, Any]		,		_snake_case					:    int		)       ->   str:


       """simple docstring"""






       return self.decoder.get(_snake_case		,		self.unk_token		)






      def      UpperCAmelCase_     (self					:    Tuple		,		_snake_case					:    Any		,		**_snake_case					:    Optional[Any]		)       ->   Union[str, Any]:


       """simple docstring"""






       return super().batch_decode(_snake_case		,		**_snake_case		)






      def      UpperCAmelCase_     (self					:    Optional[int]		,		_snake_case					:    Optional[int]		,		**_snake_case					:    List[str]		)       ->   Optional[int]:


       """simple docstring"""






       return super().decode(_snake_case		,		**_snake_case		)






      def      UpperCAmelCase_     (self					:    Optional[int]		,		_snake_case					:    List[str]		)       ->   str:


       """simple docstring"""






       lowerCamelCase_							:    List[Any]							=     self.spm_source if self._decode_use_source_tokenizer else self.spm_target
       lowerCamelCase_							:    Optional[Any]							=     []
       lowerCamelCase_							:    Optional[int]							=     ''
       for token in tokens:
        # make sure that special tokens are not decoded using sentencepiece model
        if token in self.all_special_tokens:
         out_string += sp_model.decode_pieces(_snake_case		) + token + " "
         lowerCamelCase_							:    Union[str, Any]							=     []
        else:
         current_sub_tokens.append(_snake_case		)
       out_string += sp_model.decode_pieces(_snake_case		)
       return out_string.strip()






      def      UpperCAmelCase_     (self					:    Optional[int]		,		_snake_case					:    List[Any]		,		_snake_case					:    Tuple=None		)       ->   List[int]:


       """simple docstring"""






       if token_ids_a is None:
        return token_ids_a + [self.eos_token_id]
       # We don't expect to process pairs, but leave the pair logic for API consistency
       return token_ids_a + token_ids_a + [self.eos_token_id]






      def      UpperCAmelCase_     (self					:    List[Any]		)       ->   Optional[Any]:


       """simple docstring"""






       lowerCamelCase_							:    Optional[Any]							=     self.spm_source
       lowerCamelCase_							:    Any							=     self.encoder






      def      UpperCAmelCase_     (self					:    List[str]		)       ->   Any:


       """simple docstring"""






       lowerCamelCase_							:    Optional[Any]							=     self.spm_target
       if self.separate_vocabs:
        lowerCamelCase_							:    Any							=     self.target_encoder






      @property
      def      UpperCAmelCase_     (self					:    int		)       ->   int:


       """simple docstring"""






       return len(self.encoder		)






      def      UpperCAmelCase_     (self					:    Optional[int]		,		_snake_case					:    str		,		_snake_case					:    Optional[str] = None		)       ->   Tuple[str]:


       """simple docstring"""






       if not os.path.isdir(_snake_case		):
        logger.error(f'Vocabulary path ({save_directory}) should be a directory'		)
        return
       lowerCamelCase_							:    List[str]							=     []

       if self.separate_vocabs:
        lowerCamelCase_							:    Union[str, Any]							=     os.path.join(
            _snake_case		,		(filename_prefix + '-' if filename_prefix else '') + VOCAB_FILES_NAMES['vocab']		,		)
        lowerCamelCase_							:    Tuple							=     os.path.join(
            _snake_case		,		(filename_prefix + '-' if filename_prefix else '') + VOCAB_FILES_NAMES['target_vocab_file']		,		)
        save_json(self.encoder		,		_snake_case		)
        save_json(self.target_encoder		,		_snake_case		)
        saved_files.append(_snake_case		)
        saved_files.append(_snake_case		)
       else:
        lowerCamelCase_							:    Any							=     os.path.join(
            _snake_case		,		(filename_prefix + '-' if filename_prefix else '') + VOCAB_FILES_NAMES['vocab']		)
        save_json(self.encoder		,		_snake_case		)
        saved_files.append(_snake_case		)

       for spm_save_filename, spm_orig_path, spm_model in zip(
           [VOCAB_FILES_NAMES['source_spm'], VOCAB_FILES_NAMES['target_spm']]		,		self.spm_files		,		[self.spm_source, self.spm_target]		,		):
        lowerCamelCase_							:    str							=     os.path.join(
            _snake_case		,		(filename_prefix + '-' if filename_prefix else '') + spm_save_filename		)
        if os.path.abspath(_snake_case		) != os.path.abspath(_snake_case		) and os.path.isfile(_snake_case		):
         copyfile(_snake_case		,		_snake_case		)
         saved_files.append(_snake_case		)
        elif not os.path.isfile(_snake_case		):
         with open(_snake_case		,		'wb'		) as fi:
          lowerCamelCase_							:    List[Any]							=     spm_model.serialized_model_proto()
          fi.write(_snake_case		)
         saved_files.append(_snake_case		)

       return tuple(_snake_case		)






      def      UpperCAmelCase_     (self					:    Optional[Any]		)       ->   Dict:


       """simple docstring"""






       return self.get_src_vocab()






      def      UpperCAmelCase_     (self					:    Tuple		)       ->   Tuple:


       """simple docstring"""






       return dict(self.encoder		,		**self.added_tokens_encoder		)






      def      UpperCAmelCase_     (self					:    Any		)       ->   Tuple:


       """simple docstring"""






       return dict(self.target_encoder		,		**self.added_tokens_decoder		)






      def __getstate__(self					:    Optional[Any]		)       ->   Dict:


       """simple docstring"""






       lowerCamelCase_							:    Union[str, Any]							=     self.__dict__.copy()
       state.update(
           {k: None for k in ['spm_source', 'spm_target', 'current_spm', 'punc_normalizer', 'target_vocab_file']}		)
       return state






      def __setstate__(self					:    List[str]		,		_snake_case					:    Dict		)       ->   None:


       """simple docstring"""






       lowerCamelCase_							:    Optional[Any]							=     d

       # for backward compatibility
       if not hasattr(self		,		'sp_model_kwargs'		):
        lowerCamelCase_							:    List[str]							=     {}

       lowerCamelCase_		,      lowerCamelCase_							:    List[str]							=     (load_spm(_snake_case		,		self.sp_model_kwargs		) for f in self.spm_files)
       lowerCamelCase_							:    Tuple							=     self.spm_source
       self._setup_normalizer()






      def      UpperCAmelCase_     (self					:    Any		,		*_snake_case					:    Any		,		**_snake_case					:    Tuple		)       ->   Optional[Any]:


       """simple docstring"""






       return 1






      def      UpperCAmelCase_     (self					:    List[str]		,		_snake_case					:    List[Any]		)       ->   str:


       """simple docstring"""






       lowerCamelCase_							:    Tuple							=     set(self.all_special_ids		)  # call it once instead of inside list comp
       all_special_ids.remove(self.unk_token_id		)  # <unk> is only sometimes special
       return [1 if x in all_special_ids else 0 for x in seq]






      def      UpperCAmelCase_     (self					:    Optional[int]		,		_snake_case					:    List		,		_snake_case					:    Optional[List] = None		,		_snake_case					:    bool = False		)       ->   List[int]:


       """simple docstring"""






       if already_has_special_tokens:
        return self._special_token_mask(_snake_case		)
       elif token_ids_a is None:
        return self._special_token_mask(_snake_case		) + [1]
       else:
        return self._special_token_mask(token_ids_a + token_ids_a		) + [1]




def       _a   (							lowerCamelCase__ ,			lowerCamelCase__				)  ->      sentencepiece.SentencePieceProcessor:
 lowerCamelCase_							:    List[Any]							=     sentencepiece.SentencePieceProcessor(**lowerCamelCase__				)
 spm.Load(lowerCamelCase__				)
 return spm




def       _a   (							lowerCamelCase__ ,			lowerCamelCase__				)  ->      None:
 with open(lowerCamelCase__ ,			'w'				) as f:
  json.dump(lowerCamelCase__ ,			lowerCamelCase__ ,			indent=2				)





def       _a   (							lowerCamelCase__				)  ->      Union[Dict, List]:
 with open(lowerCamelCase__ ,			'r'				) as f:
  return json.load(lowerCamelCase__				)








from functools import partial
from typing import Any, Callable, Dict, List, Type, TypeVar, Union, overload

import torch
import torch.nn as nn
import torch.types




def       _a   (							lowerCamelCase__ ,			lowerCamelCase__ ,			lowerCamelCase__				)  ->      torch.Tensor:
 # The first operation in a checkpoint can't be in-place, but it's
 # nice to have in-place addition during inference. Thus...
 if not inplace:
  lowerCamelCase_							:    Tuple							=     ma + ma
 else:
  ma += ma

 return ma




def       _a   (							lowerCamelCase__ ,			lowerCamelCase__				)  ->      torch.Tensor:
 lowerCamelCase_							:    Union[str, Any]							=     -1 * len(lowerCamelCase__				)
 lowerCamelCase_							:    Optional[Any]							=     list(range(len(tensor.shape[:zero_index]				)				)				)
 return tensor.permute(first_inds + [zero_index + i for i in inds]				)




def       _a   (							lowerCamelCase__ ,			lowerCamelCase__				)  ->      torch.Tensor:
 return t.reshape(t.shape[:-no_dims] + (-1,)				)




def       _a   (							lowerCamelCase__ ,			lowerCamelCase__ ,			lowerCamelCase__ ,			lowerCamelCase__ = 1e-4				)  ->      torch.Tensor:
 lowerCamelCase_							:    Optional[int]							=     mask.expand(*value.shape				)
 return torch.sum(mask * value ,			dim=lowerCamelCase__				) / (eps + torch.sum(lowerCamelCase__ ,			dim=lowerCamelCase__				))




def       _a   (							lowerCamelCase__ ,			lowerCamelCase__ = 2.3125 ,			lowerCamelCase__ = 21.6875 ,			lowerCamelCase__ = 64				)  ->      torch.Tensor:
 lowerCamelCase_							:    int							=     torch.linspace(lowerCamelCase__ ,			lowerCamelCase__ ,			no_bins - 1 ,			device=pts.device				)
 lowerCamelCase_							:    Tuple							=     torch.sqrt(torch.sum((pts.unsqueeze(-2				) - pts.unsqueeze(-3				)) ** 2 ,			dim=-1				)				)
 return torch.bucketize(lowerCamelCase__ ,			lowerCamelCase__				)




def       _a   (							lowerCamelCase__ ,			lowerCamelCase__				)  ->      dict:
 lowerCamelCase_							:    Optional[Any]							=     dicts[0]
 lowerCamelCase_							:    Optional[Any]							=     {}
 for k, v in first.items():
  lowerCamelCase_							:    Tuple							=     [d[k] for d in dicts]
  if isinstance(lowerCamelCase__ ,			lowerCamelCase__				):
   lowerCamelCase_							:    Optional[int]							=     dict_multimap(lowerCamelCase__ ,			lowerCamelCase__				)
  else:
   lowerCamelCase_							:    List[Any]							=     fn(lowerCamelCase__				)

 return new_dict




def       _a   (							lowerCamelCase__ ,			lowerCamelCase__				)  ->      torch.Tensor:
 lowerCamelCase_							:    str							=     v_bins.view(((1,) * len(x.shape				)) + (len(lowerCamelCase__				),)				)
 lowerCamelCase_							:    List[Any]							=     x[..., None] - reshaped_bins
 lowerCamelCase_							:    List[Any]							=     torch.argmin(torch.abs(lowerCamelCase__				) ,			dim=-1				)
 return nn.functional.one_hot(lowerCamelCase__ ,			num_classes=len(lowerCamelCase__				)				).float()




def       _a   (							lowerCamelCase__ ,			lowerCamelCase__ ,			lowerCamelCase__ = 0 ,			lowerCamelCase__ = 0				)  ->      torch.Tensor:
 lowerCamelCase_							:    List[Union[slice, torch.Tensor]]							=     []
 for i, s in enumerate(data.shape[:no_batch_dims]				):
  lowerCamelCase_							:    int							=     torch.arange(lowerCamelCase__				)
  lowerCamelCase_							:    Any							=     r.view(*(*((1,) * i), -1, *((1,) * (len(inds.shape				) - i - 1)))				)
  ranges.append(lowerCamelCase__				)

 lowerCamelCase_							:    List[Union[slice, torch.Tensor]]							=     [slice(lowerCamelCase__				) for _ in range(len(data.shape				) - no_batch_dims				)]
 lowerCamelCase_							:    Optional[int]							=     inds
 ranges.extend(lowerCamelCase__				)
 # Matt note: Editing this to get around the behaviour of using a list as an array index changing
 # in recent Numpy versions
 return data[tuple(lowerCamelCase__				)]


UpperCamelCase											=      TypeVar('''T''')




def       _a   (							lowerCamelCase__ ,			lowerCamelCase__ ,			lowerCamelCase__				)  ->      Dict[Any, Union[dict, list, tuple, Any]]:
 lowerCamelCase_							:    Dict[Any, Union[dict, list, tuple, Any]]							=     {}
 for k, v in dic.items():
  if isinstance(lowerCamelCase__ ,			lowerCamelCase__				):
   lowerCamelCase_							:    Optional[Any]							=     dict_map(lowerCamelCase__ ,			lowerCamelCase__ ,			lowerCamelCase__				)
  else:
   lowerCamelCase_							:    Dict							=     tree_map(lowerCamelCase__ ,			lowerCamelCase__ ,			lowerCamelCase__				)

 return new_dict




@overload
def       _a   (							lowerCamelCase__ ,			lowerCamelCase__ ,			lowerCamelCase__				)  ->      Any:
 ...




@overload
def       _a   (							lowerCamelCase__ ,			lowerCamelCase__ ,			lowerCamelCase__				)  ->      dict:
 ...




@overload
def       _a   (							lowerCamelCase__ ,			lowerCamelCase__ ,			lowerCamelCase__				)  ->      list:
 ...




@overload
def       _a   (							lowerCamelCase__ ,			lowerCamelCase__ ,			lowerCamelCase__				)  ->      tuple:
 ...




def       _a   (							lowerCamelCase__ ,			lowerCamelCase__ ,			lowerCamelCase__				)  ->      Any:
 if isinstance(lowerCamelCase__ ,			lowerCamelCase__				):
  return dict_map(lowerCamelCase__ ,			lowerCamelCase__ ,			lowerCamelCase__				)
 elif isinstance(lowerCamelCase__ ,			lowerCamelCase__				):
  return [tree_map(lowerCamelCase__ ,			lowerCamelCase__ ,			lowerCamelCase__				) for x in tree]
 elif isinstance(lowerCamelCase__ ,			lowerCamelCase__				):
  return tuple(tree_map(lowerCamelCase__ ,			lowerCamelCase__ ,			lowerCamelCase__				) for x in tree				)
 elif isinstance(lowerCamelCase__ ,			lowerCamelCase__				):
  return fn(lowerCamelCase__				)
 else:
  print(type(lowerCamelCase__				)				)
  raise ValueError('Not supported'				)





UpperCamelCase											=      partial(tree_map, leaf_type=torch.Tensor)
