'''simple docstring'''




import os
import pickle
import unittest
from collections import UserDict, namedtuple

import torch

from accelerate.test_utils.testing import require_cuda, require_torch_min_version
from accelerate.test_utils.training import RegressionModel
from accelerate.utils import (
    convert_outputs_to_fpaa,
    extract_model_from_parallel,
    find_device,
    listify,
    patch_environment,
    recursively_apply,
    send_to_device,
)


_UpperCamelCase				:		Dict       			=		namedtuple('ExampleNamedTuple', 'a b c')

class   _lowercase(     unittest.TestCase				):



			"""simple docstring"""





			def 			snake_case	(							self:    Optional[int]	):
				__UpperCAmelCase         =					torch.randn(5			,2	)
				__UpperCAmelCase         =					torch.device('cuda'	) if torch.cuda.is_available() else torch.device('cpu'	)

				__UpperCAmelCase         =					send_to_device(_lowerCAmelCase			,_lowerCAmelCase	)
				self.assertTrue(torch.equal(resulta.cpu()			,_lowerCAmelCase	)	)

				__UpperCAmelCase         =					send_to_device((tensor, [tensor, tensor], 1)			,_lowerCAmelCase	)
				self.assertIsInstance(_lowerCAmelCase			,_lowerCAmelCase	)
				self.assertTrue(torch.equal(resulta[0].cpu()			,_lowerCAmelCase	)	)
				self.assertIsInstance(resulta[1]			,_lowerCAmelCase	)
				self.assertTrue(torch.equal(resulta[1][0].cpu()			,_lowerCAmelCase	)	)
				self.assertTrue(torch.equal(resulta[1][1].cpu()			,_lowerCAmelCase	)	)
				self.assertEqual(resulta[2]			,1	)

				__UpperCAmelCase         =					send_to_device({'a': tensor, 'b': [tensor, tensor], 'c': 1}			,_lowerCAmelCase	)
				self.assertIsInstance(_lowerCAmelCase			,_lowerCAmelCase	)
				self.assertTrue(torch.equal(resulta['a'].cpu()			,_lowerCAmelCase	)	)
				self.assertIsInstance(resulta['b']			,_lowerCAmelCase	)
				self.assertTrue(torch.equal(resulta['b'][0].cpu()			,_lowerCAmelCase	)	)
				self.assertTrue(torch.equal(resulta['b'][1].cpu()			,_lowerCAmelCase	)	)
				self.assertEqual(resulta['c']			,1	)

				__UpperCAmelCase         =					send_to_device(ExampleNamedTuple(a=_lowerCAmelCase			,b=[tensor, tensor]			,c=1	)			,_lowerCAmelCase	)
				self.assertIsInstance(_lowerCAmelCase			,_lowerCAmelCase	)
				self.assertTrue(torch.equal(resulta.a.cpu()			,_lowerCAmelCase	)	)
				self.assertIsInstance(resulta.b			,_lowerCAmelCase	)
				self.assertTrue(torch.equal(resulta.b[0].cpu()			,_lowerCAmelCase	)	)
				self.assertTrue(torch.equal(resulta.b[1].cpu()			,_lowerCAmelCase	)	)
				self.assertEqual(resulta.c			,1	)

				__UpperCAmelCase         =					send_to_device(UserDict({'a': tensor, 'b': [tensor, tensor], 'c': 1}	)			,_lowerCAmelCase	)
				self.assertIsInstance(_lowerCAmelCase			,_lowerCAmelCase	)
				self.assertTrue(torch.equal(resulta['a'].cpu()			,_lowerCAmelCase	)	)
				self.assertIsInstance(resulta['b']			,_lowerCAmelCase	)
				self.assertTrue(torch.equal(resulta['b'][0].cpu()			,_lowerCAmelCase	)	)
				self.assertTrue(torch.equal(resulta['b'][1].cpu()			,_lowerCAmelCase	)	)
				self.assertEqual(resulta['c']			,1	)





			def 			snake_case	(							self:    int	):
				with self.assertRaises(_lowerCAmelCase	) as cm:
					__UpperCAmelCase         =					recursively_apply(torch.tensor			,(torch.tensor(1	), 1)			,error_on_other_type=_lowerCAmelCase	)
				self.assertEqual(
				    str(cm.exception	)			,'Unsupported types (<class \'int\'>) passed to `tensor`. Only nested list/tuple/dicts of objects that are valid for `is_torch_tensor` should be passed.'			,)





			def 			snake_case	(							self:    Dict	):
				__UpperCAmelCase         =					torch.tensor([1, 2, 3, 4, 5]	)
				self.assertEqual(listify(_lowerCAmelCase	)			,[1, 2, 3, 4, 5]	)

				__UpperCAmelCase         =					torch.tensor([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]]	)
				self.assertEqual(listify(_lowerCAmelCase	)			,[[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]]	)

				__UpperCAmelCase         =					torch.tensor([[[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]], [[11, 12, 13, 14, 15], [16, 17, 18, 19, 20]]]	)
				self.assertEqual(
				    listify(_lowerCAmelCase	)			,[[[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]], [[11, 12, 13, 14, 15], [16, 17, 18, 19, 20]]]	)





			def 			snake_case	(							self:    Optional[int]	):
				with patch_environment(aa=1			,BB=2	):
					self.assertEqual(os.environ.get('AA'	)			,'1'	)
					self.assertEqual(os.environ.get('BB'	)			,'2'	)

				self.assertNotIn('AA'			,os.environ	)
				self.assertNotIn('BB'			,os.environ	)





			def 			snake_case	(							self:    int	):
				__UpperCAmelCase         =					RegressionModel()
				__UpperCAmelCase         =					model.forward
				__UpperCAmelCase         =					convert_outputs_to_fpaa(model.forward	)
				__UpperCAmelCase         =					extract_model_from_parallel(_lowerCAmelCase			,keep_fpaa_wrapper=_lowerCAmelCase	)
				__UpperCAmelCase         =					pickle.dumps(_lowerCAmelCase	)





			@require_cuda
			def 			snake_case	(							self:    int	):
				__UpperCAmelCase         =					RegressionModel()
				__UpperCAmelCase         =					model.forward
				__UpperCAmelCase         =					torch.cuda.amp.autocast(dtype=torch.floataa	)(model.forward	)
				__UpperCAmelCase         =					convert_outputs_to_fpaa(model.forward	)
				__UpperCAmelCase         =					extract_model_from_parallel(_lowerCAmelCase			,keep_fpaa_wrapper=_lowerCAmelCase	)
				__UpperCAmelCase         =					pickle.dumps(_lowerCAmelCase	)





			@require_cuda
			@require_torch_min_version(version='2.0'	)
			def 			snake_case	(							self:    Optional[int]	):
				__UpperCAmelCase         =					RegressionModel()
				__UpperCAmelCase         =					model.forward
				__UpperCAmelCase         =					torch.cuda.amp.autocast(dtype=torch.floataa	)(model.forward	)
				__UpperCAmelCase         =					convert_outputs_to_fpaa(model.forward	)
				__UpperCAmelCase         =					torch.compile(model.forward			,backend='inductor'	)
				__UpperCAmelCase         =					torch.randn(4			,10	).cuda()
				__UpperCAmelCase         =					model(_lowerCAmelCase	)





			def 			snake_case	(							self:    List[str]	):
				__UpperCAmelCase         =					RegressionModel()
				# could also do a test with DistributedDataParallel, but difficult to run on CPU or single GPU
				__UpperCAmelCase         =					torch.nn.parallel.DataParallel(_lowerCAmelCase	)
				__UpperCAmelCase         =					extract_model_from_parallel(_lowerCAmelCase	)

				self.assertEqual(_lowerCAmelCase			,_lowerCAmelCase	)





			@require_torch_min_version(version='2.0'	)
			def 			snake_case	(							self:    Union[str, Any]	):
				__UpperCAmelCase         =					RegressionModel()
				__UpperCAmelCase         =					torch.compile(_lowerCAmelCase	)

				# could also do a test with DistributedDataParallel, but difficult to run on CPU or single GPU
				__UpperCAmelCase         =					torch.nn.parallel.DataParallel(_lowerCAmelCase	)
				__UpperCAmelCase         =					torch.compile(_lowerCAmelCase	)
				__UpperCAmelCase         =					extract_model_from_parallel(_lowerCAmelCase	)

				self.assertEqual(compiled_model._orig_mod			,compiled_model_unwrapped._orig_mod	)





			def 			snake_case	(							self:    int	):
				self.assertEqual(find_device([1, 'a', torch.tensor([1, 2, 3]	)]	)			,torch.device('cpu'	)	)
				self.assertEqual(find_device({'a': 1, 'b': torch.tensor([1, 2, 3]	)}	)			,torch.device('cpu'	)	)
				self.assertIsNone(find_device([1, 'a']	)	)




