


import json
from typing import List, Optional, Tuple

from tokenizers import pre_tokenizers, processors

from ...tokenization_utils_base import AddedToken, BatchEncoding
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .tokenization_longformer import LongformerTokenizer


__UpperCamelCase					:	List[str]          =				logging.get_logger(__name__)

__UpperCamelCase					:	str          =				{'vocab_file': 'vocab.json', 'merges_file': 'merges.txt', 'tokenizer_file': 'tokenizer.json'}

__UpperCamelCase					:	Union[str, Any]          =				{
    'vocab_file': {
        'allenai/longformer-base-4096': 'https://huggingface.co/allenai/longformer-base-4096/resolve/main/vocab.json',
        'allenai/longformer-large-4096': (
            'https://huggingface.co/allenai/longformer-large-4096/resolve/main/vocab.json'
        ),
        'allenai/longformer-large-4096-finetuned-triviaqa': (
            'https://huggingface.co/allenai/longformer-large-4096-finetuned-triviaqa/resolve/main/vocab.json'
        ),
        'allenai/longformer-base-4096-extra.pos.embd.only': (
            'https://huggingface.co/allenai/longformer-base-4096-extra.pos.embd.only/resolve/main/vocab.json'
        ),
        'allenai/longformer-large-4096-extra.pos.embd.only': (
            'https://huggingface.co/allenai/longformer-large-4096-extra.pos.embd.only/resolve/main/vocab.json'
        ),
    },
    'merges_file': {
        'allenai/longformer-base-4096': 'https://huggingface.co/allenai/longformer-base-4096/resolve/main/merges.txt',
        'allenai/longformer-large-4096': (
            'https://huggingface.co/allenai/longformer-large-4096/resolve/main/merges.txt'
        ),
        'allenai/longformer-large-4096-finetuned-triviaqa': (
            'https://huggingface.co/allenai/longformer-large-4096-finetuned-triviaqa/resolve/main/merges.txt'
        ),
        'allenai/longformer-base-4096-extra.pos.embd.only': (
            'https://huggingface.co/allenai/longformer-base-4096-extra.pos.embd.only/resolve/main/merges.txt'
        ),
        'allenai/longformer-large-4096-extra.pos.embd.only': (
            'https://huggingface.co/allenai/longformer-large-4096-extra.pos.embd.only/resolve/main/merges.txt'
        ),
    },
    'tokenizer_file': {
        'allenai/longformer-base-4096': (
            'https://huggingface.co/allenai/longformer-base-4096/resolve/main/tokenizer.json'
        ),
        'allenai/longformer-large-4096': (
            'https://huggingface.co/allenai/longformer-large-4096/resolve/main/tokenizer.json'
        ),
        'allenai/longformer-large-4096-finetuned-triviaqa': (
            'https://huggingface.co/allenai/longformer-large-4096-finetuned-triviaqa/resolve/main/tokenizer.json'
        ),
        'allenai/longformer-base-4096-extra.pos.embd.only': (
            'https://huggingface.co/allenai/longformer-base-4096-extra.pos.embd.only/resolve/main/tokenizer.json'
        ),
        'allenai/longformer-large-4096-extra.pos.embd.only': (
            'https://huggingface.co/allenai/longformer-large-4096-extra.pos.embd.only/resolve/main/tokenizer.json'
        ),
    },
}

__UpperCamelCase					:	Tuple          =				{
    'allenai/longformer-base-4096': 4096,
    'allenai/longformer-large-4096': 4096,
    'allenai/longformer-large-4096-finetuned-triviaqa': 4096,
    'allenai/longformer-base-4096-extra.pos.embd.only': 4096,
    'allenai/longformer-large-4096-extra.pos.embd.only': 4096,
}





