





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__(__snake_case ,__snake_case							)       ->   str:







	'''simple docstring'''







	lowerCamelCase__              =	[]
	for part_id in partition_order:
		lowerCamelCase__              =	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 lowerCAmelCase__()       ->   Dict:







	'''simple docstring'''







	lowerCamelCase__              =	pyspark.sql.SparkSession.builder.master('''local[*]'''							).appName('''pyspark'''							).getOrCreate()
	lowerCamelCase__              =	spark.range(100							).repartition(1							)
	lowerCamelCase__              =	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 lowerCAmelCase__()       ->   str:







	'''simple docstring'''







	lowerCamelCase__              =	pyspark.sql.SparkSession.builder.master('''local[*]'''							).appName('''pyspark'''							).getOrCreate()
	lowerCamelCase__              =	spark.range(10							).repartition(2							)
	lowerCamelCase__              =	[1, 0]
	lowerCamelCase__              =	_generate_iterable_examples(__snake_case ,__snake_case							)  # Reverse the partitions.
	lowerCamelCase__              =	_get_expected_row_ids_and_row_dicts_for_partition_order(__snake_case ,__snake_case							)

	for i, (row_id, row_dict) in enumerate(generate_fn()							):
		lowerCamelCase__       ,			lowerCamelCase__              =	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'''







	lowerCamelCase__              =	pyspark.sql.SparkSession.builder.master('''local[*]'''							).appName('''pyspark'''							).getOrCreate()
	lowerCamelCase__              =	spark.range(10							).repartition(1							)
	lowerCamelCase__              =	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 lowerCAmelCase__()       ->   Optional[int]:







	'''simple docstring'''







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

		lowerCamelCase__              =	SparkExamplesIterable(__snake_case							).shuffle_data_sources(__snake_case							)
		assert shuffled_it.n_shards == 3
		for i, (row_id, row_dict) in enumerate(__snake_case							):
			lowerCamelCase__       ,			lowerCamelCase__              =	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__()       ->   int:







	'''simple docstring'''







	lowerCamelCase__              =	pyspark.sql.SparkSession.builder.master('''local[*]'''							).appName('''pyspark'''							).getOrCreate()
	lowerCamelCase__              =	spark.range(20							).repartition(4							)

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

	# Partitions 1 and 3
	lowerCamelCase__              =	SparkExamplesIterable(__snake_case							).shard_data_sources(worker_id=1 ,num_workers=2							)
	assert shard_it_a.n_shards == 2
	lowerCamelCase__              =	_get_expected_row_ids_and_row_dicts_for_partition_order(__snake_case ,[1, 3]							)
	for i, (row_id, row_dict) in enumerate(__snake_case							):
		lowerCamelCase__       ,			lowerCamelCase__              =	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__()       ->   str:







	'''simple docstring'''







	lowerCamelCase__              =	pyspark.sql.SparkSession.builder.master('''local[*]'''							).appName('''pyspark'''							).getOrCreate()
	lowerCamelCase__              =	spark.range(100							).repartition(1							)
	lowerCamelCase__              =	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() == 100










from __future__ import annotations

import unittest

from transformers import is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
			import numpy
			import tensorflow as tf

			from transformers import (
			    TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
			    TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
			    TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
			    BertConfig,
			    DPRConfig,
			    TFDPRContextEncoder,
			    TFDPRQuestionEncoder,
			    TFDPRReader,
			)






class       __A					:

			'''simple docstring'''







			def __init__(					self ,						__lowerCAmelCase ,						__lowerCAmelCase=1_3 ,						__lowerCAmelCase=7 ,						__lowerCAmelCase=True ,						__lowerCAmelCase=True ,						__lowerCAmelCase=True ,						__lowerCAmelCase=True ,						__lowerCAmelCase=9_9 ,						__lowerCAmelCase=3_2 ,						__lowerCAmelCase=2 ,						__lowerCAmelCase=4 ,						__lowerCAmelCase=3_7 ,						__lowerCAmelCase="gelu" ,						__lowerCAmelCase=0.1 ,						__lowerCAmelCase=0.1 ,						__lowerCAmelCase=5_1_2 ,						__lowerCAmelCase=1_6 ,						__lowerCAmelCase=2 ,						__lowerCAmelCase=0.02 ,						__lowerCAmelCase=3 ,						__lowerCAmelCase=4 ,						__lowerCAmelCase=None ,						__lowerCAmelCase=0 ,						):



				'''simple docstring'''







				lowerCamelCase__              =	parent
				lowerCamelCase__              =	batch_size
				lowerCamelCase__              =	seq_length
				lowerCamelCase__              =	is_training
				lowerCamelCase__              =	use_input_mask
				lowerCamelCase__              =	use_token_type_ids
				lowerCamelCase__              =	use_labels
				lowerCamelCase__              =	vocab_size
				lowerCamelCase__              =	hidden_size
				lowerCamelCase__              =	num_hidden_layers
				lowerCamelCase__              =	num_attention_heads
				lowerCamelCase__              =	intermediate_size
				lowerCamelCase__              =	hidden_act
				lowerCamelCase__              =	hidden_dropout_prob
				lowerCamelCase__              =	attention_probs_dropout_prob
				lowerCamelCase__              =	max_position_embeddings
				lowerCamelCase__              =	type_vocab_size
				lowerCamelCase__              =	type_sequence_label_size
				lowerCamelCase__              =	initializer_range
				lowerCamelCase__              =	num_labels
				lowerCamelCase__              =	num_choices
				lowerCamelCase__              =	scope
				lowerCamelCase__              =	projection_dim





			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				lowerCamelCase__              =	ids_tensor([self.batch_size, self.seq_length] ,						self.vocab_size							)

				lowerCamelCase__              =	None
				if self.use_input_mask:
					# follow test_modeling_tf_ctrl.py
					lowerCamelCase__              =	random_attention_mask([self.batch_size, self.seq_length]							)

				lowerCamelCase__              =	None
				if self.use_token_type_ids:
					lowerCamelCase__              =	ids_tensor([self.batch_size, self.seq_length] ,						self.type_vocab_size							)

				lowerCamelCase__              =	None
				lowerCamelCase__              =	None
				lowerCamelCase__              =	None
				if self.use_labels:
					lowerCamelCase__              =	ids_tensor([self.batch_size] ,						self.type_sequence_label_size							)
					lowerCamelCase__              =	ids_tensor([self.batch_size, self.seq_length] ,						self.num_labels							)
					lowerCamelCase__              =	ids_tensor([self.batch_size] ,						self.num_choices							)

				lowerCamelCase__              =	BertConfig(
				    vocab_size=self.vocab_size ,						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 ,						max_position_embeddings=self.max_position_embeddings ,						type_vocab_size=self.type_vocab_size ,						is_decoder=__lowerCAmelCase ,						initializer_range=self.initializer_range ,						)
				lowerCamelCase__              =	DPRConfig(projection_dim=self.projection_dim ,						**config.to_dict()							)

				return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels





			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase							):



				'''simple docstring'''







				lowerCamelCase__              =	TFDPRContextEncoder(config=__lowerCAmelCase							)
				lowerCamelCase__              =	model(__lowerCAmelCase ,						attention_mask=__lowerCAmelCase ,						token_type_ids=__lowerCAmelCase							)
				lowerCamelCase__              =	model(__lowerCAmelCase ,						token_type_ids=__lowerCAmelCase							)
				lowerCamelCase__              =	model(__lowerCAmelCase							)
				self.parent.assertEqual(result.pooler_output.shape ,						(self.batch_size, self.projection_dim or self.hidden_size)							)





			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase							):



				'''simple docstring'''







				lowerCamelCase__              =	TFDPRQuestionEncoder(config=__lowerCAmelCase							)
				lowerCamelCase__              =	model(__lowerCAmelCase ,						attention_mask=__lowerCAmelCase ,						token_type_ids=__lowerCAmelCase							)
				lowerCamelCase__              =	model(__lowerCAmelCase ,						token_type_ids=__lowerCAmelCase							)
				lowerCamelCase__              =	model(__lowerCAmelCase							)
				self.parent.assertEqual(result.pooler_output.shape ,						(self.batch_size, self.projection_dim or self.hidden_size)							)





			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase							):



				'''simple docstring'''







				lowerCamelCase__              =	TFDPRReader(config=__lowerCAmelCase							)
				lowerCamelCase__              =	model(__lowerCAmelCase ,						attention_mask=__lowerCAmelCase							)

				self.parent.assertEqual(result.start_logits.shape ,						(self.batch_size, self.seq_length)							)
				self.parent.assertEqual(result.end_logits.shape ,						(self.batch_size, self.seq_length)							)
				self.parent.assertEqual(result.relevance_logits.shape ,						(self.batch_size,)							)





			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				lowerCamelCase__              =	self.prepare_config_and_inputs()
				(
				    (
				    lowerCamelCase__
				)       ,			(
				    lowerCamelCase__
				)       ,			(
				    lowerCamelCase__
				)       ,			(
				    lowerCamelCase__
				)       ,			(
				    lowerCamelCase__
				)       ,			(
				    lowerCamelCase__
				)       ,			(
				    lowerCamelCase__
				)       ,			
				)              =	config_and_inputs
				lowerCamelCase__              =	{'''input_ids''': input_ids}
				return config, inputs_dict






@require_tf
class       __A					(					lowerCAmelCase    ,       lowerCAmelCase    ,       unittest.TestCase				):

			'''simple docstring'''







			lowerCAmelCase_							  =		(
			    (
			        TFDPRContextEncoder,
			        TFDPRQuestionEncoder,
			        TFDPRReader,
			    )
			    if is_tf_available()
			    else ()
			)
			lowerCAmelCase_							  =		{"""feature-extraction""": TFDPRQuestionEncoder} if is_tf_available() else {}

			lowerCAmelCase_							  =		False
			lowerCAmelCase_							  =		False
			lowerCAmelCase_							  =		False
			lowerCAmelCase_							  =		False
			lowerCAmelCase_							  =		False





			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				lowerCamelCase__              =	TFDPRModelTester(self							)
				lowerCamelCase__              =	ConfigTester(self ,						config_class=__lowerCAmelCase ,						hidden_size=3_7							)





			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				self.config_tester.run_common_tests()





			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				lowerCamelCase__              =	self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_dpr_context_encoder(*__lowerCAmelCase							)





			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				lowerCamelCase__              =	self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_dpr_question_encoder(*__lowerCAmelCase							)





			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				lowerCamelCase__              =	self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_dpr_reader(*__lowerCAmelCase							)





			@slow
			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
					lowerCamelCase__              =	TFDPRContextEncoder.from_pretrained(__lowerCAmelCase							)
					self.assertIsNotNone(__lowerCAmelCase							)

				for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
					lowerCamelCase__              =	TFDPRContextEncoder.from_pretrained(__lowerCAmelCase							)
					self.assertIsNotNone(__lowerCAmelCase							)

				for model_name in TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
					lowerCamelCase__              =	TFDPRQuestionEncoder.from_pretrained(__lowerCAmelCase							)
					self.assertIsNotNone(__lowerCAmelCase							)

				for model_name in TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
					lowerCamelCase__              =	TFDPRReader.from_pretrained(__lowerCAmelCase							)
					self.assertIsNotNone(__lowerCAmelCase							)






@require_tf
class       __A					(					unittest.TestCase				):

			'''simple docstring'''







			@slow
			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				lowerCamelCase__              =	TFDPRQuestionEncoder.from_pretrained('''facebook/dpr-question_encoder-single-nq-base'''							)

				lowerCamelCase__              =	tf.constant(
				    [[1_0_1, 7_5_9_2, 1_0_1_0, 2_0_0_3, 2_0_2_6, 3_8_9_9, 1_0_1_4_0, 1_0_2_9, 1_0_2]]							)  # [CLS] hello, is my dog cute? [SEP]
				lowerCamelCase__              =	model(__lowerCAmelCase							)[0]  # embedding shape = (1, 768)
				# compare the actual values for a slice.
				lowerCamelCase__              =	tf.constant(
				    [
				        [
				            0.0323_6253,
				            0.1275_3335,
				            0.1681_8509,
				            0.0027_9786,
				            0.389_6933,
				            0.2426_4945,
				            0.217_8971,
				            -0.0233_5227,
				            -0.0848_1959,
				            -0.1432_4117,
				        ]
				    ]							)
				self.assertTrue(numpy.allclose(output[:, :1_0].numpy() ,						expected_slice.numpy() ,						atol=1E-4							)							)



