



'''simple docstring'''




import unittest

from transformers import SqueezeBertConfig, is_torch_available
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
							import torch

							from transformers import (
							    SQUEEZEBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
							    SqueezeBertForMaskedLM,
							    SqueezeBertForMultipleChoice,
							    SqueezeBertForQuestionAnswering,
							    SqueezeBertForSequenceClassification,
							    SqueezeBertForTokenClassification,
							    SqueezeBertModel,
							)

class     UpperCamelCase_				(							__magic_name__    ):





	def __init__(						self	,	A	,	A=13	,	A=7	,	A=True	,	A=True	,	A=False	,	A=True	,	A=99	,	A=32	,	A=5	,	A=4	,	A=64	,	A="gelu"	,	A=0.1	,	A=0.1	,	A=512	,	A=16	,	A=2	,	A=0.0_2	,	A=3	,	A=4	,	A=None	,	A=2	,	A=2	,	A=2	,	A=2	,	A=4	,	A=1	,	)    ->    int:
								UpperCAmelCase  :	str							  =	parent
								UpperCAmelCase  :	List[str]							  =	batch_size
								UpperCAmelCase  :	List[Any]							  =	seq_length
								UpperCAmelCase  :	int							  =	is_training
								UpperCAmelCase  :	Dict							  =	use_input_mask
								UpperCAmelCase  :	Tuple							  =	use_token_type_ids
								UpperCAmelCase  :	Dict							  =	use_labels
								UpperCAmelCase  :	Union[str, Any]							  =	vocab_size
								UpperCAmelCase  :	Dict							  =	hidden_size
								UpperCAmelCase  :	Any							  =	num_hidden_layers
								UpperCAmelCase  :	str							  =	num_attention_heads
								UpperCAmelCase  :	int							  =	intermediate_size
								UpperCAmelCase  :	List[str]							  =	hidden_act
								UpperCAmelCase  :	Optional[int]							  =	hidden_dropout_prob
								UpperCAmelCase  :	int							  =	attention_probs_dropout_prob
								UpperCAmelCase  :	Optional[int]							  =	max_position_embeddings
								UpperCAmelCase  :	Optional[int]							  =	type_vocab_size
								UpperCAmelCase  :	Union[str, Any]							  =	type_sequence_label_size
								UpperCAmelCase  :	List[Any]							  =	initializer_range
								UpperCAmelCase  :	str							  =	num_labels
								UpperCAmelCase  :	Union[str, Any]							  =	num_choices
								UpperCAmelCase  :	Dict							  =	scope
								UpperCAmelCase  :	Optional[Any]							  =	q_groups
								UpperCAmelCase  :	Dict							  =	k_groups
								UpperCAmelCase  :	Any							  =	v_groups
								UpperCAmelCase  :	Union[str, Any]							  =	post_attention_groups
								UpperCAmelCase  :	List[str]							  =	intermediate_groups
								UpperCAmelCase  :	Dict							  =	output_groups





	def 				_lowercase(						self					)    ->    Optional[Any]:
								UpperCAmelCase  :	int							  =	ids_tensor([self.batch_size, self.seq_length]	,	self.vocab_size					)

								UpperCAmelCase  :	Optional[Any]							  =	None
								if self.use_input_mask:
															UpperCAmelCase  :	Optional[Any]							  =	random_attention_mask([self.batch_size, self.seq_length]					)

								UpperCAmelCase  :	str							  =	None
								UpperCAmelCase  :	List[Any]							  =	None
								UpperCAmelCase  :	Optional[int]							  =	None
								if self.use_labels:
															UpperCAmelCase  :	Tuple							  =	ids_tensor([self.batch_size]	,	self.type_sequence_label_size					)
															UpperCAmelCase  :	Optional[Any]							  =	ids_tensor([self.batch_size, self.seq_length]	,	self.num_labels					)
															UpperCAmelCase  :	Dict							  =	ids_tensor([self.batch_size]	,	self.num_choices					)

								UpperCAmelCase  :	int							  =	self.get_config()

								return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels





	def 				_lowercase(						self					)    ->    int:
								return SqueezeBertConfig(
								    embedding_size=self.hidden_size	,	vocab_size=self.vocab_size	,	hidden_size=self.hidden_size	,	num_hidden_layers=self.num_hidden_layers	,	num_attention_heads=self.num_attention_heads	,	intermediate_size=self.intermediate_size	,	hidden_act=self.hidden_act	,	attention_probs_dropout_prob=self.hidden_dropout_prob	,	attention_dropout=self.attention_probs_dropout_prob	,	max_position_embeddings=self.max_position_embeddings	,	initializer_range=self.initializer_range	,	q_groups=self.q_groups	,	k_groups=self.k_groups	,	v_groups=self.v_groups	,	post_attention_groups=self.post_attention_groups	,	intermediate_groups=self.intermediate_groups	,	output_groups=self.output_groups	,	)





	def 				_lowercase(						self	,	A	,	A	,	A	,	A	,	A	,	A					)    ->    Tuple:
								UpperCAmelCase  :	Any							  =	SqueezeBertModel(config=A					)
								model.to(A					)
								model.eval()
								UpperCAmelCase  :	List[str]							  =	model(A	,	A					)
								UpperCAmelCase  :	Any							  =	model(A					)
								self.parent.assertEqual(result.last_hidden_state.shape	,	(self.batch_size, self.seq_length, self.hidden_size)					)





	def 				_lowercase(						self	,	A	,	A	,	A	,	A	,	A	,	A					)    ->    Dict:
								UpperCAmelCase  :	List[Any]							  =	SqueezeBertForMaskedLM(config=A					)
								model.to(A					)
								model.eval()
								UpperCAmelCase  :	int							  =	model(A	,	attention_mask=A	,	labels=A					)
								self.parent.assertEqual(result.logits.shape	,	(self.batch_size, self.seq_length, self.vocab_size)					)





	def 				_lowercase(						self	,	A	,	A	,	A	,	A	,	A	,	A					)    ->    Optional[int]:
								UpperCAmelCase  :	Optional[Any]							  =	SqueezeBertForQuestionAnswering(config=A					)
								model.to(A					)
								model.eval()
								UpperCAmelCase  :	Optional[int]							  =	model(
								    A	,	attention_mask=A	,	start_positions=A	,	end_positions=A					)
								self.parent.assertEqual(result.start_logits.shape	,	(self.batch_size, self.seq_length)					)
								self.parent.assertEqual(result.end_logits.shape	,	(self.batch_size, self.seq_length)					)





	def 				_lowercase(						self	,	A	,	A	,	A	,	A	,	A	,	A					)    ->    Optional[Any]:
								UpperCAmelCase  :	Optional[int]							  =	self.num_labels
								UpperCAmelCase  :	int							  =	SqueezeBertForSequenceClassification(A					)
								model.to(A					)
								model.eval()
								UpperCAmelCase  :	Optional[int]							  =	model(A	,	attention_mask=A	,	labels=A					)
								self.parent.assertEqual(result.logits.shape	,	(self.batch_size, self.num_labels)					)





	def 				_lowercase(						self	,	A	,	A	,	A	,	A	,	A	,	A					)    ->    Union[str, Any]:
								UpperCAmelCase  :	str							  =	self.num_labels
								UpperCAmelCase  :	Dict							  =	SqueezeBertForTokenClassification(config=A					)
								model.to(A					)
								model.eval()

								UpperCAmelCase  :	Optional[Any]							  =	model(A	,	attention_mask=A	,	labels=A					)
								self.parent.assertEqual(result.logits.shape	,	(self.batch_size, self.seq_length, self.num_labels)					)





	def 				_lowercase(						self	,	A	,	A	,	A	,	A	,	A	,	A					)    ->    List[str]:
								UpperCAmelCase  :	Union[str, Any]							  =	self.num_choices
								UpperCAmelCase  :	Any							  =	SqueezeBertForMultipleChoice(config=A					)
								model.to(A					)
								model.eval()
								UpperCAmelCase  :	Tuple							  =	input_ids.unsqueeze(1					).expand(-1	,	self.num_choices	,	-1					).contiguous()
								UpperCAmelCase  :	str							  =	input_mask.unsqueeze(1					).expand(-1	,	self.num_choices	,	-1					).contiguous()
								UpperCAmelCase  :	List[Any]							  =	model(
								    A	,	attention_mask=A	,	labels=A	,	)
								self.parent.assertEqual(result.logits.shape	,	(self.batch_size, self.num_choices)					)






	def 				_lowercase(						self					)    ->    List[str]:
								UpperCAmelCase  :	Any							  =	self.prepare_config_and_inputs()
								(UpperCAmelCase)  :	Optional[int]							  =	config_and_inputs
								UpperCAmelCase  :	Dict							  =	{"""input_ids""": input_ids, """attention_mask""": input_mask}
								return config, inputs_dict




