

import inspect
import unittest

from transformers import MobileNetVaConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

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


if is_torch_available():
     import torch

     from transformers import MobileNetVaForImageClassification, MobileNetVaForSemanticSegmentation, MobileNetVaModel
     from transformers.models.mobilenet_va.modeling_mobilenet_va import MOBILENET_V2_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
     from PIL import Image

     from transformers import MobileNetVaImageProcessor


class 			A						(							_UpperCAmelCase			):







    """simple docstring"""





    def 						snake_case__					(				self			:	str		)->   Any:





        '''simple docstring'''







        A__   				=       self.config_class(**self.inputs_dict		)
        self.parent.assertTrue(hasattr(lowercase_,'tf_padding'		)		)
        self.parent.assertTrue(hasattr(lowercase_,'depth_multiplier'		)		)




class 			A						:







    """simple docstring"""





    def __init__(				self			:	Any,lowercase_			:	Optional[int],lowercase_			:	Dict=1_3,lowercase_			:	Any=3,lowercase_			:	Any=3_2,lowercase_			:	Any=0.25,lowercase_			:	Optional[int]=8,lowercase_			:	List[str]=8,lowercase_			:	str=6,lowercase_			:	Union[str, Any]=3_2,lowercase_			:	Optional[Any]=True,lowercase_			:	Optional[int]=True,lowercase_			:	Any=True,lowercase_			:	List[Any]="relu6",lowercase_			:	Dict=1_2_8_0,lowercase_			:	Dict=0.1,lowercase_			:	Optional[int]=0.02,lowercase_			:	List[Any]=True,lowercase_			:	Union[str, Any]=True,lowercase_			:	str=1_0,lowercase_			:	str=None,)->   List[Any]:





        '''simple docstring'''







        A__   				=       parent
        A__   				=       batch_size
        A__   				=       num_channels
        A__   				=       image_size
        A__   				=       depth_multiplier
        A__   				=       depth_divisible_by
        A__   				=       min_depth
        A__   				=       expand_ratio
        A__   				=       tf_padding
        A__   				=       output_stride
        A__   				=       first_layer_is_expansion
        A__   				=       finegrained_output
        A__   				=       hidden_act
        A__   				=       last_hidden_size if finegrained_output else int(last_hidden_size * depth_multiplier		)
        A__   				=       classifier_dropout_prob
        A__   				=       use_labels
        A__   				=       is_training
        A__   				=       num_labels
        A__   				=       initializer_range
        A__   				=       scope


    def 						snake_case__					(				self			:	Optional[int]		)->   Dict:





        '''simple docstring'''







        A__   				=       floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]		)

        A__   				=       None
        A__   				=       None
        if self.use_labels:
            A__   				=       ids_tensor([self.batch_size],self.num_labels		)
            A__   				=       ids_tensor([self.batch_size, self.image_size, self.image_size],self.num_labels		)

        A__   				=       self.get_config()

        return config, pixel_values, labels, pixel_labels


    def 						snake_case__					(				self			:	Tuple		)->   List[str]:





        '''simple docstring'''







        return MobileNetVaConfig(
            num_channels=self.num_channels,image_size=self.image_size,depth_multiplier=self.depth_multiplier,depth_divisible_by=self.depth_divisible_by,min_depth=self.min_depth,expand_ratio=self.expand_ratio,output_stride=self.output_stride,first_layer_is_expansion=self.first_layer_is_expansion,finegrained_output=self.finegrained_output,hidden_act=self.hidden_act,tf_padding=self.tf_padding,classifier_dropout_prob=self.classifier_dropout_prob,initializer_range=self.initializer_range,)


    def 						snake_case__					(				self			:	Tuple,lowercase_			:	Dict,lowercase_			:	Tuple,lowercase_			:	Tuple,lowercase_			:	Union[str, Any]		)->   Optional[int]:





        '''simple docstring'''







        A__   				=       MobileNetVaModel(config=lowercase_		)
        model.to(lowercase_		)
        model.eval()
        A__   				=       model(lowercase_		)
        self.parent.assertEqual(
            result.last_hidden_state.shape,(
                self.batch_size,
                self.last_hidden_size,
                self.image_size // self.output_stride,
                self.image_size // self.output_stride,
            ),)
        self.parent.assertEqual(
            result.pooler_output.shape,(self.batch_size, self.last_hidden_size),)


    def 						snake_case__					(				self			:	List[str],lowercase_			:	Optional[Any],lowercase_			:	Dict,lowercase_			:	str,lowercase_			:	Any		)->   Union[str, Any]:





        '''simple docstring'''







        A__   				=       self.num_labels
        A__   				=       MobileNetVaForImageClassification(lowercase_		)
        model.to(lowercase_		)
        model.eval()
        A__   				=       model(lowercase_,labels=lowercase_		)
        self.parent.assertEqual(result.logits.shape,(self.batch_size, self.num_labels)		)


    def 						snake_case__					(				self			:	Dict,lowercase_			:	Optional[int],lowercase_			:	Any,lowercase_			:	List[Any],lowercase_			:	Dict		)->   Any:





        '''simple docstring'''







        A__   				=       self.num_labels
        A__   				=       MobileNetVaForSemanticSegmentation(lowercase_		)
        model.to(lowercase_		)
        model.eval()
        A__   				=       model(lowercase_		)
        self.parent.assertEqual(
            result.logits.shape,(
                self.batch_size,
                self.num_labels,
                self.image_size // self.output_stride,
                self.image_size // self.output_stride,
            ),)
        A__   				=       model(lowercase_,labels=lowercase_		)
        self.parent.assertEqual(
            result.logits.shape,(
                self.batch_size,
                self.num_labels,
                self.image_size // self.output_stride,
                self.image_size // self.output_stride,
            ),)


    def 						snake_case__					(				self			:	List[Any]		)->   int:





        '''simple docstring'''







        A__   				=       self.prepare_config_and_inputs()
        A__				,						A__				,						A__				,						A__   				=       config_and_inputs
        A__   				=       {'pixel_values': pixel_values}
        return config, inputs_dict




