






'''simple docstring'''
import inspect
import math
import tempfile
import unittest

import numpy as np

from transformers import ViTMAEConfig
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, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
 import torch
 from torch import nn

 from transformers import ViTMAEForPreTraining, ViTMAEModel
 from transformers.models.vit.modeling_vit import VIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
 from PIL import Image

 from transformers import ViTImageProcessor


class        _UpperCAmelCase						:




    def __init__(			self  :       List[str]  ,     a  :       Tuple  ,     a  :       str=1_3  ,     a  :       List[str]=3_0  ,     a  :       Union[str, Any]=2  ,     a  :       Tuple=3  ,     a  :       List[Any]=True  ,     a  :       int=True  ,     a  :       Optional[int]=3_2  ,     a  :       List[str]=5  ,     a  :       int=4  ,     a  :       int=3_7  ,     a  :       Tuple="gelu"  ,     a  :       str=0.1  ,     a  :       Tuple=0.1  ,     a  :       List[Any]=1_0  ,     a  :       int=0.02  ,     a  :       Union[str, Any]=3  ,     a  :       Dict=0.6  ,     a  :       Union[str, Any]=None  ,     ):






       '''simple docstring'''






       lowercase_	:    Dict           =     parent
       lowercase_	:    Dict           =     batch_size
       lowercase_	:    str           =     image_size
       lowercase_	:    int           =     patch_size
       lowercase_	:    Optional[Any]           =     num_channels
       lowercase_	:    Union[str, Any]           =     is_training
       lowercase_	:    List[str]           =     use_labels
       lowercase_	:    Union[str, Any]           =     hidden_size
       lowercase_	:    Tuple           =     num_hidden_layers
       lowercase_	:    Optional[int]           =     num_attention_heads
       lowercase_	:    List[Any]           =     intermediate_size
       lowercase_	:    str           =     hidden_act
       lowercase_	:    Dict           =     hidden_dropout_prob
       lowercase_	:    Dict           =     attention_probs_dropout_prob
       lowercase_	:    Union[str, Any]           =     type_sequence_label_size
       lowercase_	:    int           =     initializer_range
       lowercase_	:    List[Any]           =     mask_ratio
       lowercase_	:    Any           =     scope

       # in ViTMAE, the expected sequence length = (num_patches + 1) * (1 - config.mask_ratio), rounded above
       # (we add 1 for the [CLS] token)
       lowercase_	:    Union[str, Any]           =     (image_size // patch_size) ** 2
       lowercase_	:    int           =     int(math.ceil((1 - mask_ratio) * (num_patches + 1)			)			)




    def 				lowerCAmelCase__				(			self  :       str			):






       '''simple docstring'''






       lowercase_	:    List[str]           =     floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]			)

       lowercase_	:    int           =     None
       if self.use_labels:
          lowercase_	:    Any           =     ids_tensor([self.batch_size]  ,     self.type_sequence_label_size			)

       lowercase_	:    Optional[Any]           =     self.get_config()

       return config, pixel_values, labels




    def 				lowerCAmelCase__				(			self  :       Optional[Any]			):






       '''simple docstring'''






       return ViTMAEConfig(
           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  ,     mask_ratio=self.mask_ratio  ,     )




    def 				lowerCAmelCase__				(			self  :       Tuple  ,     a  :       Dict  ,     a  :       Union[str, Any]  ,     a  :       str			):






       '''simple docstring'''






       lowercase_	:    List[Any]           =     ViTMAEModel(config=a			)
       model.to(a			)
       model.eval()
       lowercase_	:    Any           =     model(a			)
       self.parent.assertEqual(result.last_hidden_state.shape  ,     (self.batch_size, self.seq_length, self.hidden_size)			)




    def 				lowerCAmelCase__				(			self  :       Optional[Any]  ,     a  :       Optional[int]  ,     a  :       Union[str, Any]  ,     a  :       Union[str, Any]			):






       '''simple docstring'''






       lowercase_	:    Optional[Any]           =     ViTMAEForPreTraining(a			)
       model.to(a			)
       model.eval()
       lowercase_	:    str           =     model(a			)
       lowercase_	:    Optional[int]           =     (self.image_size // self.patch_size) ** 2
       lowercase_	:    Dict           =     self.patch_size**2 * self.num_channels
       self.parent.assertEqual(result.logits.shape  ,     (self.batch_size, num_patches, expected_num_channels)			)

       # test greyscale images
       lowercase_	:    str           =     1
       lowercase_	:    Any           =     ViTMAEForPreTraining(a			)
       model.to(a			)
       model.eval()
       lowercase_	:    str           =     floats_tensor([self.batch_size, 1, self.image_size, self.image_size]			)
       lowercase_	:    Optional[int]           =     model(a			)
       lowercase_	:    List[Any]           =     self.patch_size**2
       self.parent.assertEqual(result.logits.shape  ,     (self.batch_size, num_patches, expected_num_channels)			)





    def 				lowerCAmelCase__				(			self  :       str			):






       '''simple docstring'''






       lowercase_	:    Optional[Any]           =     self.prepare_config_and_inputs()
       lowercase_    ,		lowercase_    ,		lowercase_	:    str           =     config_and_inputs
       lowercase_	:    List[str]           =     {"pixel_values": pixel_values}
       return config, inputs_dict


@require_torch
class        _UpperCAmelCase						(		snake_case				,					snake_case				,					unittest.TestCase			):
    __lowerCamelCase:   Optional[Any]		       =    (ViTMAEModel, ViTMAEForPreTraining) if is_torch_available() else ()
    __lowerCamelCase:   Optional[Any]		       =    {'feature-extraction': ViTMAEModel} if is_torch_available() else {}

    __lowerCamelCase:   Dict		       =    False
    __lowerCamelCase:   Dict		       =    False
    __lowerCamelCase:   Tuple		       =    False
    __lowerCamelCase:   str		       =    False




    def 				lowerCAmelCase__				(			self  :       int			):






       '''simple docstring'''






       lowercase_	:    Tuple           =     ViTMAEModelTester(self			)
       lowercase_	:    List[Any]           =     ConfigTester(self  ,     config_class=a  ,     has_text_modality=a  ,     hidden_size=3_7			)




    def 				lowerCAmelCase__				(			self  :       List[str]			):






       '''simple docstring'''






       self.config_tester.run_common_tests()




    @unittest.skip(reason="ViTMAE does not use inputs_embeds"			)
    def 				lowerCAmelCase__				(			self  :       Union[str, Any]			):






       '''simple docstring'''






       pass




    def 				lowerCAmelCase__				(			self  :       str			):






       '''simple docstring'''






       lowercase_    ,		lowercase_	:    List[str]           =     self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
          lowercase_	:    str           =     model_class(a			)
          self.assertIsInstance(model.get_input_embeddings()  ,     (nn.Module)			)
          lowercase_	:    Optional[int]           =     model.get_output_embeddings()
          self.assertTrue(x is None or isinstance(a  ,     nn.Linear			)			)




    def 				lowerCAmelCase__				(			self  :       List[str]			):






       '''simple docstring'''






       lowercase_    ,		lowercase_	:    int           =     self.model_tester.prepare_config_and_inputs_for_common()

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

          lowercase_	:    Tuple           =     ["pixel_values"]
          self.assertListEqual(arg_names[:1]  ,     a			)




    def 				lowerCAmelCase__				(			self  :       List[Any]			):






       '''simple docstring'''






       lowercase_	:    Optional[int]           =     self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*a			)




    def 				lowerCAmelCase__				(			self  :       Optional[int]			):






       '''simple docstring'''






       lowercase_	:    Optional[int]           =     self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_pretraining(*a			)




    def 				lowerCAmelCase__				(			self  :       Tuple  ,     a  :       Tuple  ,     a  :       List[Any]  ,     a  :       int			):






       '''simple docstring'''






       np.random.seed(2			)

       lowercase_	:    int           =     int((pt_model.config.image_size // pt_model.config.patch_size) ** 2			)
       lowercase_	:    Dict           =     np.random.uniform(size=(self.model_tester.batch_size, num_patches)			)
       lowercase_	:    List[Any]           =     torch.from_numpy(a			)

       # Add `noise` argument.
       # PT inputs will be prepared in `super().check_pt_tf_models()` with this added `noise` argument
       lowercase_	:    Tuple           =     pt_noise

       super().check_pt_tf_models(a  ,     a  ,     a			)




    def 				lowerCAmelCase__				(			self  :       List[Any]			):






       '''simple docstring'''






       lowercase_    ,		lowercase_	:    Tuple           =     self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
          lowercase_	:    Any           =     model_class(a			)
          model.to(a			)
          model.eval()
          # make random mask reproducible
          torch.manual_seed(2			)
          with torch.no_grad():
             lowercase_	:    Optional[Any]           =     model(**self._prepare_for_class(a  ,     a			)			)

          lowercase_	:    List[str]           =     outputs[0].cpu().numpy()
          lowercase_	:    List[Any]           =     0

          with tempfile.TemporaryDirectory() as tmpdirname:
             model.save_pretrained(a			)
             lowercase_	:    Tuple           =     model_class.from_pretrained(a			)
             model.to(a			)
             # make random mask reproducible
             torch.manual_seed(2			)
             with torch.no_grad():
                lowercase_	:    int           =     model(**self._prepare_for_class(a  ,     a			)			)

             # Make sure we don't have nans
             lowercase_	:    List[str]           =     after_outputs[0].cpu().numpy()
             lowercase_	:    int           =     0
             lowercase_	:    List[str]           =     np.amax(np.abs(out_a - out_a			)			)
             self.assertLessEqual(a  ,     1e-5			)




    @unittest.skip(
        reason="ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load\n    to get deterministic results."			)
    def 				lowerCAmelCase__				(			self  :       Tuple			):






       '''simple docstring'''






       pass




    @unittest.skip(
        reason="ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load\n    to get deterministic results."			)
    def 				lowerCAmelCase__				(			self  :       Union[str, Any]			):






       '''simple docstring'''






       pass




    @unittest.skip(
        reason="ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load\n    to get deterministic results."			)
    def 				lowerCAmelCase__				(			self  :       List[Any]			):






       '''simple docstring'''






       pass




    @unittest.skip(reason="ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load"			)
    def 				lowerCAmelCase__				(			self  :       Dict			):






       '''simple docstring'''






       pass




    @unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."			)
    def 				lowerCAmelCase__				(			self  :       Union[str, Any]			):






       '''simple docstring'''






       pass





    @slow
    def 				lowerCAmelCase__				(			self  :       Union[str, Any]			):






       '''simple docstring'''






       for model_name in VIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
          lowercase_	:    Tuple           =     ViTMAEModel.from_pretrained(a			)
          self.assertIsNotNone(a			)




def 		__SCREAMING_SNAKE_CASE					(  ):


   """simple docstring"""




   lowercase_	:    str           =     Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"							)
   return image





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




    @cached_property
    def 				lowerCAmelCase__				(			self  :       List[str]			):






       '''simple docstring'''






       return ViTImageProcessor.from_pretrained("facebook/vit-mae-base"			) if is_vision_available() else None





    @slow
    def 				lowerCAmelCase__				(			self  :       Any			):






       '''simple docstring'''






       np.random.seed(2			)

       lowercase_	:    Dict           =     ViTMAEForPreTraining.from_pretrained("facebook/vit-mae-base"			).to(a			)

       lowercase_	:    Optional[Any]           =     self.default_image_processor
       lowercase_	:    Union[str, Any]           =     prepare_img()
       lowercase_	:    Optional[int]           =     image_processor(images=a  ,     return_tensors="pt"			).to(a			)

       # prepare a noise vector that will be also used for testing the TF model
       # (this way we can ensure that the PT and TF models operate on the same inputs)
       lowercase_	:    Optional[int]           =     ViTMAEConfig()
       lowercase_	:    str           =     int((vit_mae_config.image_size // vit_mae_config.patch_size) ** 2			)
       lowercase_	:    Optional[Any]           =     np.random.uniform(size=(1, num_patches)			)

       # forward pass
       with torch.no_grad():
          lowercase_	:    int           =     model(**a  ,     noise=torch.from_numpy(a			).to(device=a			)			)

       # verify the logits
       lowercase_	:    int           =     torch.Size((1, 1_9_6, 7_6_8)			)
       self.assertEqual(outputs.logits.shape  ,     a			)

       lowercase_	:    Optional[Any]           =     torch.tensor(
           [[-0.0548, -1.7023, -0.9325], [0.3721, -0.5670, -0.2233], [0.8235, -1.3878, -0.3524]]			)

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












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

import fsspec

from .. import Dataset, Features, NamedSplit, config
from ..formatting import query_table
from ..packaged_modules.json.json import Json
from ..utils import logging
from ..utils.typing import NestedDataStructureLike, PathLike
from .abc import AbstractDatasetReader


class        _UpperCAmelCase						(		snake_case			):




    def __init__(			self  :       Tuple  ,     a  :       NestedDataStructureLike[PathLike]  ,     a  :       Optional[NamedSplit] = None  ,     a  :       Optional[Features] = None  ,     a  :       str = None  ,     a  :       bool = False  ,     a  :       bool = False  ,     a  :       Optional[str] = None  ,     a  :       Optional[int] = None  ,     **a  :       List[Any]  ,     ):






       '''simple docstring'''






       super().__init__(
           a  ,     split=a  ,     features=a  ,     cache_dir=a  ,     keep_in_memory=a  ,     streaming=a  ,     num_proc=a  ,     **a  ,     )
       lowercase_	:    str           =     field
       lowercase_	:    Optional[Any]           =     path_or_paths if isinstance(a  ,     a			) else {self.split: path_or_paths}
       lowercase_	:    Any           =     Json(
           cache_dir=a  ,     data_files=a  ,     features=a  ,     field=a  ,     **a  ,     )





    def 				lowerCAmelCase__				(			self  :       Optional[int]			):






       '''simple docstring'''






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

          self.builder.download_and_prepare(
              download_config=a  ,     download_mode=a  ,     verification_mode=a  ,     base_path=a  ,     num_proc=self.num_proc  ,     )
          lowercase_	:    int           =     self.builder.as_dataset(
              split=self.split  ,     verification_mode=a  ,     in_memory=self.keep_in_memory			)
       return dataset


class        _UpperCAmelCase						:




    def __init__(			self  :       str  ,     a  :       Dataset  ,     a  :       Union[PathLike, BinaryIO]  ,     a  :       Optional[int] = None  ,     a  :       Optional[int] = None  ,     **a  :       List[Any]  ,     ):






       '''simple docstring'''






       if num_proc is not None and num_proc <= 0:
          raise ValueError(f"""num_proc {num_proc} must be an integer > 0."""			)

       lowercase_	:    Dict           =     dataset
       lowercase_	:    Optional[int]           =     path_or_buf
       lowercase_	:    List[Any]           =     batch_size if batch_size else config.DEFAULT_MAX_BATCH_SIZE
       lowercase_	:    Optional[Any]           =     num_proc
       lowercase_	:    List[Any]           =     "utf-8"
       lowercase_	:    List[str]           =     to_json_kwargs




    def 				lowerCAmelCase__				(			self  :       Optional[int]			):






       '''simple docstring'''






       lowercase_	:    str           =     self.to_json_kwargs.pop("path_or_buf"  ,     a			)
       lowercase_	:    Any           =     self.to_json_kwargs.pop("orient"  ,     "records"			)
       lowercase_	:    Any           =     self.to_json_kwargs.pop("lines"  ,     True if orient == "records" else False			)
       lowercase_	:    List[str]           =     self.to_json_kwargs.pop("index"  ,     False if orient in ["split", "table"] else True			)
       lowercase_	:    int           =     self.to_json_kwargs.pop("compression"  ,     a			)

       if compression not in [None, "infer", "gzip", "bz2", "xz"]:
          raise NotImplementedError(f"""`datasets` currently does not support {compression} compression"""			)

       if isinstance(self.path_or_buf  ,     (str, bytes, os.PathLike)			):
          with fsspec.open(self.path_or_buf  ,     "wb"  ,     compression=a			) as buffer:
             lowercase_	:    Dict           =     self._write(file_obj=a  ,     orient=a  ,     lines=a  ,     index=a  ,     **self.to_json_kwargs			)
       else:
          if compression:
             raise NotImplementedError(
                 f"""The compression parameter is not supported when writing to a buffer, but compression={compression}"""
                 " was passed. Please provide a local path instead."			)
          lowercase_	:    Dict           =     self._write(
              file_obj=self.path_or_buf  ,     orient=a  ,     lines=a  ,     index=a  ,     **self.to_json_kwargs			)
       return written




    def 				lowerCAmelCase__				(			self  :       Optional[int]  ,     a  :       List[str]			):






       '''simple docstring'''






       lowercase_    ,		lowercase_    ,		lowercase_    ,		lowercase_    ,		lowercase_	:    List[Any]           =     args

       lowercase_	:    Optional[int]           =     query_table(
           table=self.dataset.data  ,     key=slice(a  ,     offset + self.batch_size			)  ,     indices=self.dataset._indices  ,     )
       lowercase_	:    Dict           =     batch.to_pandas().to_json(
           path_or_buf=a  ,     orient=a  ,     lines=a  ,     index=a  ,     **a			)
       if not json_str.endswith("\n"			):
          json_str += "\n"
       return json_str.encode(self.encoding			)





    def 				lowerCAmelCase__				(			self  :       int  ,     a  :       BinaryIO  ,     a  :       int  ,     a  :       str  ,     a  :       Union[str, Any]  ,     **a  :       str  ,     ):






       '''simple docstring'''






       lowercase_	:    Union[str, Any]           =     0

       if self.num_proc is None or self.num_proc == 1:
          for offset in logging.tqdm(
              range(0  ,     len(self.dataset			)  ,     self.batch_size			)  ,     unit="ba"  ,     disable=not logging.is_progress_bar_enabled()  ,     desc="Creating json from Arrow format"  ,     ):
             lowercase_	:    Dict           =     self._batch_json((offset, orient, lines, index, to_json_kwargs)			)
             written += file_obj.write(a			)
       else:
          lowercase_    ,		lowercase_	:    Any           =     len(self.dataset			), self.batch_size
          with multiprocessing.Pool(self.num_proc			) as pool:
             for json_str in logging.tqdm(
                 pool.imap(
                     self._batch_json  ,     [(offset, orient, lines, index, to_json_kwargs) for offset in range(0  ,     a  ,     a			)]  ,     )  ,     total=(num_rows // batch_size) + 1 if num_rows % batch_size else num_rows // batch_size  ,     unit="ba"  ,     disable=not logging.is_progress_bar_enabled()  ,     desc="Creating json from Arrow format"  ,     ):
                written += file_obj.write(a			)

       return written




