






"""simple docstring"""





import random
import unittest

import numpy as np

from transformers.testing_utils import require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin, prepare_image_inputs


if is_torch_available():
  import torch

if is_vision_available():
  import PIL

  from transformers import FlavaImageProcessor
  from transformers.image_utils import PILImageResampling
  from transformers.models.flava.image_processing_flava import (
      FLAVA_CODEBOOK_MEAN,
      FLAVA_CODEBOOK_STD,
      FLAVA_IMAGE_MEAN,
      FLAVA_IMAGE_STD,
  )
else:
  UpperCamelCase					=     UpperCamelCase					=     UpperCamelCase					=     UpperCamelCase					=     None







class 						_a  (		unittest.TestCase       ):





   '''simple docstring'''






   def __init__(       self					, UpperCAmelCase_					, UpperCAmelCase_=7					, UpperCAmelCase_=3					, UpperCAmelCase_=30					, UpperCAmelCase_=400					, UpperCAmelCase_=True					, UpperCAmelCase_=None					, UpperCAmelCase_=True					, UpperCAmelCase_=None					, UpperCAmelCase_=None					, UpperCAmelCase_=True					, UpperCAmelCase_=1 / 255					, UpperCAmelCase_=True					, UpperCAmelCase_=FLAVA_IMAGE_MEAN					, UpperCAmelCase_=FLAVA_IMAGE_STD					, UpperCAmelCase_=14					, UpperCAmelCase_=75					, UpperCAmelCase_=None					, UpperCAmelCase_=16					, UpperCAmelCase_=0.3					, UpperCAmelCase_=None					, UpperCAmelCase_=True					, UpperCAmelCase_=None					, UpperCAmelCase_=None					, UpperCAmelCase_=True					, UpperCAmelCase_=None					, UpperCAmelCase_=True					, UpperCAmelCase_=True					, UpperCAmelCase_=FLAVA_CODEBOOK_MEAN					, UpperCAmelCase_=FLAVA_CODEBOOK_STD					, )			->   Union[str, Any]:




          '''simple docstring'''




          lowercase__:							List[Any]         =     size if size is not None else {"height": 224, "width": 224}
          lowercase__:							Union[str, Any]         =     crop_size if crop_size is not None else {"height": 224, "width": 224}
          lowercase__:							int         =     codebook_size if codebook_size is not None else {"height": 112, "width": 112}
          lowercase__:							int         =     codebook_crop_size if codebook_crop_size is not None else {"height": 112, "width": 112}

          lowercase__:							List[str]         =     parent
          lowercase__:							int         =     batch_size
          lowercase__:							List[str]         =     num_channels
          lowercase__:							Any         =     do_resize
          lowercase__:							int         =     do_rescale
          lowercase__:							Optional[int]         =     rescale_factor
          lowercase__:							int         =     min_resolution
          lowercase__:							int         =     max_resolution
          lowercase__:							str         =     size
          lowercase__:							Union[str, Any]         =     resample if resample is not None else PILImageResampling.BICUBIC
          lowercase__:							Dict         =     do_normalize
          lowercase__:							Any         =     image_mean
          lowercase__:							Any         =     image_std
          lowercase__:							Optional[Any]         =     do_center_crop
          lowercase__:							Any         =     crop_size

          lowercase__:							List[str]         =     input_size_patches
          lowercase__:							List[Any]         =     total_mask_patches
          lowercase__:							Dict         =     mask_group_max_patches
          lowercase__:							Optional[int]         =     mask_group_min_patches
          lowercase__:							str         =     mask_group_min_aspect_ratio
          lowercase__:							Any         =     mask_group_max_aspect_ratio

          lowercase__:							List[Any]         =     codebook_do_resize
          lowercase__:							str         =     codebook_size
          lowercase__:							Dict         =     codebook_resample if codebook_resample is not None else PILImageResampling.LANCZOS
          lowercase__:							Optional[Any]         =     codebook_do_center_crop
          lowercase__:							Tuple         =     codebook_crop_size
          lowercase__:							List[str]         =     codebook_do_map_pixels
          lowercase__:							List[Any]         =     codebook_do_normalize
          lowercase__:							List[Any]         =     codebook_image_mean
          lowercase__:							Any         =     codebook_image_std





   def __lowercase							(       self)			->   Any:




          '''simple docstring'''




          return {
              "image_mean": self.image_mean,
              "image_std": self.image_std,
              "do_normalize": self.do_normalize,
              "do_resize": self.do_resize,
              "size": self.size,
              "resample": self.resample,
              "do_rescale": self.do_rescale,
              "rescale_factor": self.rescale_factor,
              "do_center_crop": self.do_center_crop,
              "crop_size": self.crop_size,
              "input_size_patches": self.input_size_patches,
              "total_mask_patches": self.total_mask_patches,
              "mask_group_max_patches": self.mask_group_max_patches,
              "mask_group_min_patches": self.mask_group_min_patches,
              "mask_group_min_aspect_ratio": self.mask_group_min_aspect_ratio,
              "mask_group_max_aspect_ratio": self.mask_group_min_aspect_ratio,
              "codebook_do_resize": self.codebook_do_resize,
              "codebook_size": self.codebook_size,
              "codebook_resample": self.codebook_resample,
              "codebook_do_center_crop": self.codebook_do_center_crop,
              "codebook_crop_size": self.codebook_crop_size,
              "codebook_do_map_pixels": self.codebook_do_map_pixels,
              "codebook_do_normalize": self.codebook_do_normalize,
              "codebook_image_mean": self.codebook_image_mean,
              "codebook_image_std": self.codebook_image_std,
          }





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




          '''simple docstring'''




          return (self.size["height"], self.size["width"])





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




          '''simple docstring'''




          return (
              (self.input_size_patches, self.input_size_patches)
              if not isinstance(self.input_size_patches					, UpperCAmelCase_)
              else self.input_size_patches
          )





   def __lowercase							(       self)			->   Dict:




          '''simple docstring'''




          return (self.codebook_size["height"], self.codebook_size["width"])







