




"""simple docstring"""


from __future__ import annotations

import inspect
import unittest
from math import floor

import numpy as np

from transformers import CvtConfig
from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import cached_property, is_tf_available, is_vision_available

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


if is_tf_available():
     import tensorflow as tf

     from transformers import TFCvtForImageClassification, TFCvtModel
     from transformers.models.cvt.modeling_tf_cvt import TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
     from PIL import Image

     from transformers import AutoImageProcessor






class 					lowercase(     UpperCAmelCase_ ):







 '''simple docstring'''





 def  UpperCamelCase_   (				self:		List[str]		):




   '''simple docstring'''



   _snake_case							:		Union[str, Any]								=  self.config_class(**self.inputs_dict		)
   self.parent.assertTrue(hasattr(__lowercase,      """embed_dim"""		)		)
   self.parent.assertTrue(hasattr(__lowercase,      """num_heads"""		)		)







class 					lowercase:







 '''simple docstring'''





 def __init__(				self:		Union[str, Any],      a_:		List[Any],      a_:		List[Any]=13,      a_:		Optional[int]=64,      a_:		Optional[Any]=3,      a_:		Any=[16, 48, 96],      a_:		Tuple=[1, 3, 6],      a_:		Dict=[1, 2, 10],      a_:		int=[7, 3, 3],      a_:		Optional[Any]=[4, 2, 2],      a_:		Any=[2, 1, 1],      a_:		Union[str, Any]=[2, 2, 2],      a_:		Optional[int]=[False, False, True],      a_:		List[Any]=[0.0, 0.0, 0.0],      a_:		List[Any]=0.02,      a_:		Optional[Any]=1E-12,      a_:		List[str]=True,      a_:		Dict=True,      a_:		List[Any]=2,      ):




   '''simple docstring'''



   _snake_case							:		Optional[Any]								=  parent
   _snake_case							:		Tuple								=  batch_size
   _snake_case							:		List[Any]								=  image_size
   _snake_case							:		str								=  patch_sizes
   _snake_case							:		Optional[Any]								=  patch_stride
   _snake_case							:		int								=  patch_padding
   _snake_case							:		int								=  is_training
   _snake_case							:		Optional[Any]								=  use_labels
   _snake_case							:		str								=  num_labels
   _snake_case							:		Tuple								=  num_channels
   _snake_case							:		Optional[Any]								=  embed_dim
   _snake_case							:		List[Any]								=  num_heads
   _snake_case							:		Dict								=  stride_kv
   _snake_case							:		Optional[Any]								=  depth
   _snake_case							:		Any								=  cls_token
   _snake_case							:		Dict								=  attention_drop_rate
   _snake_case							:		int								=  initializer_range
   _snake_case							:		str								=  layer_norm_eps


 def  UpperCamelCase_   (				self:		Optional[int]		):




   '''simple docstring'''



   _snake_case							:		Dict								=  floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]		)

   _snake_case							:		int								=  None
   if self.use_labels:
     # create a random int32 tensor of given shape
     _snake_case							:		Optional[Any]								=  ids_tensor([self.batch_size],      self.num_labels		)

   _snake_case							:		Dict								=  self.get_config()
   return config, pixel_values, labels


 def  UpperCamelCase_   (				self:		Dict		):




   '''simple docstring'''



   return CvtConfig(
       image_size=self.image_size,      num_labels=self.num_labels,      num_channels=self.num_channels,      embed_dim=self.embed_dim,      num_heads=self.num_heads,      patch_sizes=self.patch_sizes,      patch_padding=self.patch_padding,      patch_stride=self.patch_stride,      stride_kv=self.stride_kv,      depth=self.depth,      cls_token=self.cls_token,      attention_drop_rate=self.attention_drop_rate,      initializer_range=self.initializer_range,      )


 def  UpperCamelCase_   (				self:		Union[str, Any],      a_:		Tuple,      a_:		str,      a_:		Any		):




   '''simple docstring'''



   _snake_case							:		Any								=  TFCvtModel(config=__lowercase		)
   _snake_case							:		Optional[int]								=  model(__lowercase,      training=__lowercase		)
   _snake_case							:		List[Any]								=  (self.image_size, self.image_size)
   _snake_case							:		Tuple								=  image_size[0], image_size[1]
   for i in range(len(self.depth		)		):
     _snake_case							:		Tuple								=  floor(((height + 2 * self.patch_padding[i] - self.patch_sizes[i]) / self.patch_stride[i]) + 1		)
     _snake_case							:		int								=  floor(((width + 2 * self.patch_padding[i] - self.patch_sizes[i]) / self.patch_stride[i]) + 1		)
   self.parent.assertEqual(result.last_hidden_state.shape,      (self.batch_size, self.embed_dim[-1], height, width)		)


 def  UpperCamelCase_   (				self:		Tuple,      a_:		Optional[int],      a_:		int,      a_:		Optional[int]		):




   '''simple docstring'''



   _snake_case							:		List[str]								=  self.num_labels
   _snake_case							:		Any								=  TFCvtForImageClassification(__lowercase		)
   _snake_case							:		int								=  model(__lowercase,      labels=__lowercase,      training=__lowercase		)
   self.parent.assertEqual(result.logits.shape,      (self.batch_size, self.num_labels)		)


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




   '''simple docstring'''



   _snake_case							:		List[str]								=  self.prepare_config_and_inputs()
   _snake_case							:		List[str]								=  config_and_inputs
   _snake_case							:		List[Any]								=  {'''pixel_values''': pixel_values}
   return config, inputs_dict







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







 '''simple docstring'''





 lowercase__ 					=					(TFCvtModel, TFCvtForImageClassification) if is_tf_available() else ()
 lowercase__ 					=					(
     {"""feature-extraction""": TFCvtModel, """image-classification""": TFCvtForImageClassification}
     if is_tf_available()
     else {}
 )
 lowercase__ 					=					False
 lowercase__ 					=					False
 lowercase__ 					=					False
 lowercase__ 					=					False
 lowercase__ 					=					False


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




   '''simple docstring'''



   _snake_case							:		List[str]								=  TFCvtModelTester(self		)
   _snake_case							:		List[str]								=  TFCvtConfigTester(self,      config_class=__lowercase,      has_text_modality=__lowercase,      hidden_size=37		)


 def  UpperCamelCase_   (				self:		str		):




   '''simple docstring'''



   self.config_tester.create_and_test_config_common_properties()
   self.config_tester.create_and_test_config_to_json_string()
   self.config_tester.create_and_test_config_to_json_file()
   self.config_tester.create_and_test_config_from_and_save_pretrained()
   self.config_tester.create_and_test_config_with_num_labels()
   self.config_tester.check_config_can_be_init_without_params()
   self.config_tester.check_config_arguments_init()


 @unittest.skip(reason="""Cvt does not output attentions"""		)
 def  UpperCamelCase_   (				self:		Optional[int]		):




   '''simple docstring'''



   pass


 @unittest.skip(reason="""Cvt does not use inputs_embeds"""		)
 def  UpperCamelCase_   (				self:		List[str]		):




   '''simple docstring'''



   pass


 @unittest.skip(reason="""Cvt does not support input and output embeddings"""		)
 def  UpperCamelCase_   (				self:		Optional[int]		):




   '''simple docstring'''



   pass


 @unittest.skipIf(
     not is_tf_available() or len(tf.config.list_physical_devices("""GPU"""		)		) == 0,      reason="""TF does not support backprop for grouped convolutions on CPU.""",      )
 def  UpperCamelCase_   (				self:		str		):




   '''simple docstring'''



   super().test_dataset_conversion()


 @unittest.skipIf(
     not is_tf_available() or len(tf.config.list_physical_devices("""GPU"""		)		) == 0,      reason="""TF does not support backprop for grouped convolutions on CPU.""",      )
 @slow
 def  UpperCamelCase_   (				self:		Union[str, Any]		):




   '''simple docstring'''



   super().test_keras_fit()


 @unittest.skip(reason="""Get `Failed to determine best cudnn convolution algo.` error after using TF 2.12+cuda 11.8"""		)
 def  UpperCamelCase_   (				self:		Any		):




   '''simple docstring'''



   _snake_case							:		int								=  tf.keras.mixed_precision.Policy("""mixed_float16"""		)
   tf.keras.mixed_precision.set_global_policy(__lowercase		)
   super().test_keras_fit()
   tf.keras.mixed_precision.set_global_policy("""float32"""		)


 def  UpperCamelCase_   (				self:		Optional[int]		):




   '''simple docstring'''



   _snake_case							:		List[str]								=  self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
     _snake_case							:		Dict								=  model_class(__lowercase		)
     _snake_case							:		Dict								=  inspect.signature(model.call		)
     # signature.parameters is an OrderedDict => so arg_names order is deterministic
     _snake_case							:		Dict								=  [*signature.parameters.keys()]

     _snake_case							:		Optional[int]								=  ['''pixel_values''']
     self.assertListEqual(arg_names[:1],      __lowercase		)


 def  UpperCamelCase_   (				self:		Optional[int]		):




   '''simple docstring'''



   def check_hidden_states_output(a_:		Tuple,      a_:		Union[str, Any],      a_:		Any		):
     _snake_case							:		Tuple								=  model_class(__lowercase		)

     _snake_case							:		Optional[int]								=  model(**self._prepare_for_class(__lowercase,      __lowercase		)		)
     _snake_case							:		List[Any]								=  outputs.hidden_states

     _snake_case							:		Any								=  len(self.model_tester.depth		)
     self.assertEqual(len(__lowercase		),      __lowercase		)

     # verify the first hidden states (first block)
     self.assertListEqual(
         list(hidden_states[0].shape[-3:]		),      [
             self.model_tester.embed_dim[0],
             self.model_tester.image_size // 4,
             self.model_tester.image_size // 4,
         ],      )

   _snake_case							:		str								=  self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
     _snake_case							:		int								=  True
     check_hidden_states_output(__lowercase,      __lowercase,      __lowercase		)

     # check that output_hidden_states also work using config
     del inputs_dict["output_hidden_states"]
     _snake_case							:		Optional[Any]								=  True

     check_hidden_states_output(__lowercase,      __lowercase,      __lowercase		)


 def  UpperCamelCase_   (				self:		str		):




   '''simple docstring'''



   _snake_case							:		Optional[int]								=  self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_model(*__lowercase		)


 def  UpperCamelCase_   (				self:		str		):




   '''simple docstring'''



   _snake_case							:		str								=  self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_image_classification(*__lowercase		)


 @slow
 def  UpperCamelCase_   (				self:		Tuple		):




   '''simple docstring'''



   for model_name in TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
     _snake_case							:		Optional[int]								=  TFCvtModel.from_pretrained(__lowercase		)
     self.assertIsNotNone(__lowercase		)







