

'''simple docstring'''






import inspect
import unittest

from transformers import ViTConfig
from transformers.testing_utils import (
    require_accelerate,
    require_torch,
    require_torch_gpu,
    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 torch import nn

 from transformers import ViTForImageClassification, ViTForMaskedImageModeling, ViTModel
 from transformers.models.vit.modeling_vit import VIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
 from PIL import Image

 from transformers import ViTImageProcessor






class 			lowercase_	:




  """simple docstring"""




  def __init__(      self   :     Optional[int]					,lowercase__   :     Optional[Any]					,lowercase__   :     Any=1_3					,lowercase__   :     List[Any]=3_0					,lowercase__   :     int=2					,lowercase__   :     str=3					,lowercase__   :     Optional[Any]=True					,lowercase__   :     str=True					,lowercase__   :     Optional[Any]=3_2					,lowercase__   :     Optional[int]=5					,lowercase__   :     Any=4					,lowercase__   :     List[str]=3_7					,lowercase__   :     Tuple="gelu"					,lowercase__   :     Dict=0.1					,lowercase__   :     Dict=0.1					,lowercase__   :     List[str]=1_0					,lowercase__   :     Optional[int]=0.0_2					,lowercase__   :     int=None					,lowercase__   :     Optional[Any]=2					,):
   __lowercase							  = parent
   __lowercase							  = batch_size
   __lowercase							  = image_size
   __lowercase							  = patch_size
   __lowercase							  = num_channels
   __lowercase							  = is_training
   __lowercase							  = use_labels
   __lowercase							  = hidden_size
   __lowercase							  = num_hidden_layers
   __lowercase							  = num_attention_heads
   __lowercase							  = intermediate_size
   __lowercase							  = hidden_act
   __lowercase							  = hidden_dropout_prob
   __lowercase							  = attention_probs_dropout_prob
   __lowercase							  = type_sequence_label_size
   __lowercase							  = initializer_range
   __lowercase							  = scope
   __lowercase							  = encoder_stride

   # in ViT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
   __lowercase							  = (image_size // patch_size) ** 2
   __lowercase							  = num_patches + 1






  def 					SCREAMING_SNAKE_CASE      (      self   :     List[Any]			):
   __lowercase							  = floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]			)

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

   __lowercase							  = self.get_config()

   return config, pixel_values, labels






  def 					SCREAMING_SNAKE_CASE      (      self   :     Union[str, Any]			):
   return ViTConfig(
       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=lowercase__					,initializer_range=self.initializer_range					,encoder_stride=self.encoder_stride					,)






  def 					SCREAMING_SNAKE_CASE      (      self   :     Dict					,lowercase__   :     Optional[Any]					,lowercase__   :     Dict					,lowercase__   :     List[Any]			):
   __lowercase							  = ViTModel(config=lowercase__			)
   model.to(lowercase__			)
   model.eval()
   __lowercase							  = model(lowercase__			)
   self.parent.assertEqual(result.last_hidden_state.shape					,(self.batch_size, self.seq_length, self.hidden_size)			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     str					,lowercase__   :     Optional[int]					,lowercase__   :     Any					,lowercase__   :     Dict			):
   __lowercase							  = ViTForMaskedImageModeling(config=lowercase__			)
   model.to(lowercase__			)
   model.eval()
   __lowercase							  = model(lowercase__			)
   self.parent.assertEqual(
       result.reconstruction.shape					,(self.batch_size, self.num_channels, self.image_size, self.image_size)			)

   # test greyscale images
   __lowercase							  = 1
   __lowercase							  = ViTForMaskedImageModeling(lowercase__			)
   model.to(lowercase__			)
   model.eval()

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






  def 					SCREAMING_SNAKE_CASE      (      self   :     Optional[int]					,lowercase__   :     List[Any]					,lowercase__   :     Union[str, Any]					,lowercase__   :     List[Any]			):
   __lowercase							  = self.type_sequence_label_size
   __lowercase							  = ViTForImageClassification(lowercase__			)
   model.to(lowercase__			)
   model.eval()
   __lowercase							  = model(lowercase__					,labels=lowercase__			)
   self.parent.assertEqual(result.logits.shape					,(self.batch_size, self.type_sequence_label_size)			)

   # test greyscale images
   __lowercase							  = 1
   __lowercase							  = ViTForImageClassification(lowercase__			)
   model.to(lowercase__			)
   model.eval()

   __lowercase							  = floats_tensor([self.batch_size, 1, self.image_size, self.image_size]			)
   __lowercase							  = model(lowercase__			)
   self.parent.assertEqual(result.logits.shape					,(self.batch_size, self.type_sequence_label_size)			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     str			):
   __lowercase							  = self.prepare_config_and_inputs()
   (
       (
       __lowercase
   )  ,	(
       __lowercase
   )  ,	(
       __lowercase
   )  ,	
   )							  = config_and_inputs
   __lowercase							  = {'''pixel_values''': pixel_values}
   return config, inputs_dict






