





"""simple docstring"""





import os
import posixpath
import uuid
from dataclasses import dataclass
from typing import TYPE_CHECKING, Iterable, List, Optional, Tuple, Union

import numpy as np
import pyarrow as pa

import datasets
from datasets.arrow_writer import ArrowWriter, ParquetWriter
from datasets.config import MAX_SHARD_SIZE
from datasets.filesystems import (
    is_remote_filesystem,
    rename,
)
from datasets.iterable_dataset import _BaseExamplesIterable
from datasets.utils.py_utils import convert_file_size_to_int


lowerCAmelCase__   	=	datasets.utils.logging.get_logger(__name__)

if TYPE_CHECKING:
						import pyspark





@dataclass
class 				_lowerCamelCase					(				datasets.BuilderConfig       ):
						UpperCAmelCase_       =	None



def 					a__    (    _SCREAMING_SNAKE_CASE     ,		_SCREAMING_SNAKE_CASE     ,		):




		"""simple docstring"""


		import pyspark

		def generate_fn():
				UpperCamelCase      =		df.select("*"     ,		pyspark.sql.functions.spark_partition_id().alias("part_id"			)			)
				for partition_id in partition_order:
						UpperCamelCase      =		df_with_partition_id.select("*"			).where(F"part_id = {partition_id}"			).drop("part_id"			)
						UpperCamelCase      =		partition_df.collect()
						UpperCamelCase      =		0
						for row in rows:
								yield F"{partition_id}_{row_id}", row.asDict()
								row_id += 1

		return generate_fn





class 				_lowerCamelCase					(				_BaseExamplesIterable       ):


						def __init__(self      ,							__a      ,							__a=None      ,							)       ->      List[Any]:
								UpperCamelCase      =		df
								UpperCamelCase      =		partition_order or range(self.df.rdd.getNumPartitions()						)
								UpperCamelCase      =		_generate_iterable_examples(self.df      ,							self.partition_order						)


						def __iter__(self						)       ->      int:
								yield from self.generate_examples_fn()


						def 			snake_case_      (self      ,							__a						)       ->      Dict:
								UpperCamelCase      =		list(range(self.df.rdd.getNumPartitions()						)						)
								generator.shuffle(lowercase_						)
								return SparkExamplesIterable(self.df      ,							partition_order=lowercase_						)


						def 			snake_case_      (self      ,							__a      ,							__a						)       ->      List[Any]:
								UpperCamelCase      =		self.split_shard_indices_by_worker(lowercase_      ,							lowercase_						)
								return SparkExamplesIterable(self.df      ,							partition_order=lowercase_						)




						@property
						def 			snake_case_      (self						)       ->      Any:
								return len(self.partition_order						)