class 			lowerCAmelCase__(     snake_case__    ):






    '''simple docstring'''





    A_	:  str								=			VOCAB_FILES_NAMES
    A_	:  List[str]								=			PRETRAINED_VOCAB_FILES_MAP
    A_	:  int								=			PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
    A_	:  Dict								=			['input_ids', 'attention_mask']
    A_	:  Dict								=			LongformerTokenizer





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





       '''simple docstring'''
       super().__init__(
           __snake_case							, __snake_case							, tokenizer_file=__snake_case							, errors=__snake_case							, bos_token=__snake_case							, eos_token=__snake_case							, sep_token=__snake_case							, cls_token=__snake_case							, unk_token=__snake_case							, pad_token=__snake_case							, mask_token=__snake_case							, add_prefix_space=__snake_case							, trim_offsets=__snake_case							, **__snake_case							, )

       UpperCAmelCase_				:   Union[str, Any]							= json.loads(self.backend_tokenizer.pre_tokenizer.__getstate__()						)
       if pre_tok_state.get('''add_prefix_space'''							, __snake_case						) != add_prefix_space:
          UpperCAmelCase_				:   Any							= getattr(__snake_case							, pre_tok_state.pop('''type'''						)						)
          UpperCAmelCase_				:   str							= add_prefix_space
          UpperCAmelCase_				:   int							= pre_tok_class(**__snake_case						)

       UpperCAmelCase_				:   List[str]							= add_prefix_space

       UpperCAmelCase_				:   Union[str, Any]							= '''post_processor'''
       UpperCAmelCase_				:   List[str]							= getattr(self.backend_tokenizer							, __snake_case							, __snake_case						)
       if tokenizer_component_instance:
          UpperCAmelCase_				:   Tuple							= json.loads(tokenizer_component_instance.__getstate__()						)

          # The lists 'sep' and 'cls' must be cased in tuples for the object `post_processor_class`
          if "sep" in state:
             UpperCAmelCase_				:   Dict							= tuple(state['''sep''']						)
          if "cls" in state:
             UpperCAmelCase_				:   Union[str, Any]							= tuple(state['''cls''']						)

          UpperCAmelCase_				:   int							= False

          if state.get('''add_prefix_space'''							, __snake_case						) != add_prefix_space:
             UpperCAmelCase_				:   Dict							= add_prefix_space
             UpperCAmelCase_				:   Tuple							= True

          if state.get('''trim_offsets'''							, __snake_case						) != trim_offsets:
             UpperCAmelCase_				:   Optional[int]							= trim_offsets
             UpperCAmelCase_				:   int							= True

          if changes_to_apply:
             UpperCAmelCase_				:   Optional[Any]							= getattr(__snake_case							, state.pop('''type'''						)						)
             UpperCAmelCase_				:   Optional[int]							= component_class(**__snake_case						)
             setattr(self.backend_tokenizer							, __snake_case							, __snake_case						)





    @property
    def        _lowerCamelCase			(						self					: Tuple						):





       '''simple docstring'''
       if self._mask_token is None:
          if self.verbose:
             logger.error('''Using mask_token, but it is not set yet.'''						)
          return None
       return str(self._mask_token						)





    @mask_token.setter
    def        _lowerCamelCase			(						self					: Tuple							, __snake_case					: List[str]						):





       '''simple docstring'''
       UpperCAmelCase_				:   Optional[int]							= AddedToken(__snake_case							, lstrip=__snake_case							, rstrip=__snake_case						) if isinstance(__snake_case							, __snake_case						) else value
       UpperCAmelCase_				:   Optional[int]							= value





    def        _lowerCamelCase			(						self					: Dict							, *__snake_case					: Any							, **__snake_case					: Optional[int]						):





       '''simple docstring'''
       UpperCAmelCase_				:   Dict							= kwargs.get('''is_split_into_words'''							, __snake_case						)
       assert self.add_prefix_space or not is_split_into_words, (
           f'''You need to instantiate {self.__class__.__name__} with add_prefix_space=True '''
           "to use it with pretokenized inputs."
       )

       return super()._batch_encode_plus(*__snake_case							, **__snake_case						)





    def        _lowerCamelCase			(						self					: Optional[int]							, *__snake_case					: Union[str, Any]							, **__snake_case					: List[Any]						):





       '''simple docstring'''
       UpperCAmelCase_				:   Optional[int]							= kwargs.get('''is_split_into_words'''							, __snake_case						)

       assert self.add_prefix_space or not is_split_into_words, (
           f'''You need to instantiate {self.__class__.__name__} with add_prefix_space=True '''
           "to use it with pretokenized inputs."
       )

       return super()._encode_plus(*__snake_case							, **__snake_case						)





    def        _lowerCamelCase			(						self					: str							, __snake_case					: str							, __snake_case					: Optional[str] = None						):





       '''simple docstring'''
       UpperCAmelCase_				:   Any							= self._tokenizer.model.save(__snake_case							, name=__snake_case						)
       return tuple(__snake_case						)





    def        _lowerCamelCase			(						self					: int							, __snake_case					: List[Any]							, __snake_case					: Any=None						):





       '''simple docstring'''
       UpperCAmelCase_				:   List[str]							= [self.bos_token_id] + token_ids_a + [self.eos_token_id]
       if token_ids_a is None:
          return output

       return output + [self.eos_token_id] + token_ids_a + [self.eos_token_id]





    def        _lowerCamelCase			(						self					: str							, __snake_case					: List[int]							, __snake_case					: Optional[List[int]] = None						):





       '''simple docstring'''
       UpperCAmelCase_				:   Any							= [self.sep_token_id]
       UpperCAmelCase_				:   Optional[Any]							= [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]



