



'''simple docstring'''





import os
import re
from shutil import copyfile
from typing import List, Optional, Tuple

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


_lowerCAmelCase      						=	logging.get_logger(__name__)

_lowerCAmelCase      						=	{
    '''vocab_file''': '''vocab.txt''',
    '''merges_file''': '''bpe.codes''',
}

_lowerCAmelCase      						=	{
    '''vocab_file''': {
        '''vinai/phobert-base''': '''https://huggingface.co/vinai/phobert-base/resolve/main/vocab.txt''',
        '''vinai/phobert-large''': '''https://huggingface.co/vinai/phobert-large/resolve/main/vocab.txt''',
    },
    '''merges_file''': {
        '''vinai/phobert-base''': '''https://huggingface.co/vinai/phobert-base/resolve/main/bpe.codes''',
        '''vinai/phobert-large''': '''https://huggingface.co/vinai/phobert-large/resolve/main/bpe.codes''',
    },
}

_lowerCAmelCase      						=	{
    '''vinai/phobert-base''': 256,
    '''vinai/phobert-large''': 256,
}
def    _SCREAMING_SNAKE_CASE  (       UpperCamelCase			):



    """simple docstring"""





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

    lowerCAmelCase__ :  Any							      =					set(UpperCamelCase			)
    return pairs






