from __future__ import annotations

import unittest

from transformers import MobileBertConfig, is_tf_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
    import tensorflow as tf

    from transformers import (
        TF_MODEL_FOR_PRETRAINING_MAPPING,
        TFMobileBertForMaskedLM,
        TFMobileBertForMultipleChoice,
        TFMobileBertForNextSentencePrediction,
        TFMobileBertForPreTraining,
        TFMobileBertForQuestionAnswering,
        TFMobileBertForSequenceClassification,
        TFMobileBertForTokenClassification,
        TFMobileBertModel,
    )


@require_tf
class 	lowerCamelCase_     (   UpperCAmelCase_						,	UpperCAmelCase_						,	unittest.TestCase ):






    '''simple docstring'''






    a__	:						List[str]						=      (
        (
            TFMobileBertModel,
            TFMobileBertForMaskedLM,
            TFMobileBertForNextSentencePrediction,
            TFMobileBertForPreTraining,
            TFMobileBertForQuestionAnswering,
            TFMobileBertForSequenceClassification,
            TFMobileBertForTokenClassification,
            TFMobileBertForMultipleChoice,
        )
        if is_tf_available()
        else ()
    )
    a__	:						List[str]						=      (
        {
            """feature-extraction""": TFMobileBertModel,
            """fill-mask""": TFMobileBertForMaskedLM,
            """question-answering""": TFMobileBertForQuestionAnswering,
            """text-classification""": TFMobileBertForSequenceClassification,
            """token-classification""": TFMobileBertForTokenClassification,
            """zero-shot""": TFMobileBertForSequenceClassification,
        }
        if is_tf_available()
        else {}
    )
    a__	:						List[Any]						=      False
    a__	:						int						=      False







    def 			UpperCamelCase__   (	self , __lowercase , __lowercase , __lowercase=False)       ->	Union[str, Any]:
        __UpperCamelCase     :List[str]			  =		super()._prepare_for_class(__lowercase , __lowercase , return_labels=__lowercase)

        if return_labels:
            if model_class in get_values(__lowercase):
                __UpperCamelCase     :Union[str, Any]			  =		tf.zeros(self.model_tester.batch_size , dtype=tf.intaa)

        return inputs_dict


    class 	lowerCamelCase_     (   UpperCAmelCase_ ):






        '''simple docstring'''






        def __init__(	self , __lowercase , __lowercase=13 , __lowercase=7 , __lowercase=True , __lowercase=True , __lowercase=True , __lowercase=True , __lowercase=99 , __lowercase=32 , __lowercase=32 , __lowercase=2 , __lowercase=4 , __lowercase=37 , __lowercase="gelu" , __lowercase=0.1 , __lowercase=0.1 , __lowercase=512 , __lowercase=16 , __lowercase=2 , __lowercase=0.02 , __lowercase=3 , __lowercase=4 , __lowercase=None , )       ->	Any:
            __UpperCamelCase     :List[Any]			  =		parent
            __UpperCamelCase     :Any			  =		batch_size
            __UpperCamelCase     :Dict			  =		seq_length
            __UpperCamelCase     :List[str]			  =		is_training
            __UpperCamelCase     :List[Any]			  =		use_input_mask
            __UpperCamelCase     :Optional[int]			  =		use_token_type_ids
            __UpperCamelCase     :Any			  =		use_labels
            __UpperCamelCase     :List[str]			  =		vocab_size
            __UpperCamelCase     :Any			  =		hidden_size
            __UpperCamelCase     :Optional[int]			  =		num_hidden_layers
            __UpperCamelCase     :Union[str, Any]			  =		num_attention_heads
            __UpperCamelCase     :Dict			  =		intermediate_size
            __UpperCamelCase     :str			  =		hidden_act
            __UpperCamelCase     :List[Any]			  =		hidden_dropout_prob
            __UpperCamelCase     :Union[str, Any]			  =		attention_probs_dropout_prob
            __UpperCamelCase     :List[str]			  =		max_position_embeddings
            __UpperCamelCase     :List[str]			  =		type_vocab_size
            __UpperCamelCase     :Dict			  =		type_sequence_label_size
            __UpperCamelCase     :str			  =		initializer_range
            __UpperCamelCase     :int			  =		num_labels
            __UpperCamelCase     :int			  =		num_choices
            __UpperCamelCase     :List[str]			  =		scope
            __UpperCamelCase     :Any			  =		embedding_size







        def 			UpperCamelCase__   (	self)       ->	str:
            __UpperCamelCase     :str			  =		ids_tensor([self.batch_size, self.seq_length] , self.vocab_size)

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

            __UpperCamelCase     :Union[str, Any]			  =		None
            if self.use_token_type_ids:
                __UpperCamelCase     :Tuple			  =		ids_tensor([self.batch_size, self.seq_length] , self.type_vocab_size)

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

            __UpperCamelCase     :str			  =		MobileBertConfig(
                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 , hidden_dropout_prob=self.hidden_dropout_prob , attention_probs_dropout_prob=self.attention_probs_dropout_prob , max_position_embeddings=self.max_position_embeddings , type_vocab_size=self.type_vocab_size , initializer_range=self.initializer_range , embedding_size=self.embedding_size , )

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







        def 			UpperCamelCase__   (	self , __lowercase , __lowercase , __lowercase , __lowercase , __lowercase , __lowercase , __lowercase)       ->	str:
            __UpperCamelCase     :Dict			  =		TFMobileBertModel(config=__lowercase)
            __UpperCamelCase     :List[Any]			  =		{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
            __UpperCamelCase     :Dict			  =		model(__lowercase)

            __UpperCamelCase     :Union[str, Any]			  =		[input_ids, input_mask]
            __UpperCamelCase     :Optional[Any]			  =		model(__lowercase)

            __UpperCamelCase     :str			  =		model(__lowercase)

            self.parent.assertEqual(
                result.last_hidden_state.shape , (self.batch_size, self.seq_length, self.hidden_size))
            self.parent.assertEqual(result.pooler_output.shape , (self.batch_size, self.hidden_size))







        def 			UpperCamelCase__   (	self , __lowercase , __lowercase , __lowercase , __lowercase , __lowercase , __lowercase , __lowercase)       ->	Dict:
            __UpperCamelCase     :int			  =		TFMobileBertForMaskedLM(config=__lowercase)
            __UpperCamelCase     :Optional[int]			  =		{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
            __UpperCamelCase     :Union[str, Any]			  =		model(__lowercase)
            self.parent.assertEqual(result.logits.shape , (self.batch_size, self.seq_length, self.vocab_size))







        def 			UpperCamelCase__   (	self , __lowercase , __lowercase , __lowercase , __lowercase , __lowercase , __lowercase , __lowercase)       ->	Optional[Any]:
            __UpperCamelCase     :int			  =		TFMobileBertForNextSentencePrediction(config=__lowercase)
            __UpperCamelCase     :Any			  =		{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
            __UpperCamelCase     :Optional[int]			  =		model(__lowercase)
            self.parent.assertEqual(result.logits.shape , (self.batch_size, 2))







        def 			UpperCamelCase__   (	self , __lowercase , __lowercase , __lowercase , __lowercase , __lowercase , __lowercase , __lowercase)       ->	List[Any]:
            __UpperCamelCase     :List[str]			  =		TFMobileBertForPreTraining(config=__lowercase)
            __UpperCamelCase     :int			  =		{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
            __UpperCamelCase     :Optional[Any]			  =		model(__lowercase)
            self.parent.assertEqual(
                result.prediction_logits.shape , (self.batch_size, self.seq_length, self.vocab_size))
            self.parent.assertEqual(result.seq_relationship_logits.shape , (self.batch_size, 2))







        def 			UpperCamelCase__   (	self , __lowercase , __lowercase , __lowercase , __lowercase , __lowercase , __lowercase , __lowercase)       ->	Union[str, Any]:
            __UpperCamelCase     :Any			  =		self.num_labels
            __UpperCamelCase     :Dict			  =		TFMobileBertForSequenceClassification(config=__lowercase)
            __UpperCamelCase     :Union[str, Any]			  =		{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
            __UpperCamelCase     :List[Any]			  =		model(__lowercase)
            self.parent.assertEqual(result.logits.shape , (self.batch_size, self.num_labels))







        def 			UpperCamelCase__   (	self , __lowercase , __lowercase , __lowercase , __lowercase , __lowercase , __lowercase , __lowercase)       ->	int:
            __UpperCamelCase     :Union[str, Any]			  =		self.num_choices
            __UpperCamelCase     :Optional[Any]			  =		TFMobileBertForMultipleChoice(config=__lowercase)
            __UpperCamelCase     :Dict			  =		tf.tile(tf.expand_dims(__lowercase , 1) , (1, self.num_choices, 1))
            __UpperCamelCase     :Union[str, Any]			  =		tf.tile(tf.expand_dims(__lowercase , 1) , (1, self.num_choices, 1))
            __UpperCamelCase     :Union[str, Any]			  =		tf.tile(tf.expand_dims(__lowercase , 1) , (1, self.num_choices, 1))
            __UpperCamelCase     :Any			  =		{
                '''input_ids''': multiple_choice_inputs_ids,
                '''attention_mask''': multiple_choice_input_mask,
                '''token_type_ids''': multiple_choice_token_type_ids,
            }
            __UpperCamelCase     :Optional[int]			  =		model(__lowercase)
            self.parent.assertEqual(result.logits.shape , (self.batch_size, self.num_choices))







        def 			UpperCamelCase__   (	self , __lowercase , __lowercase , __lowercase , __lowercase , __lowercase , __lowercase , __lowercase)       ->	Tuple:
            __UpperCamelCase     :Optional[int]			  =		self.num_labels
            __UpperCamelCase     :Optional[int]			  =		TFMobileBertForTokenClassification(config=__lowercase)
            __UpperCamelCase     :Optional[Any]			  =		{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
            __UpperCamelCase     :Optional[Any]			  =		model(__lowercase)
            self.parent.assertEqual(result.logits.shape , (self.batch_size, self.seq_length, self.num_labels))







        def 			UpperCamelCase__   (	self , __lowercase , __lowercase , __lowercase , __lowercase , __lowercase , __lowercase , __lowercase)       ->	int:
            __UpperCamelCase     :int			  =		TFMobileBertForQuestionAnswering(config=__lowercase)
            __UpperCamelCase     :List[Any]			  =		{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
            __UpperCamelCase     :Any			  =		model(__lowercase)
            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)       ->	Tuple:
            __UpperCamelCase     :Optional[Any]			  =		self.prepare_config_and_inputs()
            (
                (
                __UpperCamelCase
            )  ,       (
                __UpperCamelCase
            )  ,       (
                __UpperCamelCase
            )  ,       (
                __UpperCamelCase
            )  ,       (
                __UpperCamelCase
            )  ,       (
                __UpperCamelCase
            )  ,       (
                __UpperCamelCase
            )  ,       
            )     :List[str]			  =		config_and_inputs
            __UpperCamelCase     :str			  =		{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
            return config, inputs_dict







    def 			UpperCamelCase__   (	self)       ->	Tuple:
        __UpperCamelCase     :List[Any]			  =		TFMobileBertModelTest.TFMobileBertModelTester(self)
        __UpperCamelCase     :str			  =		ConfigTester(self , config_class=__lowercase , hidden_size=37)







    def 			UpperCamelCase__   (	self)       ->	List[Any]:
        self.config_tester.run_common_tests()







    def 			UpperCamelCase__   (	self)       ->	Tuple:
        __UpperCamelCase     :Any			  =		self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_mobilebert_model(*__lowercase)







    def 			UpperCamelCase__   (	self)       ->	Optional[Any]:
        __UpperCamelCase     :Union[str, Any]			  =		self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_mobilebert_for_masked_lm(*__lowercase)







    def 			UpperCamelCase__   (	self)       ->	Optional[Any]:
        __UpperCamelCase     :Tuple			  =		self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_mobilebert_for_multiple_choice(*__lowercase)







    def 			UpperCamelCase__   (	self)       ->	Dict:
        __UpperCamelCase     :int			  =		self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_mobilebert_for_next_sequence_prediction(*__lowercase)







    def 			UpperCamelCase__   (	self)       ->	Tuple:
        __UpperCamelCase     :Optional[Any]			  =		self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_mobilebert_for_pretraining(*__lowercase)







    def 			UpperCamelCase__   (	self)       ->	List[Any]:
        __UpperCamelCase     :List[Any]			  =		self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_mobilebert_for_question_answering(*__lowercase)







    def 			UpperCamelCase__   (	self)       ->	Any:
        __UpperCamelCase     :List[str]			  =		self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_mobilebert_for_sequence_classification(*__lowercase)







    def 			UpperCamelCase__   (	self)       ->	Optional[int]:
        __UpperCamelCase     :str			  =		self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_mobilebert_for_token_classification(*__lowercase)







    @slow
    def 			UpperCamelCase__   (	self)       ->	List[Any]:
        # for model_name in TF_MOBILEBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        for model_name in ["google/mobilebert-uncased"]:
            __UpperCamelCase     :List[str]			  =		TFMobileBertModel.from_pretrained(__lowercase)
            self.assertIsNotNone(__lowercase)


@require_tf
class 	lowerCamelCase_     (   unittest.TestCase ):






    '''simple docstring'''






    @slow
    def 			UpperCamelCase__   (	self)       ->	Any:
        __UpperCamelCase     :Union[str, Any]			  =		TFMobileBertForPreTraining.from_pretrained('''google/mobilebert-uncased''')
        __UpperCamelCase     :Any			  =		tf.constant([[0, 1, 2, 3, 4, 5]])
        __UpperCamelCase     :int			  =		model(__lowercase)[0]

        __UpperCamelCase     :Optional[Any]			  =		[1, 6, 30_522]
        self.assertEqual(output.shape , __lowercase)

        __UpperCamelCase     :Dict			  =		tf.constant(
            [
                [
                    [-4.5_91_95_47, -9.24_82_95, -9.64_52_56],
                    [-6.7_30_61_75, -6.44_02_84, -6.6_05_28_37],
                    [-7.2_74_35_06, -6.7_84_79_15, -6.02_46_73],
                ]
            ])
        tf.debugging.assert_near(output[:, :3, :3] , __lowercase , atol=1E-4)






# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import torch

from ..models.auto import AutoModelForSequenceClassification, AutoTokenizer
from .base import PipelineTool


class 	lowerCamelCase_     (   UpperCAmelCase_ ):






    '''simple docstring'''






    a__	:						Optional[Any]						=      """facebook/bart-large-mnli"""
    a__	:						int						=      (
        """This is a tool that classifies an English text using provided labels. It takes two inputs: `text`, which """
        """should be the text to classify, and `labels`, which should be the list of labels to use for classification. """
        """It returns the most likely label in the list of provided `labels` for the input text."""
    )
    a__	:						Optional[Any]						=      """text_classifier"""
    a__	:						Any						=      AutoTokenizer
    a__	:						str						=      AutoModelForSequenceClassification

    a__	:						str						=      ["""text""", ["""text"""]]
    a__	:						Optional[int]						=      ["""text"""]







    def 			UpperCamelCase__   (	self)       ->	Union[str, Any]:
        super().setup()
        __UpperCamelCase     :int			  =		self.model.config
        __UpperCamelCase     :Optional[Any]			  =		-1
        for idx, label in config.idalabel.items():
            if label.lower().startswith('''entail'''):
                __UpperCamelCase     :List[Any]			  =		int(__lowercase)
        if self.entailment_id == -1:
            raise ValueError('''Could not determine the entailment ID from the model config, please pass it at init.''')







    def 			UpperCamelCase__   (	self , __lowercase , __lowercase)       ->	Union[str, Any]:
        __UpperCamelCase     :Any			  =		labels
        return self.pre_processor(
            [text] * len(__lowercase) , [f"""This example is {label}""" for label in labels] , return_tensors='''pt''' , padding='''max_length''' , )







    def 			UpperCamelCase__   (	self , __lowercase)       ->	Optional[Any]:
        __UpperCamelCase     :List[Any]			  =		outputs.logits
        __UpperCamelCase     :Any			  =		torch.argmax(logits[:, 2]).item()
        return self._labels[label_id]