@require_torch
class     UpperCamelCase_				(							__magic_name__   ,  __magic_name__   ,  unittest.TestCase    ):
	lowercase    =				(
	    (
	        SqueezeBertModel,
	        SqueezeBertForMaskedLM,
	        SqueezeBertForMultipleChoice,
	        SqueezeBertForQuestionAnswering,
	        SqueezeBertForSequenceClassification,
	        SqueezeBertForTokenClassification,
	    )
	    if is_torch_available()
	    else None
	)
	lowercase    =				(
	    {
	        'feature-extraction': SqueezeBertModel,
	        'fill-mask': SqueezeBertForMaskedLM,
	        'question-answering': SqueezeBertForQuestionAnswering,
	        'text-classification': SqueezeBertForSequenceClassification,
	        'token-classification': SqueezeBertForTokenClassification,
	        'zero-shot': SqueezeBertForSequenceClassification,
	    }
	    if is_torch_available()
	    else {}
	)
	lowercase    =				False
	lowercase    =				True
	lowercase    =				False





	def 				_lowercase(						self					)    ->    List[Any]:
								UpperCAmelCase  :	Tuple							  =	SqueezeBertModelTester(self					)
								UpperCAmelCase  :	Optional[Any]							  =	ConfigTester(self	,	config_class=A	,	dim=37					)





	def 				_lowercase(						self					)    ->    Any:
								self.config_tester.run_common_tests()





	def 				_lowercase(						self					)    ->    int:
								UpperCAmelCase  :	int							  =	self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_squeezebert_model(*A					)





	def 				_lowercase(						self					)    ->    Optional[Any]:
								UpperCAmelCase  :	Optional[int]							  =	self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_squeezebert_for_masked_lm(*A					)





	def 				_lowercase(						self					)    ->    List[str]:
								UpperCAmelCase  :	Optional[int]							  =	self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_squeezebert_for_question_answering(*A					)





	def 				_lowercase(						self					)    ->    Tuple:
								UpperCAmelCase  :	int							  =	self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_squeezebert_for_sequence_classification(*A					)





	def 				_lowercase(						self					)    ->    Optional[Any]:
								UpperCAmelCase  :	str							  =	self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_squeezebert_for_token_classification(*A					)





	def 				_lowercase(						self					)    ->    Any:
								UpperCAmelCase  :	Tuple							  =	self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_squeezebert_for_multiple_choice(*A					)






	@slow
	def 				_lowercase(						self					)    ->    Optional[int]:
								for model_name in SQUEEZEBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
															UpperCAmelCase  :	List[Any]							  =	SqueezeBertModel.from_pretrained(A					)
															self.assertIsNotNone(A					)




