
'''simple docstring'''
import os
from typing import BinaryIO, Optional, Union

import numpy as np
import pyarrow.parquet as pq

from .. import Audio, Dataset, Features, Image, NamedSplit, Value, config
from ..features.features import FeatureType, _visit
from ..formatting import query_table
from ..packaged_modules import _PACKAGED_DATASETS_MODULES
from ..packaged_modules.parquet.parquet import Parquet
from ..utils import logging
from ..utils.typing import NestedDataStructureLike, PathLike
from .abc import AbstractDatasetReader







def 		A_(     A						:		Features):
							UpperCamelCase   			=     np.inf

							def set_batch_size(A						:		FeatureType) -> None:
														nonlocal batch_size
														if isinstance(A		,     A):
																					UpperCamelCase   			=     min(A		,     config.PARQUET_ROW_GROUP_SIZE_FOR_IMAGE_DATASETS)
														elif isinstance(A		,     A):
																					UpperCamelCase   			=     min(A		,     config.PARQUET_ROW_GROUP_SIZE_FOR_AUDIO_DATASETS)
														elif isinstance(A		,     A) and feature.dtype == "binary":
																					UpperCamelCase   			=     min(A		,     config.PARQUET_ROW_GROUP_SIZE_FOR_BINARY_DATASETS)

							_visit(A		,     A)

							return None if batch_size is np.inf else batch_size


class 			SCREAMING_SNAKE_CASE__      (     snake_case_):







				def __init__(       self						,				A_						,				A_ = None						,				A_ = None						,				A_ = None						,				A_ = False						,				A_ = False						,				A_ = None						,				**A_						,				)->		Optional[Any]:







											'''simple docstring'''






											super().__init__(
											    A_						,				split=A_						,				features=A_						,				cache_dir=A_						,				keep_in_memory=A_						,				streaming=A_						,				num_proc=A_						,				**A_						,				)
											UpperCamelCase   			=     path_or_paths if isinstance(A_						,				A_     ) else {self.split: path_or_paths}
											UpperCamelCase   			=     _PACKAGED_DATASETS_MODULES['parquet'][1]
											UpperCamelCase   			=     Parquet(
											    cache_dir=A_						,				data_files=A_						,				features=A_						,				hash=A_						,				**A_						,				)







				def 						UpperCAmelCase_    (       self     )->		int:







											'''simple docstring'''






											if self.streaming:
																		UpperCamelCase   			=     self.builder.as_streaming_dataset(split=self.split     )
											# Build regular (map-style) dataset
											else:
																		UpperCamelCase   			=     None
																		UpperCamelCase   			=     None
																		UpperCamelCase   			=     None
																		UpperCamelCase   			=     None

																		self.builder.download_and_prepare(
																		    download_config=A_						,				download_mode=A_						,				verification_mode=A_						,				base_path=A_						,				num_proc=self.num_proc						,				)
																		UpperCamelCase   			=     self.builder.as_dataset(
																		    split=self.split						,				verification_mode=A_						,				in_memory=self.keep_in_memory     )
											return dataset