import os
import random
from pathlib import Path
from typing import List

import numpy as np
import torch
from torch.cuda.amp import GradScaler

from .utils import (
    MODEL_NAME,
    OPTIMIZER_NAME,
    RNG_STATE_NAME,
    SCALER_NAME,
    SCHEDULER_NAME,
    get_pretty_name,
    is_tpu_available,
    is_xpu_available,
    save,
)


if is_tpu_available(check_device=False):
       import torch_xla.core.xla_model as xm

from .logging import get_logger
from .state import PartialState


__UpperCamelCase					:	int          =				get_logger(__name__)
def   snake_case_    (			__lowercase      ,							__lowercase      ,							__lowercase      ,							__lowercase      ,							__lowercase      ,							__lowercase = None      ,							):
   # Model states
   for i, state in enumerate(__lowercase							):
      UpperCAmelCase_				:   Dict							= F'''{MODEL_NAME}.bin''' if i == 0 else F'''{MODEL_NAME}_{i}.bin'''
      UpperCAmelCase_				:   List[Any]							= os.path.join(__lowercase      ,							__lowercase							)
      save(__lowercase      ,							__lowercase							)
      logger.info(F'''Model weights saved in {output_model_file}'''							)
   # Optimizer states
   for i, opt in enumerate(__lowercase							):
      UpperCAmelCase_				:   Tuple							= opt.state_dict()
      UpperCAmelCase_				:   Any							= F'''{OPTIMIZER_NAME}.bin''' if i == 0 else F'''{OPTIMIZER_NAME}_{i}.bin'''
      UpperCAmelCase_				:   List[Any]							= os.path.join(__lowercase      ,							__lowercase							)
      save(__lowercase      ,							__lowercase							)
      logger.info(F'''Optimizer state saved in {output_optimizer_file}'''							)
   # Scheduler states
   for i, scheduler in enumerate(__lowercase							):
      UpperCAmelCase_				:   Tuple							= scheduler.state_dict()
      UpperCAmelCase_				:   Tuple							= F'''{SCHEDULER_NAME}.bin''' if i == 0 else F'''{SCHEDULER_NAME}_{i}.bin'''
      UpperCAmelCase_				:   Tuple							= os.path.join(__lowercase      ,							__lowercase							)
      save(__lowercase      ,							__lowercase							)
      logger.info(F'''Scheduler state saved in {output_scheduler_file}'''							)
   # GradScaler state
   if scaler is not None:
      UpperCAmelCase_				:   Union[str, Any]							= scaler.state_dict()
      UpperCAmelCase_				:   Union[str, Any]							= os.path.join(__lowercase      ,							__lowercase							)
      torch.save(__lowercase      ,							__lowercase							)
      logger.info(F'''Gradient scaler state saved in {output_scaler_file}'''							)
   # Random number generator states
   UpperCAmelCase_				:   Union[str, Any]							= {}
   UpperCAmelCase_				:   Tuple							= F'''{RNG_STATE_NAME}_{process_index}.pkl'''
   UpperCAmelCase_				:   List[Any]							= random.getstate()
   UpperCAmelCase_				:   Any							= np.random.get_state()
   UpperCAmelCase_				:   List[Any]							= torch.get_rng_state()
   if is_xpu_available():
      UpperCAmelCase_				:   Optional[int]							= torch.xpu.get_rng_state_all()
   else:
      UpperCAmelCase_				:   Any							= torch.cuda.get_rng_state_all()
   if is_tpu_available():
      UpperCAmelCase_				:   List[str]							= xm.get_rng_state()
   UpperCAmelCase_				:   Any							= os.path.join(__lowercase      ,							__lowercase							)
   torch.save(__lowercase      ,							__lowercase							)
   logger.info(F'''Random states saved in {output_states_file}'''							)
   return output_dir
