"""simple docstring"""




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

from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import is_sentencepiece_available, logging


if is_sentencepiece_available():
				from .tokenization_pegasus import PegasusTokenizer
else:
				__lowerCamelCase							 =     None


__lowerCamelCase							 =     logging.get_logger(__name__)


__lowerCamelCase							 =     "▁"

__lowerCamelCase							 =     {"vocab_file": "spiece.model", "tokenizer_file": "tokenizer.json"}

__lowerCamelCase							 =     {
    "vocab_file": {"google/pegasus-xsum": "https://huggingface.co/google/pegasus-xsum/resolve/main/spiece.model"},
    "tokenizer_file": {
        "google/pegasus-xsum": "https://huggingface.co/google/pegasus-xsum/resolve/main/tokenizer.json"
    },
}

__lowerCamelCase							 =     {
    "google/pegasus-xsum": 5_12,
}
class 			UpperCamelCase__(	__A       ):
						lowerCAmelCase__     :  Any  							=						VOCAB_FILES_NAMES
						lowerCAmelCase__     :  List[str]  							=						PRETRAINED_VOCAB_FILES_MAP
						lowerCAmelCase__     :  List[str]  							=						PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
						lowerCAmelCase__     :  Any  							=						PegasusTokenizer
						lowerCAmelCase__     :  Tuple  							=						['input_ids', 'attention_mask']





						def __init__(				self						,__UpperCAmelCase=None						,__UpperCAmelCase=None						,__UpperCAmelCase="<pad>"						,__UpperCAmelCase="</s>"						,__UpperCAmelCase="<unk>"						,__UpperCAmelCase="<mask_2>"						,__UpperCAmelCase="<mask_1>"						,__UpperCAmelCase=None						,__UpperCAmelCase=1_03						,**__UpperCAmelCase						,)    -> Any:
													A__		   =							offset

													if additional_special_tokens is not None:
																				if not isinstance(__UpperCAmelCase						,__UpperCAmelCase   ):
																											raise TypeError(
																											    f'''additional_special_tokens should be of type {type(__UpperCAmelCase   )}, but is'''
																											    f''' {type(__UpperCAmelCase   )}'''   )

																				A__		   =							(
																				    ([mask_token_sent] + additional_special_tokens)
																				    if mask_token_sent not in additional_special_tokens and mask_token_sent is not None
																				    else additional_special_tokens
																				)
																				# fill additional tokens with ..., <unk_token_102> in case not all additional tokens are already taken
																				additional_special_tokens_extended += [
																				    f'''<unk_{i}>''' for i in range(len(__UpperCAmelCase   )						,self.offset - 1   )
																				]

																				if len(set(__UpperCAmelCase   )   ) != len(__UpperCAmelCase   ):
																											raise ValueError(
																											    'Please make sure that the provided additional_special_tokens do not contain an incorrectly'
																											    f''' shifted list of <unk_x> tokens. Found {additional_special_tokens_extended}.'''   )
																				A__		   =							additional_special_tokens_extended
													else:
																				A__		   =							[mask_token_sent] if mask_token_sent is not None else []
																				additional_special_tokens += [f'''<unk_{i}>''' for i in range(2						,self.offset   )]

													super().__init__(
													    __UpperCAmelCase						,tokenizer_file=__UpperCAmelCase						,pad_token=__UpperCAmelCase						,eos_token=__UpperCAmelCase						,unk_token=__UpperCAmelCase						,mask_token=__UpperCAmelCase						,mask_token_sent=__UpperCAmelCase						,offset=__UpperCAmelCase						,additional_special_tokens=__UpperCAmelCase						,**__UpperCAmelCase						,)
													A__		   =							vocab_file
													A__		   =							False if not self.vocab_file else True





						def  snake_case__		(				self						,__UpperCAmelCase   )    -> str:
													A__		   =							set(self.all_special_ids   )  # call it once instead of inside list comp
													all_special_ids.remove(self.unk_token_id   )  # <unk> is only sometimes special

													if all_special_ids != set(range(len(self.additional_special_tokens   ) + 3   )   ):
																				raise ValueError(
																				    'There should be 3 special tokens: mask_token, pad_token, and eos_token +'
																				    f''' {len(self.additional_special_tokens   )} additional_special_tokens, but got {all_special_ids}'''   )

													return [1 if x in all_special_ids else 0 for x in seq]





						def  snake_case__		(				self						,__UpperCAmelCase						,__UpperCAmelCase = None						,__UpperCAmelCase = False   )    -> List[int]:
													if already_has_special_tokens:
																				return self._special_token_mask(__UpperCAmelCase   )
													elif token_ids_a is None:
																				return self._special_token_mask(__UpperCAmelCase   ) + [1]
													else:
																				return self._special_token_mask(token_ids_a + token_ids_a   ) + [1]





						def  snake_case__		(				self						,__UpperCAmelCase						,__UpperCAmelCase=None   )    -> List[int]:
													if token_ids_a is None:
																				return token_ids_a + [self.eos_token_id]
													# We don't expect to process pairs, but leave the pair logic for API consistency
													return token_ids_a + token_ids_a + [self.eos_token_id]





						def  snake_case__		(				self						,__UpperCAmelCase						,__UpperCAmelCase = None   )    -> Tuple[str]:
													if not self.can_save_slow_tokenizer:
																				raise ValueError(
																				    'Your fast tokenizer does not have the necessary information to save the vocabulary for a slow '
																				    'tokenizer.'   )

													if not os.path.isdir(__UpperCAmelCase   ):
																				logger.error(f'''Vocabulary path ({save_directory}) should be a directory'''   )
																				return
													A__		   =							os.path.join(
													    __UpperCAmelCase						,(filename_prefix + '-' if filename_prefix else '') + VOCAB_FILES_NAMES['vocab_file']   )

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

													return (out_vocab_file,)






