

import inspect
from typing import Callable, List, Optional, Set, Tuple, Union

import torch
from packaging import version
from safetensors.torch import storage_ptr, storage_size
from torch import nn

from .utils import logging


lowerCAmelCase_	     =     [nn.LayerNorm]

lowerCAmelCase_	     =     logging.get_logger(__name__)

lowerCAmelCase_	     =     version.parse(version.parse(torch.__version__).base_version)

lowerCAmelCase_	     =     parsed_torch_version_base >= version.parse('''2.0''')
lowerCAmelCase_	     =     parsed_torch_version_base >= version.parse('''1.12''')
lowerCAmelCase_	     =     parsed_torch_version_base >= version.parse('''1.11''')
lowerCAmelCase_	     =     parsed_torch_version_base < version.parse('''1.11''')


def   __SCREAMING_SNAKE_CASE (SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__  ,		self	):
						from torch import _softmax_backward_data

						if is_torch_less_than_1_11:
												return _softmax_backward_data(SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__  ,		parent.dim  ,		self	)
						else:
												return _softmax_backward_data(SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__  ,		parent.dim  ,		self.dtype	)


def   __SCREAMING_SNAKE_CASE (SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__ = 0	):
						snake_case_       =     index.to(layer.weight.device	)
						snake_case_       =     layer.weight.index_select(SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__	).clone().detach()
						if layer.bias is not None:
												if dim == 1:
																		snake_case_       =     layer.bias.clone().detach()
												else:
																		snake_case_       =     layer.bias[index].clone().detach()
						snake_case_       =     list(layer.weight.size()	)
						snake_case_       =     len(SCREAMING_SNAKE_CASE__	)
						snake_case_       =     nn.Linear(new_size[1]  ,		new_size[0]  ,		bias=layer.bias is not None	).to(layer.weight.device	)
						snake_case_       =     False
						new_layer.weight.copy_(W.contiguous()	)
						snake_case_       =     True
						if layer.bias is not None:
												snake_case_       =     False
												new_layer.bias.copy_(b.contiguous()	)
												snake_case_       =     True
						return new_layer





class      snake_case_       (		nn.Module ):







			'''simple docstring'''

			def __init__(					self			:    Optional[Any]							,       _UpperCamelCase			:    Optional[Any]							,       _UpperCamelCase			:    int							)				->Optional[Any]:
									super().__init__()
									snake_case_       =     nf
									snake_case_       =     nn.Parameter(torch.empty(_UpperCamelCase							,       _UpperCamelCase							)							)
									snake_case_       =     nn.Parameter(torch.zeros(_UpperCamelCase							)							)
									nn.init.normal_(self.weight							,       std=0.02							)







			def     snake_case__(					self			:    List[Any]							,       _UpperCamelCase			:    Optional[Any]							)				->Tuple:
									snake_case_       =     x.size()[:-1] + (self.nf,)
									snake_case_       =     torch.addmm(self.bias							,       x.view(-1							,       x.size(-1							)							)							,       self.weight							)
									snake_case_       =     x.view(_UpperCamelCase							)
									return x






def   __SCREAMING_SNAKE_CASE (SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__ = 1	):
						snake_case_       =     index.to(layer.weight.device	)
						snake_case_       =     layer.weight.index_select(SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__	).clone().detach()
						if dim == 0:
												snake_case_       =     layer.bias.clone().detach()
						else:
												snake_case_       =     layer.bias[index].clone().detach()
						snake_case_       =     list(layer.weight.size()	)
						snake_case_       =     len(SCREAMING_SNAKE_CASE__	)
						snake_case_       =     ConvaD(new_size[1]  ,		new_size[0]	).to(layer.weight.device	)
						snake_case_       =     False
						new_layer.weight.copy_(W.contiguous()	)
						snake_case_       =     True
						snake_case_       =     False
						new_layer.bias.copy_(b.contiguous()	)
						snake_case_       =     True
						return new_layer


def   __SCREAMING_SNAKE_CASE (SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__ = None	):
						if isinstance(SCREAMING_SNAKE_CASE__  ,		nn.Linear	):
												return prune_linear_layer(SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__  ,		dim=0 if dim is None else dim	)
						elif isinstance(SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__	):
												return prune_convad_layer(SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__  ,		dim=1 if dim is None else dim	)
						else:
												raise ValueError(F'''Can\'t prune layer of class {layer.__class__}'''	)