@require_torch
class 			A						(							_UpperCAmelCase   ,							_UpperCAmelCase   ,							unittest.TestCase			):







    """simple docstring"""





    lowerCamelCase		     =							(
        (MobileNetVaModel, MobileNetVaForImageClassification, MobileNetVaForSemanticSegmentation)
        if is_torch_available()
        else ()
    )
    lowerCamelCase		     =							(
        {
            'feature-extraction': MobileNetVaModel,
            'image-classification': MobileNetVaForImageClassification,
            'image-segmentation': MobileNetVaForSemanticSegmentation,
        }
        if is_torch_available()
        else {}
    )

    lowerCamelCase		     =							False
    lowerCamelCase		     =							False
    lowerCamelCase		     =							False
    lowerCamelCase		     =							False


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





        '''simple docstring'''







        A__   				=       MobileNetVaModelTester(self		)
        A__   				=       MobileNetVaConfigTester(self,config_class=lowercase_,has_text_modality=lowercase_		)


    def 						snake_case__					(				self			:	List[str]		)->   Any:





        '''simple docstring'''







        self.config_tester.run_common_tests()


    @unittest.skip(reason='MobileNetV2 does not use inputs_embeds'		)
    def 						snake_case__					(				self			:	int		)->   Any:





        '''simple docstring'''







        pass


    @unittest.skip(reason='MobileNetV2 does not support input and output embeddings'		)
    def 						snake_case__					(				self			:	List[str]		)->   str:





        '''simple docstring'''







        pass


    @unittest.skip(reason='MobileNetV2 does not output attentions'		)
    def 						snake_case__					(				self			:	Optional[int]		)->   int:





        '''simple docstring'''







        pass


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





        '''simple docstring'''







        A__				,						A__   				=       self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
            A__   				=       model_class(lowercase_		)
            A__   				=       inspect.signature(model.forward		)
            # signature.parameters is an OrderedDict => so arg_names order is deterministic
            A__   				=       [*signature.parameters.keys()]

            A__   				=       ['pixel_values']
            self.assertListEqual(arg_names[:1],lowercase_		)


    def 						snake_case__					(				self			:	Optional[Any]		)->   List[Any]:





        '''simple docstring'''







        A__   				=       self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*lowercase_		)


    def 						snake_case__					(				self			:	Tuple		)->   Optional[Any]:





        '''simple docstring'''







        def check_hidden_states_output(lowercase_			:	Optional[Any],lowercase_			:	Optional[int],lowercase_			:	Tuple		):
            A__   				=       model_class(lowercase_		)
            model.to(lowercase_		)
            model.eval()

            with torch.no_grad():
                A__   				=       model(**self._prepare_for_class(lowercase_,lowercase_		)		)

            A__   				=       outputs.hidden_states

            A__   				=       1_6
            self.assertEqual(len(lowercase_		),lowercase_		)

        A__				,						A__   				=       self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
            A__   				=       True
            check_hidden_states_output(lowercase_,lowercase_,lowercase_		)

            # check that output_hidden_states also work using config
            del inputs_dict["output_hidden_states"]
            A__   				=       True

            check_hidden_states_output(lowercase_,lowercase_,lowercase_		)


    def 						snake_case__					(				self			:	Tuple		)->   str:





        '''simple docstring'''







        A__   				=       self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_image_classification(*lowercase_		)


    def 						snake_case__					(				self			:	Tuple		)->   str:





        '''simple docstring'''







        A__   				=       self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_semantic_segmentation(*lowercase_		)


    @slow
    def 						snake_case__					(				self			:	Optional[Any]		)->   Dict:





        '''simple docstring'''







        for model_name in MOBILENET_V2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
            A__   				=       MobileNetVaModel.from_pretrained(lowercase_		)
            self.assertIsNotNone(lowercase_		)




