




"""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__					(							snake_case__							)		->   Optional[int]:
			lowerCamelCase      =   np.inf

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

			_visit(snake_case__			,					snake_case__							)

			return None if batch_size is np.inf else batch_size






class        __magic_name__				(			UpperCAmelCase__							):






				'''simple docstring'''





				def __init__(			self						,					_a						,					_a = None						,					_a = None						,					_a = None						,					_a = False						,					_a = False						,					_a = None						,					**_a						,					):





							"""simple docstring"""
							super().__init__(
							    _a						,					split=_a						,					features=_a						,					cache_dir=_a						,					keep_in_memory=_a						,					streaming=_a						,					num_proc=_a						,					**_a						,					)
							lowerCamelCase      =   path_or_paths if isinstance(_a						,					_a					) else {self.split: path_or_paths}
							lowerCamelCase      =   _PACKAGED_DATASETS_MODULES["""parquet"""][1]
							lowerCamelCase      =   Parquet(
							    cache_dir=_a						,					data_files=_a						,					features=_a						,					hash=_a						,					**_a						,					)





				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							# Build iterable dataset
							if self.streaming:
										lowerCamelCase      =   self.builder.as_streaming_dataset(split=self.split					)
							# Build regular (map-style) dataset
							else:
										lowerCamelCase      =   None
										lowerCamelCase      =   None
										lowerCamelCase      =   None
										lowerCamelCase      =   None

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






class        __magic_name__				:






				'''simple docstring'''





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





							"""simple docstring"""
							lowerCamelCase      =   dataset
							lowerCamelCase      =   path_or_buf
							lowerCamelCase      =   batch_size or get_writer_batch_size(dataset.features					)
							lowerCamelCase      =   parquet_writer_kwargs





				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							lowerCamelCase      =   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:
													lowerCamelCase      =   self._write(file_obj=_a						,					batch_size=_a						,					**self.parquet_writer_kwargs					)
							else:
										lowerCamelCase      =   self._write(file_obj=self.path_or_buf						,					batch_size=_a						,					**self.parquet_writer_kwargs					)
							return written





				def   _lowerCAmelCase							(			self						,					_a						,					_a						,					**_a					):





							"""simple docstring"""
							lowerCamelCase      =   0
							lowerCamelCase      =   parquet_writer_kwargs.pop("""path_or_buf"""						,					_a					)
							lowerCamelCase      =   self.dataset.features.arrow_schema

							lowerCamelCase      =   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"""						,					):
										lowerCamelCase      =   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 collections
import inspect
import unittest
from typing import Dict, List, Tuple

from transformers import MaskFormerSwinConfig
from transformers.testing_utils import require_torch, require_torch_multi_gpu, torch_device
from transformers.utils import is_torch_available

from ...test_backbone_common import BackboneTesterMixin
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 MaskFormerSwinBackbone
					from transformers.models.maskformer import MaskFormerSwinModel






