


import json
import os
from functools import lru_cache
from typing import List, Optional, Tuple

import regex as re

from ...tokenization_utils import AddedToken, PreTrainedTokenizer
from ...utils import logging


A      :    str		=					logging.get_logger(__name__)


A      :    Optional[int]		=					{"vocab_file": "vocab.json", "merges_file": "merges.txt"}

# See all BART models at https://huggingface.co/models?filter=bart
A      :    List[str]		=					{
    "vocab_file": {
        "facebook/bart-base": "https://huggingface.co/facebook/bart-base/resolve/main/vocab.json",
        "facebook/bart-large": "https://huggingface.co/facebook/bart-large/resolve/main/vocab.json",
        "facebook/bart-large-mnli": "https://huggingface.co/facebook/bart-large-mnli/resolve/main/vocab.json",
        "facebook/bart-large-cnn": "https://huggingface.co/facebook/bart-large-cnn/resolve/main/vocab.json",
        "facebook/bart-large-xsum": "https://huggingface.co/facebook/bart-large-xsum/resolve/main/vocab.json",
        "yjernite/bart_eli5": "https://huggingface.co/yjernite/bart_eli5/resolve/main/vocab.json",
    },
    "merges_file": {
        "facebook/bart-base": "https://huggingface.co/facebook/bart-base/resolve/main/merges.txt",
        "facebook/bart-large": "https://huggingface.co/facebook/bart-large/resolve/main/merges.txt",
        "facebook/bart-large-mnli": "https://huggingface.co/facebook/bart-large-mnli/resolve/main/merges.txt",
        "facebook/bart-large-cnn": "https://huggingface.co/facebook/bart-large-cnn/resolve/main/merges.txt",
        "facebook/bart-large-xsum": "https://huggingface.co/facebook/bart-large-xsum/resolve/main/merges.txt",
        "yjernite/bart_eli5": "https://huggingface.co/yjernite/bart_eli5/resolve/main/merges.txt",
    },
}

A      :    List[str]		=					{
    "facebook/bart-base": 1024,
    "facebook/bart-large": 1024,
    "facebook/bart-large-mnli": 1024,
    "facebook/bart-large-cnn": 1024,
    "facebook/bart-large-xsum": 1024,
    "yjernite/bart_eli5": 1024,
}

@lru_cache()
def       lowercase_		(		):

      """simple docstring"""




      lowerCamelCase__		:				List[Any]    =  (
          list(range(ord("!"						)    ,					ord("~"						) + 1						)						) + list(range(ord("¡"						)    ,					ord("¬"						) + 1						)						) + list(range(ord("®"						)    ,					ord("ÿ"						) + 1						)						)
      )
      lowerCamelCase__		:				Tuple    =  bs[:]
      lowerCamelCase__		:				Optional[Any]    =  0
      for b in range(2**8						):
            if b not in bs:
                  bs.append(_A						)
                  cs.append(2**8 + n						)
                  n += 1
      lowerCamelCase__		:				Tuple    =  [chr(_A						) for n in cs]
      return dict(zip(_A    ,					_A						)						)

def       lowercase_		(		_A			:    Optional[Any]						):

      """simple docstring"""




      lowerCamelCase__		:				List[Any]    =  set()
      lowerCamelCase__		:				Tuple    =  word[0]
      for char in word[1:]:
            pairs.add((prev_char, char)						)
            lowerCamelCase__		:				List[Any]    =  char
      return pairs