def 						_snake_case(   )							->				Optional[int]:





    '''simple docstring'''






    A__   				=       Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'				)
    return image


@require_torch
@require_vision
class 			A						(							unittest.TestCase			):







    """simple docstring"""





    @cached_property
    def 						snake_case__					(				self			:	Any		)->   Union[str, Any]:





        '''simple docstring'''







        return (
            MobileNetVaImageProcessor.from_pretrained('google/mobilenet_v2_1.0_224'		) if is_vision_available() else None
        )


    @slow
    def 						snake_case__					(				self			:	List[Any]		)->   Any:





        '''simple docstring'''







        A__   				=       MobileNetVaForImageClassification.from_pretrained('google/mobilenet_v2_1.0_224'		).to(lowercase_		)

        A__   				=       self.default_image_processor
        A__   				=       prepare_img()
        A__   				=       image_processor(images=lowercase_,return_tensors='pt'		).to(lowercase_		)

        # forward pass
        with torch.no_grad():
            A__   				=       model(**lowercase_		)

        # verify the logits
        A__   				=       torch.Size((1, 1_0_0_1)		)
        self.assertEqual(outputs.logits.shape,lowercase_		)

        A__   				=       torch.tensor([0.2_445, -1.1_993, 0.1_905]		).to(lowercase_		)

        self.assertTrue(torch.allclose(outputs.logits[0, :3],lowercase_,atol=1E-4		)		)


    @slow
    def 						snake_case__					(				self			:	Optional[Any]		)->   Optional[Any]:





        '''simple docstring'''







        A__   				=       MobileNetVaForSemanticSegmentation.from_pretrained('google/deeplabv3_mobilenet_v2_1.0_513'		)
        A__   				=       model.to(lowercase_		)

        A__   				=       MobileNetVaImageProcessor.from_pretrained('google/deeplabv3_mobilenet_v2_1.0_513'		)

        A__   				=       prepare_img()
        A__   				=       image_processor(images=lowercase_,return_tensors='pt'		).to(lowercase_		)

        # forward pass
        with torch.no_grad():
            A__   				=       model(**lowercase_		)
        A__   				=       outputs.logits

        # verify the logits
        A__   				=       torch.Size((1, 2_1, 6_5, 6_5)		)
        self.assertEqual(logits.shape,lowercase_		)

        A__   				=       torch.tensor(
            [
                [[17.5_790, 17.7_581, 18.3_355], [18.3_257, 18.4_230, 18.8_973], [18.6_169, 18.8_650, 19.2_187]],
                [[-2.1_595, -2.0_977, -2.3_741], [-2.4_226, -2.3_028, -2.6_835], [-2.7_819, -2.5_991, -2.7_706]],
                [[4.2_058, 4.8_317, 4.7_638], [4.4_136, 5.0_361, 4.9_383], [4.5_028, 4.9_644, 4.8_734]],
            ],device=lowercase_,)

        self.assertTrue(torch.allclose(logits[0, :3, :3, :3],lowercase_,atol=1E-4		)		)