class        __magic_name__				:






				'''simple docstring'''





				def __init__(			self						,					_a						,					_a=13						,					_a=32						,					_a=2						,					_a=3						,					_a=16						,					_a=[1, 2, 1]						,					_a=[2, 2, 4]						,					_a=2						,					_a=2.0						,					_a=True						,					_a=0.0						,					_a=0.0						,					_a=0.1						,					_a="gelu"						,					_a=False						,					_a=True						,					_a=0.02						,					_a=1e-5						,					_a=True						,					_a=None						,					_a=True						,					_a=10						,					_a=8						,					_a=["stage1", "stage2", "stage3"]						,					_a=[1, 2, 3]						,					):





							"""simple docstring"""
							lowerCamelCase      =   parent
							lowerCamelCase      =   batch_size
							lowerCamelCase      =   image_size
							lowerCamelCase      =   patch_size
							lowerCamelCase      =   num_channels
							lowerCamelCase      =   embed_dim
							lowerCamelCase      =   depths
							lowerCamelCase      =   num_heads
							lowerCamelCase      =   window_size
							lowerCamelCase      =   mlp_ratio
							lowerCamelCase      =   qkv_bias
							lowerCamelCase      =   hidden_dropout_prob
							lowerCamelCase      =   attention_probs_dropout_prob
							lowerCamelCase      =   drop_path_rate
							lowerCamelCase      =   hidden_act
							lowerCamelCase      =   use_absolute_embeddings
							lowerCamelCase      =   patch_norm
							lowerCamelCase      =   layer_norm_eps
							lowerCamelCase      =   initializer_range
							lowerCamelCase      =   is_training
							lowerCamelCase      =   scope
							lowerCamelCase      =   use_labels
							lowerCamelCase      =   type_sequence_label_size
							lowerCamelCase      =   encoder_stride
							lowerCamelCase      =   out_features
							lowerCamelCase      =   out_indices





				def   _lowerCAmelCase							(			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   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							return MaskFormerSwinConfig(
							    image_size=self.image_size						,					patch_size=self.patch_size						,					num_channels=self.num_channels						,					embed_dim=self.embed_dim						,					depths=self.depths						,					num_heads=self.num_heads						,					window_size=self.window_size						,					mlp_ratio=self.mlp_ratio						,					qkv_bias=self.qkv_bias						,					hidden_dropout_prob=self.hidden_dropout_prob						,					attention_probs_dropout_prob=self.attention_probs_dropout_prob						,					drop_path_rate=self.drop_path_rate						,					hidden_act=self.hidden_act						,					use_absolute_embeddings=self.use_absolute_embeddings						,					path_norm=self.patch_norm						,					layer_norm_eps=self.layer_norm_eps						,					initializer_range=self.initializer_range						,					encoder_stride=self.encoder_stride						,					out_features=self.out_features						,					out_indices=self.out_indices						,					)





				def   _lowerCAmelCase							(			self						,					_a						,					_a						,					_a					):





							"""simple docstring"""
							lowerCamelCase      =   MaskFormerSwinModel(config=_a					)
							model.to(_a					)
							model.eval()
							lowerCamelCase      =   model(_a					)

							lowerCamelCase      =   ((config.image_size // config.patch_size) ** 2) // (4 ** (len(config.depths					) - 1))
							lowerCamelCase      =   int(config.embed_dim * 2 ** (len(config.depths					) - 1)					)

							self.parent.assertEqual(result.last_hidden_state.shape						,					(self.batch_size, expected_seq_len, expected_dim)					)





				def   _lowerCAmelCase							(			self						,					_a						,					_a						,					_a					):





							"""simple docstring"""
							lowerCamelCase      =   MaskFormerSwinBackbone(config=_a					)
							model.to(_a					)
							model.eval()
							lowerCamelCase      =   model(_a					)

							# verify feature maps
							self.parent.assertEqual(len(result.feature_maps					)						,					len(config.out_features					)					)
							self.parent.assertListEqual(list(result.feature_maps[0].shape					)						,					[13, 16, 16, 16]					)

							# verify channels
							self.parent.assertEqual(len(model.channels					)						,					len(config.out_features					)					)
							self.parent.assertListEqual(model.channels						,					[16, 32, 64]					)

							# verify ValueError
							with self.parent.assertRaises(_a					):
										lowerCamelCase      =   ["""stem"""]
										lowerCamelCase      =   MaskFormerSwinBackbone(config=_a					)





				def   _lowerCAmelCase							(			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        __magic_name__				(			UpperCAmelCase__		,   UpperCAmelCase__		,   unittest.TestCase							):






				'''simple docstring'''
				__UpperCamelCase         =		(
				    (
				        MaskFormerSwinModel,
				        MaskFormerSwinBackbone,
				    )
				    if is_torch_available()
				    else ()
				)
				__UpperCamelCase         =		{"feature-extraction": MaskFormerSwinModel} if is_torch_available() else {}
				__UpperCamelCase         =		False
				__UpperCamelCase         =		False
				__UpperCamelCase         =		False
				__UpperCamelCase         =		False
				__UpperCamelCase         =		False





				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							lowerCamelCase      =   MaskFormerSwinModelTester(self					)
							lowerCamelCase      =   ConfigTester(self						,					config_class=_a						,					embed_dim=37					)





				@require_torch_multi_gpu
				@unittest.skip(
				    reason=(
				        """`MaskFormerSwinModel` outputs `hidden_states_spatial_dimensions` which doesn't work well with"""
				        """ `nn.DataParallel`"""
				    )					)
				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							pass





				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							self.create_and_test_config_common_properties()
							self.config_tester.create_and_test_config_to_json_string()
							self.config_tester.create_and_test_config_to_json_file()
							self.config_tester.create_and_test_config_from_and_save_pretrained()
							self.config_tester.create_and_test_config_with_num_labels()
							self.config_tester.check_config_can_be_init_without_params()
							self.config_tester.check_config_arguments_init()





				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							return





				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							lowerCamelCase      =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*_a					)





				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							lowerCamelCase      =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_backbone(*_a					)





				@unittest.skip("""Swin does not use inputs_embeds"""					)
				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							pass





				@unittest.skip("""Swin does not support feedforward chunking"""					)
				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							pass





				def   _lowerCAmelCase							(			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   _lowerCAmelCase							(			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					)





				@unittest.skip(reason="""MaskFormerSwin is only used as backbone and doesn't support output_attentions"""					)
				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							pass





				@unittest.skip(reason="""MaskFormerSwin is only used as an internal backbone"""					)
				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							pass





				def   _lowerCAmelCase							(			self						,					_a						,					_a						,					_a						,					_a					):





							"""simple docstring"""
							lowerCamelCase      =   model_class(_a					)
							model.to(_a					)
							model.eval()

							with torch.no_grad():
										lowerCamelCase      =   model(**self._prepare_for_class(_a						,					_a					)					)

							lowerCamelCase      =   outputs.hidden_states

							lowerCamelCase      =   getattr(
							    self.model_tester						,					"""expected_num_hidden_layers"""						,					len(self.model_tester.depths					) + 1					)
							self.assertEqual(len(_a					)						,					_a					)

							# Swin has a different seq_length
							lowerCamelCase      =   (
							    config.patch_size
							    if isinstance(config.patch_size						,					collections.abc.Iterable					)
							    else (config.patch_size, config.patch_size)
							)

							lowerCamelCase      =   (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])

							self.assertListEqual(
							    list(hidden_states[0].shape[-2:]					)						,					[num_patches, self.model_tester.embed_dim]						,					)





				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							lowerCamelCase ,  lowerCamelCase      =   self.model_tester.prepare_config_and_inputs_for_common()

							lowerCamelCase      =   (
							    self.model_tester.image_size
							    if isinstance(self.model_tester.image_size						,					collections.abc.Iterable					)
							    else (self.model_tester.image_size, self.model_tester.image_size)
							)

							for model_class in self.all_model_classes:
										lowerCamelCase      =   True
										self.check_hidden_states_output(_a						,					_a						,					_a						,					_a					)

										# check that output_hidden_states also work using config
										del inputs_dict["output_hidden_states"]
										lowerCamelCase      =   True

										self.check_hidden_states_output(_a						,					_a						,					_a						,					_a					)





				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							lowerCamelCase ,  lowerCamelCase      =   self.model_tester.prepare_config_and_inputs_for_common()
							lowerCamelCase      =   3

							lowerCamelCase      =   (
							    self.model_tester.image_size
							    if isinstance(self.model_tester.image_size						,					collections.abc.Iterable					)
							    else (self.model_tester.image_size, self.model_tester.image_size)
							)
							lowerCamelCase      =   (
							    config.patch_size
							    if isinstance(config.patch_size						,					collections.abc.Iterable					)
							    else (config.patch_size, config.patch_size)
							)

							lowerCamelCase      =   image_size[0] + patch_size[0] - (image_size[0] % patch_size[0])
							lowerCamelCase      =   image_size[1] + patch_size[1] - (image_size[1] % patch_size[1])

							for model_class in self.all_model_classes:
										lowerCamelCase      =   True
										self.check_hidden_states_output(_a						,					_a						,					_a						,					(padded_height, padded_width)					)

										# check that output_hidden_states also work using config
										del inputs_dict["output_hidden_states"]
										lowerCamelCase      =   True
										self.check_hidden_states_output(_a						,					_a						,					_a						,					(padded_height, padded_width)					)





				@unittest.skip(reason="""MaskFormerSwin doesn't have pretrained checkpoints"""					)
				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							pass





				@unittest.skip(reason="""This will be fixed once MaskFormerSwin is replaced by native Swin"""					)
				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							pass





				@unittest.skip(reason="""This will be fixed once MaskFormerSwin is replaced by native Swin"""					)
				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							pass





				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							lowerCamelCase ,  lowerCamelCase      =   self.model_tester.prepare_config_and_inputs_for_common()

							def set_nan_tensor_to_zero(_a					):
										lowerCamelCase      =   0
										return t

							def check_equivalence(_a						,					_a						,					_a						,					_a={}					):
										with torch.no_grad():
													lowerCamelCase      =   model(**_a						,					return_dict=_a						,					**_a					)
													lowerCamelCase      =   model(**_a						,					return_dict=_a						,					**_a					).to_tuple()

													def recursive_check(_a						,					_a					):
																if isinstance(_a						,					(List, Tuple)					):
																			for tuple_iterable_value, dict_iterable_value in zip(_a						,					_a					):
																						recursive_check(_a						,					_a					)
																elif isinstance(_a						,					_a					):
																			for tuple_iterable_value, dict_iterable_value in zip(
																			    tuple_object.values()						,					dict_object.values()					):
																						recursive_check(_a						,					_a					)
																elif tuple_object is None:
																			return
																else:
																			self.assertTrue(
																			    torch.allclose(
																			        set_nan_tensor_to_zero(_a					)						,					set_nan_tensor_to_zero(_a					)						,					atol=1e-5					)						,					msg=(
																			        """Tuple and dict output are not equal. Difference:"""
																			        f' {torch.max(torch.abs(tuple_object - dict_object					)					)}. Tuple has `nan`:'
																			        f' {torch.isnan(_a					).any()} and `inf`: {torch.isinf(_a					)}. Dict has'
																			        f' `nan`: {torch.isnan(_a					).any()} and `inf`: {torch.isinf(_a					)}.'
																			    )						,					)

													recursive_check(_a						,					_a					)

							for model_class in self.all_model_classes:
										lowerCamelCase      =   model_class(_a					)
										model.to(_a					)
										model.eval()

										lowerCamelCase      =   self._prepare_for_class(_a						,					_a					)
										lowerCamelCase      =   self._prepare_for_class(_a						,					_a					)
										check_equivalence(_a						,					_a						,					_a					)

										lowerCamelCase      =   self._prepare_for_class(_a						,					_a						,					return_labels=_a					)
										lowerCamelCase      =   self._prepare_for_class(_a						,					_a						,					return_labels=_a					)
										check_equivalence(_a						,					_a						,					_a					)

										lowerCamelCase      =   self._prepare_for_class(_a						,					_a					)
										lowerCamelCase      =   self._prepare_for_class(_a						,					_a					)
										check_equivalence(_a						,					_a						,					_a						,					{"""output_hidden_states""": True}					)

										lowerCamelCase      =   self._prepare_for_class(_a						,					_a						,					return_labels=_a					)
										lowerCamelCase      =   self._prepare_for_class(_a						,					_a						,					return_labels=_a					)
										check_equivalence(_a						,					_a						,					_a						,					{"""output_hidden_states""": True}					)






@require_torch
class        __magic_name__				(			unittest.TestCase		,   UpperCAmelCase__							):






				'''simple docstring'''
				__UpperCamelCase         =		(MaskFormerSwinBackbone,) if is_torch_available() else ()
				__UpperCamelCase         =		MaskFormerSwinConfig





				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							lowerCamelCase      =   MaskFormerSwinModelTester(self					)





				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							lowerCamelCase ,  lowerCamelCase      =   self.model_tester.prepare_config_and_inputs_for_common()
							lowerCamelCase      =   inputs_dict["""pixel_values"""].shape[0]

							for backbone_class in self.all_model_classes:
										lowerCamelCase      =   backbone_class(_a					)
										backbone.to(_a					)
										backbone.eval()

										lowerCamelCase      =   backbone(**_a					)

										# Test default outputs and verify feature maps
										self.assertIsInstance(outputs.feature_maps						,					_a					)
										self.assertTrue(len(outputs.feature_maps					) == len(backbone.channels					)					)
										for feature_map, n_channels in zip(outputs.feature_maps						,					backbone.channels					):
													self.assertTrue(feature_map.shape[:2]						,					(batch_size, n_channels)					)
										self.assertIsNone(outputs.hidden_states					)
										self.assertIsNone(outputs.attentions					)

										# Test output_hidden_states=True
										lowerCamelCase      =   backbone(**_a						,					output_hidden_states=_a					)
										self.assertIsNotNone(outputs.hidden_states					)
										self.assertTrue(len(outputs.hidden_states					)						,					len(backbone.stage_names					)					)
										# We skip the stem layer
										for hidden_states, n_channels in zip(outputs.hidden_states[1:]						,					backbone.channels					):
													for hidden_state in hidden_states:
																# Hidden states are in the format (batch_size, (height * width), n_channels)
																lowerCamelCase ,  lowerCamelCase ,  lowerCamelCase      =   hidden_state.shape
																self.assertTrue((h_batch_size, h_n_channels)						,					(batch_size, n_channels)					)

            # Test output_attentions=True
										if self.has_attentions:
													lowerCamelCase      =   backbone(**_a						,					output_attentions=_a					)
													self.assertIsNotNone(outputs.attentions					)