class 				_lowerCamelCase					(				datasets.DatasetBuilder       ):
						UpperCAmelCase_       =	SparkConfig


						def __init__(self      ,							__a      ,							__a = None      ,							__a = None      ,							**__a      ,							)       ->      Tuple:
								import pyspark

								UpperCamelCase      =		pyspark.sql.SparkSession.builder.getOrCreate()
								UpperCamelCase      =		df
								UpperCamelCase      =		working_dir

								super().__init__(
								    cache_dir=lowercase_      ,							config_name=str(self.df.semanticHash()						)      ,							**lowercase_      ,							)


						def 			snake_case_      (self						)       ->      Tuple:
								# Returns the path of the created file.
								def create_cache_and_write_probe(__a						):
										# makedirs with exist_ok will recursively create the directory. It will not throw an error if directories
										# already exist.
										os.makedirs(self._cache_dir      ,							exist_ok=lowercase_						)
										UpperCamelCase      =		os.path.join(self._cache_dir      ,							"fs_test" + uuid.uuida().hex						)
										# Opening the file in append mode will create a new file unless it already exists, in which case it will not
										# change the file contents.
										open(lowercase_      ,							"a"						)
										return [probe_file]

								if self._spark.conf.get("spark.master"      ,							""						).startswith("local"						):
										return

								# If the cluster is multi-node, make sure that the user provided a cache_dir and that it is on an NFS
								# accessible to the driver.
								# TODO: Stream batches to the driver using ArrowCollectSerializer instead of throwing an error.
								if self._cache_dir:
										UpperCamelCase      =		(
										    self._spark.sparkContext.parallelize(range(1						)      ,							1						).mapPartitions(lowercase_						).collect()
										)
										if os.path.isfile(probe[0]						):
												return

								raise ValueError(
								    "When using Dataset.from_spark on a multi-node cluster, the driver and all workers should be able to access cache_dir"						)


						def 			snake_case_      (self						)       ->      List[Any]:
								return datasets.DatasetInfo(features=self.config.features						)


						def 			snake_case_      (self      ,							__a						)       ->      Any:
								return [datasets.SplitGenerator(name=datasets.Split.TRAIN						)]


						def 			snake_case_      (self      ,							__a						)       ->      Tuple:
								import pyspark

								def get_arrow_batch_size(__a						):
										for batch in it:
												yield pa.RecordBatch.from_pydict({"batch_bytes": [batch.nbytes]}						)

								UpperCamelCase      =		self.df.count()
								UpperCamelCase      =		df_num_rows if df_num_rows <= 1_00 else 1_00
								# Approximate the size of each row (in Arrow format) by averaging over a max-100-row sample.
								UpperCamelCase      =		(
								    self.df.limit(lowercase_						)
								    .repartition(1						)
								    .mapInArrow(lowercase_      ,							"batch_bytes: long"						)
								    .agg(pyspark.sql.functions.sum("batch_bytes"						).alias("sample_bytes"						)						)
								    .collect()[0]
								    .sample_bytes
								    / sample_num_rows
								)
								UpperCamelCase      =		approx_bytes_per_row * df_num_rows
								if approx_total_size > max_shard_size:
										# Make sure there is at least one row per partition.
										UpperCamelCase      =		min(lowercase_      ,							int(approx_total_size / max_shard_size						)						)
										UpperCamelCase      =		self.df.repartition(lowercase_						)


						def 			snake_case_      (self      ,							__a      ,							__a      ,							__a      ,							)       ->      Union[str, Any]:
								import pyspark

								UpperCamelCase      =		ParquetWriter if file_format == "parquet" else ArrowWriter
								UpperCamelCase      =		os.path.join(self._working_dir      ,							os.path.basename(lowercase_						)						) if self._working_dir else fpath
								UpperCamelCase      =		file_format == "parquet"

								# Define these so that we don't reference self in write_arrow, which will result in a pickling error due to
								# pickling the SparkContext.
								UpperCamelCase      =		self.config.features
								UpperCamelCase      =		self._writer_batch_size
								UpperCamelCase      =		self._fs.storage_options

								def write_arrow(__a						):
										# Within the same SparkContext, no two task attempts will share the same attempt ID.
										UpperCamelCase      =		pyspark.TaskContext().taskAttemptId()
										UpperCamelCase      =		next(lowercase_      ,							lowercase_						)
										if first_batch is None:
												# Some partitions might not receive any data.
												return pa.RecordBatch.from_arrays(
												    [[task_id], [0], [0]]      ,							names=["task_id", "num_examples", "num_bytes"]      ,							)
										UpperCamelCase      =		0
										UpperCamelCase      =		writer_class(
										    features=lowercase_      ,							path=working_fpath.replace("SSSSS"      ,							F"{shard_id:05d}"						).replace("TTTTT"      ,							F"{task_id:05d}"						)      ,							writer_batch_size=lowercase_      ,							storage_options=lowercase_      ,							embed_local_files=lowercase_      ,							)
										UpperCamelCase      =		pa.Table.from_batches([first_batch]						)
										writer.write_table(lowercase_						)
										for batch in it:
												if max_shard_size is not None and writer._num_bytes >= max_shard_size:
														UpperCamelCase      =		writer.finalize()
														writer.close()
														yield pa.RecordBatch.from_arrays(
														    [[task_id], [num_examples], [num_bytes]]      ,							names=["task_id", "num_examples", "num_bytes"]      ,							)
														shard_id += 1
														UpperCamelCase      =		writer_class(
														    features=writer._features      ,							path=working_fpath.replace("SSSSS"      ,							F"{shard_id:05d}"						).replace("TTTTT"      ,							F"{task_id:05d}"						)      ,							writer_batch_size=lowercase_      ,							storage_options=lowercase_      ,							embed_local_files=lowercase_      ,							)
												UpperCamelCase      =		pa.Table.from_batches([batch]						)
												writer.write_table(lowercase_						)

										if writer._num_bytes > 0:
												UpperCamelCase      =		writer.finalize()
												writer.close()
												yield pa.RecordBatch.from_arrays(
												    [[task_id], [num_examples], [num_bytes]]      ,							names=["task_id", "num_examples", "num_bytes"]      ,							)

										if working_fpath != fpath:
												for file in os.listdir(os.path.dirname(lowercase_						)						):
														UpperCamelCase      =		os.path.join(os.path.dirname(lowercase_						)      ,							os.path.basename(lowercase_						)						)
														shutil.move(lowercase_      ,							lowercase_						)

								UpperCamelCase      =		(
								    self.df.mapInArrow(lowercase_      ,							"task_id: long, num_examples: long, num_bytes: long"						)
								    .groupBy("task_id"						)
								    .agg(
								        pyspark.sql.functions.sum("num_examples"						).alias("total_num_examples"						)      ,							pyspark.sql.functions.sum("num_bytes"						).alias("total_num_bytes"						)      ,							pyspark.sql.functions.count("num_bytes"						).alias("num_shards"						)      ,							pyspark.sql.functions.collect_list("num_examples"						).alias("shard_lengths"						)      ,							)
								    .collect()
								)
								for row in stats:
										yield row.task_id, (row.total_num_examples, row.total_num_bytes, row.num_shards, row.shard_lengths)


						def 			snake_case_      (self      ,							__a      ,							__a = "arrow"      ,							__a = None      ,							__a = None      ,							**__a      ,							)       ->      int:
								self._validate_cache_dir()

								UpperCamelCase      =		convert_file_size_to_int(max_shard_size or MAX_SHARD_SIZE						)
								self._repartition_df_if_needed(lowercase_						)
								UpperCamelCase      =		not is_remote_filesystem(self._fs						)
								UpperCamelCase      =		os.path.join if is_local else posixpath.join

								UpperCamelCase      =		"-TTTTT-SSSSS-of-NNNNN"
								UpperCamelCase      =		F"{self.name}-{split_generator.name}{SUFFIX}.{file_format}"
								UpperCamelCase      =		path_join(self._output_dir      ,							lowercase_						)

								UpperCamelCase      =		0
								UpperCamelCase      =		0
								UpperCamelCase      =		0
								UpperCamelCase      =		[]
								UpperCamelCase      =		[]

								for task_id, content in self._prepare_split_single(lowercase_      ,							lowercase_      ,							lowercase_						):
										(
										    UpperCamelCase   
										)      =		content
										if num_bytes > 0:
												total_num_examples += num_examples
												total_num_bytes += num_bytes
												total_shards += num_shards
												task_id_and_num_shards.append((task_id, num_shards)						)
												all_shard_lengths.extend(lowercase_						)

								UpperCamelCase      =		total_num_examples
								UpperCamelCase      =		total_num_bytes

								# should rename everything at the end
								logger.debug(F"Renaming {total_shards} shards."						)
								if total_shards > 1:
										UpperCamelCase      =		all_shard_lengths

										# Define fs outside of _rename_shard so that we don't reference self in the function, which will result in a
										# pickling error due to pickling the SparkContext.
										UpperCamelCase      =		self._fs

										# use the -SSSSS-of-NNNNN pattern
										def _rename_shard(
										    __a      ,							__a      ,							__a      ,							):
												rename(
												    lowercase_      ,							fpath.replace("SSSSS"      ,							F"{shard_id:05d}"						).replace("TTTTT"      ,							F"{task_id:05d}"						)      ,							fpath.replace("TTTTT-SSSSS"      ,							F"{global_shard_id:05d}"						).replace("NNNNN"      ,							F"{total_shards:05d}"						)      ,							)

										UpperCamelCase      =		[]
										UpperCamelCase      =		0
										for i in range(len(lowercase_						)						):
												UpperCamelCase      =		task_id_and_num_shards[i]
												for shard_id in range(lowercase_						):
														args.append([task_id, shard_id, global_shard_id]						)
														global_shard_id += 1
										self._spark.sparkContext.parallelize(lowercase_      ,							len(lowercase_						)						).map(lambda __a						: _rename_shard(*lowercase_						)						).collect()
								else:
										# don't use any pattern
										UpperCamelCase      =		0
										UpperCamelCase      =		task_id_and_num_shards[0][0]
										self._rename(
										    fpath.replace("SSSSS"      ,							F"{shard_id:05d}"						).replace("TTTTT"      ,							F"{task_id:05d}"						)      ,							fpath.replace(lowercase_      ,							""						)      ,							)




						def 			snake_case_      (self      ,							__a      ,							)       ->      Union[str, Any]:
								return SparkExamplesIterable(self.df						)








