



import unittest

import numpy as np

from transformers import is_flax_available
from transformers.testing_utils import require_flax

from ..test_modeling_flax_common import ids_tensor


if is_flax_available():
 import jax
 import jax.numpy as jnp

 from transformers.generation import (
     FlaxForcedBOSTokenLogitsProcessor,
     FlaxForcedEOSTokenLogitsProcessor,
     FlaxLogitsProcessorList,
     FlaxMinLengthLogitsProcessor,
     FlaxTemperatureLogitsWarper,
     FlaxTopKLogitsWarper,
     FlaxTopPLogitsWarper,
 )

@require_flax
class     _A       (	unittest.TestCase    ):



       def 		_lowerCamelCase			(		self							:     Union[str, Any]     ,		__SCREAMING_SNAKE_CASE							:     int     ,		__SCREAMING_SNAKE_CASE							:     int):


           '''simple docstring'''

           __a						       =				jnp.ones((batch_size, length)) / length
           return scores



       def 		_lowerCamelCase			(		self							:     Tuple):


           '''simple docstring'''

           __a						       =				None
           __a						       =				20

           __a						       =				self._get_uniform_logits(batch_size=2     ,		length=__SCREAMING_SNAKE_CASE)

           # tweak scores to not be uniform anymore
           __a						       =				scores.at[1, 5].set((1 / length) + 0.1)  # peak, 1st batch
           __a						       =				scores.at[1, 10].set((1 / length) - 0.4)  # valley, 1st batch

           # compute softmax
           __a						       =				jax.nn.softmax(__SCREAMING_SNAKE_CASE     ,		axis=-1)

           __a						       =				FlaxTemperatureLogitsWarper(temperature=0.5)
           __a						       =				FlaxTemperatureLogitsWarper(temperature=1.3)

           __a						       =				jax.nn.softmax(temp_dist_warper_sharper(__SCREAMING_SNAKE_CASE     ,		scores.copy()     ,		cur_len=__SCREAMING_SNAKE_CASE)     ,		axis=-1)
           __a						       =				jax.nn.softmax(temp_dist_warper_smoother(__SCREAMING_SNAKE_CASE     ,		scores.copy()     ,		cur_len=__SCREAMING_SNAKE_CASE)     ,		axis=-1)

           # uniform distribution stays uniform
           self.assertTrue(jnp.allclose(probs[0, :]     ,		warped_prob_sharp[0, :]     ,		atol=1E-3))
           self.assertTrue(jnp.allclose(probs[0, :]     ,		warped_prob_smooth[0, :]     ,		atol=1E-3))

           # sharp peaks get higher, valleys get lower
           self.assertLess(probs[1, :].max()     ,		warped_prob_sharp[1, :].max())
           self.assertGreater(probs[1, :].min()     ,		warped_prob_sharp[1, :].min())

           # smooth peaks get lower, valleys get higher
           self.assertGreater(probs[1, :].max()     ,		warped_prob_smooth[1, :].max())
           self.assertLess(probs[1, :].min()     ,		warped_prob_smooth[1, :].min())



       def 		_lowerCamelCase			(		self							:     List[Any]):


           '''simple docstring'''

           __a						       =				None
           __a						       =				10
           __a						       =				2

           # create ramp distribution
           __a						       =				np.broadcast_to(np.arange(__SCREAMING_SNAKE_CASE)[None, :]     ,		(batch_size, vocab_size)).copy()
           __a						       =				ramp_logits[1:, : vocab_size // 2] + vocab_size

           __a						       =				FlaxTopKLogitsWarper(3)

           __a						       =				top_k_warp(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE     ,		cur_len=__SCREAMING_SNAKE_CASE)

           # check that correct tokens are filtered
           self.assertListEqual(jnp.isinf(scores[0]).tolist()     ,		7 * [True] + 3 * [False])
           self.assertListEqual(jnp.isinf(scores[1]).tolist()     ,		2 * [True] + 3 * [False] + 5 * [True])

           # check special case
           __a						       =				5
           __a						       =				FlaxTopKLogitsWarper(top_k=1     ,		filter_value=0.0     ,		min_tokens_to_keep=3)

           __a						       =				np.broadcast_to(np.arange(__SCREAMING_SNAKE_CASE)[None, :]     ,		(batch_size, length)).copy()
           __a						       =				top_k_warp_safety_check(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE     ,		cur_len=__SCREAMING_SNAKE_CASE)

           # min_tokens overwrites k: 3 tokens are kept => 2 tokens are nullified
           self.assertListEqual((scores == 0.0).sum(axis=-1).tolist()     ,		[2, 2])



       def 		_lowerCamelCase			(		self							:     Dict):


           '''simple docstring'''

           __a						       =				None
           __a						       =				10
           __a						       =				2

           # create distribution and take log (inverse to Softmax as taken in TopPLogitsWarper)
           __a						       =				np.log(np.array([[0.3, 0.1, 0.1, 0.5], [0.15, 0.3, 0.3, 0.25]]))

           __a						       =				FlaxTopPLogitsWarper(0.8)
           __a						       =				np.exp(top_p_warp(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE     ,		cur_len=__SCREAMING_SNAKE_CASE))

           # dist should be filtered to keep min num values so that sum is >= top_p
           # exp (-inf) => 0
           __a						       =				np.array([[0.3, 0.0, 0.0, 0.5], [0.0, 0.3, 0.3, 0.25]])
           self.assertTrue(np.allclose(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE     ,		atol=1E-3))

           # check edge cases with negative and extreme logits
           __a						       =				np.broadcast_to(np.arange(__SCREAMING_SNAKE_CASE)[None, :]     ,		(batch_size, vocab_size)).copy() - (
               vocab_size // 2
           )

           # make ramp_logits more extreme
           __a						       =				ramp_logits[1] * 1_00.0

           # make sure at least 2 tokens are kept
           __a						       =				FlaxTopPLogitsWarper(0.9     ,		min_tokens_to_keep=2     ,		filter_value=0.0)
           __a						       =				top_p_warp(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE     ,		cur_len=__SCREAMING_SNAKE_CASE)

           # first batch should keep three tokens, second batch would keep only 1, but due to `min_tokens_to_keep=2` keeps 2.
           self.assertListEqual((filtered_dist != 0.0).sum(axis=-1).tolist()     ,		[3, 2])



       def 		_lowerCamelCase			(		self							:     Union[str, Any]):


           '''simple docstring'''

           __a						       =				20
           __a						       =				4
           __a						       =				0

           __a						       =				FlaxMinLengthLogitsProcessor(min_length=10     ,		eos_token_id=__SCREAMING_SNAKE_CASE)

           # check that min length is applied at length 5
           __a						       =				ids_tensor((batch_size, 20)     ,		vocab_size=20)
           __a						       =				5
           __a						       =				self._get_uniform_logits(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE)
           __a						       =				min_dist_processor(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE     ,		cur_len=__SCREAMING_SNAKE_CASE)
           self.assertListEqual(scores_before_min_length[:, eos_token_id].tolist()     ,		4 * [-float('''inf''')])

           # check that min length is not applied anymore at length 15
           __a						       =				self._get_uniform_logits(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE)
           __a						       =				15
           __a						       =				min_dist_processor(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE     ,		cur_len=__SCREAMING_SNAKE_CASE)
           self.assertFalse(jnp.isinf(__SCREAMING_SNAKE_CASE).any())



       def 		_lowerCamelCase			(		self							:     Optional[Any]):


           '''simple docstring'''

           __a						       =				20
           __a						       =				4
           __a						       =				0

           __a						       =				FlaxForcedBOSTokenLogitsProcessor(bos_token_id=__SCREAMING_SNAKE_CASE)

           # check that all scores are -inf except the bos_token_id score
           __a						       =				ids_tensor((batch_size, 1)     ,		vocab_size=20)
           __a						       =				1
           __a						       =				self._get_uniform_logits(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE)
           __a						       =				logits_processor(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE     ,		cur_len=__SCREAMING_SNAKE_CASE)
           self.assertTrue(jnp.isneginf(scores[:, bos_token_id + 1 :]).all())
           self.assertListEqual(scores[:, bos_token_id].tolist()     ,		4 * [0])  # score for bos_token_id shold be zero

           # check that bos_token_id is not forced if current length is greater than 1
           __a						       =				3
           __a						       =				self._get_uniform_logits(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE)
           __a						       =				logits_processor(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE     ,		cur_len=__SCREAMING_SNAKE_CASE)
           self.assertFalse(jnp.isinf(__SCREAMING_SNAKE_CASE).any())



       def 		_lowerCamelCase			(		self							:     Tuple):


           '''simple docstring'''

           __a						       =				20
           __a						       =				4
           __a						       =				0
           __a						       =				5

           __a						       =				FlaxForcedEOSTokenLogitsProcessor(max_length=__SCREAMING_SNAKE_CASE     ,		eos_token_id=__SCREAMING_SNAKE_CASE)

           # check that all scores are -inf except the eos_token_id when max_length is reached
           __a						       =				ids_tensor((batch_size, 4)     ,		vocab_size=20)
           __a						       =				4
           __a						       =				self._get_uniform_logits(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE)
           __a						       =				logits_processor(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE     ,		cur_len=__SCREAMING_SNAKE_CASE)
           self.assertTrue(jnp.isneginf(scores[:, eos_token_id + 1 :]).all())
           self.assertListEqual(scores[:, eos_token_id].tolist()     ,		4 * [0])  # score for eos_token_id should be zero

           # check that eos_token_id is not forced if max_length is not reached
           __a						       =				3
           __a						       =				self._get_uniform_logits(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE)
           __a						       =				logits_processor(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE     ,		cur_len=__SCREAMING_SNAKE_CASE)
           self.assertFalse(jnp.isinf(__SCREAMING_SNAKE_CASE).any())



       def 		_lowerCamelCase			(		self							:     str):


           '''simple docstring'''

           __a						       =				4
           __a						       =				10
           __a						       =				15
           __a						       =				2
           __a						       =				1
           __a						       =				15

           # dummy input_ids and scores
           __a						       =				ids_tensor((batch_size, sequence_length)     ,		__SCREAMING_SNAKE_CASE)
           __a						       =				input_ids.copy()

           __a						       =				self._get_uniform_logits(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE)
           __a						       =				scores.copy()

           # instantiate all dist processors
           __a						       =				FlaxTemperatureLogitsWarper(temperature=0.5)
           __a						       =				FlaxTopKLogitsWarper(3)
           __a						       =				FlaxTopPLogitsWarper(0.8)

           # instantiate all logits processors
           __a						       =				FlaxMinLengthLogitsProcessor(min_length=10     ,		eos_token_id=__SCREAMING_SNAKE_CASE)
           __a						       =				FlaxForcedBOSTokenLogitsProcessor(bos_token_id=__SCREAMING_SNAKE_CASE)
           __a						       =				FlaxForcedEOSTokenLogitsProcessor(max_length=__SCREAMING_SNAKE_CASE     ,		eos_token_id=__SCREAMING_SNAKE_CASE)

           __a						       =				10

           # no processor list
           __a						       =				temp_dist_warp(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE     ,		cur_len=__SCREAMING_SNAKE_CASE)
           __a						       =				top_k_warp(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE     ,		cur_len=__SCREAMING_SNAKE_CASE)
           __a						       =				top_p_warp(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE     ,		cur_len=__SCREAMING_SNAKE_CASE)
           __a						       =				min_dist_proc(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE     ,		cur_len=__SCREAMING_SNAKE_CASE)
           __a						       =				bos_dist_proc(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE     ,		cur_len=__SCREAMING_SNAKE_CASE)
           __a						       =				eos_dist_proc(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE     ,		cur_len=__SCREAMING_SNAKE_CASE)

           # with processor list
           __a						       =				FlaxLogitsProcessorList(
               [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc])
           __a						       =				processor(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE     ,		cur_len=__SCREAMING_SNAKE_CASE)

           # scores should be equal
           self.assertTrue(jnp.allclose(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE     ,		atol=1E-3))

           # input_ids should never be changed
           self.assertListEqual(input_ids.tolist()     ,		input_ids_comp.tolist())





       def 		_lowerCamelCase			(		self							:     Any):


           '''simple docstring'''

           __a						       =				4
           __a						       =				10
           __a						       =				15
           __a						       =				2
           __a						       =				1
           __a						       =				15

           # dummy input_ids and scores
           __a						       =				ids_tensor((batch_size, sequence_length)     ,		__SCREAMING_SNAKE_CASE)
           __a						       =				input_ids.copy()

           __a						       =				self._get_uniform_logits(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE)
           __a						       =				scores.copy()

           # instantiate all dist processors
           __a						       =				FlaxTemperatureLogitsWarper(temperature=0.5)
           __a						       =				FlaxTopKLogitsWarper(3)
           __a						       =				FlaxTopPLogitsWarper(0.8)

           # instantiate all logits processors
           __a						       =				FlaxMinLengthLogitsProcessor(min_length=10     ,		eos_token_id=__SCREAMING_SNAKE_CASE)
           __a						       =				FlaxForcedBOSTokenLogitsProcessor(bos_token_id=__SCREAMING_SNAKE_CASE)
           __a						       =				FlaxForcedEOSTokenLogitsProcessor(max_length=__SCREAMING_SNAKE_CASE     ,		eos_token_id=__SCREAMING_SNAKE_CASE)

           __a						       =				10

           # no processor list
           def run_no_processor_list(__SCREAMING_SNAKE_CASE							:     List[Any]     ,		__SCREAMING_SNAKE_CASE							:     List[Any]     ,		__SCREAMING_SNAKE_CASE							:     Optional[Any]):
               __a						       =				temp_dist_warp(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE     ,		cur_len=__SCREAMING_SNAKE_CASE)
               __a						       =				top_k_warp(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE     ,		cur_len=__SCREAMING_SNAKE_CASE)
               __a						       =				top_p_warp(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE     ,		cur_len=__SCREAMING_SNAKE_CASE)
               __a						       =				min_dist_proc(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE     ,		cur_len=__SCREAMING_SNAKE_CASE)
               __a						       =				bos_dist_proc(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE     ,		cur_len=__SCREAMING_SNAKE_CASE)
               __a						       =				eos_dist_proc(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE     ,		cur_len=__SCREAMING_SNAKE_CASE)
               return scores

           # with processor list
           def run_processor_list(__SCREAMING_SNAKE_CASE							:     str     ,		__SCREAMING_SNAKE_CASE							:     Any     ,		__SCREAMING_SNAKE_CASE							:     Optional[Any]):
               __a						       =				FlaxLogitsProcessorList(
                   [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc])
               __a						       =				processor(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE     ,		cur_len=__SCREAMING_SNAKE_CASE)
               return scores

           __a						       =				jax.jit(__SCREAMING_SNAKE_CASE)
           __a						       =				jax.jit(__SCREAMING_SNAKE_CASE)

           __a						       =				jitted_run_no_processor_list(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE)
           __a						       =				jitted_run_processor_list(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE)

           # scores should be equal
           self.assertTrue(jnp.allclose(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE     ,		atol=1E-3))

           # input_ids should never be changed
           self.assertListEqual(input_ids.tolist()     ,		input_ids_comp.tolist())









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 _lowerCamelCase(						_a   ):
       def __init__(					self,      lowerCamelCase,      lowerCamelCase=13,      lowerCamelCase=7,      lowerCamelCase=True,      lowerCamelCase=True,      lowerCamelCase=False,      lowerCamelCase=True,      lowerCamelCase=99,      lowerCamelCase=32,      lowerCamelCase=5,      lowerCamelCase=4,      lowerCamelCase=64,      lowerCamelCase="gelu",      lowerCamelCase=0.1,      lowerCamelCase=0.1,      lowerCamelCase=5_12,      lowerCamelCase=16,      lowerCamelCase=2,      lowerCamelCase=0.0_2,      lowerCamelCase=3,      lowerCamelCase=4,      lowerCamelCase=None,      lowerCamelCase=2,      lowerCamelCase=2,      lowerCamelCase=2,      lowerCamelCase=2,      lowerCamelCase=4,      lowerCamelCase=1,      )       ->   Union[str, Any]:




           """simple docstring"""

           _lowercase       :							Dict      		=		parent
           _lowercase       :							Optional[Any]      		=		batch_size
           _lowercase       :							Any      		=		seq_length
           _lowercase       :							Optional[Any]      		=		is_training
           _lowercase       :							Optional[Any]      		=		use_input_mask
           _lowercase       :							List[Any]      		=		use_token_type_ids
           _lowercase       :							List[str]      		=		use_labels
           _lowercase       :							str      		=		vocab_size
           _lowercase       :							List[str]      		=		hidden_size
           _lowercase       :							Dict      		=		num_hidden_layers
           _lowercase       :							List[str]      		=		num_attention_heads
           _lowercase       :							int      		=		intermediate_size
           _lowercase       :							Union[str, Any]      		=		hidden_act
           _lowercase       :							int      		=		hidden_dropout_prob
           _lowercase       :							List[Any]      		=		attention_probs_dropout_prob
           _lowercase       :							Dict      		=		max_position_embeddings
           _lowercase       :							Union[str, Any]      		=		type_vocab_size
           _lowercase       :							List[Any]      		=		type_sequence_label_size
           _lowercase       :							Any      		=		initializer_range
           _lowercase       :							List[str]      		=		num_labels
           _lowercase       :							Any      		=		num_choices
           _lowercase       :							Tuple      		=		scope
           _lowercase       :							Optional[Any]      		=		q_groups
           _lowercase       :							List[str]      		=		k_groups
           _lowercase       :							Optional[int]      		=		v_groups
           _lowercase       :							List[str]      		=		post_attention_groups
           _lowercase       :							Union[str, Any]      		=		intermediate_groups
           _lowercase       :							int      		=		output_groups
       def 			UpperCamelCase     (					self)       ->   Optional[int]:




           """simple docstring"""

           _lowercase       :							int      		=		ids_tensor([self.batch_size, self.seq_length],      self.vocab_size)

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

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

           _lowercase       :							Optional[Any]      		=		self.get_config()

           return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels
       def 			UpperCamelCase     (					self)       ->   Optional[int]:




           """simple docstring"""

           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 			UpperCamelCase     (					self,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase)       ->   Union[str, Any]:




           """simple docstring"""

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




           """simple docstring"""

           _lowercase       :							Dict      		=		SqueezeBertForMaskedLM(config=lowerCamelCase)
           model.to(lowerCamelCase)
           model.eval()
           _lowercase       :							Optional[Any]      		=		model(lowerCamelCase,      attention_mask=lowerCamelCase,      labels=lowerCamelCase)
           self.parent.assertEqual(result.logits.shape,      (self.batch_size, self.seq_length, self.vocab_size))
       def 			UpperCamelCase     (					self,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase)       ->   Union[str, Any]:




           """simple docstring"""

           _lowercase       :							Union[str, Any]      		=		SqueezeBertForQuestionAnswering(config=lowerCamelCase)
           model.to(lowerCamelCase)
           model.eval()
           _lowercase       :							List[Any]      		=		model(
               lowerCamelCase,      attention_mask=lowerCamelCase,      start_positions=lowerCamelCase,      end_positions=lowerCamelCase)
           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 			UpperCamelCase     (					self,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase)       ->   str:




           """simple docstring"""

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




           """simple docstring"""

           _lowercase       :							Union[str, Any]      		=		self.num_labels
           _lowercase       :							List[str]      		=		SqueezeBertForTokenClassification(config=lowerCamelCase)
           model.to(lowerCamelCase)
           model.eval()

           _lowercase       :							Union[str, Any]      		=		model(lowerCamelCase,      attention_mask=lowerCamelCase,      labels=lowerCamelCase)
           self.parent.assertEqual(result.logits.shape,      (self.batch_size, self.seq_length, self.num_labels))
       def 			UpperCamelCase     (					self,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase)       ->   Dict:




           """simple docstring"""

           _lowercase       :							str      		=		self.num_choices
           _lowercase       :							str      		=		SqueezeBertForMultipleChoice(config=lowerCamelCase)
           model.to(lowerCamelCase)
           model.eval()
           _lowercase       :							Dict      		=		input_ids.unsqueeze(1).expand(-1,      self.num_choices,      -1).contiguous()
           _lowercase       :							int      		=		input_mask.unsqueeze(1).expand(-1,      self.num_choices,      -1).contiguous()
           _lowercase       :							Optional[Any]      		=		model(
               lowerCamelCase,      attention_mask=lowerCamelCase,      labels=lowerCamelCase,      )
           self.parent.assertEqual(result.logits.shape,      (self.batch_size, self.num_choices))




       def 			UpperCamelCase     (					self)       ->   List[str]:




           """simple docstring"""

           _lowercase       :							Optional[int]      		=		self.prepare_config_and_inputs()
           ((_lowercase)		,				(_lowercase)		,				(_lowercase)		,				(_lowercase)		,				(_lowercase)		,				(_lowercase))       :							Dict      		=		config_and_inputs
           _lowercase       :							Tuple      		=		{'input_ids': input_ids, 'attention_mask': input_mask}
           return config, inputs_dict






@require_torch
class _lowerCamelCase(						_a, _a, unittest.TestCase   ):
       lowercase_ :	Union[str, Any]     =			(
           (
               SqueezeBertModel,
               SqueezeBertForMaskedLM,
               SqueezeBertForMultipleChoice,
               SqueezeBertForQuestionAnswering,
               SqueezeBertForSequenceClassification,
               SqueezeBertForTokenClassification,
           )
           if is_torch_available()
           else None
       )
       lowercase_ :	Optional[int]     =			(
           {
               """feature-extraction""": SqueezeBertModel,
               """fill-mask""": SqueezeBertForMaskedLM,
               """question-answering""": SqueezeBertForQuestionAnswering,
               """text-classification""": SqueezeBertForSequenceClassification,
               """token-classification""": SqueezeBertForTokenClassification,
               """zero-shot""": SqueezeBertForSequenceClassification,
           }
           if is_torch_available()
           else {}
       )
       lowercase_ :	Tuple     =			False
       lowercase_ :	List[str]     =			True
       lowercase_ :	int     =			False
       def 			UpperCamelCase     (					self)       ->   int:




           """simple docstring"""

           _lowercase       :							str      		=		SqueezeBertModelTester(self)
           _lowercase       :							Dict      		=		ConfigTester(self,      config_class=lowerCamelCase,      dim=37)
       def 			UpperCamelCase     (					self)       ->   List[str]:




           """simple docstring"""

           self.config_tester.run_common_tests()
       def 			UpperCamelCase     (					self)       ->   Tuple:




           """simple docstring"""

           _lowercase       :							Optional[Any]      		=		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_squeezebert_model(*lowerCamelCase)
       def 			UpperCamelCase     (					self)       ->   List[str]:




           """simple docstring"""

           _lowercase       :							Any      		=		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_squeezebert_for_masked_lm(*lowerCamelCase)
       def 			UpperCamelCase     (					self)       ->   Any:




           """simple docstring"""

           _lowercase       :							List[Any]      		=		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_squeezebert_for_question_answering(*lowerCamelCase)
       def 			UpperCamelCase     (					self)       ->   Union[str, Any]:




           """simple docstring"""

           _lowercase       :							Union[str, Any]      		=		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_squeezebert_for_sequence_classification(*lowerCamelCase)
       def 			UpperCamelCase     (					self)       ->   Any:




           """simple docstring"""

           _lowercase       :							List[Any]      		=		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_squeezebert_for_token_classification(*lowerCamelCase)
       def 			UpperCamelCase     (					self)       ->   Optional[int]:




           """simple docstring"""

           _lowercase       :							Optional[Any]      		=		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_squeezebert_for_multiple_choice(*lowerCamelCase)




       @slow
       def 			UpperCamelCase     (					self)       ->   Dict:




           """simple docstring"""

           for model_name in SQUEEZEBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
               _lowercase       :							List[Any]      		=		SqueezeBertModel.from_pretrained(lowerCamelCase)
               self.assertIsNotNone(lowerCamelCase)






@require_sentencepiece
@require_tokenizers
@require_torch
class _lowerCamelCase(						unittest.TestCase   ):
       @slow
       def 			UpperCamelCase     (					self)       ->   Optional[Any]:




           """simple docstring"""

           _lowercase       :							Union[str, Any]      		=		SqueezeBertForSequenceClassification.from_pretrained('squeezebert/squeezebert-mnli')

           _lowercase       :							Optional[int]      		=		torch.tensor([[1, 2_94_14, 2_32, 3_28, 7_40, 11_40, 1_26_95, 69, 13, 15_88, 2]])
           _lowercase       :							List[str]      		=		model(lowerCamelCase)[0]
           _lowercase       :							Union[str, Any]      		=		torch.Size((1, 3))
           self.assertEqual(output.shape,      lowerCamelCase)
           _lowercase       :							Tuple      		=		torch.tensor([[0.6_4_0_1, -0.0_3_4_9, -0.6_0_4_1]])
           self.assertTrue(torch.allclose(lowerCamelCase,      lowerCamelCase,      atol=1E-4))