class 			SCREAMING_SNAKE_CASE__      :







				def __init__(       self						,				A_						,				A_						,				A_ = None						,				**A_						,				)->		Dict:







											'''simple docstring'''






											UpperCamelCase   			=     dataset
											UpperCamelCase   			=     path_or_buf
											UpperCamelCase   			=     batch_size or get_writer_batch_size(dataset.features     )
											UpperCamelCase   			=     parquet_writer_kwargs







				def 						UpperCAmelCase_    (       self     )->		int:







											'''simple docstring'''






											UpperCamelCase   			=     self.batch_size if self.batch_size else config.DEFAULT_MAX_BATCH_SIZE

											if isinstance(self.path_or_buf						,				(str, bytes, os.PathLike)     ):
																		with open(self.path_or_buf						,				'wb+'     ) as buffer:
																									UpperCamelCase   			=     self._write(file_obj=A_						,				batch_size=A_						,				**self.parquet_writer_kwargs     )
											else:
																		UpperCamelCase   			=     self._write(file_obj=self.path_or_buf						,				batch_size=A_						,				**self.parquet_writer_kwargs     )
											return written







				def 						UpperCAmelCase_    (       self						,				A_						,				A_						,				**A_     )->		int:







											'''simple docstring'''






											UpperCamelCase   			=     0
											UpperCamelCase   			=     parquet_writer_kwargs.pop('path_or_buf'						,				A_     )
											UpperCamelCase   			=     self.dataset.features.arrow_schema

											UpperCamelCase   			=     pq.ParquetWriter(A_						,				schema=A_						,				**A_     )

											for offset in logging.tqdm(
											    range(0						,				len(self.dataset     )						,				A_     )						,				unit='ba'						,				disable=not logging.is_progress_bar_enabled()						,				desc='Creating parquet from Arrow format'						,				):
																		UpperCamelCase   			=     query_table(
																		    table=self.dataset._data						,				key=slice(A_						,				offset + batch_size     )						,				indices=self.dataset._indices if self.dataset._indices is not None else None						,				)
																		writer.write_table(A_     )
																		written += batch.nbytes
											writer.close()
											return written






'''simple docstring'''
import inspect
import unittest

from datasets import load_dataset
from packaging import version

from transformers import BeitConfig
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_torch_multi_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, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
						import torch
						from torch import nn

						from transformers import (
						    MODEL_MAPPING,
						    BeitForImageClassification,
						    BeitForMaskedImageModeling,
						    BeitForSemanticSegmentation,
						    BeitModel,
						)
						from transformers.models.beit.modeling_beit import BEIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
						import PIL
						from PIL import Image

						from transformers import BeitImageProcessor


