




"""simple docstring"""


from unittest.mock import patch

import pyspark

from datasets.packaged_modules.spark.spark import (
    Spark,
    SparkExamplesIterable,
    _generate_iterable_examples,
)

from ..utils import (
    require_dill_gt_0_3_2,
    require_not_windows,
)


def 					UpperCAmelCase__       (snake_case__							:					Any		,   snake_case__							:					str    ):
  """simple docstring"""

  _snake_case							:		Optional[Any]								=  []
  for part_id in partition_order:
    _snake_case							:		str								=  df.where(F"SPARK_PARTITION_ID() = {part_id}"    ).collect()
    for row_idx, row in enumerate(snake_case__    ):
      expected_row_ids_and_row_dicts.append((F"{part_id}_{row_idx}", row.asDict())    )
  return expected_row_ids_and_row_dicts


@require_not_windows
@require_dill_gt_0_3_2
def 					UpperCAmelCase__       ():
  """simple docstring"""

  _snake_case							:		List[str]								=  pyspark.sql.SparkSession.builder.master("""local[*]"""    ).appName("""pyspark"""    ).getOrCreate()
  _snake_case							:		Optional[Any]								=  spark.range(1_00    ).repartition(1    )
  _snake_case							:		Tuple								=  Spark(snake_case__    )
  # The id ints will be converted to Pyarrow int64s, so each row will be 8 bytes. Setting a max_shard_size of 16 means
  # that each partition can hold 2 rows.
  spark_builder._repartition_df_if_needed(max_shard_size=16    )
  # Given that the dataframe has 100 rows and each partition has 2 rows, we expect 50 partitions.
  assert spark_builder.df.rdd.getNumPartitions() == 50


@require_not_windows
@require_dill_gt_0_3_2
def 					UpperCAmelCase__       ():
  """simple docstring"""

  _snake_case							:		Tuple								=  pyspark.sql.SparkSession.builder.master("""local[*]"""    ).appName("""pyspark"""    ).getOrCreate()
  _snake_case							:		int								=  spark.range(10    ).repartition(2    )
  _snake_case							:		Tuple								=  [1, 0]
  _snake_case							:		Dict								=  _generate_iterable_examples(snake_case__		,   snake_case__    )  # Reverse the partitions.
  _snake_case							:		Optional[int]								=  _get_expected_row_ids_and_row_dicts_for_partition_order(snake_case__		,   snake_case__    )

  for i, (row_id, row_dict) in enumerate(generate_fn()    ):
    _snake_case		,					_snake_case							:		Optional[Any]								=  expected_row_ids_and_row_dicts[i]
    assert row_id == expected_row_id
    assert row_dict == expected_row_dict


@require_not_windows
@require_dill_gt_0_3_2
def 					UpperCAmelCase__       ():
  """simple docstring"""

  _snake_case							:		int								=  pyspark.sql.SparkSession.builder.master("""local[*]"""    ).appName("""pyspark"""    ).getOrCreate()
  _snake_case							:		Any								=  spark.range(10    ).repartition(1    )
  _snake_case							:		List[str]								=  SparkExamplesIterable(snake_case__    )
  assert it.n_shards == 1
  for i, (row_id, row_dict) in enumerate(snake_case__    ):
    assert row_id == F"0_{i}"
    assert row_dict == {"id": i}


@require_not_windows
@require_dill_gt_0_3_2
def 					UpperCAmelCase__       ():
  """simple docstring"""

  _snake_case							:		Tuple								=  pyspark.sql.SparkSession.builder.master("""local[*]"""    ).appName("""pyspark"""    ).getOrCreate()
  _snake_case							:		int								=  spark.range(30    ).repartition(3    )
  # Mock the generator so that shuffle reverses the partition indices.
  with patch("""numpy.random.Generator"""    ) as generator_mock:
    _snake_case							:		Tuple								=  lambda snake_case__    : x.reverse()
    _snake_case							:		List[str]								=  _get_expected_row_ids_and_row_dicts_for_partition_order(snake_case__		,   [2, 1, 0]    )

    _snake_case							:		int								=  SparkExamplesIterable(snake_case__    ).shuffle_data_sources(snake_case__    )
    assert shuffled_it.n_shards == 3
    for i, (row_id, row_dict) in enumerate(snake_case__    ):
      _snake_case		,					_snake_case							:		List[Any]								=  expected_row_ids_and_row_dicts[i]
      assert row_id == expected_row_id
      assert row_dict == expected_row_dict


