

from __future__ import annotations

import inspect
import unittest

import numpy as np

from transformers import DeiTConfig
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 (
						    TFDeiTForImageClassification,
						    TFDeiTForImageClassificationWithTeacher,
						    TFDeiTForMaskedImageModeling,
						    TFDeiTModel,
						)
						from transformers.models.deit.modeling_tf_deit import TF_DEIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
						from PIL import Image

						from transformers import DeiTImageProcessor






class 	__magic_name__    :





			def __init__(		self					,  __snake_case					,  __snake_case=13					,  __snake_case=30					,  __snake_case=2					,  __snake_case=3					,  __snake_case=True					,  __snake_case=True					,  __snake_case=32					,  __snake_case=2					,  __snake_case=4					,  __snake_case=37					,  __snake_case="gelu"					,  __snake_case=0.1					,  __snake_case=0.1					,  __snake_case=10					,  __snake_case=0.02					,  __snake_case=3					,  __snake_case=None					,  __snake_case=2					,  )   -> Optional[int]:




						'''simple docstring'''
						__a        =parent
						__a        =batch_size
						__a        =image_size
						__a        =patch_size
						__a        =num_channels
						__a        =is_training
						__a        =use_labels
						__a        =hidden_size
						__a        =num_hidden_layers
						__a        =num_attention_heads
						__a        =intermediate_size
						__a        =hidden_act
						__a        =hidden_dropout_prob
						__a        =attention_probs_dropout_prob
						__a        =type_sequence_label_size
						__a        =initializer_range
						__a        =scope
						__a        =encoder_stride

						# in DeiT, the seq length equals the number of patches + 2 (we add 2 for the [CLS] and distilation tokens)
						__a        =(image_size // patch_size) ** 2
						__a        =num_patches + 2





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




						'''simple docstring'''
						__a        =floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]					)

						__a        =None
						if self.use_labels:
									__a        =ids_tensor([self.batch_size]					,  self.type_sequence_label_size					)

						__a        =self.get_config()

						return config, pixel_values, labels





			def 		__magic_name__    (		self					)   -> str:




						'''simple docstring'''
						return DeiTConfig(
						    image_size=self.image_size					,  patch_size=self.patch_size					,  num_channels=self.num_channels					,  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					,  is_decoder=__snake_case					,  initializer_range=self.initializer_range					,  encoder_stride=self.encoder_stride					,  )





			def 		__magic_name__    (		self					,  __snake_case					,  __snake_case					,  __snake_case					)   -> Optional[int]:




						'''simple docstring'''
						__a        =TFDeiTModel(config=__snake_case					)
						__a        =model(__snake_case					)
						self.parent.assertEqual(result.last_hidden_state.shape					,  (self.batch_size, self.seq_length, self.hidden_size)					)





			def 		__magic_name__    (		self					,  __snake_case					,  __snake_case					,  __snake_case					)   -> Tuple:




						'''simple docstring'''
						__a        =TFDeiTForMaskedImageModeling(config=__snake_case					)
						__a        =model(__snake_case					)
						self.parent.assertEqual(
						    result.reconstruction.shape					,  (self.batch_size, self.num_channels, self.image_size, self.image_size)					)

						# test greyscale images
						__a        =1
						__a        =TFDeiTForMaskedImageModeling(__snake_case					)

						__a        =floats_tensor([self.batch_size, 1, self.image_size, self.image_size]					)
						__a        =model(__snake_case					)
						self.parent.assertEqual(result.reconstruction.shape					,  (self.batch_size, 1, self.image_size, self.image_size)					)





			def 		__magic_name__    (		self					,  __snake_case					,  __snake_case					,  __snake_case					)   -> Optional[int]:




						'''simple docstring'''
						__a        =self.type_sequence_label_size
						__a        =TFDeiTForImageClassification(__snake_case					)
						__a        =model(__snake_case					,  labels=__snake_case					)
						self.parent.assertEqual(result.logits.shape					,  (self.batch_size, self.type_sequence_label_size)					)

						# test greyscale images
						__a        =1
						__a        =TFDeiTForImageClassification(__snake_case					)

						__a        =floats_tensor([self.batch_size, 1, self.image_size, self.image_size]					)
						__a        =model(__snake_case					,  labels=__snake_case					)
						self.parent.assertEqual(result.logits.shape					,  (self.batch_size, self.type_sequence_label_size)					)





			def 		__magic_name__    (		self					)   -> int:




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