def   __SCREAMING_SNAKE_CASE (SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__  ,		*SCREAMING_SNAKE_CASE__	):
						assert len(SCREAMING_SNAKE_CASE__	) > 0, F'''{input_tensors} has to be a tuple/list of tensors'''

						# inspect.signature exist since python 3.5 and is a python method -> no problem with backward compatibility
						snake_case_       =     len(inspect.signature(SCREAMING_SNAKE_CASE__	).parameters	)
						if num_args_in_forward_chunk_fn != len(SCREAMING_SNAKE_CASE__	):
												raise ValueError(
												    F'''forward_chunk_fn expects {num_args_in_forward_chunk_fn} arguments, but only {len(SCREAMING_SNAKE_CASE__	)} input '''
												    '''tensors are given'''	)

						if chunk_size > 0:
												snake_case_       =     input_tensors[0].shape[chunk_dim]
												for input_tensor in input_tensors:
																		if input_tensor.shape[chunk_dim] != tensor_shape:
																								raise ValueError(
																								    F'''All input tenors have to be of the same shape: {tensor_shape}, '''
																								    F'''found shape {input_tensor.shape[chunk_dim]}'''	)

												if input_tensors[0].shape[chunk_dim] % chunk_size != 0:
																		raise ValueError(
																		    F'''The dimension to be chunked {input_tensors[0].shape[chunk_dim]} has to be a multiple of the chunk '''
																		    F'''size {chunk_size}'''	)

												snake_case_       =     input_tensors[0].shape[chunk_dim] // chunk_size

												# chunk input tensor into tuples
												snake_case_       =     tuple(input_tensor.chunk(SCREAMING_SNAKE_CASE__  ,		dim=SCREAMING_SNAKE_CASE__	) for input_tensor in input_tensors	)
												# apply forward fn to every tuple
												snake_case_       =     tuple(forward_fn(*SCREAMING_SNAKE_CASE__	) for input_tensors_chunk in zip(*SCREAMING_SNAKE_CASE__	)	)
												# concatenate output at same dimension
												return torch.cat(SCREAMING_SNAKE_CASE__  ,		dim=SCREAMING_SNAKE_CASE__	)

						return forward_fn(*SCREAMING_SNAKE_CASE__	)


def   __SCREAMING_SNAKE_CASE (SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__	):
						snake_case_       =     torch.ones(SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__	)
						snake_case_       =     set(SCREAMING_SNAKE_CASE__	) - already_pruned_heads  # Convert to set and remove already pruned heads
						for head in heads:
												# Compute how many pruned heads are before the head and move the index accordingly
												snake_case_       =     head - sum(1 if h < head else 0 for h in already_pruned_heads	)
												snake_case_       =     0
						snake_case_       =     mask.view(-1	).contiguous().eq(1	)
						snake_case_       =     torch.arange(len(SCREAMING_SNAKE_CASE__	)	)[mask].long()
						return heads, index


def   __SCREAMING_SNAKE_CASE (*SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__ = None	):
						return torch.meshgrid(*SCREAMING_SNAKE_CASE__  ,		indexing=SCREAMING_SNAKE_CASE__	)


def   __SCREAMING_SNAKE_CASE (SCREAMING_SNAKE_CASE__	):
						return tensor.device, storage_ptr(SCREAMING_SNAKE_CASE__	), storage_size(SCREAMING_SNAKE_CASE__	)


import concurrent.futures
import json
import os
import shutil
import tempfile
import unittest

from transformers import AutoTokenizer, PreTrainedTokenizerFast
from transformers.testing_utils import require_tokenizers

from ..test_tokenization_common import TokenizerTesterMixin





