




'''simple docstring'''





import collections
import os
from typing import List, Optional, Tuple

from transformers.utils import is_jieba_available, requires_backends


if is_jieba_available():
     import jieba

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


_UpperCAmelCase	:						Dict			=     logging.get_logger(__name__)

_UpperCAmelCase	:						List[str]			=     {"vocab_file": "vocab.txt"}

_UpperCAmelCase	:						int			=     {
    "vocab_file": {
        "openbmb/cpm-ant-10b": "https://huggingface.co/openbmb/cpm-ant-10b/blob/main/vocab.txt",
    },
}

_UpperCAmelCase	:						Dict			=     {
    "openbmb/cpm-ant-10b": 1024,
}



def       UpperCAmelCase__  (     lowerCamelCase   ):
  lowercase    :int      	=						collections.OrderedDict()
  with open(lowerCamelCase,			"r",			encoding="utf-8"   ) as reader:
    lowercase    :Dict      	=						reader.readlines()
  for index, token in enumerate(lowerCamelCase   ):
    lowercase    :Dict      	=						token.rstrip("\n"   )
    lowercase    :List[str]      	=						index
  return vocab




class    __lowerCAmelCase       (							lowerCAmelCase):






       def __init__(   self:					Optional[int]       ,					_lowerCAmelCase:					Optional[Any]       ,					_lowerCAmelCase:					List[str]="<unk>"       ,					_lowerCAmelCase:					Optional[Any]=2_00				):
         lowercase    :Tuple      	=						vocab
         lowercase    :str      	=						unk_token
         lowercase    :List[str]      	=						max_input_chars_per_word






       def 					SCREAMING_SNAKE_CASE    (   self:					List[str]       ,					_lowerCAmelCase:					Dict				):
         lowercase    :Tuple      	=						list(_lowerCAmelCase				)
         if len(_lowerCAmelCase				) > self.max_input_chars_per_word:
           return [self.unk_token]

         lowercase    :Union[str, Any]      	=						0
         lowercase    :Tuple      	=						[]
         while start < len(_lowerCAmelCase				):
           lowercase    :Union[str, Any]      	=						len(_lowerCAmelCase				)
           lowercase    :List[str]      	=						None
           while start < end:
             lowercase    :List[Any]      	=						"".join(chars[start:end]				)
             if substr in self.vocab:
               lowercase    :Tuple      	=						substr
               break
             end -= 1
           if cur_substr is None:
             sub_tokens.append(self.unk_token				)
             start += 1
           else:
             sub_tokens.append(_lowerCAmelCase				)
             lowercase    :Tuple      	=						end

         return sub_tokens







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






       def __init__(   self:					Union[str, Any]       ,					_lowerCAmelCase:					Optional[int]       ,					_lowerCAmelCase:					List[Any]="<d>"       ,					_lowerCAmelCase:					List[Any]="</d>"       ,					_lowerCAmelCase:					int="<s>"       ,					_lowerCAmelCase:					List[str]="</s>"       ,					_lowerCAmelCase:					int="<pad>"       ,					_lowerCAmelCase:					Tuple="<unk>"       ,					_lowerCAmelCase:					Any="</n>"       ,					_lowerCAmelCase:					Optional[Any]="</_>"       ,					_lowerCAmelCase:					str="left"       ,					**_lowerCAmelCase:					List[str]       ,					):
         requires_backends(self       ,					["jieba"]				)
         super().__init__(
             bod_token=_lowerCAmelCase       ,					eod_token=_lowerCAmelCase       ,					bos_token=_lowerCAmelCase       ,					eos_token=_lowerCAmelCase       ,					pad_token=_lowerCAmelCase       ,					unk_token=_lowerCAmelCase       ,					line_token=_lowerCAmelCase       ,					space_token=_lowerCAmelCase       ,					padding_side=_lowerCAmelCase       ,					**_lowerCAmelCase       ,					)
         lowercase    :Dict      	=						bod_token
         lowercase    :List[Any]      	=						eod_token
         lowercase    :Dict      	=						load_vocab(_lowerCAmelCase				)
         lowercase    :List[Any]      	=						self.encoder[space_token]
         lowercase    :Optional[int]      	=						self.encoder[line_token]

         del self.encoder[space_token]
         del self.encoder[line_token]

         lowercase    :Tuple      	=						collections.OrderedDict(sorted(self.encoder.items()       ,					key=lambda _lowerCAmelCase				: x[1]				)				)
         lowercase    :Optional[Any]      	=						{v: k for k, v in self.encoder.items()}

         lowercase    :Any      	=						WordpieceTokenizer(vocab=self.encoder       ,					unk_token=self.unk_token				)






       @property
       def 					SCREAMING_SNAKE_CASE    (   self:					Tuple				):
         return self.encoder[self.bod_token]






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






       @property
       def 					SCREAMING_SNAKE_CASE    (   self:					List[str]				):
         return self.encoder["\n"]






       @property
       def 					SCREAMING_SNAKE_CASE    (   self:					Union[str, Any]				):
         return len(self.encoder				)






       def 					SCREAMING_SNAKE_CASE    (   self:					Union[str, Any]				):
         return dict(self.encoder       ,					**self.added_tokens_encoder				)






       def 					SCREAMING_SNAKE_CASE    (   self:					List[Any]       ,					_lowerCAmelCase:					Optional[int]				):
         lowercase    :List[Any]      	=						[]
         for x in jieba.cut(_lowerCAmelCase       ,					cut_all=_lowerCAmelCase				):
           output_tokens.extend(self.wordpiece_tokenizer.tokenize(_lowerCAmelCase				)				)
         return output_tokens






       def 					SCREAMING_SNAKE_CASE    (   self:					Dict       ,					_lowerCAmelCase:					Tuple       ,					**_lowerCAmelCase:					List[Any]				):
         lowercase    :List[str]      	=						[i for i in token_ids if i >= 0]
         lowercase    :str      	=						[
             x for x in token_ids if x != self.pad_token_id and x != self.eos_token_id and x != self.bos_token_id
         ]
         return super()._decode(_lowerCAmelCase       ,					**_lowerCAmelCase				)






       def 					SCREAMING_SNAKE_CASE    (   self:					Any       ,					_lowerCAmelCase:					Dict				):
         return token in self.encoder






       def 					SCREAMING_SNAKE_CASE    (   self:					Tuple       ,					_lowerCAmelCase:					List[str]				):
         return "".join(_lowerCAmelCase				)






       def 					SCREAMING_SNAKE_CASE    (   self:					Any       ,					_lowerCAmelCase:					List[Any]				):
         return self.encoder.get(_lowerCAmelCase       ,					self.encoder.get(self.unk_token				)				)






       def 					SCREAMING_SNAKE_CASE    (   self:					List[str]       ,					_lowerCAmelCase:					Tuple				):
         return self.decoder.get(_lowerCAmelCase       ,					self.unk_token				)






       def 					SCREAMING_SNAKE_CASE    (   self:					Optional[Any]       ,					_lowerCAmelCase:					str       ,					_lowerCAmelCase:					Optional[str] = None				):
         if os.path.isdir(_lowerCAmelCase				):
           lowercase    :Optional[Any]      	=						os.path.join(
               _lowerCAmelCase       ,					(filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]				)
         else:
           lowercase    :Tuple      	=						(filename_prefix + "-" if filename_prefix else "") + save_directory
         lowercase    :int      	=						0
         if " " in self.encoder:
           lowercase    :Tuple      	=						self.encoder[" "]
           del self.encoder[" "]
         if "\n" in self.encoder:
           lowercase    :List[str]      	=						self.encoder["\n"]
           del self.encoder["\n"]
         lowercase    :List[Any]      	=						collections.OrderedDict(sorted(self.encoder.items()       ,					key=lambda _lowerCAmelCase				: x[1]				)				)
         with open(_lowerCAmelCase       ,					"w"       ,					encoding="utf-8"				) as writer:
           for token, token_index in self.encoder.items():
             if index != token_index:
               logger.warning(
                   F"Saving vocabulary to {vocab_file}: vocabulary indices are not consecutive."
                   " Please check that the vocabulary is not corrupted!"				)
               lowercase    :Optional[int]      	=						token_index
             writer.write(token + "\n"				)
             index += 1
         return (vocab_file,)






       def 					SCREAMING_SNAKE_CASE    (   self:					List[str]       ,					_lowerCAmelCase:					List[int]       ,					_lowerCAmelCase:					List[int] = None				):
         if token_ids_a is None:
           return [self.bos_token_id] + token_ids_a
         return [self.bos_token_id] + token_ids_a + [self.bos_token_id] + token_ids_a






       def 					SCREAMING_SNAKE_CASE    (   self:					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				)

         if token_ids_a is not None:
           return [1] + ([0] * len(_lowerCAmelCase				)) + [1] + ([0] * len(_lowerCAmelCase				))
         return [1] + ([0] * len(_lowerCAmelCase				))









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   )