@require_torch
@require_vision
class 						_a  (		lowercase_							,	unittest.TestCase       ):





   '''simple docstring'''






   UpperCamelCase__     	=							FlavaImageProcessor if is_vision_available() else None
   UpperCamelCase__     	=							None





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




          '''simple docstring'''




          lowercase__:							int         =     FlavaImageProcessingTester(self)





   @property
   def __lowercase							(       self)			->   Union[str, Any]:




          '''simple docstring'''




          return self.image_processor_tester.prepare_image_processor_dict()





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




          '''simple docstring'''




          lowercase__:							Optional[Any]         =     self.image_processing_class(**self.image_processor_dict)
          self.assertTrue(hasattr(UpperCAmelCase_					, "image_mean"))
          self.assertTrue(hasattr(UpperCAmelCase_					, "image_std"))
          self.assertTrue(hasattr(UpperCAmelCase_					, "do_normalize"))
          self.assertTrue(hasattr(UpperCAmelCase_					, "do_resize"))
          self.assertTrue(hasattr(UpperCAmelCase_					, "resample"))
          self.assertTrue(hasattr(UpperCAmelCase_					, "crop_size"))
          self.assertTrue(hasattr(UpperCAmelCase_					, "do_center_crop"))
          self.assertTrue(hasattr(UpperCAmelCase_					, "do_rescale"))
          self.assertTrue(hasattr(UpperCAmelCase_					, "rescale_factor"))
          self.assertTrue(hasattr(UpperCAmelCase_					, "masking_generator"))
          self.assertTrue(hasattr(UpperCAmelCase_					, "codebook_do_resize"))
          self.assertTrue(hasattr(UpperCAmelCase_					, "codebook_size"))
          self.assertTrue(hasattr(UpperCAmelCase_					, "codebook_resample"))
          self.assertTrue(hasattr(UpperCAmelCase_					, "codebook_do_center_crop"))
          self.assertTrue(hasattr(UpperCAmelCase_					, "codebook_crop_size"))
          self.assertTrue(hasattr(UpperCAmelCase_					, "codebook_do_map_pixels"))
          self.assertTrue(hasattr(UpperCAmelCase_					, "codebook_do_normalize"))
          self.assertTrue(hasattr(UpperCAmelCase_					, "codebook_image_mean"))
          self.assertTrue(hasattr(UpperCAmelCase_					, "codebook_image_std"))





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




          '''simple docstring'''




          lowercase__:							Optional[int]         =     self.image_processing_class.from_dict(self.image_processor_dict)
          self.assertEqual(image_processor.size					, {"height": 224, "width": 224})
          self.assertEqual(image_processor.crop_size					, {"height": 224, "width": 224})
          self.assertEqual(image_processor.codebook_size					, {"height": 112, "width": 112})
          self.assertEqual(image_processor.codebook_crop_size					, {"height": 112, "width": 112})

          lowercase__:							Dict         =     self.image_processing_class.from_dict(
              self.image_processor_dict					, size=42					, crop_size=84					, codebook_size=33					, codebook_crop_size=66)
          self.assertEqual(image_processor.size					, {"height": 42, "width": 42})
          self.assertEqual(image_processor.crop_size					, {"height": 84, "width": 84})
          self.assertEqual(image_processor.codebook_size					, {"height": 33, "width": 33})
          self.assertEqual(image_processor.codebook_crop_size					, {"height": 66, "width": 66})





   def __lowercase							(       self)			->   Tuple:




          '''simple docstring'''




          pass





   def __lowercase							(       self)			->   int:




          '''simple docstring'''




          lowercase__:							Dict         =     self.image_processing_class(**self.image_processor_dict)
          # create random PIL images
          lowercase__:							Any         =     prepare_image_inputs(self.image_processor_tester					, equal_resolution=UpperCAmelCase_)
          for image in image_inputs:
                 self.assertIsInstance(UpperCAmelCase_					, PIL.Image.Image)

          # Test not batched input
          lowercase__:							Tuple         =     image_processing(image_inputs[0]					, return_tensors="pt")

          # Test no bool masked pos
          self.assertFalse("bool_masked_pos" in encoded_images)

          lowercase__    ,      lowercase__:							Optional[Any]         =     self.image_processor_tester.get_expected_image_size()

          self.assertEqual(
              encoded_images.pixel_values.shape					, (1, self.image_processor_tester.num_channels, expected_height, expected_width)					, )

          # Test batched
          lowercase__:							List[Any]         =     image_processing(UpperCAmelCase_					, return_tensors="pt")
          lowercase__    ,      lowercase__:							Optional[Any]         =     self.image_processor_tester.get_expected_image_size()

          # Test no bool masked pos
          self.assertFalse("bool_masked_pos" in encoded_images)

          self.assertEqual(
              encoded_images.pixel_values.shape					, (
                  self.image_processor_tester.batch_size,
                  self.image_processor_tester.num_channels,
                  expected_height,
                  expected_width,
              )					, )





   def __lowercase							(       self					, UpperCAmelCase_					, UpperCAmelCase_)			->   int:




          '''simple docstring'''




          lowercase__:							Dict         =     self.image_processing_class(**self.image_processor_dict)
          # create random tensors
          lowercase__:							Any         =     prepare_image_inputs(self.image_processor_tester					, equal_resolution=UpperCAmelCase_					, **UpperCAmelCase_)
          for image in image_inputs:
                 self.assertIsInstance(UpperCAmelCase_					, UpperCAmelCase_)

          # Test not batched input
          lowercase__:							Any         =     image_processing(image_inputs[0]					, return_tensors="pt")

          lowercase__    ,      lowercase__:							Optional[int]         =     self.image_processor_tester.get_expected_image_size()
          self.assertEqual(
              encoded_images.pixel_values.shape					, (1, self.image_processor_tester.num_channels, expected_height, expected_width)					, )

          lowercase__:							List[str]         =     image_processing(UpperCAmelCase_					, return_image_mask=UpperCAmelCase_					, return_tensors="pt")

          lowercase__    ,      lowercase__:							int         =     self.image_processor_tester.get_expected_image_size()
          self.assertEqual(
              encoded_images.pixel_values.shape					, (
                  self.image_processor_tester.batch_size,
                  self.image_processor_tester.num_channels,
                  expected_height,
                  expected_width,
              )					, )

          lowercase__    ,      lowercase__:							Union[str, Any]         =     self.image_processor_tester.get_expected_mask_size()
          self.assertEqual(
              encoded_images.bool_masked_pos.shape					, (
                  self.image_processor_tester.batch_size,
                  expected_height,
                  expected_width,
              )					, )

          # Test batched
          lowercase__:							Any         =     image_processing(UpperCAmelCase_					, return_tensors="pt").pixel_values

          lowercase__    ,      lowercase__:							Any         =     self.image_processor_tester.get_expected_image_size()
          self.assertEqual(
              encoded_images.shape					, (
                  self.image_processor_tester.batch_size,
                  self.image_processor_tester.num_channels,
                  expected_height,
                  expected_width,
              )					, )

          # Test masking
          lowercase__:							str         =     image_processing(UpperCAmelCase_					, return_image_mask=UpperCAmelCase_					, return_tensors="pt")

          lowercase__    ,      lowercase__:							Dict         =     self.image_processor_tester.get_expected_image_size()
          self.assertEqual(
              encoded_images.pixel_values.shape					, (
                  self.image_processor_tester.batch_size,
                  self.image_processor_tester.num_channels,
                  expected_height,
                  expected_width,
              )					, )

          lowercase__    ,      lowercase__:							str         =     self.image_processor_tester.get_expected_mask_size()
          self.assertEqual(
              encoded_images.bool_masked_pos.shape					, (
                  self.image_processor_tester.batch_size,
                  expected_height,
                  expected_width,
              )					, )





   def __lowercase							(       self)			->   Dict:




          '''simple docstring'''




          self._test_call_framework(np.ndarray					, prepare_kwargs={"numpify": True})





   def __lowercase							(       self)			->   int:




          '''simple docstring'''




          self._test_call_framework(torch.Tensor					, prepare_kwargs={"torchify": True})





   def __lowercase							(       self)			->   str:




          '''simple docstring'''




          random.seed(1_234)
          lowercase__:							List[Any]         =     self.image_processing_class(**self.image_processor_dict)
          lowercase__:							List[str]         =     prepare_image_inputs(self.image_processor_tester					, equal_resolution=UpperCAmelCase_					, torchify=UpperCAmelCase_)

          # Test not batched input
          lowercase__:							List[Any]         =     image_processing(image_inputs[0]					, return_image_mask=UpperCAmelCase_					, return_tensors="pt")
          self.assertEqual(encoded_images.bool_masked_pos.sum().item()					, 75)





   def __lowercase							(       self)			->   Optional[int]:




          '''simple docstring'''




          lowercase__:							Any         =     self.image_processing_class(**self.image_processor_dict)
          # create random PIL images
          lowercase__:							Union[str, Any]         =     prepare_image_inputs(self.image_processor_tester					, equal_resolution=UpperCAmelCase_)
          for image in image_inputs:
                 self.assertIsInstance(UpperCAmelCase_					, PIL.Image.Image)

          # Test not batched input
          lowercase__:							int         =     image_processing(image_inputs[0]					, return_codebook_pixels=UpperCAmelCase_					, return_tensors="pt")
          lowercase__    ,      lowercase__:							Dict         =     self.image_processor_tester.get_expected_codebook_image_size()
          self.assertEqual(
              encoded_images.codebook_pixel_values.shape					, (1, self.image_processor_tester.num_channels, expected_height, expected_width)					, )

          # Test batched
          lowercase__:							Union[str, Any]         =     image_processing(UpperCAmelCase_					, return_codebook_pixels=UpperCAmelCase_					, return_tensors="pt")
          lowercase__    ,      lowercase__:							int         =     self.image_processor_tester.get_expected_codebook_image_size()
          self.assertEqual(
              encoded_images.codebook_pixel_values.shape					, (
                  self.image_processor_tester.batch_size,
                  self.image_processor_tester.num_channels,
                  expected_height,
                  expected_width,
              )					, )