def   snake_case_    (			__lowercase      ,							__lowercase      ,							__lowercase      ,							__lowercase      ,							__lowercase      ,							__lowercase=None      ,							__lowercase=None      ,							**__lowercase      ,							):
   if map_location not in [None, "cpu", "on_device"]:
      raise TypeError(
          '''Unsupported optimizer map location passed, please choose one of `None`, `\'cpu\'`, or `\'on_device\'`'''							)
   if map_location is None:
      UpperCAmelCase_				:   Optional[Any]							= '''cpu'''
   elif map_location == "on_device":
      UpperCAmelCase_				:   Optional[int]							= PartialState().device
   # Model states
   for i, model in enumerate(__lowercase							):
      UpperCAmelCase_				:   List[Any]							= F'''{MODEL_NAME}.bin''' if i == 0 else F'''{MODEL_NAME}_{i}.bin'''
      UpperCAmelCase_				:   int							= os.path.join(__lowercase      ,							__lowercase							)
      models[i].load_state_dict(torch.load(__lowercase      ,							map_location=__lowercase							)      ,							**__lowercase							)
   logger.info('''All model weights loaded successfully'''							)

   # Optimizer states
   for i, opt in enumerate(__lowercase							):
      UpperCAmelCase_				:   List[str]							= F'''{OPTIMIZER_NAME}.bin''' if i == 0 else F'''{OPTIMIZER_NAME}_{i}.bin'''
      UpperCAmelCase_				:   List[Any]							= os.path.join(__lowercase      ,							__lowercase							)
      UpperCAmelCase_				:   Tuple							= torch.load(__lowercase      ,							map_location=__lowercase							)
      optimizers[i].load_state_dict(__lowercase							)
   logger.info('''All optimizer states loaded successfully'''							)

   # Scheduler states
   for i, scheduler in enumerate(__lowercase							):
      UpperCAmelCase_				:   Dict							= F'''{SCHEDULER_NAME}.bin''' if i == 0 else F'''{SCHEDULER_NAME}_{i}.bin'''
      UpperCAmelCase_				:   int							= os.path.join(__lowercase      ,							__lowercase							)
      scheduler.load_state_dict(torch.load(__lowercase							)							)
   logger.info('''All scheduler states loaded successfully'''							)

   # GradScaler state
   if scaler is not None:
      UpperCAmelCase_				:   List[str]							= os.path.join(__lowercase      ,							__lowercase							)
      scaler.load_state_dict(torch.load(__lowercase							)							)
      logger.info('''GradScaler state loaded successfully'''							)

   # Random states
   try:
      UpperCAmelCase_				:   Optional[Any]							= torch.load(os.path.join(__lowercase      ,							F'''{RNG_STATE_NAME}_{process_index}.pkl'''							)							)
      random.setstate(states['''random_state''']							)
      np.random.set_state(states['''numpy_random_seed''']							)
      torch.set_rng_state(states['''torch_manual_seed''']							)
      if is_xpu_available():
         torch.xpu.set_rng_state_all(states['''torch_xpu_manual_seed''']							)
      else:
         torch.cuda.set_rng_state_all(states['''torch_cuda_manual_seed''']							)
      if is_tpu_available():
         xm.set_rng_state(states['''xm_seed''']							)
      logger.info('''All random states loaded successfully'''							)
   except Exception:
      logger.info('''Could not load random states'''							)
def   snake_case_    (			__lowercase      ,							__lowercase      ,							__lowercase = 0							):
   UpperCAmelCase_				:   Optional[int]							= Path(__lowercase							) / F'''custom_checkpoint_{index}.pkl'''
   logger.info(F'''Saving the state of {get_pretty_name(__lowercase							)} to {save_location}'''							)
   torch.save(obj.state_dict()      ,							__lowercase							)


def   snake_case_    (			__lowercase      ,							__lowercase      ,							__lowercase = 0							):
   UpperCAmelCase_				:   Dict							= F'''{path}/custom_checkpoint_{index}.pkl'''
   logger.info(F'''Loading the state of {get_pretty_name(__lowercase							)} from {load_location}'''							)
   obj.load_state_dict(torch.load(__lowercase      ,							map_location='''cpu'''							)							)