@require_not_windows
@require_dill_gt_0_3_2
def 					UpperCAmelCase__       ():
  """simple docstring"""

  _snake_case							:		int								=  pyspark.sql.SparkSession.builder.master("""local[*]"""    ).appName("""pyspark"""    ).getOrCreate()
  _snake_case							:		Dict								=  spark.range(20    ).repartition(4    )

  # Partitions 0 and 2
  _snake_case							:		List[Any]								=  SparkExamplesIterable(snake_case__    ).shard_data_sources(worker_id=0		,   num_workers=2    )
  assert shard_it_a.n_shards == 2
  _snake_case							:		Any								=  _get_expected_row_ids_and_row_dicts_for_partition_order(snake_case__		,   [0, 2]    )
  for i, (row_id, row_dict) in enumerate(snake_case__    ):
    _snake_case		,					_snake_case							:		List[str]								=  expected_row_ids_and_row_dicts_a[i]
    assert row_id == expected_row_id
    assert row_dict == expected_row_dict

  # Partitions 1 and 3
  _snake_case							:		List[str]								=  SparkExamplesIterable(snake_case__    ).shard_data_sources(worker_id=1		,   num_workers=2    )
  assert shard_it_a.n_shards == 2
  _snake_case							:		Union[str, Any]								=  _get_expected_row_ids_and_row_dicts_for_partition_order(snake_case__		,   [1, 3]    )
  for i, (row_id, row_dict) in enumerate(snake_case__    ):
    _snake_case		,					_snake_case							:		Any								=  expected_row_ids_and_row_dicts_a[i]
    assert row_id == expected_row_id
    assert row_dict == expected_row_dict






@require_not_windows
@require_dill_gt_0_3_2
def 					UpperCAmelCase__       ():
  """simple docstring"""

  _snake_case							:		str								=  pyspark.sql.SparkSession.builder.master("""local[*]"""    ).appName("""pyspark"""    ).getOrCreate()
  _snake_case							:		Dict								=  spark.range(1_00    ).repartition(1    )
  _snake_case							:		List[Any]								=  Spark(snake_case__    )
  # Choose a small max_shard_size for maximum partitioning.
  spark_builder._repartition_df_if_needed(max_shard_size=1    )
  # The new number of partitions should not be greater than the number of rows.
  assert spark_builder.df.rdd.getNumPartitions() == 1_00





'''simple docstring'''



import json
import os
import shutil
import tempfile
import unittest

import numpy as np
import pytest

from transformers import CLIPTokenizer, CLIPTokenizerFast
from transformers.models.clip.tokenization_clip import VOCAB_FILES_NAMES
from transformers.testing_utils import require_vision
from transformers.utils import IMAGE_PROCESSOR_NAME, is_vision_available


if is_vision_available():
  from PIL import Image

  from transformers import CLIPImageProcessor, CLIPProcessor







