




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


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

if TYPE_CHECKING:
					import pyspark
@dataclass
class 			a__    (					datasets.BuilderConfig      ):
						_a							:    Optional[datasets.Features]					=    None







def 	_a      (						SCREAMING_SNAKE_CASE_							:				"pyspark.sql.DataFrame"      ,      SCREAMING_SNAKE_CASE_							:				List[int]      ,      ):
					import pyspark

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

					return generate_fn
class 			a__    (					_BaseExamplesIterable      ):






						def __init__(   self		,			_A		,			_A=None		,			):




											"""simple docstring"""



											__lowerCAmelCase				   =       df
											__lowerCAmelCase				   =       partition_order or range(self.df.rdd.getNumPartitions()			)
											__lowerCAmelCase				   =       _generate_iterable_examples(self.df		,			self.partition_order			)






						def __iter__(   self			):




											"""simple docstring"""



											yield from self.generate_examples_fn()






						def 				__SCREAMING_SNAKE_CASE(   self		,			_A			):




											"""simple docstring"""



											__lowerCAmelCase				   =       list(range(self.df.rdd.getNumPartitions()			)			)
											generator.shuffle(_A			)
											return SparkExamplesIterable(self.df		,			partition_order=_A			)






						def 				__SCREAMING_SNAKE_CASE(   self		,			_A		,			_A			):




											"""simple docstring"""



											__lowerCAmelCase				   =       self.split_shard_indices_by_worker(_A		,			_A			)
											return SparkExamplesIterable(self.df		,			partition_order=_A			)






						@property
						def 				__SCREAMING_SNAKE_CASE(   self			):




											"""simple docstring"""



											return len(self.partition_order			)






