





'''simple docstring'''






import time
from dataclasses import dataclass
from multiprocessing import Pool
from unittest import TestCase
from unittest.mock import patch

import multiprocess
import numpy as np
import pytest

from datasets.utils.py_utils import (
    NestedDataStructure,
    asdict,
    iflatmap_unordered,
    map_nested,
    temp_seed,
    temporary_assignment,
    zip_dict,
)

from .utils import require_tf, require_torch



def        UpperCamelCase__		(						__magic_name__  :   int      )				->			List[Any]:  # picklable for multiprocessing
 '''simple docstring'''






 return x.sum()



def        UpperCamelCase__		(						__magic_name__  :   Any      )				->			Optional[int]:  # picklable for multiprocessing
 '''simple docstring'''






 return i + 1
@dataclass
class 					__snake_case  :






   '''simple docstring'''





   lowerCamelCase__    							=      42
   lowerCamelCase__    							=      42



class 					__snake_case  (					__SCREAMING_SNAKE_CASE							):






   '''simple docstring'''





   def 		__UpperCamelCase						(							self       ):
    snake_case__     :					str    =  {}
    snake_case__     :					Tuple    =  []
    snake_case__     :					List[Any]    =  1
    snake_case__     :					Union[str, Any]    =  [1, 2]
    snake_case__     :					str    =  {"""a""": 1, """b""": 2}
    snake_case__     :					Dict    =  {"""a""": [1, 2], """b""": [3, 4]}
    snake_case__     :					Dict    =  {"""a""": {"""1""": 1}, """b""": 2}
    snake_case__     :					Optional[int]    =  {"""a""": 1, """b""": 2, """c""": 3, """d""": 4}
    snake_case__     :					int    =  {}
    snake_case__     :					Tuple    =  []
    snake_case__     :					Optional[Any]    =  2
    snake_case__     :					Tuple    =  [2, 3]
    snake_case__     :					str    =  {"""a""": 2, """b""": 3}
    snake_case__     :					List[Any]    =  {"""a""": [2, 3], """b""": [4, 5]}
    snake_case__     :					Union[str, Any]    =  {"""a""": {"""1""": 2}, """b""": 3}
    snake_case__     :					Optional[int]    =  {"""a""": 2, """b""": 3, """c""": 4, """d""": 5}
    self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE       )	,		__SCREAMING_SNAKE_CASE       )
    self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE       )	,		__SCREAMING_SNAKE_CASE       )
    self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE       )	,		__SCREAMING_SNAKE_CASE       )
    self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE       )	,		__SCREAMING_SNAKE_CASE       )
    self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE       )	,		__SCREAMING_SNAKE_CASE       )
    self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE       )	,		__SCREAMING_SNAKE_CASE       )
    self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE       )	,		__SCREAMING_SNAKE_CASE       )
    self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE       )	,		__SCREAMING_SNAKE_CASE       )

    snake_case__     :					Optional[int]    =  2
    self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		num_proc=__SCREAMING_SNAKE_CASE       )	,		__SCREAMING_SNAKE_CASE       )
    self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		num_proc=__SCREAMING_SNAKE_CASE       )	,		__SCREAMING_SNAKE_CASE       )
    self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		num_proc=__SCREAMING_SNAKE_CASE       )	,		__SCREAMING_SNAKE_CASE       )
    self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		num_proc=__SCREAMING_SNAKE_CASE       )	,		__SCREAMING_SNAKE_CASE       )
    self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		num_proc=__SCREAMING_SNAKE_CASE       )	,		__SCREAMING_SNAKE_CASE       )
    self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		num_proc=__SCREAMING_SNAKE_CASE       )	,		__SCREAMING_SNAKE_CASE       )
    self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		num_proc=__SCREAMING_SNAKE_CASE       )	,		__SCREAMING_SNAKE_CASE       )
    self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		num_proc=__SCREAMING_SNAKE_CASE       )	,		__SCREAMING_SNAKE_CASE       )

    snake_case__     :					int    =  {"""a""": np.eye(2       ), """b""": np.zeros(3       ), """c""": np.ones(2       )}
    snake_case__     :					int    =  {"""a""": 2, """b""": 0, """c""": 2}
    snake_case__     :					Union[str, Any]    =  {
        """a""": np.eye(2       ).astype(__SCREAMING_SNAKE_CASE       ),
        """b""": np.zeros(3       ).astype(__SCREAMING_SNAKE_CASE       ),
        """c""": np.ones(2       ).astype(__SCREAMING_SNAKE_CASE       ),
    }
    self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		map_numpy=__SCREAMING_SNAKE_CASE       )	,		__SCREAMING_SNAKE_CASE       )
    self.assertEqual(
        {k: v.tolist() for k, v in map_nested(__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		map_numpy=__SCREAMING_SNAKE_CASE       ).items()}	,		{k: v.tolist() for k, v in expected_map_nested_sna_int.items()}	,		)
    self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		map_numpy=__SCREAMING_SNAKE_CASE	,		num_proc=__SCREAMING_SNAKE_CASE       )	,		__SCREAMING_SNAKE_CASE       )
    self.assertEqual(
        {k: v.tolist() for k, v in map_nested(__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		map_numpy=__SCREAMING_SNAKE_CASE	,		num_proc=__SCREAMING_SNAKE_CASE       ).items()}	,		{k: v.tolist() for k, v in expected_map_nested_sna_int.items()}	,		)
    with self.assertRaises(__SCREAMING_SNAKE_CASE       ):  # can't pickle a local lambda
     map_nested(lambda __SCREAMING_SNAKE_CASE       : x + 1	,		__SCREAMING_SNAKE_CASE	,		num_proc=__SCREAMING_SNAKE_CASE       )

   def 		__UpperCamelCase						(							self       ):
    snake_case__     :					Union[str, Any]    =  {"""a""": 1, """b""": 2}
    snake_case__     :					Optional[int]    =  {"""a""": 3, """b""": 4}
    snake_case__     :					List[Any]    =  {"""a""": 5, """b""": 6}
    snake_case__     :					Optional[Any]    =  sorted([("""a""", (1, 3, 5)), ("""b""", (2, 4, 6))]       )
    self.assertEqual(sorted(zip_dict(__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE       )       )	,		__SCREAMING_SNAKE_CASE       )


   def 		__UpperCamelCase						(							self       ):
    class 					__snake_case  :






       '''simple docstring'''





       lowerCamelCase__    							=      '''bar'''



    snake_case__     :					str    =  Foo()
    self.assertEqual(foo.my_attr	,		"""bar"""       )
    with temporary_assignment(__SCREAMING_SNAKE_CASE	,		"""my_attr"""	,		"""BAR"""       ):
     self.assertEqual(foo.my_attr	,		"""BAR"""       )
    self.assertEqual(foo.my_attr	,		"""bar"""       )