@require_vision
class 						lowerCAmelCase__   (      unittest.TestCase					):



 """simple docstring"""



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


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      tempfile.mkdtemp()

      # fmt: off
      __SCREAMING_SNAKE_CASE     =      ["""l""", """o""", """w""", """e""", """r""", """s""", """t""", """i""", """d""", """n""", """lo""", """l</w>""", """w</w>""", """r</w>""", """t</w>""", """low</w>""", """er</w>""", """lowest</w>""", """newer</w>""", """wider""", """<unk>""", """<|startoftext|>""", """<|endoftext|>"""]
      # fmt: on
      __SCREAMING_SNAKE_CASE     =      dict(zip(__SCREAMING_SNAKE_CASE					,							range(len(__SCREAMING_SNAKE_CASE			)			)			)			)
      __SCREAMING_SNAKE_CASE     =      ["""#version: 0.2""", """l o""", """lo w</w>""", """e r</w>""", """"""]
      __SCREAMING_SNAKE_CASE     =      {"""unk_token""": """<unk>"""}

      __SCREAMING_SNAKE_CASE     =      os.path.join(self.tmpdirname					,							VOCAB_FILES_NAMES["""vocab_file"""]			)
      __SCREAMING_SNAKE_CASE     =      os.path.join(self.tmpdirname					,							VOCAB_FILES_NAMES["""merges_file"""]			)
      with open(self.vocab_file					,							"""w"""					,							encoding="""utf-8"""			) as fp:
           fp.write(json.dumps(__SCREAMING_SNAKE_CASE			) + """\n"""			)
      with open(self.merges_file					,							"""w"""					,							encoding="""utf-8"""			) as fp:
           fp.write("""\n""".join(__SCREAMING_SNAKE_CASE			)			)

      __SCREAMING_SNAKE_CASE     =      {
          """do_resize""": True,
          """size""": 20,
          """do_center_crop""": True,
          """crop_size""": 18,
          """do_normalize""": True,
          """image_mean""": [0.48145466, 0.4578275, 0.40821073],
          """image_std""": [0.26862954, 0.26130258, 0.27577711],
      }
      __SCREAMING_SNAKE_CASE     =      os.path.join(self.tmpdirname					,							__SCREAMING_SNAKE_CASE			)
      with open(self.image_processor_file					,							"""w"""					,							encoding="""utf-8"""			) as fp:
           json.dump(__SCREAMING_SNAKE_CASE					,							__SCREAMING_SNAKE_CASE			)
 def       UpperCAmelCase__		(		self					:       Optional[int]					,							**__SCREAMING_SNAKE_CASE					:       Optional[int]			)  -> str:


      """simple docstring"""




      return CLIPTokenizer.from_pretrained(self.tmpdirname					,							**__SCREAMING_SNAKE_CASE			)
 def       UpperCAmelCase__		(		self					:       Tuple					,							**__SCREAMING_SNAKE_CASE					:       Any			)  -> int:


      """simple docstring"""




      return CLIPTokenizerFast.from_pretrained(self.tmpdirname					,							**__SCREAMING_SNAKE_CASE			)
 def       UpperCAmelCase__		(		self					:       Optional[Any]					,							**__SCREAMING_SNAKE_CASE					:       Optional[int]			)  -> List[str]:


      """simple docstring"""




      return CLIPImageProcessor.from_pretrained(self.tmpdirname					,							**__SCREAMING_SNAKE_CASE			)
 def       UpperCAmelCase__		(		self					:       Tuple			)  -> str:


      """simple docstring"""




      shutil.rmtree(self.tmpdirname			)
 def       UpperCAmelCase__		(		self					:       Optional[Any]			)  -> str:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      [np.random.randint(255					,							size=(3, 30, 400)					,							dtype=np.uinta			)]

      __SCREAMING_SNAKE_CASE     =      [Image.fromarray(np.moveaxis(__SCREAMING_SNAKE_CASE					,							0					,							-1			)			) for x in image_inputs]

      return image_inputs
 def       UpperCAmelCase__		(		self					:       Optional[int]			)  -> Any:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      self.get_tokenizer()
      __SCREAMING_SNAKE_CASE     =      self.get_rust_tokenizer()
      __SCREAMING_SNAKE_CASE     =      self.get_image_processor()

      __SCREAMING_SNAKE_CASE     =      CLIPProcessor(tokenizer=__SCREAMING_SNAKE_CASE					,							image_processor=__SCREAMING_SNAKE_CASE			)
      processor_slow.save_pretrained(self.tmpdirname			)
      __SCREAMING_SNAKE_CASE     =      CLIPProcessor.from_pretrained(self.tmpdirname					,							use_fast=__SCREAMING_SNAKE_CASE			)

      __SCREAMING_SNAKE_CASE     =      CLIPProcessor(tokenizer=__SCREAMING_SNAKE_CASE					,							image_processor=__SCREAMING_SNAKE_CASE			)
      processor_fast.save_pretrained(self.tmpdirname			)
      __SCREAMING_SNAKE_CASE     =      CLIPProcessor.from_pretrained(self.tmpdirname			)

      self.assertEqual(processor_slow.tokenizer.get_vocab()					,							tokenizer_slow.get_vocab()			)
      self.assertEqual(processor_fast.tokenizer.get_vocab()					,							tokenizer_fast.get_vocab()			)
      self.assertEqual(tokenizer_slow.get_vocab()					,							tokenizer_fast.get_vocab()			)
      self.assertIsInstance(processor_slow.tokenizer					,							__SCREAMING_SNAKE_CASE			)
      self.assertIsInstance(processor_fast.tokenizer					,							__SCREAMING_SNAKE_CASE			)

      self.assertEqual(processor_slow.image_processor.to_json_string()					,							image_processor.to_json_string()			)
      self.assertEqual(processor_fast.image_processor.to_json_string()					,							image_processor.to_json_string()			)
      self.assertIsInstance(processor_slow.image_processor					,							__SCREAMING_SNAKE_CASE			)
      self.assertIsInstance(processor_fast.image_processor					,							__SCREAMING_SNAKE_CASE			)
 def       UpperCAmelCase__		(		self					:       Optional[int]			)  -> Dict:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      CLIPProcessor(tokenizer=self.get_tokenizer()					,							image_processor=self.get_image_processor()			)
      processor.save_pretrained(self.tmpdirname			)

      __SCREAMING_SNAKE_CASE     =      self.get_tokenizer(bos_token="""(BOS)"""					,							eos_token="""(EOS)"""			)
      __SCREAMING_SNAKE_CASE     =      self.get_image_processor(do_normalize=__SCREAMING_SNAKE_CASE					,							padding_value=1.0			)

      __SCREAMING_SNAKE_CASE     =      CLIPProcessor.from_pretrained(
          self.tmpdirname					,							bos_token="""(BOS)"""					,							eos_token="""(EOS)"""					,							do_normalize=__SCREAMING_SNAKE_CASE					,							padding_value=1.0			)

      self.assertEqual(processor.tokenizer.get_vocab()					,							tokenizer_add_kwargs.get_vocab()			)
      self.assertIsInstance(processor.tokenizer					,							__SCREAMING_SNAKE_CASE			)

      self.assertEqual(processor.image_processor.to_json_string()					,							image_processor_add_kwargs.to_json_string()			)
      self.assertIsInstance(processor.image_processor					,							__SCREAMING_SNAKE_CASE			)
 def       UpperCAmelCase__		(		self					:       List[str]			)  -> Dict:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      self.get_image_processor()
      __SCREAMING_SNAKE_CASE     =      self.get_tokenizer()

      __SCREAMING_SNAKE_CASE     =      CLIPProcessor(tokenizer=__SCREAMING_SNAKE_CASE					,							image_processor=__SCREAMING_SNAKE_CASE			)

      __SCREAMING_SNAKE_CASE     =      self.prepare_image_inputs()

      __SCREAMING_SNAKE_CASE     =      image_processor(__SCREAMING_SNAKE_CASE					,							return_tensors="""np"""			)
      __SCREAMING_SNAKE_CASE     =      processor(images=__SCREAMING_SNAKE_CASE					,							return_tensors="""np"""			)

      for key in input_image_proc.keys():
           self.assertAlmostEqual(input_image_proc[key].sum()					,							input_processor[key].sum()					,							delta=1E-2			)
 def       UpperCAmelCase__		(		self					:       List[Any]			)  -> int:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      self.get_image_processor()
      __SCREAMING_SNAKE_CASE     =      self.get_tokenizer()

      __SCREAMING_SNAKE_CASE     =      CLIPProcessor(tokenizer=__SCREAMING_SNAKE_CASE					,							image_processor=__SCREAMING_SNAKE_CASE			)

      __SCREAMING_SNAKE_CASE     =      """lower newer"""

      __SCREAMING_SNAKE_CASE     =      processor(text=__SCREAMING_SNAKE_CASE			)

      __SCREAMING_SNAKE_CASE     =      tokenizer(__SCREAMING_SNAKE_CASE			)

      for key in encoded_tok.keys():
           self.assertListEqual(encoded_tok[key]					,							encoded_processor[key]			)
 def       UpperCAmelCase__		(		self					:       Dict			)  -> Union[str, Any]:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      self.get_image_processor()
      __SCREAMING_SNAKE_CASE     =      self.get_tokenizer()

      __SCREAMING_SNAKE_CASE     =      CLIPProcessor(tokenizer=__SCREAMING_SNAKE_CASE					,							image_processor=__SCREAMING_SNAKE_CASE			)

      __SCREAMING_SNAKE_CASE     =      """lower newer"""
      __SCREAMING_SNAKE_CASE     =      self.prepare_image_inputs()

      __SCREAMING_SNAKE_CASE     =      processor(text=__SCREAMING_SNAKE_CASE					,							images=__SCREAMING_SNAKE_CASE			)

      self.assertListEqual(list(inputs.keys()			)					,							["""input_ids""", """attention_mask""", """pixel_values"""]			)

      # test if it raises when no input is passed
      with pytest.raises(__SCREAMING_SNAKE_CASE			):
           processor()
 def       UpperCAmelCase__		(		self					:       List[str]			)  -> Union[str, Any]:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      self.get_image_processor()
      __SCREAMING_SNAKE_CASE     =      self.get_tokenizer()

      __SCREAMING_SNAKE_CASE     =      CLIPProcessor(tokenizer=__SCREAMING_SNAKE_CASE					,							image_processor=__SCREAMING_SNAKE_CASE			)

      __SCREAMING_SNAKE_CASE     =      [[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

      __SCREAMING_SNAKE_CASE     =      processor.batch_decode(__SCREAMING_SNAKE_CASE			)
      __SCREAMING_SNAKE_CASE     =      tokenizer.batch_decode(__SCREAMING_SNAKE_CASE			)

      self.assertListEqual(__SCREAMING_SNAKE_CASE					,							__SCREAMING_SNAKE_CASE			)


 def       UpperCAmelCase__		(		self					:       int			)  -> str:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE     =      self.get_image_processor()
      __SCREAMING_SNAKE_CASE     =      self.get_tokenizer()

      __SCREAMING_SNAKE_CASE     =      CLIPProcessor(tokenizer=__SCREAMING_SNAKE_CASE					,							image_processor=__SCREAMING_SNAKE_CASE			)

      __SCREAMING_SNAKE_CASE     =      """lower newer"""
      __SCREAMING_SNAKE_CASE     =      self.prepare_image_inputs()

      __SCREAMING_SNAKE_CASE     =      processor(text=__SCREAMING_SNAKE_CASE					,							images=__SCREAMING_SNAKE_CASE			)

      self.assertListEqual(list(inputs.keys()			)					,							processor.model_input_names			)