@require_tf
class 	__magic_name__    (			lowerCAmelCase_  ,    lowerCAmelCase_  ,    unittest.TestCase      ):
			SCREAMING_SNAKE_CASE							   =						(
			    (
			        TFDeiTModel,
			        TFDeiTForImageClassification,
			        TFDeiTForImageClassificationWithTeacher,
			        TFDeiTForMaskedImageModeling,
			    )
			    if is_tf_available()
			    else ()
			)
			SCREAMING_SNAKE_CASE							   =						(
			    {
			        'feature-extraction': TFDeiTModel,
			        'image-classification': (TFDeiTForImageClassification, TFDeiTForImageClassificationWithTeacher),
			    }
			    if is_tf_available()
			    else {}
			)

			SCREAMING_SNAKE_CASE							   =						False
			SCREAMING_SNAKE_CASE							   =						False
			SCREAMING_SNAKE_CASE							   =						False
			SCREAMING_SNAKE_CASE							   =						False





			def 		__magic_name__    (		self					)   -> Optional[int]:




						'''simple docstring'''
						__a        =TFDeiTModelTester(self					)
						__a        =ConfigTester(self					,  config_class=__snake_case					,  has_text_modality=__snake_case					,  hidden_size=37					)





			def 		__magic_name__    (		self					)   -> int:




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





			@unittest.skip(reason='DeiT does not use inputs_embeds'					)
			def 		__magic_name__    (		self					)   -> Union[str, Any]:




						'''simple docstring'''
						pass





			def 		__magic_name__    (		self					)   -> int:




						'''simple docstring'''
						__a      ,    __a        =self.model_tester.prepare_config_and_inputs_for_common()

						for model_class in self.all_model_classes:
									__a        =model_class(__snake_case					)
									self.assertIsInstance(model.get_input_embeddings()					,  (tf.keras.layers.Layer)					)
									__a        =model.get_output_embeddings()
									self.assertTrue(x is None or isinstance(__snake_case					,  tf.keras.layers.Dense					)					)





			def 		__magic_name__    (		self					)   -> List[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(__snake_case					)
									__a        =inspect.signature(model.call					)
									# signature.parameters is an OrderedDict => so arg_names order is deterministic
									__a        =[*signature.parameters.keys()]

									__a        =['pixel_values']
									self.assertListEqual(arg_names[:1]					,  __snake_case					)





			def 		__magic_name__    (		self					)   -> Any:




						'''simple docstring'''
						__a        =self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*__snake_case					)





			def 		__magic_name__    (		self					)   -> Tuple:




						'''simple docstring'''
						__a        =self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_masked_image_modeling(*__snake_case					)





			def 		__magic_name__    (		self					)   -> Tuple:




						'''simple docstring'''
						__a        =self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_image_classification(*__snake_case					)





			def 		__magic_name__    (		self					,  __snake_case					,  __snake_case					,  __snake_case=False					)   -> Any:




						'''simple docstring'''
						__a        =super()._prepare_for_class(__snake_case					,  __snake_case					,  return_labels=__snake_case					)

						if return_labels:
									if "labels" in inputs_dict and "labels" not in inspect.signature(model_class.call					).parameters:
												del inputs_dict["labels"]

						return inputs_dict





			@slow
			def 		__magic_name__    (		self					)   -> List[str]:




						'''simple docstring'''
						for model_name in TF_DEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									__a        =TFDeiTModel.from_pretrained(__snake_case					)
									self.assertIsNotNone(__snake_case					)