@pytest.mark.parametrize(
    """iterable_length, num_proc, expected_num_proc"""		,       [
        (1, None, 1),
        (1, 1, 1),
        (2, None, 1),
        (2, 1, 1),
        (2, 2, 1),
        (2, 3, 1),
        (3, 2, 1),
        (16, 16, 16),
        (16, 17, 16),
        (17, 16, 16),
    ]		,       )
def        UpperCamelCase__		(						__magic_name__  :   List[Any]		,       __magic_name__  :   List[Any]		,       __magic_name__  :   Any      )				->			Optional[int]:
 '''simple docstring'''






 with patch("""datasets.utils.py_utils._single_map_nested"""      ) as mock_single_map_nested, patch(
     """datasets.parallel.parallel.Pool"""      ) as mock_multiprocessing_pool:
  snake_case__     :					Optional[int]    =  {f"{i}": i for i in range(__magic_name__      )}
  snake_case__     :					List[str]    =  map_nested(lambda __magic_name__      : x + 10		,       __magic_name__		,       num_proc=__magic_name__		,       parallel_min_length=16      )
  if expected_num_proc == 1:
   assert mock_single_map_nested.called
   assert not mock_multiprocessing_pool.called
  else:
   assert not mock_single_map_nested.called
   assert mock_multiprocessing_pool.called
   assert mock_multiprocessing_pool.call_args[0][0] == expected_num_proc