@require_sentencepiece
@require_tokenizers
@require_torch
class     UpperCamelCase_				(							unittest.TestCase    ):





	@slow
	def 				_lowercase(						self					)    ->    List[Any]:
								UpperCAmelCase  :	Optional[int]							  =	SqueezeBertForSequenceClassification.from_pretrained("""squeezebert/squeezebert-mnli"""					)

								UpperCAmelCase  :	Any							  =	torch.tensor([[1, 29414, 232, 328, 740, 1140, 12695, 69, 13, 1588, 2]]					)
								UpperCAmelCase  :	Dict							  =	model(A					)[0]
								UpperCAmelCase  :	Dict							  =	torch.Size((1, 3)					)
								self.assertEqual(output.shape	,	A					)
								UpperCAmelCase  :	Union[str, Any]							  =	torch.tensor([[0.6_4_0_1, -0.0_3_4_9, -0.6_0_4_1]]					)
								self.assertTrue(torch.allclose(A	,	A	,	atol=1e-4					)					)






'''simple docstring'''




# Lint as: python3
import sys
from collections.abc import Mapping
from typing import TYPE_CHECKING, Dict, Optional

import numpy as np
import pyarrow as pa

from .. import config
from ..utils.logging import get_logger
from ..utils.py_utils import map_nested
from .formatting import TensorFormatter