def 					UpperCAmelCase__       ():
  """simple docstring"""

  _snake_case							:		str								=  Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""    )
  return image






@require_tf
@require_vision
class 					lowercase(     unittest.TestCase ):







 '''simple docstring'''





 @cached_property
 def  UpperCamelCase_   (				self:		List[str]		):




   '''simple docstring'''



   return AutoImageProcessor.from_pretrained(TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST[0]		)


 @slow
 def  UpperCamelCase_   (				self:		int		):




   '''simple docstring'''



   _snake_case							:		List[str]								=  TFCvtForImageClassification.from_pretrained(TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST[0]		)

   _snake_case							:		Optional[Any]								=  self.default_image_processor
   _snake_case							:		Tuple								=  prepare_img()
   _snake_case							:		Dict								=  image_processor(images=__lowercase,      return_tensors="""tf"""		)

   # forward pass
   _snake_case							:		str								=  model(**__lowercase		)

   # verify the logits
   _snake_case							:		List[str]								=  tf.TensorShape((1, 1_000)		)
   self.assertEqual(outputs.logits.shape,      __lowercase		)

   _snake_case							:		Dict								=  tf.constant([0.9_285, 0.9_015, -0.3_150]		)
   self.assertTrue(np.allclose(outputs.logits[0, :3].numpy(),      __lowercase,      atol=1E-4		)		)