"""simple docstring"""




import os
import socket
from contextlib import contextmanager

import torch

from ..commands.config.default import write_basic_config  # noqa: F401
from ..state import PartialState
from .dataclasses import DistributedType
from .imports import is_deepspeed_available, is_tpu_available
from .transformer_engine import convert_model
from .versions import is_torch_version


if is_deepspeed_available():
				from deepspeed import DeepSpeedEngine

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






def      UpperCAmelCase    (		UpperCamelCase__ ):






							"""simple docstring"""





							if is_torch_version('<'				,			'2.0.0' ) or not hasattr(UpperCamelCase__				,			'_dynamo' ):
														return False
							return isinstance(UpperCamelCase__				,			torch._dynamo.eval_frame.OptimizedModule )






def      UpperCAmelCase    (		UpperCamelCase__				,			UpperCamelCase__ = True ):






							"""simple docstring"""





							A__		   =							(torch.nn.parallel.DistributedDataParallel, torch.nn.DataParallel)

							A__		   =							is_compiled_module(UpperCamelCase__ )
							if is_compiled:
														A__		   =							model
														A__		   =							model._orig_mod

							if is_deepspeed_available():
														options += (DeepSpeedEngine,)

							while isinstance(UpperCamelCase__				,			UpperCamelCase__ ):
														A__		   =							model.module

							if not keep_fpaa_wrapper:
														A__		   =							getattr(UpperCamelCase__				,			'forward' )
														A__		   =							model.__dict__.pop('_original_forward'				,			UpperCamelCase__ )
														if original_forward is not None:
																					while hasattr(UpperCamelCase__				,			'__wrapped__' ):
																												A__		   =							forward.__wrapped__
																												if forward == original_forward:
																																			break
																					A__		   =							forward
														if getattr(UpperCamelCase__				,			'_converted_to_transformer_engine'				,			UpperCamelCase__ ):
																					convert_model(UpperCamelCase__				,			to_transformer_engine=UpperCamelCase__ )

							if is_compiled:
														A__		   =							model
														A__		   =							compiled_model

							return model






def      UpperCAmelCase    (		):






							"""simple docstring"""





							PartialState().wait_for_everyone()






def      UpperCAmelCase    (		UpperCamelCase__				,			UpperCamelCase__ ):






							"""simple docstring"""





							if PartialState().distributed_type == DistributedType.TPU:
														xm.save(UpperCamelCase__				,			UpperCamelCase__ )
							elif PartialState().local_process_index == 0:
														torch.save(UpperCamelCase__				,			UpperCamelCase__ )






@contextmanager
def      UpperCAmelCase    (		**UpperCamelCase__ ):






							"""simple docstring"""





							for key, value in kwargs.items():
														A__		   =							str(UpperCamelCase__ )

							yield

							for key in kwargs:
														if key.upper() in os.environ:
																					del os.environ[key.upper()]






def      UpperCAmelCase    (		UpperCamelCase__ ):






							"""simple docstring"""





							if not hasattr(UpperCamelCase__				,			'__qualname__' ) and not hasattr(UpperCamelCase__				,			'__name__' ):
														A__		   =							getattr(UpperCamelCase__				,			'__class__'				,			UpperCamelCase__ )
							if hasattr(UpperCamelCase__				,			'__qualname__' ):
														return obj.__qualname__
							if hasattr(UpperCamelCase__				,			'__name__' ):
														return obj.__name__
							return str(UpperCamelCase__ )






def      UpperCAmelCase    (		UpperCamelCase__				,			UpperCamelCase__ ):






							"""simple docstring"""





							for key, value in source.items():
														if isinstance(UpperCamelCase__				,			UpperCamelCase__ ):
																					A__		   =							destination.setdefault(UpperCamelCase__				,			{} )
																					merge_dicts(UpperCamelCase__				,			UpperCamelCase__ )
														else:
																					A__		   =							value

							return destination






def      UpperCAmelCase    (		UpperCamelCase__ = None ):






							"""simple docstring"""





							if port is None:
														A__		   =							29_500
							with socket.socket(socket.AF_INET				,			socket.SOCK_STREAM ) as s:
														return s.connect_ex(('localhost', port) ) == 0