'''simple docstring'''





import os
import sys
from pathlib import Path

import pytest

from datasets import Dataset, IterableDataset
from datasets.distributed import split_dataset_by_node

from .utils import execute_subprocess_async, get_torch_dist_unique_port, require_torch






def  lowerCamelCase_  (     ):
 lowerCamelCase__	=  Dataset.from_dict({"i": range(17)})
 lowerCamelCase__	=  len(lowercase__)
 lowerCamelCase__	=  3
 lowerCamelCase__	=  [
     split_dataset_by_node(lowercase__					,  rank=lowercase__					,  world_size=lowercase__) for rank in range(lowercase__)
 ]
 assert sum(len(lowercase__) for ds in datasets_per_rank) == full_size
 assert len({tuple(x.values()) for ds in datasets_per_rank for x in ds}) == full_size






def  lowerCamelCase_  (     ):
 def gen():
  return ({"i": i} for i in range(17))

 lowerCamelCase__	=  3
 lowerCamelCase__	=  IterableDataset.from_generator(lowercase__)
 lowerCamelCase__	=  len(list(lowercase__))
 lowerCamelCase__	=  [
     split_dataset_by_node(lowercase__					,  rank=lowercase__					,  world_size=lowercase__) for rank in range(lowercase__)
 ]
 assert sum(len(list(lowercase__)) for ds in datasets_per_rank) == full_size
 assert len({tuple(x.values()) for ds in datasets_per_rank for x in ds}) == full_size