import argparse
import collections

import torch
from flax import traverse_util
from tax import checkpoints

from transformers import TaConfig, TaEncoderModel, TaForConditionalGeneration
from transformers.utils import logging


logging.set_verbosity_info()
def 						_snake_case(   SCREAMING_SNAKE_CASE__      :     List[Any]   ,      SCREAMING_SNAKE_CASE__      :     List[str]   ,      SCREAMING_SNAKE_CASE__      :     Optional[int]   ,      SCREAMING_SNAKE_CASE__      :     Optional[int]="attention"				)							->				Union[str, Any]:





    '''simple docstring'''






    A__   				=       params[f'{prefix}/layers_{i}/{layer_name}/key/kernel']
    A__   				=       params[f'{prefix}/layers_{i}/{layer_name}/out/kernel']
    A__   				=       params[f'{prefix}/layers_{i}/{layer_name}/query/kernel']
    A__   				=       params[f'{prefix}/layers_{i}/{layer_name}/value/kernel']
    return k, o, q, v
def 						_snake_case(   SCREAMING_SNAKE_CASE__      :     Any   ,      SCREAMING_SNAKE_CASE__      :     Tuple   ,      SCREAMING_SNAKE_CASE__      :     Optional[Any]   ,      SCREAMING_SNAKE_CASE__      :     Dict=False				)							->				str:





    '''simple docstring'''






    if split_mlp_wi:
        A__   				=       params[f'{prefix}/layers_{i}/mlp/wi_0/kernel']
        A__   				=       params[f'{prefix}/layers_{i}/mlp/wi_1/kernel']
        A__   				=       (wi_a, wi_a)
    else:
        A__   				=       params[f'{prefix}/layers_{i}/mlp/wi/kernel']

    A__   				=       params[f'{prefix}/layers_{i}/mlp/wo/kernel']
    return wi, wo
def 						_snake_case(   SCREAMING_SNAKE_CASE__      :     Optional[Any]   ,      SCREAMING_SNAKE_CASE__      :     Tuple   ,      SCREAMING_SNAKE_CASE__      :     Dict   ,      SCREAMING_SNAKE_CASE__      :     int				)							->				str:





    '''simple docstring'''






    return params[f'{prefix}/layers_{i}/{layer_name}/scale']