def        UpperCamelCase_(							):




			"""simple docstring"""






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






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





			@cached_property
			def 		__magic_name__    (		self					)   -> Dict:




						'''simple docstring'''
						return (
						    DeiTImageProcessor.from_pretrained('facebook/deit-base-distilled-patch16-224'					)
						    if is_vision_available()
						    else None
						)





			@slow
			def 		__magic_name__    (		self					)   -> Any:




						'''simple docstring'''
						__a        =TFDeiTForImageClassificationWithTeacher.from_pretrained('facebook/deit-base-distilled-patch16-224'					)

						__a        =self.default_image_processor
						__a        =prepare_img()
						__a        =image_processor(images=__snake_case					,  return_tensors='tf'					)

						# forward pass
						__a        =model(**__snake_case					)

						# verify the logits
						__a        =tf.TensorShape((1, 1000)					)
						self.assertEqual(outputs.logits.shape					,  __snake_case					)

						__a        =tf.constant([-1.0266, 0.1912, -1.2861]					)

						self.assertTrue(np.allclose(outputs.logits[0, :3]					,  __snake_case					,  atol=1e-4					)					)









import argparse
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.local_sgd import LocalSGD


########################################################################
# This is a fully working simple example to use Accelerate
# with LocalSGD, which is a method to synchronize model
# parameters every K batches. It is different, but complementary
# to gradient accumulation.
#
# This example trains a Bert base model on GLUE MRPC
# in any of the following settings (with the same script):
#   - single CPU or single GPU
#   - multi GPUS (using PyTorch distributed mode)
#   - (multi) TPUs
#   - fp16 (mixed-precision) or fp32 (normal precision)
#
# To run it in each of these various modes, follow the instructions
# in the readme for examples:
# https://github.com/huggingface/accelerate/tree/main/examples
#
########################################################################


_lowerCAmelCase					:			Union[str, Any]  				=			16
_lowerCAmelCase					:			List[str]  				=			32






def        UpperCamelCase_(							_snake_case   :  Accelerator       ,     _snake_case   :  int = 16	):




			"""simple docstring"""






			__a        =AutoTokenizer.from_pretrained('bert-base-cased'	)
			__a        =load_dataset('glue'       ,     'mrpc'	)

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

			# Apply the method we just defined to all the examples in all the splits of the dataset
			# starting with the main process first:
			with accelerator.main_process_first():
						__a        =datasets.map(
						    _snake_case       ,     batched=_snake_case       ,     remove_columns=['idx', 'sentence1', 'sentence2']       ,     )

			# We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
			# transformers library
			__a        =tokenized_datasets.rename_column('label'       ,     'labels'	)

			def collate_fn(_snake_case   :  List[Any]	):
						# On TPU it's best to pad everything to the same length or training will be very slow.
						__a        =128 if accelerator.distributed_type == DistributedType.TPU else None
						# When using mixed precision we want round multiples of 8/16
						if accelerator.mixed_precision == "fp8":
									__a        =16
						elif accelerator.mixed_precision != "no":
									__a        =8
						else:
									__a        =None

						return tokenizer.pad(
						    _snake_case       ,     padding='longest'       ,     max_length=_snake_case       ,     pad_to_multiple_of=_snake_case       ,     return_tensors='pt'       ,     )

			# Instantiate dataloaders.
			__a        =DataLoader(
			    tokenized_datasets['train']       ,     shuffle=_snake_case       ,     collate_fn=_snake_case       ,     batch_size=_snake_case	)
			__a        =DataLoader(
			    tokenized_datasets['validation']       ,     shuffle=_snake_case       ,     collate_fn=_snake_case       ,     batch_size=_snake_case	)

			return train_dataloader, eval_dataloader