class 					__snake_case  (					__SCREAMING_SNAKE_CASE							):






   '''simple docstring'''





   @require_tf
   def 		__UpperCamelCase						(							self       ):
    import tensorflow as tf
    from tensorflow.keras import layers

    snake_case__     :					Any    =  layers.Dense(2       )

    def gen_random_output():
     snake_case__     :					Union[str, Any]    =  tf.random.uniform((1, 3)       )
     return model(__SCREAMING_SNAKE_CASE       ).numpy()

    with temp_seed(4_2	,		set_tensorflow=__SCREAMING_SNAKE_CASE       ):
     snake_case__     :					Optional[Any]    =  gen_random_output()
    with temp_seed(4_2	,		set_tensorflow=__SCREAMING_SNAKE_CASE       ):
     snake_case__     :					Dict    =  gen_random_output()
    snake_case__     :					Optional[Any]    =  gen_random_output()

    np.testing.assert_equal(__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE       )
    self.assertGreater(np.abs(outa - outa       ).sum()	,		0       )

   @require_torch
   def 		__UpperCamelCase						(							self       ):
    import torch

    def gen_random_output():
     snake_case__     :					List[Any]    =  torch.nn.Linear(3	,		2       )
     snake_case__     :					Dict    =  torch.rand(1	,		3       )
     return model(__SCREAMING_SNAKE_CASE       ).detach().numpy()

    with temp_seed(4_2	,		set_pytorch=__SCREAMING_SNAKE_CASE       ):
     snake_case__     :					int    =  gen_random_output()
    with temp_seed(4_2	,		set_pytorch=__SCREAMING_SNAKE_CASE       ):
     snake_case__     :					str    =  gen_random_output()
    snake_case__     :					Any    =  gen_random_output()

    np.testing.assert_equal(__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE       )
    self.assertGreater(np.abs(outa - outa       ).sum()	,		0       )


   def 		__UpperCamelCase						(							self       ):
    def gen_random_output():
     return np.random.rand(1	,		3       )

    with temp_seed(4_2       ):
     snake_case__     :					Any    =  gen_random_output()
    with temp_seed(4_2       ):
     snake_case__     :					Optional[Any]    =  gen_random_output()
    snake_case__     :					int    =  gen_random_output()

    np.testing.assert_equal(__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE       )
    self.assertGreater(np.abs(outa - outa       ).sum()	,		0       )



@pytest.mark.parametrize("""input_data"""		,       [{}]      )
def        UpperCamelCase__		(						__magic_name__  :   List[Any]      )				->			Dict:
 '''simple docstring'''






 snake_case__     :					Union[str, Any]    =  NestedDataStructure(__magic_name__      ).data
 assert output_data == input_data



@pytest.mark.parametrize(
    """data, expected_output"""		,       [
        ({}, []),
        ([], []),
        ("""foo""", ["""foo"""]),
        (["""foo""", """bar"""], ["""foo""", """bar"""]),
        ([["""foo""", """bar"""]], ["""foo""", """bar"""]),
        ([[["""foo"""], ["""bar"""]]], ["""foo""", """bar"""]),
        ([[["""foo"""], """bar"""]], ["""foo""", """bar"""]),
        ({"""a""": 1, """b""": 2}, [1, 2]),
        ({"""a""": [1, 2], """b""": [3, 4]}, [1, 2, 3, 4]),
        ({"""a""": [[1, 2]], """b""": [[3, 4]]}, [1, 2, 3, 4]),
        ({"""a""": [[1, 2]], """b""": [3, 4]}, [1, 2, 3, 4]),
        ({"""a""": [[[1], [2]]], """b""": [[[3], [4]]]}, [1, 2, 3, 4]),
        ({"""a""": [[[1], [2]]], """b""": [[3, 4]]}, [1, 2, 3, 4]),
        ({"""a""": [[[1], [2]]], """b""": [3, 4]}, [1, 2, 3, 4]),
        ({"""a""": [[[1], [2]]], """b""": [3, [4]]}, [1, 2, 3, 4]),
        ({"""a""": {"""1""": 1}, """b""": 2}, [1, 2]),
        ({"""a""": {"""1""": [1]}, """b""": 2}, [1, 2]),
        ({"""a""": {"""1""": [1]}, """b""": [2]}, [1, 2]),
    ]		,       )
