






'''simple docstring'''




import warnings
from typing import List

import numpy as np

from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding
from ...utils import is_flax_available, is_tf_available, is_torch_available

class    _UpperCAmelCase			(     snake_case_						):




   """simple docstring"""

   snake_case    				=						['''image_processor''', '''tokenizer''']
   snake_case    				=						'''OwlViTImageProcessor'''
   snake_case    				=						('''CLIPTokenizer''', '''CLIPTokenizerFast''')

   def __init__(		self :	List[Any]      ,      __UpperCAmelCase :	int=None      ,      __UpperCAmelCase :	str=None      ,      **__UpperCAmelCase :	Tuple		):





         '''simple docstring'''





         _A			     =	None
         if "feature_extractor" in kwargs:
               warnings.warn(
                   "The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`"
                   " instead."      ,      __UpperCAmelCase      ,      )
               _A			     =	kwargs.pop("feature_extractor"		)

         _A			     =	image_processor if image_processor is not None else feature_extractor
         if image_processor is None:
               raise ValueError("You need to specify an `image_processor`."		)
         if tokenizer is None:
               raise ValueError("You need to specify a `tokenizer`."		)

         super().__init__(__UpperCAmelCase      ,      __UpperCAmelCase		)

   def __call__(		self :	Optional[Any]      ,      __UpperCAmelCase :	Any=None      ,      __UpperCAmelCase :	List[Any]=None      ,      __UpperCAmelCase :	Union[str, Any]=None      ,      __UpperCAmelCase :	Any="max_length"      ,      __UpperCAmelCase :	List[Any]="np"      ,      **__UpperCAmelCase :	str		):





         '''simple docstring'''





         if text is None and query_images is None and images is None:
               raise ValueError(
                   "You have to specify at least one text or query image or image. All three cannot be none."		)

         if text is not None:
               if isinstance(__UpperCAmelCase      ,      __UpperCAmelCase		) or (isinstance(__UpperCAmelCase      ,      __UpperCAmelCase		) and not isinstance(text[0]      ,      __UpperCAmelCase		)):
                     _A			     =	[self.tokenizer(__UpperCAmelCase      ,      padding=__UpperCAmelCase      ,      return_tensors=__UpperCAmelCase      ,      **__UpperCAmelCase		)]

               elif isinstance(__UpperCAmelCase      ,      __UpperCAmelCase		) and isinstance(text[0]      ,      __UpperCAmelCase		):
                     _A			     =	[]

                     # Maximum number of queries across batch
                     _A			     =	max([len(__UpperCAmelCase		) for t in text]		)

                     # Pad all batch samples to max number of text queries
                     for t in text:
                           if len(__UpperCAmelCase		) != max_num_queries:
                                 _A			     =	t + [" "] * (max_num_queries - len(__UpperCAmelCase		))

                           _A			     =	self.tokenizer(__UpperCAmelCase      ,      padding=__UpperCAmelCase      ,      return_tensors=__UpperCAmelCase      ,      **__UpperCAmelCase		)
                           encodings.append(__UpperCAmelCase		)
               else:
                     raise TypeError("Input text should be a string, a list of strings or a nested list of strings"		)

               if return_tensors == "np":
                     _A			     =	np.concatenate([encoding["input_ids"] for encoding in encodings]      ,      axis=0		)
                     _A			     =	np.concatenate([encoding["attention_mask"] for encoding in encodings]      ,      axis=0		)

               elif return_tensors == "jax" and is_flax_available():
                     import jax.numpy as jnp

                     _A			     =	jnp.concatenate([encoding["input_ids"] for encoding in encodings]      ,      axis=0		)
                     _A			     =	jnp.concatenate([encoding["attention_mask"] for encoding in encodings]      ,      axis=0		)

               elif return_tensors == "pt" and is_torch_available():
                     import torch

                     _A			     =	torch.cat([encoding["input_ids"] for encoding in encodings]      ,      dim=0		)
                     _A			     =	torch.cat([encoding["attention_mask"] for encoding in encodings]      ,      dim=0		)

               elif return_tensors == "tf" and is_tf_available():
                     import tensorflow as tf

                     _A			     =	tf.stack([encoding["input_ids"] for encoding in encodings]      ,      axis=0		)
                     _A			     =	tf.stack([encoding["attention_mask"] for encoding in encodings]      ,      axis=0		)

               else:
                     raise ValueError("Target return tensor type could not be returned"		)

               _A			     =	BatchEncoding()
               _A			     =	input_ids
               _A			     =	attention_mask

         if query_images is not None:
               _A			     =	BatchEncoding()
               _A			     =	self.image_processor(
                   __UpperCAmelCase      ,      return_tensors=__UpperCAmelCase      ,      **__UpperCAmelCase		).pixel_values
               _A			     =	query_pixel_values

         if images is not None:
               _A			     =	self.image_processor(__UpperCAmelCase      ,      return_tensors=__UpperCAmelCase      ,      **__UpperCAmelCase		)

         if text is not None and images is not None:
               _A			     =	image_features.pixel_values
               return encoding
         elif query_images is not None and images is not None:
               _A			     =	image_features.pixel_values
               return encoding
         elif text is not None or query_images is not None:
               return encoding
         else:
               return BatchEncoding(data=dict(**__UpperCAmelCase		)      ,      tensor_type=__UpperCAmelCase		)

   def 						lowerCAmelCase (		self :	Tuple      ,      *__UpperCAmelCase :	Dict      ,      **__UpperCAmelCase :	Dict		):





         '''simple docstring'''





         return self.image_processor.post_process(*__UpperCAmelCase      ,      **__UpperCAmelCase		)

   def 						lowerCAmelCase (		self :	int      ,      *__UpperCAmelCase :	int      ,      **__UpperCAmelCase :	int		):





         '''simple docstring'''





         return self.image_processor.post_process_object_detection(*__UpperCAmelCase      ,      **__UpperCAmelCase		)

   def 						lowerCAmelCase (		self :	Dict      ,      *__UpperCAmelCase :	str      ,      **__UpperCAmelCase :	Optional[Any]		):





         '''simple docstring'''





         return self.image_processor.post_process_image_guided_detection(*__UpperCAmelCase      ,      **__UpperCAmelCase		)

   def 						lowerCAmelCase (		self :	str      ,      *__UpperCAmelCase :	Any      ,      **__UpperCAmelCase :	Union[str, Any]		):





         '''simple docstring'''





         return self.tokenizer.batch_decode(*__UpperCAmelCase      ,      **__UpperCAmelCase		)

   def 						lowerCAmelCase (		self :	int      ,      *__UpperCAmelCase :	Any      ,      **__UpperCAmelCase :	int		):





         '''simple docstring'''





         return self.tokenizer.decode(*__UpperCAmelCase      ,      **__UpperCAmelCase		)

   @property
   def 						lowerCAmelCase (		self :	Union[str, Any]		):





         '''simple docstring'''





         warnings.warn(
             "`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead."      ,      __UpperCAmelCase      ,      )
         return self.image_processor_class

   @property
   def 						lowerCAmelCase (		self :	int		):





         '''simple docstring'''





         warnings.warn(
             "`feature_extractor` is deprecated and will be removed in v5. Use `image_processor` instead."      ,      __UpperCAmelCase      ,      )
         return self.image_processor










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 			_lowerCAmelCase       (					__lowerCAmelCase  ,  __lowerCAmelCase      )   ->					Tuple:


				"""simple docstring"""

				snake_case__	:      Union[str, Any]     		=	[]
				for part_id in partition_order:
								snake_case__	:      Any     		=	df.where(f"""SPARK_PARTITION_ID() = {part_id}"""      ).collect()
								for row_idx, row in enumerate(__lowerCAmelCase      ):
												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 			_lowerCAmelCase       (					)   ->					Tuple:


				"""simple docstring"""

				snake_case__	:      Any     		=	pyspark.sql.SparkSession.builder.master('''local[*]'''      ).appName('''pyspark'''      ).getOrCreate()
				snake_case__	:      Optional[Any]     		=	spark.range(100      ).repartition(1      )
				snake_case__	:      Optional[int]     		=	Spark(__lowerCAmelCase      )
				# 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 			_lowerCAmelCase       (					)   ->					Optional[int]:


				"""simple docstring"""

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

				for i, (row_id, row_dict) in enumerate(generate_fn()      ):
								snake_case__		, snake_case__	:      Union[str, 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 			_lowerCAmelCase       (					)   ->					Any:


				"""simple docstring"""

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



@require_not_windows
@require_dill_gt_0_3_2
def 			_lowerCAmelCase       (					)   ->					Dict:


				"""simple docstring"""

				snake_case__	:      List[str]     		=	pyspark.sql.SparkSession.builder.master('''local[*]'''      ).appName('''pyspark'''      ).getOrCreate()
				snake_case__	:      Tuple     		=	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__	:      Union[str, Any]     		=	lambda __lowerCAmelCase      : x.reverse()
								snake_case__	:      List[str]     		=	_get_expected_row_ids_and_row_dicts_for_partition_order(__lowerCAmelCase  ,  [2, 1, 0]      )

								snake_case__	:      List[str]     		=	SparkExamplesIterable(__lowerCAmelCase      ).shuffle_data_sources(__lowerCAmelCase      )
								assert shuffled_it.n_shards == 3
								for i, (row_id, row_dict) in enumerate(__lowerCAmelCase      ):
												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 			_lowerCAmelCase       (					)   ->					List[Any]:


				"""simple docstring"""

				snake_case__	:      Union[str, Any]     		=	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(__lowerCAmelCase      ).shard_data_sources(worker_id=0  ,  num_workers=2      )
				assert shard_it_a.n_shards == 2
				snake_case__	:      int     		=	_get_expected_row_ids_and_row_dicts_for_partition_order(__lowerCAmelCase  ,  [0, 2]      )
				for i, (row_id, row_dict) in enumerate(__lowerCAmelCase      ):
								snake_case__		, snake_case__	:      Dict     		=	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(__lowerCAmelCase      ).shard_data_sources(worker_id=1  ,  num_workers=2      )
				assert shard_it_a.n_shards == 2
				snake_case__	:      List[Any]     		=	_get_expected_row_ids_and_row_dicts_for_partition_order(__lowerCAmelCase  ,  [1, 3]      )
				for i, (row_id, row_dict) in enumerate(__lowerCAmelCase      ):
								snake_case__		, snake_case__	:      Optional[int]     		=	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 			_lowerCAmelCase       (					)   ->					Dict:


				"""simple docstring"""

				snake_case__	:      int     		=	pyspark.sql.SparkSession.builder.master('''local[*]'''      ).appName('''pyspark'''      ).getOrCreate()
				snake_case__	:      Dict     		=	spark.range(100      ).repartition(1      )
				snake_case__	:      Tuple     		=	Spark(__lowerCAmelCase      )
				# 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() == 100