import argparse
import json
import os

import evaluate
import torch
from datasets import load_dataset
from torch.optim import AdamW
from torch.utils.data import DataLoader
from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed

from accelerate import Accelerator, DistributedType
from accelerate.utils.deepspeed import DummyOptim, DummyScheduler


__lowercase							 =					16
__lowercase							 =					32



def lowerCamelCase  (   SCREAMING_SNAKE_CASE  ,			SCREAMING_SNAKE_CASE = 16  ,			SCREAMING_SNAKE_CASE = "bert-base-cased"	):
    '''simple docstring'''







    __UpperCamelCase     :List[str]			  =		AutoTokenizer.from_pretrained(SCREAMING_SNAKE_CASE	)
    __UpperCamelCase     :Any			  =		load_dataset('''glue'''  ,			'''mrpc'''	)

    def tokenize_function(SCREAMING_SNAKE_CASE	):
        # max_length=None => use the model max length (it's actually the default)
        __UpperCamelCase     :int			  =		tokenizer(examples['''sentence1''']  ,			examples['''sentence2''']  ,			truncation=SCREAMING_SNAKE_CASE  ,			max_length=SCREAMING_SNAKE_CASE	)
        return outputs

    # Apply the method we just defined to all the examples in all the splits of the dataset
    __UpperCamelCase     :Tuple			  =		datasets.map(
        SCREAMING_SNAKE_CASE  ,			batched=SCREAMING_SNAKE_CASE  ,			remove_columns=['''idx''', '''sentence1''', '''sentence2''']  ,			load_from_cache_file=SCREAMING_SNAKE_CASE	)

    # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
    # transformers library
    __UpperCamelCase     :List[str]			  =		tokenized_datasets.rename_column('''label'''  ,			'''labels'''	)

    def collate_fn(SCREAMING_SNAKE_CASE	):
        # On TPU it's best to pad everything to the same length or training will be very slow.
        if accelerator.distributed_type == DistributedType.TPU:
            return tokenizer.pad(SCREAMING_SNAKE_CASE  ,			padding='''max_length'''  ,			max_length=128  ,			return_tensors='''pt'''	)
        return tokenizer.pad(SCREAMING_SNAKE_CASE  ,			padding='''longest'''  ,			return_tensors='''pt'''	)

    # Instantiate dataloaders.
    __UpperCamelCase     :Union[str, Any]			  =		DataLoader(
        tokenized_datasets['''train''']  ,			shuffle=SCREAMING_SNAKE_CASE  ,			collate_fn=SCREAMING_SNAKE_CASE  ,			batch_size=SCREAMING_SNAKE_CASE	)
    __UpperCamelCase     :Dict			  =		DataLoader(
        tokenized_datasets['''validation''']  ,			shuffle=SCREAMING_SNAKE_CASE  ,			collate_fn=SCREAMING_SNAKE_CASE  ,			batch_size=SCREAMING_SNAKE_CASE	)

    return train_dataloader, eval_dataloader



