






import inspect
import unittest

import numpy as np

from transformers import ViTConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor


if is_flax_available():
	import jax

	from transformers.models.vit.modeling_flax_vit import FlaxViTForImageClassification, FlaxViTModel




class  UpperCamelCase_      (			unittest.TestCase				):



	'''simple docstring'''




	def __init__(       self       ,			a       ,			a=13       ,			a=30       ,			a=2       ,			a=3       ,			a=True       ,			a=True       ,			a=32       ,			a=5       ,			a=4       ,			a=37       ,			a="gelu"       ,			a=0.1       ,			a=0.1       ,			a=10       ,			a=0.02       ,			)       ->    int:
								snake_case_				=							parent
								snake_case_				=							batch_size
								snake_case_				=							image_size
								snake_case_				=							patch_size
								snake_case_				=							num_channels
								snake_case_				=							is_training
								snake_case_				=							use_labels
								snake_case_				=							hidden_size
								snake_case_				=							num_hidden_layers
								snake_case_				=							num_attention_heads
								snake_case_				=							intermediate_size
								snake_case_				=							hidden_act
								snake_case_				=							hidden_dropout_prob
								snake_case_				=							attention_probs_dropout_prob
								snake_case_				=							type_sequence_label_size
								snake_case_				=							initializer_range

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




	def 							_UpperCamelCase			(       self					)       ->    int:
								snake_case_				=							floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]					)

								snake_case_				=							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=a       ,			initializer_range=self.initializer_range       ,			)

								return config, pixel_values




	def 							_UpperCamelCase			(       self       ,			a       ,			a					)       ->    Any:
								snake_case_				=							FlaxViTModel(config=a					)
								snake_case_				=							model(a					)
								# expected sequence length = num_patches + 1 (we add 1 for the [CLS] token)
								snake_case_				=							(self.image_size, self.image_size)
								snake_case_				=							(self.patch_size, self.patch_size)
								snake_case_				=							(image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
								self.parent.assertEqual(result.last_hidden_state.shape       ,			(self.batch_size, num_patches + 1, self.hidden_size)					)




	def 							_UpperCamelCase			(       self       ,			a       ,			a					)       ->    Dict:
								snake_case_				=							self.type_sequence_label_size
								snake_case_				=							FlaxViTForImageClassification(config=a					)
								snake_case_				=							model(a					)
								self.parent.assertEqual(result.logits.shape       ,			(self.batch_size, self.type_sequence_label_size)					)

								# test greyscale images
								snake_case_				=							1
								snake_case_				=							FlaxViTForImageClassification(a					)

								snake_case_				=							floats_tensor([self.batch_size, 1, self.image_size, self.image_size]					)
								snake_case_				=							model(a					)






	def 							_UpperCamelCase			(       self					)       ->    int:
								snake_case_				=							self.prepare_config_and_inputs()
								(
								    (
								    snake_case_
								)			,		(
								    snake_case_
								)			,		
								)				=							config_and_inputs
								snake_case_				=							{'pixel_values': pixel_values}
								return config, inputs_dict






@require_flax
class  UpperCamelCase_      (			snake_case_       ,					unittest.TestCase				):



	'''simple docstring'''




	lowerCAmelCase													= (FlaxViTModel, FlaxViTForImageClassification) if is_flax_available() else ()




	def 							_UpperCamelCase			(       self					)       ->    None:
								snake_case_				=							FlaxViTModelTester(self					)
								snake_case_				=							ConfigTester(self       ,			config_class=a       ,			has_text_modality=a       ,			hidden_size=37					)




	def 							_UpperCamelCase			(       self					)       ->    int:
								self.config_tester.run_common_tests()




	def 							_UpperCamelCase			(       self					)       ->    Optional[Any]:
								snake_case_				=							self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*a					)




	def 							_UpperCamelCase			(       self					)       ->    Tuple:
								snake_case_				=							self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_image_classification(*a					)




	def 							_UpperCamelCase			(       self					)       ->    Union[str, Any]:
								snake_case_			,		snake_case_				=							self.model_tester.prepare_config_and_inputs_for_common()

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

															snake_case_				=							['pixel_values']
															self.assertListEqual(arg_names[:1]       ,			a					)




	def 							_UpperCamelCase			(       self					)       ->    Union[str, Any]:
								snake_case_			,		snake_case_				=							self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
															with self.subTest(model_class.__name__					):
																						snake_case_				=							self._prepare_for_class(a       ,			a					)
																						snake_case_				=							model_class(a					)

																						@jax.jit
																						def model_jitted(a       ,			**a					):
																													return model(pixel_values=a       ,			**a					)

																						with self.subTest('JIT Enabled'					):
																													snake_case_				=							model_jitted(**a					).to_tuple()

																						with self.subTest('JIT Disabled'					):
																													with jax.disable_jit():
																																				snake_case_				=							model_jitted(**a					).to_tuple()

																						self.assertEqual(len(a					)       ,			len(a					)					)
																						for jitted_output, output in zip(a       ,			a					):
																													self.assertEqual(jitted_output.shape       ,			output.shape					)






	@slow
	def 							_UpperCamelCase			(       self					)       ->    Dict:
								for model_class_name in self.all_model_classes:
															snake_case_				=							model_class_name.from_pretrained('google/vit-base-patch16-224'					)
															snake_case_				=							model(np.ones((1, 3, 2_24, 2_24)					)					)
															self.assertIsNotNone(a					)