class 			a__    (					datasets.DatasetBuilder      ):
						_a							:    Optional[Any]					=    SparkConfig






						def __init__(   self		,			_A		,			_A = None		,			_A = None		,			**_A		,			):




											"""simple docstring"""



											import pyspark

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

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






						def 				__SCREAMING_SNAKE_CASE(   self			):




											"""simple docstring"""



											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=_A			)
																__lowerCAmelCase				   =       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(_A		,			"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:
																__lowerCAmelCase				   =       (
																    self._spark.sparkContext.parallelize(range(1			)		,			1			).mapPartitions(_A			).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 				__SCREAMING_SNAKE_CASE(   self			):




											"""simple docstring"""



											return datasets.DatasetInfo(features=self.config.features			)






						def 				__SCREAMING_SNAKE_CASE(   self		,			_A			):




											"""simple docstring"""



											return [datasets.SplitGenerator(name=datasets.Split.TRAIN			)]






						def 				__SCREAMING_SNAKE_CASE(   self		,			_A			):




											"""simple docstring"""



											import pyspark

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

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






						def 				__SCREAMING_SNAKE_CASE(   self		,			_A		,			_A		,			_A		,			):




											"""simple docstring"""



											import pyspark

											__lowerCAmelCase				   =       ParquetWriter if file_format == "parquet" else ArrowWriter
											__lowerCAmelCase				   =       os.path.join(self._working_dir		,			os.path.basename(_A			)			) if self._working_dir else fpath
											__lowerCAmelCase				   =       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.
											__lowerCAmelCase				   =       self.config.features
											__lowerCAmelCase				   =       self._writer_batch_size
											__lowerCAmelCase				   =       self._fs.storage_options

											def write_arrow(_A			):
																# Within the same SparkContext, no two task attempts will share the same attempt ID.
																__lowerCAmelCase				   =       pyspark.TaskContext().taskAttemptId()
																__lowerCAmelCase				   =       next(_A		,			_A			)
																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"]		,			)
																__lowerCAmelCase				   =       0
																__lowerCAmelCase				   =       writer_class(
																    features=_A		,			path=working_fpath.replace("SSSSS"		,			f"""{shard_id:05d}"""			).replace("TTTTT"		,			f"""{task_id:05d}"""			)		,			writer_batch_size=_A		,			storage_options=_A		,			embed_local_files=_A		,			)
																__lowerCAmelCase				   =       pa.Table.from_batches([first_batch]			)
																writer.write_table(_A			)
																for batch in it:
																					if max_shard_size is not None and writer._num_bytes >= max_shard_size:
																										__lowerCAmelCase   ,							__lowerCAmelCase				   =       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
																										__lowerCAmelCase				   =       writer_class(
																										    features=writer._features		,			path=working_fpath.replace("SSSSS"		,			f"""{shard_id:05d}"""			).replace("TTTTT"		,			f"""{task_id:05d}"""			)		,			writer_batch_size=_A		,			storage_options=_A		,			embed_local_files=_A		,			)
																					__lowerCAmelCase				   =       pa.Table.from_batches([batch]			)
																					writer.write_table(_A			)

																if writer._num_bytes > 0:
																					__lowerCAmelCase   ,							__lowerCAmelCase				   =       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(_A			)			):
																										__lowerCAmelCase				   =       os.path.join(os.path.dirname(_A			)		,			os.path.basename(_A			)			)
																										shutil.move(_A		,			_A			)

											__lowerCAmelCase				   =       (
											    self.df.mapInArrow(_A		,			"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 				__SCREAMING_SNAKE_CASE(   self		,			_A		,			_A = "arrow"		,			_A = None		,			_A = None		,			**_A		,			):




											"""simple docstring"""



											self._validate_cache_dir()

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

											__lowerCAmelCase				   =       "-TTTTT-SSSSS-of-NNNNN"
											__lowerCAmelCase				   =       f"""{self.name}-{split_generator.name}{SUFFIX}.{file_format}"""
											__lowerCAmelCase				   =       path_join(self._output_dir		,			_A			)

											__lowerCAmelCase				   =       0
											__lowerCAmelCase				   =       0
											__lowerCAmelCase				   =       0
											__lowerCAmelCase				   =       []
											__lowerCAmelCase				   =       []

											for task_id, content in self._prepare_split_single(_A		,			_A		,			_A			):
																(
																    (
																    __lowerCAmelCase
																)   ,							(
																    __lowerCAmelCase
																)   ,							(
																    __lowerCAmelCase
																)   ,							(
																    __lowerCAmelCase
																)   ,							
																)				   =       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(_A			)

											__lowerCAmelCase				   =       total_num_examples
											__lowerCAmelCase				   =       total_num_bytes

											# should rename everything at the end
											logger.debug(f"""Renaming {total_shards} shards."""			)
											if total_shards > 1:
																__lowerCAmelCase				   =       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.
																__lowerCAmelCase				   =       self._fs

																# use the -SSSSS-of-NNNNN pattern
																def _rename_shard(
																    _A		,			_A		,			_A		,			):
																					rename(
																					    _A		,			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}"""			)		,			)

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






						def 				__SCREAMING_SNAKE_CASE(   self		,			_A		,			):




											"""simple docstring"""



											return SparkExamplesIterable(self.df			)










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__    :






						def __init__(   self		,			_A		,			_A=1_3		,			_A=3_0		,			_A=2		,			_A=3		,			_A=True		,			_A=True		,			_A=3_2		,			_A=5		,			_A=4		,			_A=3_7		,			_A="gelu"		,			_A=0.1		,			_A=0.1		,			_A=1_0		,			_A=0.02		,			_A=3		,			_A=None		,			_A=2		,			):




											"""simple docstring"""



											__lowerCAmelCase				   =       parent
											__lowerCAmelCase				   =       batch_size
											__lowerCAmelCase				   =       image_size
											__lowerCAmelCase				   =       patch_size
											__lowerCAmelCase				   =       num_channels
											__lowerCAmelCase				   =       is_training
											__lowerCAmelCase				   =       use_labels
											__lowerCAmelCase				   =       hidden_size
											__lowerCAmelCase				   =       num_hidden_layers
											__lowerCAmelCase				   =       num_attention_heads
											__lowerCAmelCase				   =       intermediate_size
											__lowerCAmelCase				   =       hidden_act
											__lowerCAmelCase				   =       hidden_dropout_prob
											__lowerCAmelCase				   =       attention_probs_dropout_prob
											__lowerCAmelCase				   =       type_sequence_label_size
											__lowerCAmelCase				   =       initializer_range
											__lowerCAmelCase				   =       scope
											__lowerCAmelCase				   =       encoder_stride

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






						def 				__SCREAMING_SNAKE_CASE(   self			):




											"""simple docstring"""



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

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

											__lowerCAmelCase				   =       self.get_config()

											return config, pixel_values, labels






						def 				__SCREAMING_SNAKE_CASE(   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=_A		,			initializer_range=self.initializer_range		,			encoder_stride=self.encoder_stride		,			)






						def 				__SCREAMING_SNAKE_CASE(   self		,			_A		,			_A		,			_A			):




											"""simple docstring"""



											__lowerCAmelCase				   =       DeiTModel(config=_A			)
											model.to(_A			)
											model.eval()
											__lowerCAmelCase				   =       model(_A			)
											self.parent.assertEqual(result.last_hidden_state.shape		,			(self.batch_size, self.seq_length, self.hidden_size)			)






						def 				__SCREAMING_SNAKE_CASE(   self		,			_A		,			_A		,			_A			):




											"""simple docstring"""



											__lowerCAmelCase				   =       DeiTForMaskedImageModeling(config=_A			)
											model.to(_A			)
											model.eval()
											__lowerCAmelCase				   =       model(_A			)
											self.parent.assertEqual(
											    result.reconstruction.shape		,			(self.batch_size, self.num_channels, self.image_size, self.image_size)			)

											# test greyscale images
											__lowerCAmelCase				   =       1
											__lowerCAmelCase				   =       DeiTForMaskedImageModeling(_A			)
											model.to(_A			)
											model.eval()

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






						def 				__SCREAMING_SNAKE_CASE(   self		,			_A		,			_A		,			_A			):




											"""simple docstring"""



											__lowerCAmelCase				   =       self.type_sequence_label_size
											__lowerCAmelCase				   =       DeiTForImageClassification(_A			)
											model.to(_A			)
											model.eval()
											__lowerCAmelCase				   =       model(_A		,			labels=_A			)
											self.parent.assertEqual(result.logits.shape		,			(self.batch_size, self.type_sequence_label_size)			)

											# test greyscale images
											__lowerCAmelCase				   =       1
											__lowerCAmelCase				   =       DeiTForImageClassification(_A			)
											model.to(_A			)
											model.eval()

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






						def 				__SCREAMING_SNAKE_CASE(   self			):




											"""simple docstring"""



											__lowerCAmelCase				   =       self.prepare_config_and_inputs()
											(
											    (
											    __lowerCAmelCase
											)   ,							(
											    __lowerCAmelCase
											)   ,							(
											    __lowerCAmelCase
											)   ,							
											)				   =       config_and_inputs
											__lowerCAmelCase				   =       {"pixel_values": pixel_values}
											return config, inputs_dict






@require_torch
class 			a__    (					snake_case__ ,				snake_case__ ,				unittest.TestCase      ):
						_a							:    Optional[Any]					=    (
						    (
						        DeiTModel,
						        DeiTForImageClassification,
						        DeiTForImageClassificationWithTeacher,
						        DeiTForMaskedImageModeling,
						    )
						    if is_torch_available()
						    else ()
						)
						_a							:    int					=    (
						    {
						        """feature-extraction""": DeiTModel,
						        """image-classification""": (DeiTForImageClassification, DeiTForImageClassificationWithTeacher),
						    }
						    if is_torch_available()
						    else {}
						)

						_a							:    Optional[Any]					=    False
						_a							:    Tuple					=    False
						_a							:    Tuple					=    False






						def 				__SCREAMING_SNAKE_CASE(   self			):




											"""simple docstring"""



											__lowerCAmelCase				   =       DeiTModelTester(self			)
											__lowerCAmelCase				   =       ConfigTester(self		,			config_class=_A		,			has_text_modality=_A		,			hidden_size=3_7			)






						def 				__SCREAMING_SNAKE_CASE(   self			):




											"""simple docstring"""



											self.config_tester.run_common_tests()






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




											"""simple docstring"""



											pass






						def 				__SCREAMING_SNAKE_CASE(   self			):




											"""simple docstring"""



											__lowerCAmelCase   ,							__lowerCAmelCase				   =       self.model_tester.prepare_config_and_inputs_for_common()

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






						def 				__SCREAMING_SNAKE_CASE(   self			):




											"""simple docstring"""



											__lowerCAmelCase   ,							__lowerCAmelCase				   =       self.model_tester.prepare_config_and_inputs_for_common()

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

																__lowerCAmelCase				   =       ["pixel_values"]
																self.assertListEqual(arg_names[:1]		,			_A			)






						def 				__SCREAMING_SNAKE_CASE(   self			):




											"""simple docstring"""



											__lowerCAmelCase				   =       self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*_A			)






						def 				__SCREAMING_SNAKE_CASE(   self			):




											"""simple docstring"""



											__lowerCAmelCase				   =       self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_masked_image_modeling(*_A			)






						def 				__SCREAMING_SNAKE_CASE(   self			):




											"""simple docstring"""



											__lowerCAmelCase				   =       self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_image_classification(*_A			)






						def 				__SCREAMING_SNAKE_CASE(   self		,			_A		,			_A		,			_A=False			):




											"""simple docstring"""



											__lowerCAmelCase				   =       super()._prepare_for_class(_A		,			_A		,			return_labels=_A			)

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

											return inputs_dict






						def 				__SCREAMING_SNAKE_CASE(   self			):




											"""simple docstring"""



											if not self.model_tester.is_training:
																return

											__lowerCAmelCase   ,							__lowerCAmelCase				   =       self.model_tester.prepare_config_and_inputs_for_common()
											__lowerCAmelCase				   =       True

											for model_class in self.all_model_classes:
																# DeiTForImageClassificationWithTeacher supports inference-only
																if (
																    model_class in get_values(_A			)
																    or model_class.__name__ == "DeiTForImageClassificationWithTeacher"
																):
																					continue
																__lowerCAmelCase				   =       model_class(_A			)
																model.to(_A			)
																model.train()
																__lowerCAmelCase				   =       self._prepare_for_class(_A		,			_A		,			return_labels=_A			)
																__lowerCAmelCase				   =       model(**_A			).loss
																loss.backward()






						def 				__SCREAMING_SNAKE_CASE(   self			):




											"""simple docstring"""



											__lowerCAmelCase   ,							__lowerCAmelCase				   =       self.model_tester.prepare_config_and_inputs_for_common()
											if not self.model_tester.is_training:
																return

											__lowerCAmelCase				   =       False
											__lowerCAmelCase				   =       True

											for model_class in self.all_model_classes:
																if model_class in get_values(_A			) or not model_class.supports_gradient_checkpointing:
																					continue
																# DeiTForImageClassificationWithTeacher supports inference-only
																if model_class.__name__ == "DeiTForImageClassificationWithTeacher":
																					continue
																__lowerCAmelCase				   =       model_class(_A			)
																model.gradient_checkpointing_enable()
																model.to(_A			)
																model.train()
																__lowerCAmelCase				   =       self._prepare_for_class(_A		,			_A		,			return_labels=_A			)
																__lowerCAmelCase				   =       model(**_A			).loss
																loss.backward()






						def 				__SCREAMING_SNAKE_CASE(   self			):




											"""simple docstring"""



											__lowerCAmelCase   ,							__lowerCAmelCase				   =       self.model_tester.prepare_config_and_inputs_for_common()

											__lowerCAmelCase				   =       [
											    {"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(_A			),
																        *get_values(_A			),
																    ]
																    or model_class.__name__ == "DeiTForImageClassificationWithTeacher"
																):
																					continue

																for problem_type in problem_types:
																					with self.subTest(msg=f"""Testing {model_class} with {problem_type['title']}"""			):
																										__lowerCAmelCase				   =       problem_type["title"]
																										__lowerCAmelCase				   =       problem_type["num_labels"]

																										__lowerCAmelCase				   =       model_class(_A			)
																										model.to(_A			)
																										model.train()

																										__lowerCAmelCase				   =       self._prepare_for_class(_A		,			_A		,			return_labels=_A			)

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

																										__lowerCAmelCase				   =       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=_A			) as warning_list:
																															__lowerCAmelCase				   =       model(**_A			).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 				__SCREAMING_SNAKE_CASE(   self			):




											"""simple docstring"""



											for model_name in DEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																__lowerCAmelCase				   =       DeiTModel.from_pretrained(_A			)
																self.assertIsNotNone(_A			)







def 	_a      (						):
					__lowerCAmelCase				   =       Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"     )
					return image


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






						@cached_property
						def 				__SCREAMING_SNAKE_CASE(   self			):




											"""simple docstring"""



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






						@slow
						def 				__SCREAMING_SNAKE_CASE(   self			):




											"""simple docstring"""



											__lowerCAmelCase				   =       DeiTForImageClassificationWithTeacher.from_pretrained("facebook/deit-base-distilled-patch16-224"			).to(
											    _A			)

											__lowerCAmelCase				   =       self.default_image_processor
											__lowerCAmelCase				   =       prepare_img()
											__lowerCAmelCase				   =       image_processor(images=_A		,			return_tensors="pt"			).to(_A			)

											# forward pass
											with torch.no_grad():
																__lowerCAmelCase				   =       model(**_A			)

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

											__lowerCAmelCase				   =       torch.tensor([-1.02_66, 0.19_12, -1.28_61]			).to(_A			)

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






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




											"""simple docstring"""



											__lowerCAmelCase				   =       DeiTModel.from_pretrained(
											    "facebook/deit-base-distilled-patch16-224"		,			torch_dtype=torch.floataa		,			device_map="auto"			)
											__lowerCAmelCase				   =       self.default_image_processor

											__lowerCAmelCase				   =       prepare_img()
											__lowerCAmelCase				   =       image_processor(images=_A		,			return_tensors="pt"			)
											__lowerCAmelCase				   =       inputs.pixel_values.to(_A			)

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