def lowerCamelCase  (   SCREAMING_SNAKE_CASE  ,			SCREAMING_SNAKE_CASE	):
    '''simple docstring'''







    __UpperCamelCase     :List[Any]			  =		Accelerator()

    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
    __UpperCamelCase     :int			  =		config['''lr''']
    __UpperCamelCase     :str			  =		int(config['''num_epochs''']	)
    __UpperCamelCase     :Any			  =		int(config['''seed''']	)
    __UpperCamelCase     :Dict			  =		int(config['''batch_size''']	)
    __UpperCamelCase     :Optional[Any]			  =		args.model_name_or_path

    set_seed(SCREAMING_SNAKE_CASE	)
    __UpperCamelCase  ,       __UpperCamelCase     :Dict			  =		get_dataloaders(SCREAMING_SNAKE_CASE  ,			SCREAMING_SNAKE_CASE  ,			SCREAMING_SNAKE_CASE	)

    # Instantiate the model (we build the model here so that the seed also control new weights initialization)
    __UpperCamelCase     :Any			  =		AutoModelForSequenceClassification.from_pretrained(SCREAMING_SNAKE_CASE  ,			return_dict=SCREAMING_SNAKE_CASE	)

    # Instantiate optimizer
    __UpperCamelCase     :List[str]			  =		(
        AdamW
        if accelerator.state.deepspeed_plugin is None
        or '''optimizer''' not in accelerator.state.deepspeed_plugin.deepspeed_config
        else DummyOptim
    )
    __UpperCamelCase     :Optional[Any]			  =		optimizer_cls(params=model.parameters()  ,			lr=SCREAMING_SNAKE_CASE	)

    if accelerator.state.deepspeed_plugin is not None:
        __UpperCamelCase     :Dict			  =		accelerator.state.deepspeed_plugin.deepspeed_config[
            '''gradient_accumulation_steps'''
        ]
    else:
        __UpperCamelCase     :Dict			  =		1
    __UpperCamelCase     :Tuple			  =		(len(SCREAMING_SNAKE_CASE	) * num_epochs) // gradient_accumulation_steps

    # Instantiate scheduler
    if (
        accelerator.state.deepspeed_plugin is None
        or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config
    ):
        __UpperCamelCase     :str			  =		get_linear_schedule_with_warmup(
            optimizer=SCREAMING_SNAKE_CASE  ,			num_warmup_steps=0  ,			num_training_steps=SCREAMING_SNAKE_CASE  ,			)
    else:
        __UpperCamelCase     :Dict			  =		DummyScheduler(SCREAMING_SNAKE_CASE  ,			total_num_steps=SCREAMING_SNAKE_CASE  ,			warmup_num_steps=0	)

    # Prepare everything
    # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
    # prepare method.
    __UpperCamelCase  ,       __UpperCamelCase  ,       __UpperCamelCase  ,       __UpperCamelCase  ,       __UpperCamelCase     :int			  =		accelerator.prepare(
        SCREAMING_SNAKE_CASE  ,			SCREAMING_SNAKE_CASE  ,			SCREAMING_SNAKE_CASE  ,			SCREAMING_SNAKE_CASE  ,			SCREAMING_SNAKE_CASE	)

    # We need to keep track of how many total steps we have iterated over
    __UpperCamelCase     :List[Any]			  =		0
    # We also need to keep track of the stating epoch so files are named properly
    __UpperCamelCase     :Dict			  =		0

    # Now we train the model
    __UpperCamelCase     :Any			  =		evaluate.load('''glue'''  ,			'''mrpc'''	)
    __UpperCamelCase     :Union[str, Any]			  =		0
    __UpperCamelCase     :Optional[int]			  =		{}
    for epoch in range(SCREAMING_SNAKE_CASE  ,			SCREAMING_SNAKE_CASE	):
        model.train()
        for step, batch in enumerate(SCREAMING_SNAKE_CASE	):
            __UpperCamelCase     :Optional[Any]			  =		model(**SCREAMING_SNAKE_CASE	)
            __UpperCamelCase     :Tuple			  =		outputs.loss
            __UpperCamelCase     :str			  =		loss / gradient_accumulation_steps
            accelerator.backward(SCREAMING_SNAKE_CASE	)
            if step % gradient_accumulation_steps == 0:
                optimizer.step()
                lr_scheduler.step()
                optimizer.zero_grad()

            overall_step += 1

        model.eval()
        __UpperCamelCase     :Any			  =		0
        for step, batch in enumerate(SCREAMING_SNAKE_CASE	):
            # We could avoid this line since we set the accelerator with `device_placement=True`.
            batch.to(accelerator.device	)
            with torch.no_grad():
                __UpperCamelCase     :Any			  =		model(**SCREAMING_SNAKE_CASE	)
            __UpperCamelCase     :Optional[int]			  =		outputs.logits.argmax(dim=-1	)
            # It is slightly faster to call this once, than multiple times
            __UpperCamelCase  ,       __UpperCamelCase     :List[Any]			  =		accelerator.gather(
                (predictions, batch['''labels'''])	)  # If we are in a multiprocess environment, the last batch has duplicates
            if accelerator.use_distributed:
                if step == len(SCREAMING_SNAKE_CASE	) - 1:
                    __UpperCamelCase     :List[str]			  =		predictions[: len(eval_dataloader.dataset	) - samples_seen]
                    __UpperCamelCase     :Optional[int]			  =		references[: len(eval_dataloader.dataset	) - samples_seen]
                else:
                    samples_seen += references.shape[0]
            metric.add_batch(
                predictions=SCREAMING_SNAKE_CASE  ,			references=SCREAMING_SNAKE_CASE  ,			)

        __UpperCamelCase     :Dict			  =		metric.compute()
        # Use accelerator.print to print only on the main process.
        accelerator.print(f"""epoch {epoch}:"""  ,			SCREAMING_SNAKE_CASE	)
        __UpperCamelCase     :str			  =		eval_metric['''accuracy''']

        if best_performance < eval_metric["accuracy"]:
            __UpperCamelCase     :int			  =		eval_metric['''accuracy''']

    if args.performance_lower_bound is not None:
        assert (
            args.performance_lower_bound <= best_performance
        ), f"""Best performance metric {best_performance} is lower than the lower bound {args.performance_lower_bound}"""

    accelerator.wait_for_everyone()
    if accelerator.is_main_process:
        with open(os.path.join(args.output_dir  ,			'''all_results.json'''	)  ,			'''w'''	) as f:
            json.dump(SCREAMING_SNAKE_CASE  ,			SCREAMING_SNAKE_CASE	)