@pytest.mark.parametrize("shards_per_node"					,  [1, 2, 3])
def  lowerCamelCase_  (     lowercase__):
 def gen(lowercase__):
  for shard in shards:
   yield from ({"i": i, "shard": shard} for i in range(17))

 lowerCamelCase__	=  3
 lowerCamelCase__	=  shards_per_node * world_size
 lowerCamelCase__	=  {"shards": [F'''shard_{shard_idx}.txt''' for shard_idx in range(lowercase__)]}
 lowerCamelCase__	=  IterableDataset.from_generator(lowercase__					,  gen_kwargs=lowercase__)
 lowerCamelCase__	=  len(list(lowercase__))
 assert full_ds.n_shards == world_size * shards_per_node
 lowerCamelCase__	=  [
     split_dataset_by_node(lowercase__					,  rank=lowercase__					,  world_size=lowercase__) for rank in range(lowercase__)
 ]
 assert [ds.n_shards for ds in datasets_per_rank] == [shards_per_node] * world_size
 assert sum(len(list(lowercase__)) for ds in datasets_per_rank) == full_size
 assert len({tuple(x.values()) for ds in datasets_per_rank for x in ds}) == full_size






def  lowerCamelCase_  (     ):
 def gen():
  return ({"i": i} for i in range(17))

 lowerCamelCase__	=  2
 lowerCamelCase__	=  IterableDataset.from_generator(lowercase__)
 lowerCamelCase__	=  len(list(lowercase__))

 lowerCamelCase__	=  split_dataset_by_node(lowercase__					,  rank=0					,  world_size=lowercase__).shuffle(seed=42)
 assert len(list(lowercase__)) == 1 + full_size // world_size
 with pytest.raises(lowercase__):
  split_dataset_by_node(lowercase__					,  rank=0					,  world_size=lowercase__).shuffle()

 lowerCamelCase__	=  split_dataset_by_node(full_ds.shuffle(seed=42)					,  rank=0					,  world_size=lowercase__)
 assert len(list(lowercase__)) == 1 + full_size // world_size
 with pytest.raises(lowercase__):
  split_dataset_by_node(full_ds.shuffle()					,  rank=0					,  world_size=lowercase__)