@require_tokenizers
class      snake_case_       (		__A     ,	unittest.TestCase ):







			'''simple docstring'''

			SCREAMING_SNAKE_CASE					:       str						      =    PreTrainedTokenizerFast
			SCREAMING_SNAKE_CASE					:       List[str]						      =    False
			SCREAMING_SNAKE_CASE					:       Tuple						      =    True
			SCREAMING_SNAKE_CASE					:       str						      =    "tokenizer_file"







			def     snake_case__(					self			:    Dict							)				->int:
									snake_case_       =     False  # because we don't have pretrained_vocab_files_map
									super().setUp()
									snake_case_       =     True

									snake_case_       =     ['''robot-test/dummy-tokenizer-fast''', '''robot-test/dummy-tokenizer-wordlevel''']
									snake_case_       =     '''SaulLu/dummy-tokenizer-bytelevel-bpe'''

									# Inclusion of 2 tokenizers to test different types of models (Unigram and WordLevel for the moment)
									snake_case_       =     [(PreTrainedTokenizerFast, model_path, {}) for model_path in model_paths]

									snake_case_       =     PreTrainedTokenizerFast.from_pretrained(model_paths[0]							)
									tokenizer.save_pretrained(self.tmpdirname							)







			def     snake_case__(					self			:    Union[str, Any]							)				->Dict:
									# We disable this test for PreTrainedTokenizerFast because it is the only tokenizer that is not linked to any
									# model
									pass







			def     snake_case__(					self			:    str							)				->Any:
									# We disable this test for PreTrainedTokenizerFast because it is the only tokenizer that is not linked to any
									# model
									pass







			def     snake_case__(					self			:    List[Any]							)				->str:
									# We disable this test for PreTrainedTokenizerFast because it is the only tokenizer that is not linked to any
									# model
									pass







			def     snake_case__(					self			:    str							)				->List[Any]:
									# PreTrainedTokenizerFast doesn't have tokenizer_file in its signature
									pass







			def     snake_case__(					self			:    Any							)				->Union[str, Any]:
									snake_case_       =     self.tmpdirname
									# Here we want to test the 2 available tokenizers that use 2 different types of models: Unigram and WordLevel.
									for tokenizer, pretrained_name, kwargs in self.tokenizers_list:
															with self.subTest(f'''{tokenizer.__class__.__name__} ({pretrained_name})'''							):
																					try:
																											snake_case_       =     tempfile.mkdtemp()
																											snake_case_       =     self.rust_tokenizer_class.from_pretrained(_UpperCamelCase							,       **_UpperCamelCase							)

																											tokenizer.save_pretrained(self.tmpdirname							)
																											super().test_training_new_tokenizer()
																					finally:
																											# Even if the test fails, we must be sure that the folder is deleted and that the default tokenizer
																											# is restored
																											shutil.rmtree(self.tmpdirname							)
																											snake_case_       =     tmpdirname_orig







			def     snake_case__(					self			:    List[str]							)				->List[Any]:
									snake_case_       =     self.tmpdirname
									# Here we want to test the 2 available tokenizers that use 2 different types of models: Unigram and WordLevel.
									for tokenizer, pretrained_name, kwargs in self.tokenizers_list:
															with self.subTest(f'''{tokenizer.__class__.__name__} ({pretrained_name})'''							):
																					try:
																											snake_case_       =     tempfile.mkdtemp()
																											snake_case_       =     self.rust_tokenizer_class.from_pretrained(_UpperCamelCase							,       **_UpperCamelCase							)

																											tokenizer.save_pretrained(self.tmpdirname							)
																											super().test_training_new_tokenizer_with_special_tokens_change()
																					finally:
																											# Even if the test fails, we must be sure that the folder is deleted and that the default tokenizer
																											# is restored
																											shutil.rmtree(self.tmpdirname							)
																											snake_case_       =     tmpdirname_orig







			def     snake_case__(					self			:    Optional[int]							)				->Optional[Any]:
									snake_case_       =     self.rust_tokenizer_class.from_pretrained(self.bytelevel_bpe_model_name							)

									snake_case_       =     ('''a''' for _ in range(1_0_0_0							))
									snake_case_       =     tokenizer.train_new_from_iterator(text_iterator=_UpperCamelCase							,       length=1_0_0_0							,       vocab_size=5_0							)

									snake_case_       =     new_tokenizer.encode('''a🤗'''							)
									self.assertEqual(_UpperCamelCase							,       [6_4, 1_7_2, 2_5_3, 9_7, 2_4_5]							)