class        lowerCAmelCase_(				SCREAMING_SNAKE_CASE_     ):
  '''simple docstring'''




  __lowercase							:      Union[str, Any]				=      VOCAB_FILES_NAMES
  __lowercase							:      int				=      PRETRAINED_VOCAB_FILES_MAP
  __lowercase							:      List[Any]				=      PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES




  def __init__(    self  ,__UpperCAmelCase  ,__UpperCAmelCase  ,__UpperCAmelCase="<s>"  ,__UpperCAmelCase="</s>"  ,__UpperCAmelCase="</s>"  ,__UpperCAmelCase="<s>"  ,__UpperCAmelCase="<unk>"  ,__UpperCAmelCase="<pad>"  ,__UpperCAmelCase="<mask>"  ,**__UpperCAmelCase  ,)    ->		List[str]:
      super().__init__(
          bos_token=__UpperCAmelCase  ,eos_token=__UpperCAmelCase  ,unk_token=__UpperCAmelCase  ,sep_token=__UpperCAmelCase  ,cls_token=__UpperCAmelCase  ,pad_token=__UpperCAmelCase  ,mask_token=__UpperCAmelCase  ,**__UpperCAmelCase  ,)

      lowerCAmelCase__ :  List[Any]							      =					vocab_file
      lowerCAmelCase__ :  Dict							      =					merges_file

      lowerCAmelCase__ :  int							      =					{}
      lowerCAmelCase__ :  str							      =					0
      lowerCAmelCase__ :  Optional[Any]							      =					1
      lowerCAmelCase__ :  Union[str, Any]							      =					2
      lowerCAmelCase__ :  Optional[int]							      =					3

      self.add_from_file(__UpperCAmelCase      )

      lowerCAmelCase__ :  Optional[int]							      =					{v: k for k, v in self.encoder.items()}

      with open(__UpperCAmelCase  ,encoding="""utf-8"""      ) as merges_handle:
          lowerCAmelCase__ :  int							      =					merges_handle.read().split("""\n"""      )[:-1]
      lowerCAmelCase__ :  Optional[int]							      =					[tuple(merge.split()[:-1]      ) for merge in merges]
      lowerCAmelCase__ :  Dict							      =					dict(zip(__UpperCAmelCase  ,range(len(__UpperCAmelCase      )      )      )      )
      lowerCAmelCase__ :  Dict							      =					{}




  def 							UpperCAmelCase_	(    self  ,__UpperCAmelCase  ,__UpperCAmelCase = None      )    ->		List[int]:

      if token_ids_a is None:
          return [self.cls_token_id] + token_ids_a + [self.sep_token_id]
      lowerCAmelCase__ :  Optional[int]							      =					[self.cls_token_id]
      lowerCAmelCase__ :  Tuple							      =					[self.sep_token_id]
      return cls + token_ids_a + sep + sep + token_ids_a + sep




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

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

      if token_ids_a is None:
          return [1] + ([0] * len(__UpperCAmelCase      )) + [1]
      return [1] + ([0] * len(__UpperCAmelCase      )) + [1, 1] + ([0] * len(__UpperCAmelCase      )) + [1]




  def 							UpperCAmelCase_	(    self  ,__UpperCAmelCase  ,__UpperCAmelCase = None      )    ->		List[int]:
      lowerCAmelCase__ :  Optional[Any]							      =					[self.sep_token_id]
      lowerCAmelCase__ :  Dict							      =					[self.cls_token_id]

      if token_ids_a is None:
          return len(cls + token_ids_a + sep      ) * [0]
      return len(cls + token_ids_a + sep + sep + token_ids_a + sep      ) * [0]




  @property
  def 							UpperCAmelCase_	(    self      )    ->		Tuple:
      return len(self.encoder      )




  def 							UpperCAmelCase_	(    self      )    ->		Optional[Any]:
      return dict(self.encoder  ,**self.added_tokens_encoder      )




  def 							UpperCAmelCase_	(    self  ,__UpperCAmelCase      )    ->		List[str]:
      if token in self.cache:
          return self.cache[token]
      lowerCAmelCase__ :  List[Any]							      =					tuple(__UpperCAmelCase      )
      lowerCAmelCase__ :  List[str]							      =					tuple(list(word[:-1]      ) + [word[-1] + """</w>"""]      )
      lowerCAmelCase__ :  Optional[int]							      =					get_pairs(__UpperCAmelCase      )

      if not pairs:
          return token

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

              if word[i] == first and i < len(__UpperCAmelCase      ) - 1 and word[i + 1] == second:
                  new_word.append(first + second      )
                  i += 2
              else:
                  new_word.append(word[i]      )
                  i += 1
          lowerCAmelCase__ :  str							      =					tuple(__UpperCAmelCase      )
          lowerCAmelCase__ :  int							      =					new_word
          if len(__UpperCAmelCase      ) == 1:
              break
          else:
              lowerCAmelCase__ :  Any							      =					get_pairs(__UpperCAmelCase      )
      lowerCAmelCase__ :  List[str]							      =					"""@@ """.join(__UpperCAmelCase      )
      lowerCAmelCase__ :  Union[str, Any]							      =					word[:-4]
      lowerCAmelCase__ :  Optional[Any]							      =					word
      return word




  def 							UpperCAmelCase_	(    self  ,__UpperCAmelCase      )    ->		Optional[int]:
      lowerCAmelCase__ :  Dict							      =					[]

      lowerCAmelCase__ :  str							      =					re.findall(R"""\S+\n?"""  ,__UpperCAmelCase      )

      for token in words:
          split_tokens.extend(list(self.bpe(__UpperCAmelCase      ).split(""" """      )      )      )
      return split_tokens




  def 							UpperCAmelCase_	(    self  ,__UpperCAmelCase      )    ->		Union[str, Any]:
      return self.encoder.get(__UpperCAmelCase  ,self.encoder.get(self.unk_token      )      )




  def 							UpperCAmelCase_	(    self  ,__UpperCAmelCase      )    ->		Tuple:
      return self.decoder.get(__UpperCAmelCase  ,self.unk_token      )




  def 							UpperCAmelCase_	(    self  ,__UpperCAmelCase      )    ->		List[str]:
      lowerCAmelCase__ :  List[str]							      =					""" """.join(__UpperCAmelCase      ).replace("""@@ """  ,""""""      ).strip()
      return out_string




  def 							UpperCAmelCase_	(    self  ,__UpperCAmelCase  ,__UpperCAmelCase = None      )    ->		Tuple[str]:
      if not os.path.isdir(__UpperCAmelCase      ):
          logger.error(F"""Vocabulary path ({save_directory}) should be a directory"""      )
          return
      lowerCAmelCase__ :  List[str]							      =					os.path.join(
          __UpperCAmelCase  ,(filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""vocab_file"""]      )
      lowerCAmelCase__ :  List[str]							      =					os.path.join(
          __UpperCAmelCase  ,(filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""merges_file"""]      )

      if os.path.abspath(self.vocab_file      ) != os.path.abspath(__UpperCAmelCase      ):
          copyfile(self.vocab_file  ,__UpperCAmelCase      )

      if os.path.abspath(self.merges_file      ) != os.path.abspath(__UpperCAmelCase      ):
          copyfile(self.merges_file  ,__UpperCAmelCase      )

      return out_vocab_file, out_merge_file






  def 							UpperCAmelCase_	(    self  ,__UpperCAmelCase      )    ->		Tuple:
      if isinstance(__UpperCAmelCase  ,__UpperCAmelCase      ):
          try:
              with open(__UpperCAmelCase  ,"""r"""  ,encoding="""utf-8"""      ) as fd:
                  self.add_from_file(__UpperCAmelCase      )
          except FileNotFoundError as fnfe:
              raise fnfe
          except UnicodeError:
              raise Exception(F"""Incorrect encoding detected in {f}, please rebuild the dataset"""      )
          return

      lowerCAmelCase__ :  Tuple							      =					f.readlines()
      for lineTmp in lines:
          lowerCAmelCase__ :  Optional[int]							      =					lineTmp.strip()
          lowerCAmelCase__ :  List[str]							      =					line.rfind(""" """      )
          if idx == -1:
              raise ValueError("""Incorrect dictionary format, expected '<token> <cnt>'"""      )
          lowerCAmelCase__ :  Optional[Any]							      =					line[:idx]
          lowerCAmelCase__ :  int							      =					len(self.encoder      )