@require_torch
class 			lowercase_	(lowerCamelCase__				,       lowerCamelCase__				,       unittest.TestCase      ):




  """simple docstring"""




  SCREAMING_SNAKE_CASE    :  Dict							=			(
      (
          ViTModel,
          ViTForImageClassification,
          ViTForMaskedImageModeling,
      )
      if is_torch_available()
      else ()
  )
  SCREAMING_SNAKE_CASE    :  Dict							=			(
      {'feature-extraction': ViTModel, 'image-classification': ViTForImageClassification}
      if is_torch_available()
      else {}
  )
  SCREAMING_SNAKE_CASE    :  Union[str, Any]							=			True

  SCREAMING_SNAKE_CASE    :  Optional[Any]							=			False
  SCREAMING_SNAKE_CASE    :  Tuple							=			False
  SCREAMING_SNAKE_CASE    :  Dict							=			False






  def 					SCREAMING_SNAKE_CASE      (      self   :     int			):
   __lowercase							  = ViTModelTester(self			)
   __lowercase							  = ConfigTester(self					,config_class=lowercase__					,has_text_modality=lowercase__					,hidden_size=3_7			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     Union[str, Any]			):
   self.config_tester.run_common_tests()






  @unittest.skip(reason='''ViT does not use inputs_embeds'''			)
  def 					SCREAMING_SNAKE_CASE      (      self   :     Dict			):
   pass






  def 					SCREAMING_SNAKE_CASE      (      self   :     Dict			):
   __lowercase  ,	__lowercase							  = self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
    __lowercase							  = model_class(lowercase__			)
    self.assertIsInstance(model.get_input_embeddings()					,(nn.Module)			)
    __lowercase							  = model.get_output_embeddings()
    self.assertTrue(x is None or isinstance(lowercase__					,nn.Linear			)			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     Any			):
   __lowercase  ,	__lowercase							  = self.model_tester.prepare_config_and_inputs_for_common()

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

    __lowercase							  = ['''pixel_values''']
    self.assertListEqual(arg_names[:1]					,lowercase__			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     Optional[int]			):
   __lowercase							  = self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_model(*lowercase__			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     List[str]			):
   __lowercase							  = self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_masked_image_modeling(*lowercase__			)






  def 					SCREAMING_SNAKE_CASE      (      self   :     Tuple			):
   __lowercase							  = self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_image_classification(*lowercase__			)






  @slow
  def 					SCREAMING_SNAKE_CASE      (      self   :     List[str]			):
   for model_name in VIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
    __lowercase							  = ViTModel.from_pretrained(lowercase__			)
    self.assertIsNotNone(lowercase__			)


