
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 __snake_case		(				lowerCAmelCase					):
    def __init__(  self		,snake_case		,snake_case=13		,snake_case=7		,snake_case=True		,snake_case=True		,snake_case=False		,snake_case=True		,snake_case=99		,snake_case=32		,snake_case=5		,snake_case=4		,snake_case=64		,snake_case="gelu"		,snake_case=0.1		,snake_case=0.1		,snake_case=512		,snake_case=16		,snake_case=2		,snake_case=0.02		,snake_case=3		,snake_case=4		,snake_case=None		,snake_case=2		,snake_case=2		,snake_case=2		,snake_case=2		,snake_case=4		,snake_case=1		,):



           '''simple docstring'''







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



           '''simple docstring'''







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

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

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

           lowercase							:					Union[str, Any]	     =				self.get_config()

           return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels
    def 				_SCREAMING_SNAKE_CASE	(  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 				_SCREAMING_SNAKE_CASE	(  self		,snake_case		,snake_case		,snake_case		,snake_case		,snake_case		,snake_case				):



           '''simple docstring'''







           lowercase							:					Optional[Any]	     =				SqueezeBertModel(config=snake_case				)
           model.to(snake_case				)
           model.eval()
           lowercase							:					Tuple	     =				model(snake_case		,snake_case				)
           lowercase							:					int	     =				model(snake_case				)
           self.parent.assertEqual(result.last_hidden_state.shape		,(self.batch_size, self.seq_length, self.hidden_size)				)
    def 				_SCREAMING_SNAKE_CASE	(  self		,snake_case		,snake_case		,snake_case		,snake_case		,snake_case		,snake_case				):



           '''simple docstring'''







           lowercase							:					List[str]	     =				SqueezeBertForMaskedLM(config=snake_case				)
           model.to(snake_case				)
           model.eval()
           lowercase							:					str	     =				model(snake_case		,attention_mask=snake_case		,labels=snake_case				)
           self.parent.assertEqual(result.logits.shape		,(self.batch_size, self.seq_length, self.vocab_size)				)
    def 				_SCREAMING_SNAKE_CASE	(  self		,snake_case		,snake_case		,snake_case		,snake_case		,snake_case		,snake_case				):



           '''simple docstring'''







           lowercase							:					Any	     =				SqueezeBertForQuestionAnswering(config=snake_case				)
           model.to(snake_case				)
           model.eval()
           lowercase							:					List[Any]	     =				model(
               snake_case		,attention_mask=snake_case		,start_positions=snake_case		,end_positions=snake_case				)
           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 				_SCREAMING_SNAKE_CASE	(  self		,snake_case		,snake_case		,snake_case		,snake_case		,snake_case		,snake_case				):



           '''simple docstring'''







           lowercase							:					Dict	     =				self.num_labels
           lowercase							:					List[Any]	     =				SqueezeBertForSequenceClassification(snake_case				)
           model.to(snake_case				)
           model.eval()
           lowercase							:					str	     =				model(snake_case		,attention_mask=snake_case		,labels=snake_case				)
           self.parent.assertEqual(result.logits.shape		,(self.batch_size, self.num_labels)				)
    def 				_SCREAMING_SNAKE_CASE	(  self		,snake_case		,snake_case		,snake_case		,snake_case		,snake_case		,snake_case				):



           '''simple docstring'''







           lowercase							:					List[Any]	     =				self.num_labels
           lowercase							:					List[Any]	     =				SqueezeBertForTokenClassification(config=snake_case				)
           model.to(snake_case				)
           model.eval()

           lowercase							:					str	     =				model(snake_case		,attention_mask=snake_case		,labels=snake_case				)
           self.parent.assertEqual(result.logits.shape		,(self.batch_size, self.seq_length, self.num_labels)				)
    def 				_SCREAMING_SNAKE_CASE	(  self		,snake_case		,snake_case		,snake_case		,snake_case		,snake_case		,snake_case				):



           '''simple docstring'''







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



           '''simple docstring'''







           lowercase							:					Any	     =				self.prepare_config_and_inputs()
           ((lowercase)     ,      (lowercase)     ,      (lowercase)     ,      (lowercase)     ,      (lowercase)     ,      (lowercase))							:					Any	     =				config_and_inputs
           lowercase							:					List[Any]	     =				{"""input_ids""": input_ids, """attention_mask""": input_mask}
           return config, inputs_dict






@require_torch
class __snake_case		(				lowerCAmelCase ,     lowerCAmelCase ,     unittest.TestCase					):
    _a		:						Union[str, Any]=    (
        (
            SqueezeBertModel,
            SqueezeBertForMaskedLM,
            SqueezeBertForMultipleChoice,
            SqueezeBertForQuestionAnswering,
            SqueezeBertForSequenceClassification,
            SqueezeBertForTokenClassification,
        )
        if is_torch_available()
        else None
    )
    _a		:						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 {}
    )
    _a		:						List[str]=    False
    _a		:						List[str]=    True
    _a		:						Optional[int]=    False
    def 				_SCREAMING_SNAKE_CASE	(  self				):



           '''simple docstring'''







           lowercase							:					List[str]	     =				SqueezeBertModelTester(self				)
           lowercase							:					Optional[Any]	     =				ConfigTester(self		,config_class=snake_case		,dim=37				)
    def 				_SCREAMING_SNAKE_CASE	(  self				):



           '''simple docstring'''







           self.config_tester.run_common_tests()
    def 				_SCREAMING_SNAKE_CASE	(  self				):



           '''simple docstring'''







           lowercase							:					Optional[Any]	     =				self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_squeezebert_model(*snake_case				)
    def 				_SCREAMING_SNAKE_CASE	(  self				):



           '''simple docstring'''







           lowercase							:					Union[str, Any]	     =				self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_squeezebert_for_masked_lm(*snake_case				)
    def 				_SCREAMING_SNAKE_CASE	(  self				):



           '''simple docstring'''







           lowercase							:					Tuple	     =				self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_squeezebert_for_question_answering(*snake_case				)
    def 				_SCREAMING_SNAKE_CASE	(  self				):



           '''simple docstring'''







           lowercase							:					Union[str, Any]	     =				self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_squeezebert_for_sequence_classification(*snake_case				)
    def 				_SCREAMING_SNAKE_CASE	(  self				):



           '''simple docstring'''







           lowercase							:					Optional[int]	     =				self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_squeezebert_for_token_classification(*snake_case				)
    def 				_SCREAMING_SNAKE_CASE	(  self				):



           '''simple docstring'''







           lowercase							:					Tuple	     =				self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_squeezebert_for_multiple_choice(*snake_case				)
    @slow
    def 				_SCREAMING_SNAKE_CASE	(  self				):



           '''simple docstring'''







           for model_name in SQUEEZEBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                  lowercase							:					Optional[Any]	     =				SqueezeBertModel.from_pretrained(snake_case				)
                  self.assertIsNotNone(snake_case				)