@pytest.mark.parametrize("streaming"					,  [False, True])
@require_torch
@pytest.mark.skipif(os.name == "nt"					,  reason="execute_subprocess_async doesn't support windows")
@pytest.mark.integration
def  lowerCamelCase_  (     lowercase__):
 lowerCamelCase__	=  2
 lowerCamelCase__	=  get_torch_dist_unique_port()
 lowerCamelCase__	=  Path(lowercase__).resolve().parent / "distributed_scripts" / "launch_torch_distributed.py"
 lowerCamelCase__	=  F'''
        -m torch.distributed.launch
        --nproc_per_node={nproc_per_node}
        --master_port={master_port}
        {test_script}
    '''.split()
 lowerCamelCase__	=  F'''
        --streaming={streaming}
    '''.split()
 lowerCamelCase__	=  [sys.executable] + distributed_args + args
 execute_subprocess_async(lowercase__					,  env=os.environ.copy())






@pytest.mark.parametrize(
    "nproc_per_node, num_workers"					,  [
        (2, 2),  # each node has 2 shards and each worker has 1 shards
        (3, 2),  # each node uses all the shards but skips examples, and each worker has 2 shards
    ]					,  )
@require_torch
@pytest.mark.skipif(os.name == "nt"					,  reason="execute_subprocess_async doesn't support windows")
@pytest.mark.integration
def  lowerCamelCase_  (     lowercase__					,  lowercase__):
 lowerCamelCase__	=  True
 lowerCamelCase__	=  get_torch_dist_unique_port()
 lowerCamelCase__	=  Path(lowercase__).resolve().parent / "distributed_scripts" / "launch_torch_distributed.py"
 lowerCamelCase__	=  F'''
        -m torch.distributed.launch
        --nproc_per_node={nproc_per_node}
        --master_port={master_port}
        {test_script}
    '''.split()
 lowerCamelCase__	=  F'''
        --streaming={streaming}
        --num_workers={num_workers}
    '''.split()
 lowerCamelCase__	=  [sys.executable] + distributed_args + args
 execute_subprocess_async(lowercase__					,  env=os.environ.copy())