def 		_A     (    ):






 """simple docstring"""



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






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




  """simple docstring"""




  @cached_property
  def 					SCREAMING_SNAKE_CASE      (      self   :     Union[str, Any]			):
   return ViTImageProcessor.from_pretrained('''google/vit-base-patch16-224'''			) if is_vision_available() else None






  @slow
  def 					SCREAMING_SNAKE_CASE      (      self   :     Union[str, Any]			):
   __lowercase							  = ViTForImageClassification.from_pretrained('''google/vit-base-patch16-224'''			).to(lowercase__			)

   __lowercase							  = self.default_image_processor
   __lowercase							  = prepare_img()
   __lowercase							  = image_processor(images=lowercase__					,return_tensors='''pt'''			).to(lowercase__			)

   # forward pass
   with torch.no_grad():
    __lowercase							  = model(**lowercase__			)

   # verify the logits
   __lowercase							  = torch.Size((1, 1_0_0_0)			)
   self.assertEqual(outputs.logits.shape					,lowercase__			)

   __lowercase							  = torch.tensor([-0.2_7_4_4, 0.8_2_1_5, -0.0_8_3_6]			).to(lowercase__			)

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






  @slow
  def 					SCREAMING_SNAKE_CASE      (      self   :     List[str]			):
   # ViT models have an `interpolate_pos_encoding` argument in their forward method,
   # allowing to interpolate the pre-trained position embeddings in order to use
   # the model on higher resolutions. The DINO model by Facebook AI leverages this
   # to visualize self-attention on higher resolution images.
   __lowercase							  = ViTModel.from_pretrained('''facebook/dino-vits8'''			).to(lowercase__			)

   __lowercase							  = ViTImageProcessor.from_pretrained('''facebook/dino-vits8'''					,size=4_8_0			)
   __lowercase							  = prepare_img()
   __lowercase							  = image_processor(images=lowercase__					,return_tensors='''pt'''			)
   __lowercase							  = inputs.pixel_values.to(lowercase__			)

   # forward pass
   with torch.no_grad():
    __lowercase							  = model(lowercase__					,interpolate_pos_encoding=lowercase__			)

   # verify the logits
   __lowercase							  = torch.Size((1, 3_6_0_1, 3_8_4)			)
   self.assertEqual(outputs.last_hidden_state.shape					,lowercase__			)

   __lowercase							  = torch.tensor(
       [[4.2_3_4_0, 4.3_9_0_6, -6.6_6_9_2], [4.5_4_6_3, 1.8_9_2_8, -6.7_2_5_7], [4.4_4_2_9, 0.8_4_9_6, -5.8_5_8_5]]			).to(lowercase__			)

   self.assertTrue(torch.allclose(outputs.last_hidden_state[0, :3, :3]					,lowercase__					,atol=1e-4			)			)






  @slow
  @require_accelerate
  @require_torch_gpu
  def 					SCREAMING_SNAKE_CASE      (      self   :     Tuple			):
   __lowercase							  = ViTModel.from_pretrained('''facebook/dino-vits8'''					,torch_dtype=torch.floataa					,device_map='''auto'''			)
   __lowercase							  = self.default_image_processor

   __lowercase							  = prepare_img()
   __lowercase							  = image_processor(images=lowercase__					,return_tensors='''pt'''			)
   __lowercase							  = inputs.pixel_values.to(lowercase__			)

   # forward pass to make sure inference works in fp16
   with torch.no_grad():
    __lowercase							  = model(lowercase__			)




'''simple docstring'''






# Copyright 2022 The HuggingFace 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 argparse
import os
import subprocess

from packaging.version import Version, parse

from accelerate.commands.config.config_args import default_config_file, load_config_from_file


lowerCAmelCase__											=	'''Run commands across TPU VMs for initial setup before running `accelerate launch`.'''