# For testing only
if os.environ.get("TESTING_MOCKED_DATALOADERS", None) == "1":
						from accelerate.test_utils.training import mocked_dataloaders

						_lowerCAmelCase					:			List[Any]  				=			mocked_dataloaders  # noqa: F811






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




			"""simple docstring"""






			if os.environ.get('TESTING_MOCKED_DATALOADERS'       ,     _snake_case	) == "1":
						__a        =2
			# New Code #
			__a        =int(args.gradient_accumulation_steps	)
			__a        =int(args.local_sgd_steps	)
			# Initialize accelerator
			__a        =Accelerator(
			    cpu=args.cpu       ,     mixed_precision=args.mixed_precision       ,     gradient_accumulation_steps=_snake_case	)
			if accelerator.distributed_type not in [DistributedType.NO, DistributedType.MULTI_CPU, DistributedType.MULTI_GPU]:
						raise NotImplementedError('LocalSGD is supported only for CPUs and GPUs (no DeepSpeed or MegatronLM)'	)
			# Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
			__a        =config['lr']
			__a        =int(config['num_epochs']	)
			__a        =int(config['seed']	)
			__a        =int(config['batch_size']	)

			__a        =evaluate.load('glue'       ,     'mrpc'	)

			set_seed(_snake_case	)
			__a      ,    __a        =get_dataloaders(_snake_case       ,     _snake_case	)
			# Instantiate the model (we build the model here so that the seed also control new weights initialization)
			__a        =AutoModelForSequenceClassification.from_pretrained('bert-base-cased'       ,     return_dict=_snake_case	)

			# We could avoid this line since the accelerator is set with `device_placement=True` (default value).
			# Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
			# creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
			__a        =model.to(accelerator.device	)

			# Instantiate optimizer
			__a        =AdamW(params=model.parameters()       ,     lr=_snake_case	)

			# Instantiate scheduler
			__a        =get_linear_schedule_with_warmup(
			    optimizer=_snake_case       ,     num_warmup_steps=100       ,     num_training_steps=(len(_snake_case	) * num_epochs)       ,     )

			# 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.
			__a      ,    __a      ,    __a      ,    __a      ,    __a        =accelerator.prepare(
			    _snake_case       ,     _snake_case       ,     _snake_case       ,     _snake_case       ,     _snake_case	)

			# Now we train the model
			for epoch in range(_snake_case	):
						model.train()
						with LocalSGD(
						    accelerator=_snake_case       ,     model=_snake_case       ,     local_sgd_steps=_snake_case       ,     enabled=local_sgd_steps is not None	) as local_sgd:
									for step, batch in enumerate(_snake_case	):
												# We could avoid this line since we set the accelerator with `device_placement=True`.
												batch.to(accelerator.device	)
												# New code #
												# We use the new `accumulate` context manager to perform gradient accumulation
												# We also currently do not support TPUs nor advise it as bugs were found on the XLA side when running our tests.
												with accelerator.accumulate(_snake_case	):
															__a        =model(**_snake_case	)
															__a        =output.loss
															accelerator.backward(_snake_case	)
															optimizer.step()
															lr_scheduler.step()
															optimizer.zero_grad()
															# LocalSGD-specific line
															local_sgd.step()

						model.eval()
						for step, batch in enumerate(_snake_case	):
									# We could avoid this line since we set the accelerator with `device_placement=True`.
									batch.to(accelerator.device	)
									with torch.no_grad():
												__a        =model(**_snake_case	)
									__a        =outputs.logits.argmax(dim=-1	)
									__a      ,    __a        =accelerator.gather_for_metrics((predictions, batch['labels'])	)
									metric.add_batch(
									    predictions=_snake_case       ,     references=_snake_case       ,     )

						__a        =metric.compute()
						# Use accelerator.print to print only on the main process.
						accelerator.print(F'epoch {epoch}:'       ,     _snake_case	)






def        UpperCamelCase_(							):




			"""simple docstring"""






			__a        =argparse.ArgumentParser(description='Simple example of training script.'	)
			parser.add_argument(
			    '--mixed_precision'       ,     type=_snake_case       ,     default=_snake_case       ,     choices=['no', 'fp16', 'bf16', 'fp8']       ,     help='Whether to use mixed precision. Choose'
			    'between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10.'
			    'and an Nvidia Ampere GPU.'       ,     )
			# New Code #
			parser.add_argument(
			    '--gradient_accumulation_steps'       ,     type=_snake_case       ,     default=1       ,     help='The number of minibatches to be ran before gradients are accumulated.'       ,     )
			parser.add_argument(
			    '--local_sgd_steps'       ,     type=_snake_case       ,     default=8       ,     help='Number of local SGD steps or None to disable local SGD'	)
			parser.add_argument('--cpu'       ,     action='store_true'       ,     help='If passed, will train on the CPU.'	)
			__a        =parser.parse_args()
			__a        ={'lr': 2e-5, 'num_epochs': 3, 'seed': 42, 'batch_size': 16}
			training_function(_snake_case       ,     _snake_case	)


if __name__ == "__main__":
						main()