def        UpperCamelCase__		(						__magic_name__  :   Tuple		,       __magic_name__  :   Union[str, Any]      )				->			Optional[int]:
 '''simple docstring'''






 snake_case__     :					Optional[int]    =  NestedDataStructure(__magic_name__      ).flatten()
 assert output == expected_output



def        UpperCamelCase__		(						)				->			Union[str, Any]:
 '''simple docstring'''






 snake_case__     :					int    =  A(x=1		,       y="""foobar"""      )
 snake_case__     :					Optional[int]    =  {"""x""": 1, """y""": """foobar"""}
 assert asdict(__magic_name__      ) == expected_output

 snake_case__     :					Tuple    =  {"""a""": {"""b""": A(x=10		,       y="""foo"""      )}, """c""": [A(x=20		,       y="""bar"""      )]}
 snake_case__     :					str    =  {"""a""": {"""b""": {"""x""": 10, """y""": """foo"""}}, """c""": [{"""x""": 20, """y""": """bar"""}]}
 assert asdict(__magic_name__      ) == expected_output

 with pytest.raises(__magic_name__      ):
  asdict([1, A(x=10		,       y="""foo"""      )]      )



def        UpperCamelCase__		(						__magic_name__  :   str      )				->			str:
 '''simple docstring'''






 return text.split()



def        UpperCamelCase__		(						__magic_name__  :   Union[str, Any]      )				->			Optional[int]:
 '''simple docstring'''






 yield (time.time(), content)
 time.sleep(2      )
 yield (time.time(), content)