def 		_A     (    A__=None      ):






 """simple docstring"""



 if subparsers is not None:
  __lowercase							  = subparsers.add_parser('''tpu-config'''     ,				description=_description      )
 else:
  __lowercase							  = argparse.ArgumentParser('''Accelerate tpu-config command'''     ,				description=_description      )
 # Core arguments
 __lowercase							  = parser.add_argument_group(
     '''Config Arguments'''     ,				'''Arguments that can be configured through `accelerate config`.'''      )
 config_args.add_argument(
     '''--config_file'''     ,				type=A__     ,				default=A__     ,				help='''Path to the config file to use for accelerate.'''     ,				)
 config_args.add_argument(
     '''--tpu_name'''     ,				default=A__     ,				help='''The name of the TPU to use. If not specified, will use the TPU specified in the config file.'''     ,				)
 config_args.add_argument(
     '''--tpu_zone'''     ,				default=A__     ,				help='''The zone of the TPU to use. If not specified, will use the zone specified in the config file.'''     ,				)
 __lowercase							  = parser.add_argument_group('''TPU Arguments'''     ,				'''Arguments for options ran inside the TPU.'''      )
 pod_args.add_argument(
     '''--use_alpha'''     ,				action='''store_true'''     ,				help='''Whether to use `gcloud alpha` when running the TPU training script instead of `gcloud`.'''     ,				)
 pod_args.add_argument(
     '''--command_file'''     ,				default=A__     ,				help='''The path to the file containing the commands to run on the pod on startup.'''     ,				)
 pod_args.add_argument(
     '''--command'''     ,				action='''append'''     ,				nargs='''+'''     ,				help='''A command to run on the pod. Can be passed multiple times.'''     ,				)
 pod_args.add_argument(
     '''--install_accelerate'''     ,				action='''store_true'''     ,				help='''Whether to install accelerate on the pod. Defaults to False.'''     ,				)
 pod_args.add_argument(
     '''--accelerate_version'''     ,				default='''latest'''     ,				help='''The version of accelerate to install on the pod. If not specified, will use the latest pypi version. Specify \'dev\' to install from GitHub.'''     ,				)
 pod_args.add_argument(
     '''--debug'''     ,				action='''store_true'''     ,				help='''If set, will print the command that would be run instead of running it.'''      )

 if subparsers is not None:
  parser.set_defaults(func=A__      )
 return parser


def 		_A     (    A__      ):






 """simple docstring"""



 __lowercase							  = None

 # Get the default from the config file if it exists.
 if args.config_file is not None or os.path.isfile(A__      ):
  __lowercase							  = load_config_from_file(args.config_file      )
  if not args.command_file and defaults.command_file is not None and not args.command:
   __lowercase							  = defaults.command_file
  if not args.command and defaults.commands is not None:
   __lowercase							  = defaults.commands
  if not args.tpu_name:
   __lowercase							  = defaults.tpu_name
  if not args.tpu_zone:
   __lowercase							  = defaults.tpu_zone
 if args.accelerate_version == "dev":
  __lowercase							  = '''git+https://github.com/huggingface/accelerate.git'''
 elif args.accelerate_version == "latest":
  __lowercase							  = '''accelerate -U'''
 elif isinstance(parse(args.accelerate_version      )     ,				A__      ):
  __lowercase							  = F"accelerate=={args.accelerate_version}"

 if not args.command_file and not args.command:
  raise ValueError('''You must specify either a command file or a command to run on the pod.'''      )

 if args.command_file:
  with open(args.command_file     ,				'''r'''      ) as f:
   __lowercase							  = [f.read().splitlines()]

    # To turn list of lists into list of strings
 if isinstance(args.command[0]     ,				A__      ):
  __lowercase							  = [line for cmd in args.command for line in cmd]
 # Default to the shared folder and install accelerate
 __lowercase							  = ['''cd /usr/share''']
 if args.install_accelerate:
  new_cmd += [F"pip install {args.accelerate_version}"]
 new_cmd += args.command
 __lowercase							  = '''; '''.join(A__      )

 # Then send it to gcloud
 # Eventually try to use google-api-core to do this instead of subprocess
 __lowercase							  = ['''gcloud''']
 if args.use_alpha:
  cmd += ["alpha"]
 cmd += [
     "compute",
     "tpus",
     "tpu-vm",
     "ssh",
     args.tpu_name,
     "--zone",
     args.tpu_zone,
     "--command",
     args.command,
     "--worker",
     "all",
 ]
 if args.debug:
  print(F"Running {' '.join(A__      )}"      )
  return
 subprocess.run(A__      )
 print('''Successfully setup pod.'''      )






def 		_A     (    ):






 """simple docstring"""



 __lowercase							  = tpu_command_parser()
 __lowercase							  = parser.parse_args()

 tpu_command_launcher(A__      )