def lowerCamelCase  (   ):
    '''simple docstring'''







    __UpperCamelCase     :Tuple			  =		argparse.ArgumentParser(description='''Simple example of training script tracking peak GPU memory usage.'''	)
    parser.add_argument(
        '''--model_name_or_path'''  ,			type=SCREAMING_SNAKE_CASE  ,			default='''bert-base-cased'''  ,			help='''Path to pretrained model or model identifier from huggingface.co/models.'''  ,			required=SCREAMING_SNAKE_CASE  ,			)
    parser.add_argument(
        '''--output_dir'''  ,			type=SCREAMING_SNAKE_CASE  ,			default='''.'''  ,			help='''Optional save directory where all checkpoint folders will be stored. Default is the current working directory.'''  ,			)
    parser.add_argument(
        '''--performance_lower_bound'''  ,			type=SCREAMING_SNAKE_CASE  ,			default=SCREAMING_SNAKE_CASE  ,			help='''Optional lower bound for the performance metric. If set, the training will throw error when the performance metric drops below this value.'''  ,			)
    parser.add_argument(
        '''--num_epochs'''  ,			type=SCREAMING_SNAKE_CASE  ,			default=3  ,			help='''Number of train epochs.'''  ,			)
    __UpperCamelCase     :List[str]			  =		parser.parse_args()
    __UpperCamelCase     :Tuple			  =		{'''lr''': 2e-5, '''num_epochs''': args.num_epochs, '''seed''': 42, '''batch_size''': 16}
    training_function(SCREAMING_SNAKE_CASE  ,			SCREAMING_SNAKE_CASE	)


if __name__ == "__main__":
    main()