def        UpperCamelCase__		(						)				->			Optional[int]:
 '''simple docstring'''






 with Pool(2      ) as pool:
  snake_case__     :					Optional[Any]    =  list(iflatmap_unordered(__magic_name__		,       _split_text		,       kwargs_iterable=[{"""text""": """hello there"""}] * 10      )      )
  assert out.count("""hello"""      ) == 10
  assert out.count("""there"""      ) == 10
  assert len(__magic_name__      ) == 20

 # check multiprocess from pathos (uses dill for pickling)
 with multiprocess.Pool(2      ) as pool:
  snake_case__     :					Dict    =  list(iflatmap_unordered(__magic_name__		,       _split_text		,       kwargs_iterable=[{"""text""": """hello there"""}] * 10      )      )
  assert out.count("""hello"""      ) == 10
  assert out.count("""there"""      ) == 10
  assert len(__magic_name__      ) == 20

 # check that we get items as fast as possible
 with Pool(2      ) as pool:
  snake_case__     :					Any    =  []
  for yield_time, content in iflatmap_unordered(
      __magic_name__		,       _aseconds_generator_of_aitems_with_timing		,       kwargs_iterable=[{"""content""": """a"""}, {"""content""": """b"""}]      ):
   assert yield_time < time.time() + 0.1, "we should each item directly after it was yielded"
   out.append(__magic_name__      )
  assert out.count("""a"""      ) == 2
  assert out.count("""b"""      ) == 2
  assert len(__magic_name__      ) == 4









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  (						lowerCAmelCase__ ):




 def __init__(    self						,					__UpperCAmelCase						,					__UpperCAmelCase=13						,					__UpperCAmelCase=7						,					__UpperCAmelCase=True						,					__UpperCAmelCase=True						,					__UpperCAmelCase=False						,					__UpperCAmelCase=True						,					__UpperCAmelCase=99						,					__UpperCAmelCase=32						,					__UpperCAmelCase=5						,					__UpperCAmelCase=4						,					__UpperCAmelCase=64						,					__UpperCAmelCase="gelu"						,					__UpperCAmelCase=0.1						,					__UpperCAmelCase=0.1						,					__UpperCAmelCase=512						,					__UpperCAmelCase=16						,					__UpperCAmelCase=2						,					__UpperCAmelCase=0.02						,					__UpperCAmelCase=3						,					__UpperCAmelCase=4						,					__UpperCAmelCase=None						,					__UpperCAmelCase=2						,					__UpperCAmelCase=2						,					__UpperCAmelCase=2						,					__UpperCAmelCase=2						,					__UpperCAmelCase=4						,					__UpperCAmelCase=1						,					):



    '''simple docstring'''
    __lowerCamelCase    				=     parent
    __lowerCamelCase    				=     batch_size
    __lowerCamelCase    				=     seq_length
    __lowerCamelCase    				=     is_training
    __lowerCamelCase    				=     use_input_mask
    __lowerCamelCase    				=     use_token_type_ids
    __lowerCamelCase    				=     use_labels
    __lowerCamelCase    				=     vocab_size
    __lowerCamelCase    				=     hidden_size
    __lowerCamelCase    				=     num_hidden_layers
    __lowerCamelCase    				=     num_attention_heads
    __lowerCamelCase    				=     intermediate_size
    __lowerCamelCase    				=     hidden_act
    __lowerCamelCase    				=     hidden_dropout_prob
    __lowerCamelCase    				=     attention_probs_dropout_prob
    __lowerCamelCase    				=     max_position_embeddings
    __lowerCamelCase    				=     type_vocab_size
    __lowerCamelCase    				=     type_sequence_label_size
    __lowerCamelCase    				=     initializer_range
    __lowerCamelCase    				=     num_labels
    __lowerCamelCase    				=     num_choices
    __lowerCamelCase    				=     scope
    __lowerCamelCase    				=     q_groups
    __lowerCamelCase    				=     k_groups
    __lowerCamelCase    				=     v_groups
    __lowerCamelCase    				=     post_attention_groups
    __lowerCamelCase    				=     intermediate_groups
    __lowerCamelCase    				=     output_groups




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     ids_tensor([self.batch_size, self.seq_length]						,					self.vocab_size					)

    __lowerCamelCase    				=     None
    if self.use_input_mask:
       __lowerCamelCase    				=     random_attention_mask([self.batch_size, self.seq_length]					)

    __lowerCamelCase    				=     None
    __lowerCamelCase    				=     None
    __lowerCamelCase    				=     None
    if self.use_labels:
       __lowerCamelCase    				=     ids_tensor([self.batch_size]						,					self.type_sequence_label_size					)
       __lowerCamelCase    				=     ids_tensor([self.batch_size, self.seq_length]						,					self.num_labels					)
       __lowerCamelCase    				=     ids_tensor([self.batch_size]						,					self.num_choices					)

    __lowerCamelCase    				=     self.get_config()

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




 def 							lowerCamelCase   (    self					):



    '''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 							lowerCamelCase   (    self						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase					):



    '''simple docstring'''
    __lowerCamelCase    				=     SqueezeBertModel(config=__UpperCAmelCase					)
    model.to(__UpperCAmelCase					)
    model.eval()
    __lowerCamelCase    				=     model(__UpperCAmelCase						,					__UpperCAmelCase					)
    __lowerCamelCase    				=     model(__UpperCAmelCase					)
    self.parent.assertEqual(result.last_hidden_state.shape						,					(self.batch_size, self.seq_length, self.hidden_size)					)




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase					):



    '''simple docstring'''
    __lowerCamelCase    				=     SqueezeBertForMaskedLM(config=__UpperCAmelCase					)
    model.to(__UpperCAmelCase					)
    model.eval()
    __lowerCamelCase    				=     model(__UpperCAmelCase						,					attention_mask=__UpperCAmelCase						,					labels=__UpperCAmelCase					)
    self.parent.assertEqual(result.logits.shape						,					(self.batch_size, self.seq_length, self.vocab_size)					)




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase					):



    '''simple docstring'''
    __lowerCamelCase    				=     SqueezeBertForQuestionAnswering(config=__UpperCAmelCase					)
    model.to(__UpperCAmelCase					)
    model.eval()
    __lowerCamelCase    				=     model(
        __UpperCAmelCase						,					attention_mask=__UpperCAmelCase						,					start_positions=__UpperCAmelCase						,					end_positions=__UpperCAmelCase					)
    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 							lowerCamelCase   (    self						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase					):



    '''simple docstring'''
    __lowerCamelCase    				=     self.num_labels
    __lowerCamelCase    				=     SqueezeBertForSequenceClassification(__UpperCAmelCase					)
    model.to(__UpperCAmelCase					)
    model.eval()
    __lowerCamelCase    				=     model(__UpperCAmelCase						,					attention_mask=__UpperCAmelCase						,					labels=__UpperCAmelCase					)
    self.parent.assertEqual(result.logits.shape						,					(self.batch_size, self.num_labels)					)




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase					):



    '''simple docstring'''
    __lowerCamelCase    				=     self.num_labels
    __lowerCamelCase    				=     SqueezeBertForTokenClassification(config=__UpperCAmelCase					)
    model.to(__UpperCAmelCase					)
    model.eval()

    __lowerCamelCase    				=     model(__UpperCAmelCase						,					attention_mask=__UpperCAmelCase						,					labels=__UpperCAmelCase					)
    self.parent.assertEqual(result.logits.shape						,					(self.batch_size, self.seq_length, self.num_labels)					)




 def 							lowerCamelCase   (    self						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase						,					__UpperCAmelCase					):



    '''simple docstring'''
    __lowerCamelCase    				=     self.num_choices
    __lowerCamelCase    				=     SqueezeBertForMultipleChoice(config=__UpperCAmelCase					)
    model.to(__UpperCAmelCase					)
    model.eval()
    __lowerCamelCase    				=     input_ids.unsqueeze(1					).expand(-1						,					self.num_choices						,					-1					).contiguous()
    __lowerCamelCase    				=     input_mask.unsqueeze(1					).expand(-1						,					self.num_choices						,					-1					).contiguous()
    __lowerCamelCase    				=     model(
        __UpperCAmelCase						,					attention_mask=__UpperCAmelCase						,					labels=__UpperCAmelCase						,					)
    self.parent.assertEqual(result.logits.shape						,					(self.batch_size, self.num_choices)					)




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     self.prepare_config_and_inputs()
    ((__lowerCamelCase)				,(__lowerCamelCase)				,(__lowerCamelCase)				,(__lowerCamelCase)				,(__lowerCamelCase)				,(__lowerCamelCase))    				=     config_and_inputs
    __lowerCamelCase    				=     {'''input_ids''': input_ids, '''attention_mask''': input_mask}
    return config, inputs_dict



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




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     SqueezeBertModelTester(self					)
    __lowerCamelCase    				=     ConfigTester(self						,					config_class=__UpperCAmelCase						,					dim=37					)




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    self.config_tester.run_common_tests()




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_squeezebert_model(*__UpperCAmelCase					)




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_squeezebert_for_masked_lm(*__UpperCAmelCase					)




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_squeezebert_for_question_answering(*__UpperCAmelCase					)




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_squeezebert_for_sequence_classification(*__UpperCAmelCase					)




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_squeezebert_for_token_classification(*__UpperCAmelCase					)




 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_squeezebert_for_multiple_choice(*__UpperCAmelCase					)




 @slow
 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    for model_name in SQUEEZEBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
       __lowerCamelCase    				=     SqueezeBertModel.from_pretrained(__UpperCAmelCase					)
       self.assertIsNotNone(__UpperCAmelCase					)



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




 @slow
 def 							lowerCamelCase   (    self					):



    '''simple docstring'''
    __lowerCamelCase    				=     SqueezeBertForSequenceClassification.from_pretrained('''squeezebert/squeezebert-mnli'''					)

    __lowerCamelCase    				=     torch.tensor([[1, 29414, 232, 328, 740, 1140, 12695, 69, 13, 1588, 2]]					)
    __lowerCamelCase    				=     model(__UpperCAmelCase					)[0]
    __lowerCamelCase    				=     torch.Size((1, 3)					)
    self.assertEqual(output.shape						,					__UpperCAmelCase					)
    __lowerCamelCase    				=     torch.tensor([[0.6_401, -0.0_349, -0.6_041]]					)
    self.assertTrue(torch.allclose(__UpperCAmelCase						,					__UpperCAmelCase						,					atol=1E-4					)					)