def 						_snake_case(   SCREAMING_SNAKE_CASE__      :     dict   ,      *, SCREAMING_SNAKE_CASE__      :     int   ,      SCREAMING_SNAKE_CASE__      :     bool				)							->				int:





    '''simple docstring'''






    A__   				=       traverse_util.flatten_dict(variables['target']				)
    A__   				=       {'/'.join(SCREAMING_SNAKE_CASE__				): v for k, v in old.items()}

    # v1.1 models have a gated GeLU with wi_0 and wi_1 instead of wi
    A__   				=       'encoder/layers_0/mlp/wi_0/kernel' in old
    print('Split MLP:'   ,      SCREAMING_SNAKE_CASE__				)

    A__   				=       collections.OrderedDict()

    # Shared embeddings.
    A__   				=       old['token_embedder/embedding']

    # Encoder.
    for i in range(SCREAMING_SNAKE_CASE__				):
        # Block i, layer 0 (Self Attention).
        A__   				=       tax_layer_norm_lookup(SCREAMING_SNAKE_CASE__   ,      SCREAMING_SNAKE_CASE__   ,      'encoder'   ,      'pre_attention_layer_norm'				)
        A__				,						A__				,						A__				,						A__   				=       tax_attention_lookup(SCREAMING_SNAKE_CASE__   ,      SCREAMING_SNAKE_CASE__   ,      'encoder'   ,      'attention'				)
        A__   				=       layer_norm
        A__   				=       k.T
        A__   				=       o.T
        A__   				=       q.T
        A__   				=       v.T

        # Block i, layer 1 (MLP).
        A__   				=       tax_layer_norm_lookup(SCREAMING_SNAKE_CASE__   ,      SCREAMING_SNAKE_CASE__   ,      'encoder'   ,      'pre_mlp_layer_norm'				)
        A__				,						A__   				=       tax_mlp_lookup(SCREAMING_SNAKE_CASE__   ,      SCREAMING_SNAKE_CASE__   ,      'encoder'   ,      SCREAMING_SNAKE_CASE__				)
        A__   				=       layer_norm
        if split_mlp_wi:
            A__   				=       wi[0].T
            A__   				=       wi[1].T
        else:
            A__   				=       wi.T
        A__   				=       wo.T

    A__   				=       old[
        'encoder/relpos_bias/rel_embedding'
    ].T
    A__   				=       old['encoder/encoder_norm/scale']

    if not is_encoder_only:
        # Decoder.
        for i in range(SCREAMING_SNAKE_CASE__				):
            # Block i, layer 0 (Self Attention).
            A__   				=       tax_layer_norm_lookup(SCREAMING_SNAKE_CASE__   ,      SCREAMING_SNAKE_CASE__   ,      'decoder'   ,      'pre_self_attention_layer_norm'				)
            A__				,						A__				,						A__				,						A__   				=       tax_attention_lookup(SCREAMING_SNAKE_CASE__   ,      SCREAMING_SNAKE_CASE__   ,      'decoder'   ,      'self_attention'				)
            A__   				=       layer_norm
            A__   				=       k.T
            A__   				=       o.T
            A__   				=       q.T
            A__   				=       v.T

            # Block i, layer 1 (Cross Attention).
            A__   				=       tax_layer_norm_lookup(SCREAMING_SNAKE_CASE__   ,      SCREAMING_SNAKE_CASE__   ,      'decoder'   ,      'pre_cross_attention_layer_norm'				)
            A__				,						A__				,						A__				,						A__   				=       tax_attention_lookup(SCREAMING_SNAKE_CASE__   ,      SCREAMING_SNAKE_CASE__   ,      'decoder'   ,      'encoder_decoder_attention'				)
            A__   				=       layer_norm
            A__   				=       k.T
            A__   				=       o.T
            A__   				=       q.T
            A__   				=       v.T

            # Block i, layer 2 (MLP).
            A__   				=       tax_layer_norm_lookup(SCREAMING_SNAKE_CASE__   ,      SCREAMING_SNAKE_CASE__   ,      'decoder'   ,      'pre_mlp_layer_norm'				)
            A__				,						A__   				=       tax_mlp_lookup(SCREAMING_SNAKE_CASE__   ,      SCREAMING_SNAKE_CASE__   ,      'decoder'   ,      SCREAMING_SNAKE_CASE__				)
            A__   				=       layer_norm
            if split_mlp_wi:
                A__   				=       wi[0].T
                A__   				=       wi[1].T
            else:
                A__   				=       wi.T
            A__   				=       wo.T

        A__   				=       old['decoder/decoder_norm/scale']
        A__   				=       old[
            'decoder/relpos_bias/rel_embedding'
        ].T

        # LM Head (only in v1.1 checkpoints, in v1.0 embeddings are used instead)
        if "decoder/logits_dense/kernel" in old:
            A__   				=       old['decoder/logits_dense/kernel'].T

    return new