'''simple docstring'''





from pickle import UnpicklingError

import jax
import jax.numpy as jnp
import numpy as np
from flax.serialization import from_bytes
from flax.traverse_util import flatten_dict

from ..utils import logging


_lowerCAmelCase      						=	logging.get_logger(__name__)
def    _SCREAMING_SNAKE_CASE  (       UpperCamelCase							,    UpperCamelCase			):



    """simple docstring"""





    try:
        with open(UpperCamelCase							,    """rb"""			) as flax_state_f:
            lowerCAmelCase__ :  Union[str, Any]							      =					from_bytes(UpperCamelCase							,    flax_state_f.read()			)
    except UnpicklingError as e:
        try:
            with open(UpperCamelCase			) as f:
                if f.read().startswith("""version"""			):
                    raise OSError(
                        """You seem to have cloned a repository without having git-lfs installed. Please"""
                        """ install git-lfs and run `git lfs install` followed by `git lfs pull` in the"""
                        """ folder you cloned."""			)
                else:
                    raise ValueError from e
        except (UnicodeDecodeError, ValueError):
            raise EnvironmentError(f"""Unable to convert {model_file} to Flax deserializable object. """			)

    return load_flax_weights_in_pytorch_model(UpperCamelCase							,    UpperCamelCase			)