@require_sentencepiece
@require_tokenizers
@require_torch
class __snake_case		(				unittest.TestCase					):
    @slow
    def 				_SCREAMING_SNAKE_CASE	(  self				):



           '''simple docstring'''







           lowercase							:					List[str]	     =				SqueezeBertForSequenceClassification.from_pretrained("""squeezebert/squeezebert-mnli"""				)

           lowercase							:					Dict	     =				torch.tensor([[1, 29414, 232, 328, 740, 1140, 12695, 69, 13, 1588, 2]]				)
           lowercase							:					List[str]	     =				model(snake_case				)[0]
           lowercase							:					Tuple	     =				torch.Size((1, 3)				)
           self.assertEqual(output.shape		,snake_case				)
           lowercase							:					Any	     =				torch.tensor([[0.6_401, -0.0_349, -0.6_041]]				)
           self.assertTrue(torch.allclose(snake_case		,snake_case		,atol=1e-4				)				)


from bisect import bisect
from itertools import accumulate







def 						_snake_case(  SCREAMING_SNAKE_CASE__ ,     SCREAMING_SNAKE_CASE__ ,     SCREAMING_SNAKE_CASE__ ,     SCREAMING_SNAKE_CASE__    )			->   Optional[int]:
       lowercase							:					Dict	     =				sorted(zip(SCREAMING_SNAKE_CASE__ ,     SCREAMING_SNAKE_CASE__    ) ,     key=lambda SCREAMING_SNAKE_CASE__    : x[0] / x[1] ,     reverse=SCREAMING_SNAKE_CASE__    )
       lowercase     ,      lowercase							:					Optional[Any]	     =				[i[0] for i in r], [i[1] for i in r]
       lowercase							:					Any	     =				list(accumulate(SCREAMING_SNAKE_CASE__    )    )
       lowercase							:					int	     =				bisect(SCREAMING_SNAKE_CASE__ ,     SCREAMING_SNAKE_CASE__    )
       return (
           0
           if k == 0
           else sum(vl[:k]    ) + (w - acc[k - 1]) * (vl[k]) / (wt[k])
           if k != n
           else sum(vl[:k]    )
       )


if __name__ == "__main__":
  import doctest

  doctest.testmod()