def 						_snake_case(   SCREAMING_SNAKE_CASE__      :     Optional[Any]   ,      SCREAMING_SNAKE_CASE__      :     bool				)							->				Dict:





    '''simple docstring'''






    A__   				=       collections.OrderedDict([(k, torch.from_numpy(v.copy()				)) for (k, v) in converted_params.items()]				)

    # Add what is missing.
    if "encoder.embed_tokens.weight" not in state_dict:
        A__   				=       state_dict['shared.weight']

    if not is_encoder_only:
        if "decoder.embed_tokens.weight" not in state_dict:
            A__   				=       state_dict['shared.weight']

        if "lm_head.weight" not in state_dict:  # For old 1.0 models.
            print('Using shared word embeddings as lm_head.'				)
            A__   				=       state_dict['shared.weight']

    return state_dict
def 						_snake_case(   SCREAMING_SNAKE_CASE__      :     Union[str, Any]   ,      SCREAMING_SNAKE_CASE__      :     Any   ,      SCREAMING_SNAKE_CASE__      :     Any   ,      SCREAMING_SNAKE_CASE__      :     Tuple				)							->				int:





    '''simple docstring'''






    A__   				=       checkpoints.load_tax_checkpoint(SCREAMING_SNAKE_CASE__				)
    A__   				=       convert_tax_to_pytorch(SCREAMING_SNAKE_CASE__   ,      num_layers=config.num_layers   ,      is_encoder_only=SCREAMING_SNAKE_CASE__				)
    A__   				=       make_state_dict(SCREAMING_SNAKE_CASE__   ,      SCREAMING_SNAKE_CASE__				)
    model.load_state_dict(SCREAMING_SNAKE_CASE__   ,      strict=SCREAMING_SNAKE_CASE__				)
def 						_snake_case(   SCREAMING_SNAKE_CASE__      :     Optional[int]   ,      SCREAMING_SNAKE_CASE__      :     List[Any]   ,      SCREAMING_SNAKE_CASE__      :     Tuple   ,      SCREAMING_SNAKE_CASE__      :     bool = False				)							->				Any:





    '''simple docstring'''






    A__   				=       TaConfig.from_json_file(SCREAMING_SNAKE_CASE__				)
    print(f'Building PyTorch model from configuration: {config}'				)
    # Non-v1.1 checkpoints could also use T5Model, but this works for all.
    # The v1.0 checkpoints will simply have an LM head that is the word embeddings.
    if is_encoder_only:
        A__   				=       TaEncoderModel(SCREAMING_SNAKE_CASE__				)
    else:
        A__   				=       TaForConditionalGeneration(SCREAMING_SNAKE_CASE__				)

    # Load weights from tf checkpoint
    load_tax_weights_in_ta(SCREAMING_SNAKE_CASE__   ,      SCREAMING_SNAKE_CASE__   ,      SCREAMING_SNAKE_CASE__   ,      SCREAMING_SNAKE_CASE__				)

    # Save pytorch-model
    print(f'Save PyTorch model to {pytorch_dump_path}'				)
    model.save_pretrained(SCREAMING_SNAKE_CASE__				)

    # Verify that we can load the checkpoint.
    model.from_pretrained(SCREAMING_SNAKE_CASE__				)
    print('Done'				)


if __name__ == "__main__":
     lowercase_			=  argparse.ArgumentParser(description="Converts a native T5X checkpoint into a PyTorch checkpoint.")
     # Required parameters
     parser.add_argument(
         "--t5x_checkpoint_path", default=None, type=str, required=True, help="Path to the T5X checkpoint."
     )
     parser.add_argument(
         "--config_file",
         default=None,
         type=str,
         required=True,
         help="The config json file corresponding to the pre-trained T5 model.\nThis specifies the model architecture.",
     )
     parser.add_argument(
         "--pytorch_dump_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
     )
     parser.add_argument(
         "--is_encoder_only", action="store_true", help="Check if the model is encoder-decoder model", default=False
     )
     lowercase_			=  parser.parse_args()
     convert_tax_checkpoint_to_pytorch(
         args.tax_checkpoint_path, args.config_file, args.pytorch_dump_path, args.is_encoder_only
     )