if TYPE_CHECKING:
							import jax
							import jaxlib

a     :    Any      						=					get_logger()

a     :    Optional[dict]      						=					None

class     UpperCamelCase_				(							TensorFormatter[Mapping, 'jax.Array', Mapping]    ):





	def __init__(						self	,	A=None	,	A=None	,	**A					)    ->    str:
								super().__init__(features=A					)
								import jax
								from jaxlib.xla_client import Device

								if isinstance(A	,	A					):
															raise ValueError(
															    f'''Expected {device} to be a `str` not {type(A					)}, as `jaxlib.xla_extension.Device` '''
															    """is not serializable neither with `pickle` nor with `dill`. Instead you can surround """
															    """the device with `str()` to get its string identifier that will be internally mapped """
															    """to the actual `jaxlib.xla_extension.Device`."""					)
								UpperCAmelCase  :	Optional[int]							  =	device if isinstance(A	,	A					) else str(jax.devices()[0]					)
								# using global variable since `jaxlib.xla_extension.Device` is not serializable neither
								# with `pickle` nor with `dill`, so we need to use a global variable instead
								global DEVICE_MAPPING
								if DEVICE_MAPPING is None:
															UpperCAmelCase  :	Any							  =	self._map_devices_to_str()
								if self.device not in list(DEVICE_MAPPING.keys()					):
															logger.warning(
															    f'''Device with string identifier {self.device} not listed among the available '''
															    f'''devices: {list(DEVICE_MAPPING.keys()					)}, so falling back to the default '''
															    f'''device: {str(jax.devices()[0]					)}.'''					)
															UpperCAmelCase  :	List[Any]							  =	str(jax.devices()[0]					)
								UpperCAmelCase  :	Union[str, Any]							  =	jnp_array_kwargs





	@staticmethod
	def 				_lowercase(						)    ->    Dict[str, "jaxlib.xla_extension.Device"]:
								import jax

								return {str(A					): device for device in jax.devices()}





	def 				_lowercase(						self	,	A					)    ->    str:
								import jax
								import jax.numpy as jnp

								if isinstance(A	,	A					) and column:
															if all(
															    isinstance(A	,	jax.Array					) and x.shape == column[0].shape and x.dtype == column[0].dtype for x in column					):
																						return jnp.stack(A	,	axis=0					)
								return column





	def 				_lowercase(						self	,	A					)    ->    Tuple:
								import jax
								import jax.numpy as jnp

								if isinstance(A	,	(str, bytes, type(A					))					):
															return value
								elif isinstance(A	,	(np.character, np.ndarray)					) and np.issubdtype(value.dtype	,	np.character					):
															return value.tolist()

								UpperCAmelCase  :	List[str]							  =	{}

								if isinstance(A	,	(np.number, np.ndarray)					) and np.issubdtype(value.dtype	,	np.integer					):
															# the default int precision depends on the jax config
															# see https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#double-64bit-precision
															if jax.config.jax_enable_xaa:
																						UpperCAmelCase  :	str							  =	{"""dtype""": jnp.intaa}
															else:
																						UpperCAmelCase  :	int							  =	{"""dtype""": jnp.intaa}
								elif isinstance(A	,	(np.number, np.ndarray)					) and np.issubdtype(value.dtype	,	np.floating					):
															UpperCAmelCase  :	Any							  =	{"""dtype""": jnp.floataa}
								elif config.PIL_AVAILABLE and "PIL" in sys.modules:
															import PIL.Image

															if isinstance(A	,	PIL.Image.Image					):
																						UpperCAmelCase  :	List[str]							  =	np.asarray(A					)

        # using global variable since `jaxlib.xla_extension.Device` is not serializable neither
        # with `pickle` nor with `dill`, so we need to use a global variable instead
								global DEVICE_MAPPING
								if DEVICE_MAPPING is None:
															UpperCAmelCase  :	Dict							  =	self._map_devices_to_str()

								with jax.default_device(DEVICE_MAPPING[self.device]					):
															# calling jnp.array on a np.ndarray does copy the data
															# see https://github.com/google/jax/issues/4486
															return jnp.array(A	,	**{**default_dtype, **self.jnp_array_kwargs}					)





	def 				_lowercase(						self	,	A					)    ->    Tuple:
								import jax

								# support for torch, tf, jax etc.
								if config.TORCH_AVAILABLE and "torch" in sys.modules:
															import torch

															if isinstance(A	,	torch.Tensor					):
																						return self._tensorize(data_struct.detach().cpu().numpy()[()]					)
								if hasattr(A	,	"""__array__"""					) and not isinstance(A	,	jax.Array					):
															UpperCAmelCase  :	Optional[int]							  =	data_struct.__array__()
								# support for nested types like struct of list of struct
								if isinstance(A	,	np.ndarray					):
															if data_struct.dtype == object:  # jax arrays cannot be instantied from an array of objects
																						return self._consolidate([self.recursive_tensorize(A					) for substruct in data_struct]					)
								elif isinstance(A	,	(list, tuple)					):
															return self._consolidate([self.recursive_tensorize(A					) for substruct in data_struct]					)
								return self._tensorize(A					)





	def 				_lowercase(						self	,	A					)    ->    Dict:
								return map_nested(self._recursive_tensorize	,	A	,	map_list=A					)





	def 				_lowercase(						self	,	A					)    ->    Mapping:
								UpperCAmelCase  :	Union[str, Any]							  =	self.numpy_arrow_extractor().extract_row(A					)
								UpperCAmelCase  :	Dict							  =	self.python_features_decoder.decode_row(A					)
								return self.recursive_tensorize(A					)





	def 				_lowercase(						self	,	A					)    ->    "jax.Array":
								UpperCAmelCase  :	int							  =	self.numpy_arrow_extractor().extract_column(A					)
								UpperCAmelCase  :	Optional[Any]							  =	self.python_features_decoder.decode_column(A	,	pa_table.column_names[0]					)
								UpperCAmelCase  :	Optional[int]							  =	self.recursive_tensorize(A					)
								UpperCAmelCase  :	Any							  =	self._consolidate(A					)
								return column






	def 				_lowercase(						self	,	A					)    ->    Mapping:
								UpperCAmelCase  :	Optional[int]							  =	self.numpy_arrow_extractor().extract_batch(A					)
								UpperCAmelCase  :	List[str]							  =	self.python_features_decoder.decode_batch(A					)
								UpperCAmelCase  :	Union[str, Any]							  =	self.recursive_tensorize(A					)
								for column_name in batch:
															UpperCAmelCase  :	Optional[Any]							  =	self._consolidate(batch[column_name]					)
								return batch

