



"""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				):




					"""simple docstring"""
					snake_case_    :str										=    np.inf

					def set_batch_size(_A				) -> None:
										nonlocal batch_size
										if isinstance(_A,  _A				):
															snake_case_    :Any										=    min(_A,  config.PARQUET_ROW_GROUP_SIZE_FOR_IMAGE_DATASETS				)
										elif isinstance(_A,  _A				):
															snake_case_    :Dict										=    min(_A,  config.PARQUET_ROW_GROUP_SIZE_FOR_AUDIO_DATASETS				)
										elif isinstance(_A,  _A				) and feature.dtype == "binary":
															snake_case_    :str										=    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 						__lowerCAmelCase				(_lowercase	):


					'''simple docstring'''


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




										"""simple docstring"""





										super().__init__(
										    __SCREAMING_SNAKE_CASE     ,							split=__SCREAMING_SNAKE_CASE     ,							features=__SCREAMING_SNAKE_CASE     ,							cache_dir=__SCREAMING_SNAKE_CASE     ,							keep_in_memory=__SCREAMING_SNAKE_CASE     ,							streaming=__SCREAMING_SNAKE_CASE     ,							num_proc=__SCREAMING_SNAKE_CASE     ,							**__SCREAMING_SNAKE_CASE     ,							)
										snake_case_    :Optional[Any]										=    path_or_paths if isinstance(__SCREAMING_SNAKE_CASE     ,							__SCREAMING_SNAKE_CASE ) else {self.split: path_or_paths}
										snake_case_    :List[str]										=    _PACKAGED_DATASETS_MODULES["parquet"][1]
										snake_case_    :List[Any]										=    Parquet(
										    cache_dir=__SCREAMING_SNAKE_CASE     ,							data_files=__SCREAMING_SNAKE_CASE     ,							features=__SCREAMING_SNAKE_CASE     ,							hash=__SCREAMING_SNAKE_CASE     ,							**__SCREAMING_SNAKE_CASE     ,							)







					def 		_a			(   self ):




										"""simple docstring"""





										if self.streaming:
															snake_case_    :Optional[int]										=    self.builder.as_streaming_dataset(split=self.split )
										# Build regular (map-style) dataset
										else:
															snake_case_    :List[str]										=    None
															snake_case_    :Optional[int]										=    None
															snake_case_    :Optional[int]										=    None
															snake_case_    :List[str]										=    None

															self.builder.download_and_prepare(
															    download_config=__SCREAMING_SNAKE_CASE     ,							download_mode=__SCREAMING_SNAKE_CASE     ,							verification_mode=__SCREAMING_SNAKE_CASE     ,							base_path=__SCREAMING_SNAKE_CASE     ,							num_proc=self.num_proc     ,							)
															snake_case_    :Optional[Any]										=    self.builder.as_dataset(
															    split=self.split     ,							verification_mode=__SCREAMING_SNAKE_CASE     ,							in_memory=self.keep_in_memory )
										return dataset



class 						__lowerCAmelCase				:


					'''simple docstring'''


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




										"""simple docstring"""





										snake_case_    :str										=    dataset
										snake_case_    :int										=    path_or_buf
										snake_case_    :Optional[Any]										=    batch_size or get_writer_batch_size(dataset.features )
										snake_case_    :Tuple										=    parquet_writer_kwargs







					def 		_a			(   self ):




										"""simple docstring"""





										snake_case_    :Any										=    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:
																				snake_case_    :List[Any]										=    self._write(file_obj=__SCREAMING_SNAKE_CASE     ,							batch_size=__SCREAMING_SNAKE_CASE     ,							**self.parquet_writer_kwargs )
										else:
															snake_case_    :Tuple										=    self._write(file_obj=self.path_or_buf     ,							batch_size=__SCREAMING_SNAKE_CASE     ,							**self.parquet_writer_kwargs )
										return written







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




										"""simple docstring"""





										snake_case_    :str										=    0
										snake_case_    :Tuple										=    parquet_writer_kwargs.pop("path_or_buf"     ,							__SCREAMING_SNAKE_CASE )
										snake_case_    :Union[str, Any]										=    self.dataset.features.arrow_schema

										snake_case_    :List[Any]										=    pq.ParquetWriter(__SCREAMING_SNAKE_CASE     ,							schema=__SCREAMING_SNAKE_CASE     ,							**__SCREAMING_SNAKE_CASE )

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



"""simple docstring"""




import inspect
import unittest

from transformers import YolosConfig
from transformers.testing_utils import require_torch, 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
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
    import torch
    from torch import nn

    from transformers import YolosForObjectDetection, YolosModel
    from transformers.models.yolos.modeling_yolos import YOLOS_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
    from PIL import Image

    from transformers import AutoImageProcessor







class     lowercase:



     def __init__(							self     ,				__SCREAMING_SNAKE_CASE     ,				__SCREAMING_SNAKE_CASE=1_3     ,				__SCREAMING_SNAKE_CASE=[3_0, 3_0]     ,				__SCREAMING_SNAKE_CASE=2     ,				__SCREAMING_SNAKE_CASE=3     ,				__SCREAMING_SNAKE_CASE=True     ,				__SCREAMING_SNAKE_CASE=True     ,				__SCREAMING_SNAKE_CASE=3_2     ,				__SCREAMING_SNAKE_CASE=5     ,				__SCREAMING_SNAKE_CASE=4     ,				__SCREAMING_SNAKE_CASE=3_7     ,				__SCREAMING_SNAKE_CASE="gelu"     ,				__SCREAMING_SNAKE_CASE=0.1     ,				__SCREAMING_SNAKE_CASE=0.1     ,				__SCREAMING_SNAKE_CASE=1_0     ,				__SCREAMING_SNAKE_CASE=0.02     ,				__SCREAMING_SNAKE_CASE=3     ,				__SCREAMING_SNAKE_CASE=None     ,				__SCREAMING_SNAKE_CASE=8     ,				__SCREAMING_SNAKE_CASE=1_0     ,				)		->     str:







         """simple docstring"""



         a__		=			parent
         a__		=			batch_size
         a__		=			image_size
         a__		=			patch_size
         a__		=			num_channels
         a__		=			is_training
         a__		=			use_labels
         a__		=			hidden_size
         a__		=			num_hidden_layers
         a__		=			num_attention_heads
         a__		=			intermediate_size
         a__		=			hidden_act
         a__		=			hidden_dropout_prob
         a__		=			attention_probs_dropout_prob
         a__		=			type_sequence_label_size
         a__		=			initializer_range
         a__		=			num_labels
         a__		=			scope
         a__		=			n_targets
         a__		=			num_detection_tokens
         # we set the expected sequence length (which is used in several tests)
         # expected sequence length = num_patches + 1 (we add 1 for the [CLS] token) + num_detection_tokens
         a__		=			(image_size[1] // patch_size) * (image_size[0] // patch_size)
         a__		=			num_patches + 1 + self.num_detection_tokens



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







         """simple docstring"""



         a__		=			floats_tensor([self.batch_size, self.num_channels, self.image_size[0], self.image_size[1]]   )

         a__		=			None
         if self.use_labels:
             # labels is a list of Dict (each Dict being the labels for a given example in the batch)
             a__		=			[]
             for i in range(self.batch_size   ):
                 a__		=			{}
                 a__		=			torch.randint(
                     high=self.num_labels     ,				size=(self.n_targets,)     ,				device=__SCREAMING_SNAKE_CASE   )
                 a__		=			torch.rand(self.n_targets     ,				4     ,				device=__SCREAMING_SNAKE_CASE   )
                 labels.append(__SCREAMING_SNAKE_CASE   )

         a__		=			self.get_config()

         return config, pixel_values, labels



     def 			lowercase__  (							self   )		->     Tuple:







         """simple docstring"""



         return YolosConfig(
             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=__SCREAMING_SNAKE_CASE     ,				initializer_range=self.initializer_range     ,				num_detection_tokens=self.num_detection_tokens     ,				num_labels=self.num_labels     ,				)



     def 			lowercase__  (							self     ,				__SCREAMING_SNAKE_CASE     ,				__SCREAMING_SNAKE_CASE     ,				__SCREAMING_SNAKE_CASE   )		->     Optional[Any]:







         """simple docstring"""



         a__		=			YolosModel(config=__SCREAMING_SNAKE_CASE   )
         model.to(__SCREAMING_SNAKE_CASE   )
         model.eval()
         a__		=			model(__SCREAMING_SNAKE_CASE   )
         self.parent.assertEqual(
             result.last_hidden_state.shape     ,				(self.batch_size, self.expected_seq_len, self.hidden_size)   )



     def 			lowercase__  (							self     ,				__SCREAMING_SNAKE_CASE     ,				__SCREAMING_SNAKE_CASE     ,				__SCREAMING_SNAKE_CASE   )		->     Any:







         """simple docstring"""



         a__		=			YolosForObjectDetection(__SCREAMING_SNAKE_CASE   )
         model.to(__SCREAMING_SNAKE_CASE   )
         model.eval()

         a__		=			model(pixel_values=__SCREAMING_SNAKE_CASE   )
         a__		=			model(__SCREAMING_SNAKE_CASE   )

         self.parent.assertEqual(result.logits.shape     ,				(self.batch_size, self.num_detection_tokens, self.num_labels + 1)   )
         self.parent.assertEqual(result.pred_boxes.shape     ,				(self.batch_size, self.num_detection_tokens, 4)   )

         a__		=			model(pixel_values=__SCREAMING_SNAKE_CASE     ,				labels=__SCREAMING_SNAKE_CASE   )

         self.parent.assertEqual(result.loss.shape     ,				()   )
         self.parent.assertEqual(result.logits.shape     ,				(self.batch_size, self.num_detection_tokens, self.num_labels + 1)   )
         self.parent.assertEqual(result.pred_boxes.shape     ,				(self.batch_size, self.num_detection_tokens, 4)   )




     def 			lowercase__  (							self   )		->     int:







         """simple docstring"""



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







@require_torch
class     lowercase(_lowercase    ,						_lowercase    ,						unittest.TestCase       ):
     __snake_case:	List[str]						=     (YolosModel, YolosForObjectDetection) if is_torch_available() else ()
     __snake_case:	Tuple						=     (
         {'feature-extraction': YolosModel, 'object-detection': YolosForObjectDetection} if is_torch_available() else {}
     )

     __snake_case:	Any						=     False
     __snake_case:	List[Any]						=     False
     __snake_case:	Optional[Any]						=     False
     __snake_case:	int						=     False



     def 			lowercase__  (							self     ,				__SCREAMING_SNAKE_CASE     ,				__SCREAMING_SNAKE_CASE     ,				__SCREAMING_SNAKE_CASE=False   )		->     Optional[int]:







         """simple docstring"""



         a__		=			super()._prepare_for_class(__SCREAMING_SNAKE_CASE     ,				__SCREAMING_SNAKE_CASE     ,				return_labels=__SCREAMING_SNAKE_CASE   )

         if return_labels:
             if model_class.__name__ == "YolosForObjectDetection":
                 a__		=			[]
                 for i in range(self.model_tester.batch_size   ):
                     a__		=			{}
                     a__		=			torch.ones(
                         size=(self.model_tester.n_targets,)     ,				device=__SCREAMING_SNAKE_CASE     ,				dtype=torch.long   )
                     a__		=			torch.ones(
                         self.model_tester.n_targets     ,				4     ,				device=__SCREAMING_SNAKE_CASE     ,				dtype=torch.float   )
                     labels.append(__SCREAMING_SNAKE_CASE   )
                 a__		=			labels

         return inputs_dict



     def 			lowercase__  (							self   )		->     Optional[Any]:







         """simple docstring"""



         a__		=			YolosModelTester(self   )
         a__		=			ConfigTester(self     ,				config_class=__SCREAMING_SNAKE_CASE     ,				has_text_modality=__SCREAMING_SNAKE_CASE     ,				hidden_size=3_7   )



     def 			lowercase__  (							self   )		->     Tuple:







         """simple docstring"""



         self.config_tester.run_common_tests()



     def 			lowercase__  (							self   )		->     int:







         """simple docstring"""



         pass



     def 			lowercase__  (							self   )		->     str:







         """simple docstring"""



         a__	,     a__		=			self.model_tester.prepare_config_and_inputs_for_common()

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



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







         """simple docstring"""



         a__	,     a__		=			self.model_tester.prepare_config_and_inputs_for_common()

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

             a__		=			['pixel_values']
             self.assertListEqual(arg_names[:1]     ,				__SCREAMING_SNAKE_CASE   )



     def 			lowercase__  (							self   )		->     Tuple:







         """simple docstring"""



         a__		=			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*__SCREAMING_SNAKE_CASE   )



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







         """simple docstring"""



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

         # in YOLOS, the seq_len is different
         a__		=			self.model_tester.expected_seq_len
         for model_class in self.all_model_classes:
             a__		=			True
             a__		=			False
             a__		=			True
             a__		=			model_class(__SCREAMING_SNAKE_CASE   )
             model.to(__SCREAMING_SNAKE_CASE   )
             model.eval()
             with torch.no_grad():
                 a__		=			model(**self._prepare_for_class(__SCREAMING_SNAKE_CASE     ,				__SCREAMING_SNAKE_CASE   )   )
             a__		=			outputs.attentions
             self.assertEqual(len(__SCREAMING_SNAKE_CASE   )     ,				self.model_tester.num_hidden_layers   )

             # check that output_attentions also work using config
             del inputs_dict["output_attentions"]
             a__		=			True
             a__		=			model_class(__SCREAMING_SNAKE_CASE   )
             model.to(__SCREAMING_SNAKE_CASE   )
             model.eval()
             with torch.no_grad():
                 a__		=			model(**self._prepare_for_class(__SCREAMING_SNAKE_CASE     ,				__SCREAMING_SNAKE_CASE   )   )
             a__		=			outputs.attentions
             self.assertEqual(len(__SCREAMING_SNAKE_CASE   )     ,				self.model_tester.num_hidden_layers   )

             self.assertListEqual(
                 list(attentions[0].shape[-3:]   )     ,				[self.model_tester.num_attention_heads, seq_len, seq_len]     ,				)
             a__		=			len(__SCREAMING_SNAKE_CASE   )

             # Check attention is always last and order is fine
             a__		=			True
             a__		=			True
             a__		=			model_class(__SCREAMING_SNAKE_CASE   )
             model.to(__SCREAMING_SNAKE_CASE   )
             model.eval()
             with torch.no_grad():
                 a__		=			model(**self._prepare_for_class(__SCREAMING_SNAKE_CASE     ,				__SCREAMING_SNAKE_CASE   )   )

             a__		=			1
             self.assertEqual(out_len + added_hidden_states     ,				len(__SCREAMING_SNAKE_CASE   )   )

             a__		=			outputs.attentions

             self.assertEqual(len(__SCREAMING_SNAKE_CASE   )     ,				self.model_tester.num_hidden_layers   )
             self.assertListEqual(
                 list(self_attentions[0].shape[-3:]   )     ,				[self.model_tester.num_attention_heads, seq_len, seq_len]     ,				)



     def 			lowercase__  (							self   )		->     Dict:







         """simple docstring"""



         def check_hidden_states_output(__SCREAMING_SNAKE_CASE     ,				__SCREAMING_SNAKE_CASE     ,				__SCREAMING_SNAKE_CASE   ):
             a__		=			model_class(__SCREAMING_SNAKE_CASE   )
             model.to(__SCREAMING_SNAKE_CASE   )
             model.eval()

             with torch.no_grad():
                 a__		=			model(**self._prepare_for_class(__SCREAMING_SNAKE_CASE     ,				__SCREAMING_SNAKE_CASE   )   )

             a__		=			outputs.hidden_states

             a__		=			getattr(
                 self.model_tester     ,				'expected_num_hidden_layers'     ,				self.model_tester.num_hidden_layers + 1   )
             self.assertEqual(len(__SCREAMING_SNAKE_CASE   )     ,				__SCREAMING_SNAKE_CASE   )

             # YOLOS has a different seq_length
             a__		=			self.model_tester.expected_seq_len

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

         a__	,     a__		=			self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
             a__		=			True
             check_hidden_states_output(__SCREAMING_SNAKE_CASE     ,				__SCREAMING_SNAKE_CASE     ,				__SCREAMING_SNAKE_CASE   )

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

             check_hidden_states_output(__SCREAMING_SNAKE_CASE     ,				__SCREAMING_SNAKE_CASE     ,				__SCREAMING_SNAKE_CASE   )



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







         """simple docstring"""



         a__		=			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_object_detection(*__SCREAMING_SNAKE_CASE   )




     @slow
     def 			lowercase__  (							self   )		->     List[str]:







         """simple docstring"""



         for model_name in YOLOS_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             a__		=			YolosModel.from_pretrained(__SCREAMING_SNAKE_CASE   )
             self.assertIsNotNone(__SCREAMING_SNAKE_CASE   )




def     __magic_name__		(					)     ->						Union[str, Any]:
    a__		=			Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'     )
    return image







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



     @cached_property
     def 			lowercase__  (							self   )		->     List[Any]:







         """simple docstring"""



         return AutoImageProcessor.from_pretrained('hustvl/yolos-small'   ) if is_vision_available() else None




     @slow
     def 			lowercase__  (							self   )		->     int:







         """simple docstring"""



         a__		=			YolosForObjectDetection.from_pretrained('hustvl/yolos-small'   ).to(__SCREAMING_SNAKE_CASE   )

         a__		=			self.default_image_processor
         a__		=			prepare_img()
         a__		=			image_processor(images=__SCREAMING_SNAKE_CASE     ,				return_tensors='pt'   ).to(__SCREAMING_SNAKE_CASE   )

         # forward pass
         with torch.no_grad():
             a__		=			model(inputs.pixel_values   )

         # verify outputs
         a__		=			torch.Size((1, 1_0_0, 9_2)   )
         self.assertEqual(outputs.logits.shape     ,				__SCREAMING_SNAKE_CASE   )

         a__		=			torch.tensor(
             [[-24.02_48, -10.30_24, -14.82_90], [-42.03_92, -16.82_00, -27.43_34], [-27.27_43, -11.81_54, -18.71_48]]     ,				device=__SCREAMING_SNAKE_CASE     ,				)
         a__		=			torch.tensor(
             [[0.25_59, 0.54_55, 0.47_06], [0.29_89, 0.72_79, 0.18_75], [0.77_32, 0.40_17, 0.44_62]]     ,				device=__SCREAMING_SNAKE_CASE   )
         self.assertTrue(torch.allclose(outputs.logits[0, :3, :3]     ,				__SCREAMING_SNAKE_CASE     ,				atol=1e-4   )   )
         self.assertTrue(torch.allclose(outputs.pred_boxes[0, :3, :3]     ,				__SCREAMING_SNAKE_CASE     ,				atol=1e-4   )   )

         # verify postprocessing
         a__		=			image_processor.post_process_object_detection(
             __SCREAMING_SNAKE_CASE     ,				threshold=0.3     ,				target_sizes=[image.size[::-1]]   )[0]
         a__		=			torch.tensor([0.99_94, 0.97_90, 0.99_64, 0.99_72, 0.98_61]   ).to(__SCREAMING_SNAKE_CASE   )
         a__		=			[7_5, 7_5, 1_7, 6_3, 1_7]
         a__		=			torch.tensor([3_35.06_09, 79.38_48, 3_75.42_16, 1_87.24_95]   ).to(__SCREAMING_SNAKE_CASE   )

         self.assertEqual(len(results['scores']   )     ,				5   )
         self.assertTrue(torch.allclose(results['scores']     ,				__SCREAMING_SNAKE_CASE     ,				atol=1e-4   )   )
         self.assertSequenceEqual(results['labels'].tolist()     ,				__SCREAMING_SNAKE_CASE   )
         self.assertTrue(torch.allclose(results['boxes'][0, :]     ,				__SCREAMING_SNAKE_CASE   )   )