import os
from pathlib import Path
from unittest.mock import patch

import pytest
import zstandard as zstd

from datasets.download.download_config import DownloadConfig
from datasets.utils.file_utils import (
    OfflineModeIsEnabled,
    cached_path,
    fsspec_get,
    fsspec_head,
    ftp_get,
    ftp_head,
    get_from_cache,
    http_get,
    http_head,
)


lowercase  							=					"\\n    Text data.\n    Second line of data."

lowercase  							=					"file"





@pytest.fixture(scope='session')
def 						__UpperCAmelCase			(     a_):
							snake_case_				=							tmp_path_factory.mktemp('data') / (FILE_PATH + '.zstd')
							snake_case_				=							bytes(a_							,     'utf-8')
							with zstd.open(a_							,     'wb') as f:
														f.write(a_)
							return path





@pytest.fixture
def 						__UpperCAmelCase			(     a_):
							with open(os.path.join(tmpfs.local_root_dir							,     a_)							,     'w') as f:
														f.write(a_)
							return FILE_PATH





@pytest.mark.parametrize('compression_format'							,     ['gzip', 'xz', 'zstd'])
def 						__UpperCAmelCase			(     a_							,     a_							,     a_							,     a_							,     a_							,     a_):
							snake_case_				=							{'gzip': gz_file, 'xz': xz_file, 'zstd': zstd_path}
							snake_case_				=							input_paths[compression_format]
							snake_case_				=							tmp_path / 'cache'
							snake_case_				=							DownloadConfig(cache_dir=a_							,     extract_compressed_file=a_)
							snake_case_				=							cached_path(a_							,     download_config=a_)
							with open(a_) as f:
														snake_case_				=							f.read()
							with open(a_) as f:
														snake_case_				=							f.read()
							assert extracted_file_content == expected_file_content





@pytest.mark.parametrize('default_extracted'							,     [True, False])
@pytest.mark.parametrize('default_cache_dir'							,     [True, False])
def 						__UpperCAmelCase			(     a_							,     a_							,     a_							,     a_							,     a_):
							snake_case_				=							'custom_cache'
							snake_case_				=							'custom_extracted_dir'
							snake_case_				=							tmp_path / 'custom_extracted_path'
							if default_extracted:
														snake_case_				=							('downloads' if default_cache_dir else custom_cache_dir, 'extracted')
							else:
														monkeypatch.setattr('datasets.config.EXTRACTED_DATASETS_DIR'							,     a_)
														monkeypatch.setattr('datasets.config.EXTRACTED_DATASETS_PATH'							,     str(a_))
														snake_case_				=							custom_extracted_path.parts[-2:] if default_cache_dir else (custom_cache_dir, custom_extracted_dir)

							snake_case_				=							xz_file
							snake_case_				=							(
							    DownloadConfig(extract_compressed_file=a_)
							    if default_cache_dir
							    else DownloadConfig(cache_dir=tmp_path / custom_cache_dir							,     extract_compressed_file=a_)
							)
							snake_case_				=							cached_path(a_							,     download_config=a_)
							assert Path(a_).parent.parts[-2:] == expected





def 						__UpperCAmelCase			(     a_):
							# absolute path
							snake_case_				=							str(Path(a_).resolve())
							assert cached_path(a_) == text_file
							# relative path
							snake_case_				=							str(Path(a_).resolve().relative_to(Path(os.getcwd())))
							assert cached_path(a_) == text_file





def 						__UpperCAmelCase			(     a_):
							# absolute path
							snake_case_				=							str(tmp_path.resolve() / '__missing_file__.txt')
							with pytest.raises(a_):
														cached_path(a_)
							# relative path
							snake_case_				=							'./__missing_file__.txt'
							with pytest.raises(a_):
														cached_path(a_)





def 						__UpperCAmelCase			(     a_):
							snake_case_				=							get_from_cache(f'''tmp://{tmpfs_file}''')
							with open(a_) as f:
														snake_case_				=							f.read()
							assert output_file_content == FILE_CONTENT





@patch('datasets.config.HF_DATASETS_OFFLINE'							,     a_)
def 						__UpperCAmelCase			(     ):
							with pytest.raises(a_):
														cached_path('https://huggingface.co')





@patch('datasets.config.HF_DATASETS_OFFLINE'							,     a_)
def 						__UpperCAmelCase			(     a_):
							snake_case_				=							tmp_path_factory.mktemp('data') / 'file.html'
							with pytest.raises(a_):
														http_get('https://huggingface.co'							,     temp_file=a_)
							with pytest.raises(a_):
														http_head('https://huggingface.co')





@patch('datasets.config.HF_DATASETS_OFFLINE'							,     a_)
def 						__UpperCAmelCase			(     a_):
							snake_case_				=							tmp_path_factory.mktemp('data') / 'file.html'
							with pytest.raises(a_):
														ftp_get('ftp://huggingface.co'							,     temp_file=a_)
							with pytest.raises(a_):
														ftp_head('ftp://huggingface.co')





@patch('datasets.config.HF_DATASETS_OFFLINE'							,     a_)
def 						__UpperCAmelCase			(     a_):
							snake_case_				=							tmp_path_factory.mktemp('data') / 'file.html'
							with pytest.raises(a_):
														fsspec_get('s3://huggingface.co'							,     temp_file=a_)
							with pytest.raises(a_):
														fsspec_head('s3://huggingface.co')