'''simple docstring'''




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				:		int       			=		logging.get_logger(__name__)

_UpperCamelCase				:		Tuple       			=		{'vocab_file': 'vocab.json', 'merges_file': 'merges.txt', 'tokenizer_file': 'tokenizer.json'}

_UpperCamelCase				:		int       			=		{
    '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				:		List[str]       			=		{
    'allenai/longformer-base-4096': 40_96,
    'allenai/longformer-large-4096': 40_96,
    'allenai/longformer-large-4096-finetuned-triviaqa': 40_96,
    'allenai/longformer-base-4096-extra.pos.embd.only': 40_96,
    'allenai/longformer-large-4096-extra.pos.embd.only': 40_96,
}

class   _lowercase(     _lowerCamelCase				):



   """simple docstring"""





   __lowerCamelCase    =  VOCAB_FILES_NAMES
   __lowerCamelCase    =  PRETRAINED_VOCAB_FILES_MAP
   __lowerCamelCase    =  PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
   __lowerCamelCase    =  ['''input_ids''', '''attention_mask''']
   __lowerCamelCase    =  LongformerTokenizer





   def __init__(							self:    Any			,a:    Optional[int]=None			,a:    Optional[int]=None			,a:    int=None			,a:    Any="replace"			,a:    Optional[int]="<s>"			,a:    Any="</s>"			,a:    Optional[int]="</s>"			,a:    Dict="<s>"			,a:    Dict="<unk>"			,a:    int="<pad>"			,a:    List[str]="<mask>"			,a:    Union[str, Any]=False			,a:    Any=True			,**a:    str			,):
    super().__init__(
        a			,a			,tokenizer_file=a			,errors=a			,bos_token=a			,eos_token=a			,sep_token=a			,cls_token=a			,unk_token=a			,pad_token=a			,mask_token=a			,add_prefix_space=a			,trim_offsets=a			,**a			,)

    __UpperCAmelCase         =					json.loads(self.backend_tokenizer.pre_tokenizer.__getstate__()	)
    if pre_tok_state.get('add_prefix_space'			,a	) != add_prefix_space:
     __UpperCAmelCase         =					getattr(a			,pre_tok_state.pop('type'	)	)
     __UpperCAmelCase         =					add_prefix_space
     __UpperCAmelCase         =					pre_tok_class(**a	)

    __UpperCAmelCase         =					add_prefix_space

    __UpperCAmelCase         =					'post_processor'
    __UpperCAmelCase         =					getattr(self.backend_tokenizer			,a			,a	)
    if tokenizer_component_instance:
     __UpperCAmelCase         =					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         =					tuple(state['sep']	)
     if "cls" in state:
      __UpperCAmelCase         =					tuple(state['cls']	)

     __UpperCAmelCase         =					False

     if state.get('add_prefix_space'			,a	) != add_prefix_space:
      __UpperCAmelCase         =					add_prefix_space
      __UpperCAmelCase         =					True

     if state.get('trim_offsets'			,a	) != trim_offsets:
      __UpperCAmelCase         =					trim_offsets
      __UpperCAmelCase         =					True

     if changes_to_apply:
      __UpperCAmelCase         =					getattr(a			,state.pop('type'	)	)
      __UpperCAmelCase         =					component_class(**a	)
      setattr(self.backend_tokenizer			,a			,a	)





   @property
   def 			snake_case	(							self:    List[Any]	):
    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 			snake_case	(							self:    int			,a:    Optional[int]	):
    __UpperCAmelCase         =					AddedToken(a			,lstrip=a			,rstrip=a	) if isinstance(a			,a	) else value
    __UpperCAmelCase         =					value





   def 			snake_case	(							self:    Optional[int]			,*a:    Dict			,**a:    Dict	):
    __UpperCAmelCase         =					kwargs.get('is_split_into_words'			,a	)
    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(*a			,**a	)





   def 			snake_case	(							self:    str			,*a:    Dict			,**a:    str	):
    __UpperCAmelCase         =					kwargs.get('is_split_into_words'			,a	)

    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(*a			,**a	)





   def 			snake_case	(							self:    List[str]			,a:    str			,a:    Optional[str] = None	):
    __UpperCAmelCase         =					self._tokenizer.model.save(a			,name=a	)
    return tuple(a	)





   def 			snake_case	(							self:    Tuple			,a:    Dict			,a:    Tuple=None	):
    __UpperCAmelCase         =					[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 			snake_case	(							self:    int			,a:    List[int]			,a:    Optional[List[int]] = None	):
    __UpperCAmelCase         =					[self.sep_token_id]
    __UpperCAmelCase         =					[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]