"""simple docstring"""





import inspect
import unittest
import warnings

from transformers import DeiTConfig
from transformers.models.auto import get_values
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 (
        MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING,
        MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
        MODEL_MAPPING,
        DeiTForImageClassification,
        DeiTForImageClassificationWithTeacher,
        DeiTForMaskedImageModeling,
        DeiTModel,
    )
    from transformers.models.deit.modeling_deit import DEIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
    from PIL import Image

    from transformers import DeiTImageProcessor


class A_	:




  '''simple docstring'''







  def __init__(		self				,   lowercase_				,   lowercase_=13				,   lowercase_=30				,   lowercase_=2				,   lowercase_=3				,   lowercase_=True				,   lowercase_=True				,   lowercase_=32				,   lowercase_=5				,   lowercase_=4				,   lowercase_=37				,   lowercase_="gelu"				,   lowercase_=0.1				,   lowercase_=0.1				,   lowercase_=10				,   lowercase_=0.02				,   lowercase_=3				,   lowercase_=None				,   lowercase_=2				,   ):

     """simple docstring"""
     UpperCAmelCase_					:					List[str]			     =    parent
     UpperCAmelCase_					:					int			     =    batch_size
     UpperCAmelCase_					:					int			     =    image_size
     UpperCAmelCase_					:					List[Any]			     =    patch_size
     UpperCAmelCase_					:					Any			     =    num_channels
     UpperCAmelCase_					:					Optional[int]			     =    is_training
     UpperCAmelCase_					:					Union[str, Any]			     =    use_labels
     UpperCAmelCase_					:					Union[str, Any]			     =    hidden_size
     UpperCAmelCase_					:					str			     =    num_hidden_layers
     UpperCAmelCase_					:					List[str]			     =    num_attention_heads
     UpperCAmelCase_					:					str			     =    intermediate_size
     UpperCAmelCase_					:					str			     =    hidden_act
     UpperCAmelCase_					:					List[Any]			     =    hidden_dropout_prob
     UpperCAmelCase_					:					Union[str, Any]			     =    attention_probs_dropout_prob
     UpperCAmelCase_					:					str			     =    type_sequence_label_size
     UpperCAmelCase_					:					str			     =    initializer_range
     UpperCAmelCase_					:					Union[str, Any]			     =    scope
     UpperCAmelCase_					:					str			     =    encoder_stride

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







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					Any			     =    floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]  )

     UpperCAmelCase_					:					Tuple			     =    None
     if self.use_labels:
        UpperCAmelCase_					:					Any			     =    ids_tensor([self.batch_size]				,   self.type_sequence_label_size  )

     UpperCAmelCase_					:					Union[str, Any]			     =    self.get_config()

     return config, pixel_values, labels







  def        UpperCamelCase__					(		self  ):

     """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=lowercase_				,   initializer_range=self.initializer_range				,   encoder_stride=self.encoder_stride				,   )







  def        UpperCamelCase__					(		self				,   lowercase_				,   lowercase_				,   lowercase_  ):

     """simple docstring"""
     UpperCAmelCase_					:					Optional[int]			     =    DeiTModel(config=lowercase_  )
     model.to(lowercase_  )
     model.eval()
     UpperCAmelCase_					:					List[Any]			     =    model(lowercase_  )
     self.parent.assertEqual(result.last_hidden_state.shape				,   (self.batch_size, self.seq_length, self.hidden_size)  )







  def        UpperCamelCase__					(		self				,   lowercase_				,   lowercase_				,   lowercase_  ):

     """simple docstring"""
     UpperCAmelCase_					:					Optional[int]			     =    DeiTForMaskedImageModeling(config=lowercase_  )
     model.to(lowercase_  )
     model.eval()
     UpperCAmelCase_					:					List[Any]			     =    model(lowercase_  )
     self.parent.assertEqual(
         result.reconstruction.shape				,   (self.batch_size, self.num_channels, self.image_size, self.image_size)  )

     # test greyscale images
     UpperCAmelCase_					:					List[str]			     =    1
     UpperCAmelCase_					:					Optional[Any]			     =    DeiTForMaskedImageModeling(lowercase_  )
     model.to(lowercase_  )
     model.eval()

     UpperCAmelCase_					:					Dict			     =    floats_tensor([self.batch_size, 1, self.image_size, self.image_size]  )
     UpperCAmelCase_					:					Optional[int]			     =    model(lowercase_  )
     self.parent.assertEqual(result.reconstruction.shape				,   (self.batch_size, 1, self.image_size, self.image_size)  )







  def        UpperCamelCase__					(		self				,   lowercase_				,   lowercase_				,   lowercase_  ):

     """simple docstring"""
     UpperCAmelCase_					:					Tuple			     =    self.type_sequence_label_size
     UpperCAmelCase_					:					Union[str, Any]			     =    DeiTForImageClassification(lowercase_  )
     model.to(lowercase_  )
     model.eval()
     UpperCAmelCase_					:					List[str]			     =    model(lowercase_				,   labels=lowercase_  )
     self.parent.assertEqual(result.logits.shape				,   (self.batch_size, self.type_sequence_label_size)  )

     # test greyscale images
     UpperCAmelCase_					:					Union[str, Any]			     =    1
     UpperCAmelCase_					:					Optional[int]			     =    DeiTForImageClassification(lowercase_  )
     model.to(lowercase_  )
     model.eval()

     UpperCAmelCase_					:					Dict			     =    floats_tensor([self.batch_size, 1, self.image_size, self.image_size]  )
     UpperCAmelCase_					:					List[Any]			     =    model(lowercase_				,   labels=lowercase_  )
     self.parent.assertEqual(result.logits.shape				,   (self.batch_size, self.type_sequence_label_size)  )







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					List[str]			     =    self.prepare_config_and_inputs()
     (
         (
         UpperCAmelCase_
     )     ,   (
         UpperCAmelCase_
     )     ,   (
         UpperCAmelCase_
     )     ,   
     )					:					Dict			     =    config_and_inputs
     UpperCAmelCase_					:					Optional[int]			     =    {"pixel_values": pixel_values}
     return config, inputs_dict




@require_torch
class A_	(lowercase__				,lowercase__				,unittest.TestCase      ):




  '''simple docstring'''
  SCREAMING_SNAKE_CASE__			:						Any       =		(
      (
          DeiTModel,
          DeiTForImageClassification,
          DeiTForImageClassificationWithTeacher,
          DeiTForMaskedImageModeling,
      )
      if is_torch_available()
      else ()
  )
  SCREAMING_SNAKE_CASE__			:						Tuple       =		(
      {
          """feature-extraction""": DeiTModel,
          """image-classification""": (DeiTForImageClassification, DeiTForImageClassificationWithTeacher),
      }
      if is_torch_available()
      else {}
  )

  SCREAMING_SNAKE_CASE__			:						List[Any]       =		False
  SCREAMING_SNAKE_CASE__			:						Optional[Any]       =		False
  SCREAMING_SNAKE_CASE__			:						List[str]       =		False







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					Any			     =    DeiTModelTester(self  )
     UpperCAmelCase_					:					Optional[int]			     =    ConfigTester(self				,   config_class=lowercase_				,   has_text_modality=lowercase_				,   hidden_size=37  )







  def        UpperCamelCase__					(		self  ):

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







  @unittest.skip(reason="DeiT does not use inputs_embeds"  )
  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     pass







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_     ,   UpperCAmelCase_					:					Any			     =    self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
        UpperCAmelCase_					:					List[Any]			     =    model_class(lowercase_  )
        self.assertIsInstance(model.get_input_embeddings()				,   (nn.Module)  )
        UpperCAmelCase_					:					Any			     =    model.get_output_embeddings()
        self.assertTrue(x is None or isinstance(lowercase_				,   nn.Linear  )  )







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_     ,   UpperCAmelCase_					:					Dict			     =    self.model_tester.prepare_config_and_inputs_for_common()

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

        UpperCAmelCase_					:					Optional[int]			     =    ["pixel_values"]
        self.assertListEqual(arg_names[:1]				,   lowercase_  )







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					int			     =    self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*lowercase_  )







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					Dict			     =    self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_masked_image_modeling(*lowercase_  )







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					Optional[int]			     =    self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_image_classification(*lowercase_  )







  def        UpperCamelCase__					(		self				,   lowercase_				,   lowercase_				,   lowercase_=False  ):

     """simple docstring"""
     UpperCAmelCase_					:					Tuple			     =    super()._prepare_for_class(lowercase_				,   lowercase_				,   return_labels=lowercase_  )

     if return_labels:
        if model_class.__name__ == "DeiTForImageClassificationWithTeacher":
           del inputs_dict["labels"]

     return inputs_dict







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     if not self.model_tester.is_training:
        return

     UpperCAmelCase_     ,   UpperCAmelCase_					:					Union[str, Any]			     =    self.model_tester.prepare_config_and_inputs_for_common()
     UpperCAmelCase_					:					Union[str, Any]			     =    True

     for model_class in self.all_model_classes:
        # DeiTForImageClassificationWithTeacher supports inference-only
        if (
            model_class in get_values(lowercase_  )
            or model_class.__name__ == "DeiTForImageClassificationWithTeacher"
        ):
           continue
        UpperCAmelCase_					:					Optional[int]			     =    model_class(lowercase_  )
        model.to(lowercase_  )
        model.train()
        UpperCAmelCase_					:					List[Any]			     =    self._prepare_for_class(lowercase_				,   lowercase_				,   return_labels=lowercase_  )
        UpperCAmelCase_					:					Dict			     =    model(**lowercase_  ).loss
        loss.backward()







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_     ,   UpperCAmelCase_					:					int			     =    self.model_tester.prepare_config_and_inputs_for_common()
     if not self.model_tester.is_training:
        return

     UpperCAmelCase_					:					Dict			     =    False
     UpperCAmelCase_					:					Optional[int]			     =    True

     for model_class in self.all_model_classes:
        if model_class in get_values(lowercase_  ) or not model_class.supports_gradient_checkpointing:
           continue
        # DeiTForImageClassificationWithTeacher supports inference-only
        if model_class.__name__ == "DeiTForImageClassificationWithTeacher":
           continue
        UpperCAmelCase_					:					List[str]			     =    model_class(lowercase_  )
        model.gradient_checkpointing_enable()
        model.to(lowercase_  )
        model.train()
        UpperCAmelCase_					:					Optional[int]			     =    self._prepare_for_class(lowercase_				,   lowercase_				,   return_labels=lowercase_  )
        UpperCAmelCase_					:					Any			     =    model(**lowercase_  ).loss
        loss.backward()







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_     ,   UpperCAmelCase_					:					Tuple			     =    self.model_tester.prepare_config_and_inputs_for_common()

     UpperCAmelCase_					:					Dict			     =    [
         {"title": "multi_label_classification", "num_labels": 2, "dtype": torch.float},
         {"title": "single_label_classification", "num_labels": 1, "dtype": torch.long},
         {"title": "regression", "num_labels": 1, "dtype": torch.float},
     ]

     for model_class in self.all_model_classes:
        if (
            model_class
            not in [
                *get_values(lowercase_  ),
                *get_values(lowercase_  ),
            ]
            or model_class.__name__ == "DeiTForImageClassificationWithTeacher"
        ):
           continue

        for problem_type in problem_types:
           with self.subTest(msg=F"""Testing {model_class} with {problem_type["title"]}"""  ):
              UpperCAmelCase_					:					str			     =    problem_type["title"]
              UpperCAmelCase_					:					List[Any]			     =    problem_type["num_labels"]

              UpperCAmelCase_					:					Union[str, Any]			     =    model_class(lowercase_  )
              model.to(lowercase_  )
              model.train()

              UpperCAmelCase_					:					int			     =    self._prepare_for_class(lowercase_				,   lowercase_				,   return_labels=lowercase_  )

              if problem_type["num_labels"] > 1:
                 UpperCAmelCase_					:					List[Any]			     =    inputs["labels"].unsqueeze(1  ).repeat(1				,   problem_type["num_labels"]  )

              UpperCAmelCase_					:					Tuple			     =    inputs["labels"].to(problem_type["dtype"]  )

              # This tests that we do not trigger the warning form PyTorch "Using a target size that is different
              # to the input size. This will likely lead to incorrect results due to broadcasting. Please ensure
              # they have the same size." which is a symptom something in wrong for the regression problem.
              # See https://github.com/huggingface/transformers/issues/11780
              with warnings.catch_warnings(record=lowercase_  ) as warning_list:
                 UpperCAmelCase_					:					List[str]			     =    model(**lowercase_  ).loss
              for w in warning_list:
                 if "Using a target size that is different to the input size" in str(w.message  ):
                    raise ValueError(
                        F"""Something is going wrong in the regression problem: intercepted {w.message}"""  )

              loss.backward()







  @slow
  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     for model_name in DEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        UpperCAmelCase_					:					Union[str, Any]			     =    DeiTModel.from_pretrained(lowercase_  )
        self.assertIsNotNone(lowercase_  )







def 	__a     (		):
   UpperCAmelCase_					:					Any			     =    Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"  )
   return image


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




  '''simple docstring'''







  @cached_property
  def        UpperCamelCase__					(		self  ):

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







  @slow
  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					Tuple			     =    DeiTForImageClassificationWithTeacher.from_pretrained("facebook/deit-base-distilled-patch16-224"  ).to(
         lowercase_  )

     UpperCAmelCase_					:					List[str]			     =    self.default_image_processor
     UpperCAmelCase_					:					List[str]			     =    prepare_img()
     UpperCAmelCase_					:					int			     =    image_processor(images=lowercase_				,   return_tensors="pt"  ).to(lowercase_  )

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

     # verify the logits
     UpperCAmelCase_					:					List[str]			     =    torch.Size((1, 1000)  )
     self.assertEqual(outputs.logits.shape				,   lowercase_  )

     UpperCAmelCase_					:					str			     =    torch.tensor([-1.02_66, 0.19_12, -1.28_61]  ).to(lowercase_  )

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







  @slow
  @require_accelerate
  @require_torch_gpu
  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					List[str]			     =    DeiTModel.from_pretrained(
         "facebook/deit-base-distilled-patch16-224"				,   torch_dtype=torch.floataa				,   device_map="auto"  )
     UpperCAmelCase_					:					str			     =    self.default_image_processor

     UpperCAmelCase_					:					Union[str, Any]			     =    prepare_img()
     UpperCAmelCase_					:					List[Any]			     =    image_processor(images=lowercase_				,   return_tensors="pt"  )
     UpperCAmelCase_					:					List[str]			     =    inputs.pixel_values.to(lowercase_  )

     # forward pass to make sure inference works in fp16
     with torch.no_grad():
        UpperCAmelCase_					:					int			     =    model(lowercase_  )