@require_tokenizers
class      snake_case_       (		unittest.TestCase ):







			'''simple docstring'''

			def     snake_case__(					self			:    Optional[int]							)				->int:
									snake_case_       =     AutoTokenizer.from_pretrained('''bert-base-cased'''							)
									snake_case_       =     json.loads(tokenizer._tokenizer.to_str()							)
									snake_case_       =     len(_UpperCamelCase							)

									with tempfile.TemporaryDirectory() as tmp_dir:
															# Hack to save this in the tokenizer_config.json
															snake_case_       =     ['''tokenizer.4.0.0.json''']
															tokenizer.save_pretrained(_UpperCamelCase							)
															json.dump(_UpperCamelCase							,       open(os.path.join(_UpperCamelCase							,       '''tokenizer.4.0.0.json'''							)							,       '''w'''							)							)

															# This should pick the new tokenizer file as the version of Transformers is > 4.0.0
															snake_case_       =     AutoTokenizer.from_pretrained(_UpperCamelCase							)
															self.assertEqual(len(_UpperCamelCase							)							,       len(_UpperCamelCase							) + 1							)
															snake_case_       =     json.loads(new_tokenizer._tokenizer.to_str()							)
															self.assertIn('''huggingface'''							,       json_tokenizer['''model''']['''vocab''']							)

															# Will need to be adjusted if we reach v42 and this test is still here.
															# Should pick the old tokenizer file as the version of Transformers is < 4.0.0
															shutil.move(os.path.join(_UpperCamelCase							,       '''tokenizer.4.0.0.json'''							)							,       os.path.join(_UpperCamelCase							,       '''tokenizer.42.0.0.json'''							)							)
															snake_case_       =     ['''tokenizer.42.0.0.json''']
															tokenizer.save_pretrained(_UpperCamelCase							)
															snake_case_       =     AutoTokenizer.from_pretrained(_UpperCamelCase							)
															self.assertEqual(len(_UpperCamelCase							)							,       len(_UpperCamelCase							)							)
															snake_case_       =     json.loads(new_tokenizer._tokenizer.to_str()							)
															self.assertNotIn('''huggingface'''							,       json_tokenizer['''model''']['''vocab''']							)







			def     snake_case__(					self			:    List[str]							)				->Optional[Any]:
									# This repo has two tokenizer files, one for v4.0.0 and above with an added token, one for versions lower.
									snake_case_       =     '''hf-internal-testing/test-two-tokenizers'''

									# This should pick the new tokenizer file as the version of Transformers is > 4.0.0
									snake_case_       =     AutoTokenizer.from_pretrained(_UpperCamelCase							)
									self.assertEqual(len(_UpperCamelCase							)							,       2_8_9_9_7							)
									snake_case_       =     json.loads(tokenizer._tokenizer.to_str()							)
									self.assertIn('''huggingface'''							,       json_tokenizer['''model''']['''vocab''']							)

									# Testing an older version by monkey-patching the version in the module it's used.
									import transformers as old_transformers

									snake_case_       =     '''3.0.0'''
									snake_case_       =     old_transformers.models.auto.AutoTokenizer.from_pretrained(_UpperCamelCase							)
									self.assertEqual(len(_UpperCamelCase							)							,       2_8_9_9_6							)
									snake_case_       =     json.loads(old_tokenizer._tokenizer.to_str()							)
									self.assertNotIn('''huggingface'''							,       json_tokenizer['''model''']['''vocab''']							)






@require_tokenizers
class      snake_case_       (		unittest.TestCase ):







			'''simple docstring'''

			def     snake_case__(					self			:    str							)				->Optional[int]:
									# See https://github.com/huggingface/transformers/pull/12550
									# and https://github.com/huggingface/tokenizers/issues/537
									snake_case_       =     PreTrainedTokenizerFast.from_pretrained('''robot-test/dummy-tokenizer-wordlevel'''							)
									snake_case_       =     '''The Matrix is a 1999 science fiction action film.'''

									with concurrent.futures.ThreadPoolExecutor() as executor:
															snake_case_       =     [executor.submit(self.fetch							,       _UpperCamelCase							,       _UpperCamelCase							) for i in range(1_0							)]
															snake_case_       =     [future.result() for future in futures]
															self.assertEqual(_UpperCamelCase							,       [[1, 1_0, 0, 8, 0, 1_8, 0, 0, 0, 2] for i in range(1_0							)]							)







			def     snake_case__(					self			:    Optional[int]							,       _UpperCamelCase			:    Dict							,       _UpperCamelCase			:    str							)				->Tuple:
									return tokenizer.encode(_UpperCamelCase							,       truncation='''longest_first'''							,       padding='''longest'''							)