class 			SCREAMING_SNAKE_CASE__      :







				def __init__(       self						,				A_						,				A_=100						,				A_=13						,				A_=30						,				A_=2						,				A_=3						,				A_=True						,				A_=True						,				A_=32						,				A_=4						,				A_=4						,				A_=37						,				A_="gelu"						,				A_=0.1						,				A_=0.1						,				A_=10						,				A_=0.02						,				A_=3						,				A_=None						,				A_=[0, 1, 2, 3]						,				)->		Any:







											'''simple docstring'''






											UpperCamelCase   			=     parent
											UpperCamelCase   			=     100
											UpperCamelCase   			=     batch_size
											UpperCamelCase   			=     image_size
											UpperCamelCase   			=     patch_size
											UpperCamelCase   			=     num_channels
											UpperCamelCase   			=     is_training
											UpperCamelCase   			=     use_labels
											UpperCamelCase   			=     hidden_size
											UpperCamelCase   			=     num_hidden_layers
											UpperCamelCase   			=     num_attention_heads
											UpperCamelCase   			=     intermediate_size
											UpperCamelCase   			=     hidden_act
											UpperCamelCase   			=     hidden_dropout_prob
											UpperCamelCase   			=     attention_probs_dropout_prob
											UpperCamelCase   			=     type_sequence_label_size
											UpperCamelCase   			=     initializer_range
											UpperCamelCase   			=     scope
											UpperCamelCase   			=     out_indices
											UpperCamelCase   			=     num_labels

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







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







											'''simple docstring'''






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

											UpperCamelCase   			=     None
											UpperCamelCase   			=     None
											if self.use_labels:
																		UpperCamelCase   			=     ids_tensor([self.batch_size]						,				self.type_sequence_label_size     )
																		UpperCamelCase   			=     ids_tensor([self.batch_size, self.image_size, self.image_size]						,				self.num_labels     )

											UpperCamelCase   			=     self.get_config()

											return config, pixel_values, labels, pixel_labels







				def 						UpperCAmelCase_    (       self     )->		Dict:







											'''simple docstring'''






											return BeitConfig(
											    vocab_size=self.vocab_size						,				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						,				out_indices=self.out_indices						,				)







				def 						UpperCAmelCase_    (       self						,				A_						,				A_						,				A_						,				A_     )->		List[str]:







											'''simple docstring'''






											UpperCamelCase   			=     BeitModel(config=A_     )
											model.to(A_     )
											model.eval()
											UpperCamelCase   			=     model(A_     )
											self.parent.assertEqual(result.last_hidden_state.shape						,				(self.batch_size, self.seq_length, self.hidden_size)     )







				def 						UpperCAmelCase_    (       self						,				A_						,				A_						,				A_						,				A_     )->		Any:







											'''simple docstring'''






											UpperCamelCase   			=     BeitForMaskedImageModeling(config=A_     )
											model.to(A_     )
											model.eval()
											UpperCamelCase   			=     model(A_     )
											self.parent.assertEqual(result.logits.shape						,				(self.batch_size, self.seq_length - 1, self.vocab_size)     )







				def 						UpperCAmelCase_    (       self						,				A_						,				A_						,				A_						,				A_     )->		Optional[int]:







											'''simple docstring'''






											UpperCamelCase   			=     self.type_sequence_label_size
											UpperCamelCase   			=     BeitForImageClassification(A_     )
											model.to(A_     )
											model.eval()
											UpperCamelCase   			=     model(A_						,				labels=A_     )
											self.parent.assertEqual(result.logits.shape						,				(self.batch_size, self.type_sequence_label_size)     )

											# test greyscale images
											UpperCamelCase   			=     1
											UpperCamelCase   			=     BeitForImageClassification(A_     )
											model.to(A_     )
											model.eval()

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







				def 						UpperCAmelCase_    (       self						,				A_						,				A_						,				A_						,				A_     )->		Optional[Any]:







											'''simple docstring'''






											UpperCamelCase   			=     self.num_labels
											UpperCamelCase   			=     BeitForSemanticSegmentation(A_     )
											model.to(A_     )
											model.eval()
											UpperCamelCase   			=     model(A_     )
											self.parent.assertEqual(
											    result.logits.shape						,				(self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)     )
											UpperCamelCase   			=     model(A_						,				labels=A_     )
											self.parent.assertEqual(
											    result.logits.shape						,				(self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)     )







				def 						UpperCAmelCase_    (       self     )->		int:







											'''simple docstring'''






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