def    _SCREAMING_SNAKE_CASE  (       UpperCamelCase							,    UpperCamelCase			):



    """simple docstring"""





    try:
        import torch  # noqa: F401
    except ImportError:
        logger.error(
            """Loading Flax weights in PyTorch requires both PyTorch and Flax to be installed. Please see"""
            """ https://pytorch.org/ and https://flax.readthedocs.io/en/latest/installation.html for installation"""
            """ instructions."""			)
        raise

    # check if we have bf16 weights
    lowerCAmelCase__ :  str							      =					flatten_dict(jax.tree_util.tree_map(lambda UpperCamelCase			: x.dtype == jnp.bfloataa							,    UpperCamelCase			)			).values()
    if any(UpperCamelCase			):
        # convert all weights to fp32 if they are bf16 since torch.from_numpy can-not handle bf16

        # and bf16 is not fully supported in PT yet.
        logger.warning(
            """Found ``bfloat16`` weights in Flax model. Casting all ``bfloat16`` weights to ``float32`` """
            """before loading those in PyTorch model."""			)
        lowerCAmelCase__ :  Dict							      =					jax.tree_util.tree_map(
            lambda UpperCamelCase			: params.astype(np.floataa			) if params.dtype == jnp.bfloataa else params							,    UpperCamelCase			)

    lowerCAmelCase__ :  Any							      =					""""""

    lowerCAmelCase__ :  Any							      =					flatten_dict(UpperCamelCase							,    sep="""."""			)
    lowerCAmelCase__ :  Optional[int]							      =					pt_model.state_dict()

    # keep track of unexpected & missing keys
    lowerCAmelCase__ :  Optional[Any]							      =					[]
    lowerCAmelCase__ :  int							      =					set(pt_model_dict.keys()			)

    for flax_key_tuple, flax_tensor in flax_state_dict.items():
        lowerCAmelCase__ :  Union[str, Any]							      =					flax_key_tuple.split("""."""			)

        if flax_key_tuple_array[-1] == "kernel" and flax_tensor.ndim == 4:
            lowerCAmelCase__ :  Optional[int]							      =					flax_key_tuple_array[:-1] + ["""weight"""]
            lowerCAmelCase__ :  Any							      =					jnp.transpose(UpperCamelCase							,    (3, 2, 0, 1)			)
        elif flax_key_tuple_array[-1] == "kernel":
            lowerCAmelCase__ :  str							      =					flax_key_tuple_array[:-1] + ["""weight"""]
            lowerCAmelCase__ :  Any							      =					flax_tensor.T
        elif flax_key_tuple_array[-1] == "scale":
            lowerCAmelCase__ :  int							      =					flax_key_tuple_array[:-1] + ["""weight"""]

        if "time_embedding" not in flax_key_tuple_array:
            for i, flax_key_tuple_string in enumerate(UpperCamelCase			):
                lowerCAmelCase__ :  List[str]							      =					(
                    flax_key_tuple_string.replace("""_0"""							,    """.0"""			)
                    .replace("""_1"""							,    """.1"""			)
                    .replace("""_2"""							,    """.2"""			)
                    .replace("""_3"""							,    """.3"""			)
                    .replace("""_4"""							,    """.4"""			)
                    .replace("""_5"""							,    """.5"""			)
                    .replace("""_6"""							,    """.6"""			)
                    .replace("""_7"""							,    """.7"""			)
                    .replace("""_8"""							,    """.8"""			)
                    .replace("""_9"""							,    """.9"""			)
                )

        lowerCAmelCase__ :  Union[str, Any]							      =					""".""".join(UpperCamelCase			)

        if flax_key in pt_model_dict:
            if flax_tensor.shape != pt_model_dict[flax_key].shape:
                raise ValueError(
                    f"""Flax checkpoint seems to be incorrect. Weight {flax_key_tuple} was expected """
                    f"""to be of shape {pt_model_dict[flax_key].shape}, but is {flax_tensor.shape}."""			)
            else:
                # add weight to pytorch dict
                lowerCAmelCase__ :  int							      =					np.asarray(UpperCamelCase			) if not isinstance(UpperCamelCase							,    np.ndarray			) else flax_tensor
                lowerCAmelCase__ :  int							      =					torch.from_numpy(UpperCamelCase			)
                # remove from missing keys
                missing_keys.remove(UpperCamelCase			)
        else:
            # weight is not expected by PyTorch model
            unexpected_keys.append(UpperCamelCase			)

    pt_model.load_state_dict(UpperCamelCase			)

    # re-transform missing_keys to list
    lowerCAmelCase__ :  Optional[int]							      =					list(UpperCamelCase			)

    if len(UpperCamelCase			) > 0:
        logger.warning(
            """Some weights of the Flax model were not used when initializing the PyTorch model"""
            f""" {pt_model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are initializing"""
            f""" {pt_model.__class__.__name__} from a Flax model trained on another task or with another architecture"""
            """ (e.g. initializing a BertForSequenceClassification model from a FlaxBertForPreTraining model).\n- This"""
            f""" IS NOT expected if you are initializing {pt_model.__class__.__name__} from a Flax model that you expect"""
            """ to be exactly identical (e.g. initializing a BertForSequenceClassification model from a"""
            """ FlaxBertForSequenceClassification model)."""			)
    if len(UpperCamelCase			) > 0:
        logger.warning(
            f"""Some weights of {pt_model.__class__.__name__} were not initialized from the Flax model and are newly"""
            f""" initialized: {missing_keys}\nYou should probably TRAIN this model on a down-stream task to be able to"""
            """ use it for predictions and inference."""			)

    return pt_model