class 			_lowercase (							lowercase__):


       """simple docstring"""




       A__											=			VOCAB_FILES_NAMES
       A__											=			PRETRAINED_VOCAB_FILES_MAP
       A__											=			PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
       A__											=			["input_ids", "attention_mask"]





       def __init__(   self							:			Union[str, Any]					,      __lowerCamelCase							:			Optional[int]					,      __lowerCamelCase							:			Any					,      __lowerCamelCase							:			Optional[int]="replace"					,      __lowerCamelCase							:			Optional[Any]="<s>"					,      __lowerCamelCase							:			Union[str, Any]="</s>"					,      __lowerCamelCase							:			Optional[int]="</s>"					,      __lowerCamelCase							:			Tuple="<s>"					,      __lowerCamelCase							:			int="<unk>"					,      __lowerCamelCase							:			str="<pad>"					,      __lowerCamelCase							:			Tuple="<mask>"					,      __lowerCamelCase							:			List[Any]=False					,      **__lowerCamelCase							:			Any					,      ):


             '''simple docstring'''



             lowerCamelCase__		:				Union[str, Any]    =  AddedToken(__lowerCamelCase					,      lstrip=__lowerCamelCase					,      rstrip=__lowerCamelCase       ) if isinstance(__lowerCamelCase					,      __lowerCamelCase       ) else bos_token
             lowerCamelCase__		:				Optional[Any]    =  AddedToken(__lowerCamelCase					,      lstrip=__lowerCamelCase					,      rstrip=__lowerCamelCase       ) if isinstance(__lowerCamelCase					,      __lowerCamelCase       ) else eos_token
             lowerCamelCase__		:				List[str]    =  AddedToken(__lowerCamelCase					,      lstrip=__lowerCamelCase					,      rstrip=__lowerCamelCase       ) if isinstance(__lowerCamelCase					,      __lowerCamelCase       ) else sep_token
             lowerCamelCase__		:				Optional[int]    =  AddedToken(__lowerCamelCase					,      lstrip=__lowerCamelCase					,      rstrip=__lowerCamelCase       ) if isinstance(__lowerCamelCase					,      __lowerCamelCase       ) else cls_token
             lowerCamelCase__		:				Tuple    =  AddedToken(__lowerCamelCase					,      lstrip=__lowerCamelCase					,      rstrip=__lowerCamelCase       ) if isinstance(__lowerCamelCase					,      __lowerCamelCase       ) else unk_token
             lowerCamelCase__		:				Union[str, Any]    =  AddedToken(__lowerCamelCase					,      lstrip=__lowerCamelCase					,      rstrip=__lowerCamelCase       ) if isinstance(__lowerCamelCase					,      __lowerCamelCase       ) else pad_token

             # Mask token behave like a normal word, i.e. include the space before it
             lowerCamelCase__		:				Optional[Any]    =  AddedToken(__lowerCamelCase					,      lstrip=__lowerCamelCase					,      rstrip=__lowerCamelCase       ) if isinstance(__lowerCamelCase					,      __lowerCamelCase       ) else mask_token

             super().__init__(
                 errors=__lowerCamelCase					,      bos_token=__lowerCamelCase					,      eos_token=__lowerCamelCase					,      unk_token=__lowerCamelCase					,      sep_token=__lowerCamelCase					,      cls_token=__lowerCamelCase					,      pad_token=__lowerCamelCase					,      mask_token=__lowerCamelCase					,      add_prefix_space=__lowerCamelCase					,      **__lowerCamelCase					,      )

             with open(__lowerCamelCase					,      encoding="utf-8"       ) as vocab_handle:
                   lowerCamelCase__		:				Dict    =  json.load(__lowerCamelCase       )
             lowerCamelCase__		:				str    =  {v: k for k, v in self.encoder.items()}
             lowerCamelCase__		:				str    =  errors  # how to handle errors in decoding
             lowerCamelCase__		:				Any    =  bytes_to_unicode()
             lowerCamelCase__		:				Tuple    =  {v: k for k, v in self.byte_encoder.items()}
             with open(__lowerCamelCase					,      encoding="utf-8"       ) as merges_handle:
                   lowerCamelCase__		:				Optional[Any]    =  merges_handle.read().split("\n"       )[1:-1]
             lowerCamelCase__		:				str    =  [tuple(merge.split()       ) for merge in bpe_merges]
             lowerCamelCase__		:				str    =  dict(zip(__lowerCamelCase					,      range(len(__lowerCamelCase       )       )       )       )
             lowerCamelCase__		:				Optional[Any]    =  {}
             lowerCamelCase__		:				str    =  add_prefix_space

             # Should have added re.IGNORECASE so BPE merges can happen for capitalized versions of contractions
             lowerCamelCase__		:				Any    =  re.compile(R"'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"       )





       @property
       def 						lowerCAmelCase     (   self							:			Any       ):


             '''simple docstring'''



             return len(self.encoder       )





       def 						lowerCAmelCase     (   self							:			str       ):


             '''simple docstring'''



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





       def 						lowerCAmelCase     (   self							:			Optional[Any]					,      __lowerCamelCase							:			str       ):


             '''simple docstring'''



             if token in self.cache:
                   return self.cache[token]
             lowerCamelCase__		:				Tuple    =  tuple(__lowerCamelCase       )
             lowerCamelCase__		:				Dict    =  get_pairs(__lowerCamelCase       )

             if not pairs:
                   return token

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

                         if word[i] == first and i < len(__lowerCamelCase       ) - 1 and word[i + 1] == second:
                               new_word.append(first + second       )
                               i += 2
                         else:
                               new_word.append(word[i]       )
                               i += 1
                   lowerCamelCase__		:				Optional[int]    =  tuple(__lowerCamelCase       )
                   lowerCamelCase__		:				Any    =  new_word
                   if len(__lowerCamelCase       ) == 1:
                         break
                   else:
                         lowerCamelCase__		:				Tuple    =  get_pairs(__lowerCamelCase       )
             lowerCamelCase__		:				Optional[int]    =  " ".join(__lowerCamelCase       )
             lowerCamelCase__		:				Optional[Any]    =  word
             return word





       def 						lowerCAmelCase     (   self							:			str					,      __lowerCamelCase							:			List[str]       ):


             '''simple docstring'''



             lowerCamelCase__		:				Tuple    =  []
             for token in re.findall(self.pat					,      __lowerCamelCase       ):
                   lowerCamelCase__		:				Union[str, Any]    =  "".join(
                       self.byte_encoder[b] for b in token.encode("utf-8"       )       )  # Maps all our bytes to unicode strings, avoiding control tokens of the BPE (spaces in our case)
                   bpe_tokens.extend(bpe_token for bpe_token in self.bpe(__lowerCamelCase       ).split(" "       )       )
             return bpe_tokens





       def 						lowerCAmelCase     (   self							:			int					,      __lowerCamelCase							:			Union[str, Any]       ):


             '''simple docstring'''



             return self.encoder.get(__lowerCamelCase					,      self.encoder.get(self.unk_token       )       )





       def 						lowerCAmelCase     (   self							:			Tuple					,      __lowerCamelCase							:			Optional[Any]       ):


             '''simple docstring'''



             return self.decoder.get(__lowerCamelCase       )





       def 						lowerCAmelCase     (   self							:			List[str]					,      __lowerCamelCase							:			int       ):


             '''simple docstring'''



             lowerCamelCase__		:				List[Any]    =  "".join(__lowerCamelCase       )
             lowerCamelCase__		:				Tuple    =  bytearray([self.byte_decoder[c] for c in text]       ).decode("utf-8"					,      errors=self.errors       )
             return text





       def 						lowerCAmelCase     (   self							:			Any					,      __lowerCamelCase							:			str					,      __lowerCamelCase							:			Optional[str] = None       ):


             '''simple docstring'''



             if not os.path.isdir(__lowerCamelCase       ):
                   logger.error(f"Vocabulary path ({save_directory}) should be a directory"       )
                   return
             lowerCamelCase__		:				str    =  os.path.join(
                 __lowerCamelCase					,      (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]       )
             lowerCamelCase__		:				Optional[Any]    =  os.path.join(
                 __lowerCamelCase					,      (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["merges_file"]       )

             with open(__lowerCamelCase					,      "w"					,      encoding="utf-8"       ) as f:
                   f.write(json.dumps(self.encoder					,      indent=2					,      sort_keys=__lowerCamelCase					,      ensure_ascii=__lowerCamelCase       ) + "\n"       )

             lowerCamelCase__		:				Optional[Any]    =  0
             with open(__lowerCamelCase					,      "w"					,      encoding="utf-8"       ) as writer:
                   writer.write("#version: 0.2\n"       )
                   for bpe_tokens, token_index in sorted(self.bpe_ranks.items()					,      key=lambda __lowerCamelCase       : kv[1]       ):
                         if index != token_index:
                               logger.warning(
                                   f"Saving vocabulary to {merge_file}: BPE merge indices are not consecutive."
                                   " Please check that the tokenizer is not corrupted!"       )
                               lowerCamelCase__		:				Union[str, Any]    =  token_index
                         writer.write(" ".join(__lowerCamelCase       ) + "\n"       )
                         index += 1

             return vocab_file, merge_file





       def 						lowerCAmelCase     (   self							:			Any					,      __lowerCamelCase							:			List[int]					,      __lowerCamelCase							:			Optional[List[int]] = None       ):


             '''simple docstring'''



             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__		:				Optional[Any]    =  [self.sep_token_id]
             return cls + token_ids_a + sep + sep + token_ids_a + sep





       def 						lowerCAmelCase     (   self							:			List[str]					,      __lowerCamelCase							:			List[int]					,      __lowerCamelCase							:			Optional[List[int]] = None					,      __lowerCamelCase							:			bool = False       ):


             '''simple docstring'''



             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 None:
                   return [1] + ([0] * len(__lowerCamelCase       )) + [1]
             return [1] + ([0] * len(__lowerCamelCase       )) + [1, 1] + ([0] * len(__lowerCamelCase       )) + [1]





       def 						lowerCAmelCase     (   self							:			Any					,      __lowerCamelCase							:			List[int]					,      __lowerCamelCase							:			Optional[List[int]] = None       ):


             '''simple docstring'''



             lowerCamelCase__		:				Dict    =  [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]





       def 						lowerCAmelCase     (   self							:			Optional[Any]					,      __lowerCamelCase							:			List[str]					,      __lowerCamelCase							:			Union[str, Any]=False					,      **__lowerCamelCase							:			Union[str, Any]       ):


             '''simple docstring'''



             lowerCamelCase__		:				List[str]    =  kwargs.pop("add_prefix_space"					,      self.add_prefix_space       )
             if (is_split_into_words or add_prefix_space) and (len(__lowerCamelCase       ) > 0 and not text[0].isspace()):
                   lowerCamelCase__		:				List[Any]    =  " " + text
             return (text, kwargs)






from .constants import (
    MODEL_NAME,
    OPTIMIZER_NAME,
    RNG_STATE_NAME,
    SAFE_WEIGHTS_INDEX_NAME,
    SAFE_WEIGHTS_NAME,
    SCALER_NAME,
    SCHEDULER_NAME,
    TORCH_LAUNCH_PARAMS,
    WEIGHTS_INDEX_NAME,
    WEIGHTS_NAME,
)
from .dataclasses import (
    BnbQuantizationConfig,
    ComputeEnvironment,
    CustomDtype,
    DeepSpeedPlugin,
    DistributedDataParallelKwargs,
    DistributedType,
    DynamoBackend,
    FPaRecipeKwargs,
    FullyShardedDataParallelPlugin,
    GradientAccumulationPlugin,
    GradScalerKwargs,
    InitProcessGroupKwargs,
    KwargsHandler,
    LoggerType,
    MegatronLMPlugin,
    PrecisionType,
    ProjectConfiguration,
    RNGType,
    SageMakerDistributedType,
    TensorInformation,
    TorchDynamoPlugin,
)
from .environment import get_int_from_env, parse_choice_from_env, parse_flag_from_env
from .imports import (
    get_ccl_version,
    is_abit_bnb_available,
    is_abit_bnb_available,
    is_aim_available,
    is_bfaa_available,
    is_bnb_available,
    is_botoa_available,
    is_ccl_available,
    is_comet_ml_available,
    is_datasets_available,
    is_deepspeed_available,
    is_fpa_available,
    is_ipex_available,
    is_megatron_lm_available,
    is_mlflow_available,
    is_mps_available,
    is_npu_available,
    is_rich_available,
    is_safetensors_available,
    is_sagemaker_available,
    is_tensorboard_available,
    is_tpu_available,
    is_transformers_available,
    is_wandb_available,
    is_xpu_available,
)
from .modeling import (
    check_device_map,
    check_tied_parameters_in_config,
    check_tied_parameters_on_same_device,
    compute_module_sizes,
    convert_file_size_to_int,
    dtype_byte_size,
    find_tied_parameters,
    get_balanced_memory,
    get_max_layer_size,
    get_max_memory,
    get_mixed_precision_context_manager,
    id_tensor_storage,
    infer_auto_device_map,
    load_checkpoint_in_model,
    load_offloaded_weights,
    load_state_dict,
    named_module_tensors,
    retie_parameters,
    set_module_tensor_to_device,
    shard_checkpoint,
)
from .offload import (
    OffloadedWeightsLoader,
    PrefixedDataset,
    extract_submodules_state_dict,
    load_offloaded_weight,
    offload_state_dict,
    offload_weight,
    save_offload_index,
)
from .operations import (
    broadcast,
    broadcast_object_list,
    concatenate,
    convert_outputs_to_fpaa,
    convert_to_fpaa,
    find_batch_size,
    find_device,
    gather,
    gather_object,
    get_data_structure,
    honor_type,
    initialize_tensors,
    is_namedtuple,
    is_tensor_information,
    is_torch_tensor,
    listify,
    pad_across_processes,
    recursively_apply,
    reduce,
    send_to_device,
    slice_tensors,
)
from .versions import compare_versions, is_torch_version


if is_deepspeed_available():
    from .deepspeed import (
        DeepSpeedEngineWrapper,
        DeepSpeedOptimizerWrapper,
        DeepSpeedSchedulerWrapper,
        DummyOptim,
        DummyScheduler,
        HfDeepSpeedConfig,
    )

from .bnb import has_abit_bnb_layers, load_and_quantize_model
from .fsdp_utils import load_fsdp_model, load_fsdp_optimizer, save_fsdp_model, save_fsdp_optimizer
from .launch import (
    PrepareForLaunch,
    _filter_args,
    prepare_deepspeed_cmd_env,
    prepare_multi_gpu_env,
    prepare_sagemager_args_inputs,
    prepare_simple_launcher_cmd_env,
    prepare_tpu,
)
from .megatron_lm import (
    AbstractTrainStep,
    BertTrainStep,
    GPTTrainStep,
    MegatronEngine,
    MegatronLMDummyDataLoader,
    MegatronLMDummyScheduler,
    MegatronLMOptimizerWrapper,
    MegatronLMSchedulerWrapper,
    TaTrainStep,
    avg_losses_across_data_parallel_group,
    gather_across_data_parallel_groups,
)
from .megatron_lm import initialize as megatron_lm_initialize
from .megatron_lm import prepare_data_loader as megatron_lm_prepare_data_loader
from .megatron_lm import prepare_model as megatron_lm_prepare_model
from .megatron_lm import prepare_optimizer as megatron_lm_prepare_optimizer
from .megatron_lm import prepare_scheduler as megatron_lm_prepare_scheduler
from .memory import find_executable_batch_size, release_memory
from .other import (
    extract_model_from_parallel,
    get_pretty_name,
    is_port_in_use,
    merge_dicts,
    patch_environment,
    save,
    wait_for_everyone,
    write_basic_config,
)
from .random import set_seed, synchronize_rng_state, synchronize_rng_states
from .torch_xla import install_xla
from .tqdm import tqdm
from .transformer_engine import convert_model, has_transformer_engine_layers