@require_torch
class 			SCREAMING_SNAKE_CASE__      (     snake_case_							,       snake_case_							,       unittest.TestCase):
				lowerCAmelCase_     			=							(
				    (BeitModel, BeitForImageClassification, BeitForMaskedImageModeling, BeitForSemanticSegmentation)
				    if is_torch_available()
				    else ()
				)
				lowerCAmelCase_     			=							(
				    {
				        """feature-extraction""": BeitModel,
				        """image-classification""": BeitForImageClassification,
				        """image-segmentation""": BeitForSemanticSegmentation,
				    }
				    if is_torch_available()
				    else {}
				)

				lowerCAmelCase_     			=							False
				lowerCAmelCase_     			=							False
				lowerCAmelCase_     			=							False







				def 						UpperCAmelCase_    (       self     )->		Any:







											'''simple docstring'''






											UpperCamelCase   			=     BeitModelTester(self     )
											UpperCamelCase   			=     ConfigTester(self						,				config_class=A_						,				has_text_modality=A_						,				hidden_size=37     )







				def 						UpperCAmelCase_    (       self     )->		Union[str, Any]:







											'''simple docstring'''






											self.config_tester.run_common_tests()







				@unittest.skip(reason='BEiT does not use inputs_embeds'     )
				def 						UpperCAmelCase_    (       self     )->		Optional[int]:







											'''simple docstring'''






											pass







				@require_torch_multi_gpu
				@unittest.skip(reason='BEiT has some layers using `add_module` which doesn\'t work well with `nn.DataParallel`'     )
				def 						UpperCAmelCase_    (       self     )->		Optional[Any]:







											'''simple docstring'''






											pass







				def 						UpperCAmelCase_    (       self     )->		Tuple:







											'''simple docstring'''






											UpperCamelCase			,	UpperCamelCase   			=     self.model_tester.prepare_config_and_inputs_for_common()

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







				def 						UpperCAmelCase_    (       self     )->		List[Any]:







											'''simple docstring'''






											UpperCamelCase			,	UpperCamelCase   			=     self.model_tester.prepare_config_and_inputs_for_common()

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

																		UpperCamelCase   			=     ['pixel_values']
																		self.assertListEqual(arg_names[:1]						,				A_     )







				def 						UpperCAmelCase_    (       self     )->		Union[str, Any]:







											'''simple docstring'''






											UpperCamelCase   			=     self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*A_     )







				def 						UpperCAmelCase_    (       self     )->		List[Any]:







											'''simple docstring'''






											UpperCamelCase   			=     self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_masked_lm(*A_     )







				def 						UpperCAmelCase_    (       self     )->		Any:







											'''simple docstring'''






											UpperCamelCase   			=     self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_image_classification(*A_     )







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







											'''simple docstring'''






											UpperCamelCase   			=     self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_semantic_segmentation(*A_     )







				def 						UpperCAmelCase_    (       self     )->		int:







											'''simple docstring'''






											if not self.model_tester.is_training:
																		return

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

											for model_class in self.all_model_classes:
																		# we don't test BeitForMaskedImageModeling
																		if model_class in [*get_values(A_     ), BeitForMaskedImageModeling]:
																									continue

																		UpperCamelCase   			=     model_class(A_     )
																		model.to(A_     )
																		model.train()
																		UpperCamelCase   			=     self._prepare_for_class(A_						,				A_						,				return_labels=A_     )
																		UpperCamelCase   			=     model(**A_     ).loss
																		loss.backward()







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







											'''simple docstring'''






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

											UpperCamelCase   			=     False
											UpperCamelCase   			=     True

											for model_class in self.all_model_classes:
																		# we don't test BeitForMaskedImageModeling
																		if (
																		    model_class in [*get_values(A_     ), BeitForMaskedImageModeling]
																		    or not model_class.supports_gradient_checkpointing
																		):
																									continue

																		UpperCamelCase   			=     model_class(A_     )
																		model.gradient_checkpointing_enable()
																		model.to(A_     )
																		model.train()
																		UpperCamelCase   			=     self._prepare_for_class(A_						,				A_						,				return_labels=A_     )
																		UpperCamelCase   			=     model(**A_     ).loss
																		loss.backward()







				def 						UpperCAmelCase_    (       self     )->		Union[str, Any]:







											'''simple docstring'''






											UpperCamelCase			,	UpperCamelCase   			=     self.model_tester.prepare_config_and_inputs_for_common()

											UpperCamelCase   			=     _config_zero_init(A_     )
											for model_class in self.all_model_classes:
																		UpperCamelCase   			=     model_class(config=A_     )
																		for name, param in model.named_parameters():
																									# we skip lambda parameters as these require special initial values
																									# determined by config.layer_scale_init_value
																									if "lambda" in name:
																																continue
																									if param.requires_grad:
																																self.assertIn(
																																    ((param.data.mean() * 1e9).round() / 1e9).item()						,				[0.0, 1.0]						,				msg=F'''Parameter {name} of model {model_class} seems not properly initialized'''						,				)







				@slow
				def 						UpperCAmelCase_    (       self     )->		Dict:







											'''simple docstring'''






											for model_name in BEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																		UpperCamelCase   			=     BeitModel.from_pretrained(A_     )
																		self.assertIsNotNone(A_     )







def 		A_(     ):
							UpperCamelCase   			=     Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png')
							return image





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







				@cached_property
				def 						UpperCAmelCase_    (       self     )->		Optional[int]:







											'''simple docstring'''






											return BeitImageProcessor.from_pretrained('microsoft/beit-base-patch16-224'     ) if is_vision_available() else None







				@slow
				def 						UpperCAmelCase_    (       self     )->		Optional[Any]:







											'''simple docstring'''






											UpperCamelCase   			=     BeitForMaskedImageModeling.from_pretrained('microsoft/beit-base-patch16-224-pt22k'     ).to(A_     )

											UpperCamelCase   			=     self.default_image_processor
											UpperCamelCase   			=     prepare_img()
											UpperCamelCase   			=     image_processor(images=A_						,				return_tensors='pt'     ).pixel_values.to(A_     )

											# prepare bool_masked_pos
											UpperCamelCase   			=     torch.ones((1, 196)						,				dtype=torch.bool     ).to(A_     )

											# forward pass
											with torch.no_grad():
																		UpperCamelCase   			=     model(pixel_values=A_						,				bool_masked_pos=A_     )
											UpperCamelCase   			=     outputs.logits

											# verify the logits
											UpperCamelCase   			=     torch.Size((1, 196, 8192)     )
											self.assertEqual(logits.shape						,				A_     )

											UpperCamelCase   			=     torch.tensor(
											    [[-3.2_437, 0.5_072, -13.9_174], [-3.2_456, 0.4_948, -13.9_401], [-3.2_033, 0.5_121, -13.8_550]]     ).to(A_     )

											self.assertTrue(torch.allclose(logits[bool_masked_pos][:3, :3]						,				A_						,				atol=1e-2     )     )







				@slow
				def 						UpperCAmelCase_    (       self     )->		Union[str, Any]:







											'''simple docstring'''






											UpperCamelCase   			=     BeitForImageClassification.from_pretrained('microsoft/beit-base-patch16-224'     ).to(A_     )

											UpperCamelCase   			=     self.default_image_processor
											UpperCamelCase   			=     prepare_img()
											UpperCamelCase   			=     image_processor(images=A_						,				return_tensors='pt'     ).to(A_     )

											# forward pass
											with torch.no_grad():
																		UpperCamelCase   			=     model(**A_     )
											UpperCamelCase   			=     outputs.logits

											# verify the logits
											UpperCamelCase   			=     torch.Size((1, 1000)     )
											self.assertEqual(logits.shape						,				A_     )

											UpperCamelCase   			=     torch.tensor([-1.2_385, -1.0_987, -1.0_108]     ).to(A_     )

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

											UpperCamelCase   			=     281
											self.assertEqual(logits.argmax(-1     ).item()						,				A_     )







				@slow
				def 						UpperCAmelCase_    (       self     )->		Optional[Any]:







											'''simple docstring'''






											UpperCamelCase   			=     BeitForImageClassification.from_pretrained('microsoft/beit-large-patch16-224-pt22k-ft22k'     ).to(
											    A_     )

											UpperCamelCase   			=     self.default_image_processor
											UpperCamelCase   			=     prepare_img()
											UpperCamelCase   			=     image_processor(images=A_						,				return_tensors='pt'     ).to(A_     )

											# forward pass
											with torch.no_grad():
																		UpperCamelCase   			=     model(**A_     )
											UpperCamelCase   			=     outputs.logits

											# verify the logits
											UpperCamelCase   			=     torch.Size((1, 21841)     )
											self.assertEqual(logits.shape						,				A_     )

											UpperCamelCase   			=     torch.tensor([1.6_881, -0.2_787, 0.5_901]     ).to(A_     )

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

											UpperCamelCase   			=     2396
											self.assertEqual(logits.argmax(-1     ).item()						,				A_     )







				@slow
				def 						UpperCAmelCase_    (       self     )->		Any:







											'''simple docstring'''






											UpperCamelCase   			=     BeitForSemanticSegmentation.from_pretrained('microsoft/beit-base-finetuned-ade-640-640'     )
											UpperCamelCase   			=     model.to(A_     )

											UpperCamelCase   			=     BeitImageProcessor(do_resize=A_						,				size=640						,				do_center_crop=A_     )

											UpperCamelCase   			=     load_dataset('hf-internal-testing/fixtures_ade20k'						,				split='test'     )
											UpperCamelCase   			=     Image.open(ds[0]['file']     )
											UpperCamelCase   			=     image_processor(images=A_						,				return_tensors='pt'     ).to(A_     )

											# forward pass
											with torch.no_grad():
																		UpperCamelCase   			=     model(**A_     )
											UpperCamelCase   			=     outputs.logits

											# verify the logits
											UpperCamelCase   			=     torch.Size((1, 150, 160, 160)     )
											self.assertEqual(logits.shape						,				A_     )

											UpperCamelCase   			=     version.parse(PIL.__version__     ) < version.parse('9.0.0'     )

											if is_pillow_less_than_a:
																		UpperCamelCase   			=     torch.tensor(
																		    [
																		        [[-4.9_225, -2.3_954, -3.0_522], [-2.8_822, -1.0_046, -1.7_561], [-2.9_549, -1.3_228, -2.1_347]],
																		        [[-5.8_168, -3.4_129, -4.0_778], [-3.8_651, -2.2_214, -3.0_277], [-3.8_356, -2.4_643, -3.3_535]],
																		        [[-0.0_078, 3.9_952, 4.0_754], [2.9_856, 4.6_944, 5.0_035], [3.2_413, 4.7_813, 4.9_969]],
																		    ]						,				device=A_						,				)
											else:
																		UpperCamelCase   			=     torch.tensor(
																		    [
																		        [[-4.8_960, -2.3_688, -3.0_355], [-2.8_478, -0.9_836, -1.7_418], [-2.9_449, -1.3_332, -2.1_456]],
																		        [[-5.8_081, -3.4_124, -4.1_006], [-3.8_561, -2.2_081, -3.0_323], [-3.8_365, -2.4_601, -3.3_669]],
																		        [[-0.0_309, 3.9_868, 4.0_540], [2.9_640, 4.6_877, 4.9_976], [3.2_081, 4.7_690, 4.9_942]],
																		    ]						,				device=A_						,				)

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







				@slow
				def 						UpperCAmelCase_    (       self     )->		Tuple:







											'''simple docstring'''






											UpperCamelCase   			=     BeitForSemanticSegmentation.from_pretrained('microsoft/beit-base-finetuned-ade-640-640'     )
											UpperCamelCase   			=     model.to(A_     )

											UpperCamelCase   			=     BeitImageProcessor(do_resize=A_						,				size=640						,				do_center_crop=A_     )

											UpperCamelCase   			=     load_dataset('hf-internal-testing/fixtures_ade20k'						,				split='test'     )
											UpperCamelCase   			=     Image.open(ds[0]['file']     )
											UpperCamelCase   			=     image_processor(images=A_						,				return_tensors='pt'     ).to(A_     )

											# forward pass
											with torch.no_grad():
																		UpperCamelCase   			=     model(**A_     )

											UpperCamelCase   			=     outputs.logits.detach().cpu()

											UpperCamelCase   			=     image_processor.post_process_semantic_segmentation(outputs=A_						,				target_sizes=[(500, 300)]     )
											UpperCamelCase   			=     torch.Size((500, 300)     )
											self.assertEqual(segmentation[0].shape						,				A_     )

											UpperCamelCase   			=     image_processor.post_process_semantic_segmentation(outputs=A_     )
											UpperCamelCase   			=     torch.Size((160, 160)     )
											self.assertEqual(segmentation[0].shape						,				A_     )




