


'''simple docstring'''



import math
import unittest

from transformers import BioGptConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
	import torch

	from transformers import (
	    BioGptForCausalLM,
	    BioGptForSequenceClassification,
	    BioGptForTokenClassification,
	    BioGptModel,
	    BioGptTokenizer,
	)
	from transformers.models.biogpt.modeling_biogpt import BIOGPT_PRETRAINED_MODEL_ARCHIVE_LIST





class 				__UpperCamelCase							:




	def __init__(  self		,							__a		,							__a=13		,							__a=7		,							__a=True		,							__a=True		,							__a=False		,							__a=True		,							__a=99		,							__a=32		,							__a=5		,							__a=4		,							__a=37		,							__a="gelu"		,							__a=0.1		,							__a=0.1		,							__a=512		,							__a=16		,							__a=2		,							__a=0.02		,							__a=3		,							__a=4		,							__a=None		,							):

						'''simple docstring'''





						__a     :			Union[str, Any]          =    parent
						__a     :			Optional[Any]          =    batch_size
						__a     :			Optional[int]          =    seq_length
						__a     :			Any          =    is_training
						__a     :			Optional[int]          =    use_input_mask
						__a     :			Union[str, Any]          =    use_token_type_ids
						__a     :			str          =    use_labels
						__a     :			int          =    vocab_size
						__a     :			Tuple          =    hidden_size
						__a     :			List[Any]          =    num_hidden_layers
						__a     :			List[str]          =    num_attention_heads
						__a     :			List[Any]          =    intermediate_size
						__a     :			Optional[Any]          =    hidden_act
						__a     :			Optional[int]          =    hidden_dropout_prob
						__a     :			Any          =    attention_probs_dropout_prob
						__a     :			Union[str, Any]          =    max_position_embeddings
						__a     :			Tuple          =    type_vocab_size
						__a     :			int          =    type_sequence_label_size
						__a     :			Optional[Any]          =    initializer_range
						__a     :			Optional[int]          =    num_labels
						__a     :			List[Any]          =    num_choices
						__a     :			List[str]          =    scope




	def  __UpperCAmelCase				(  self				):

						'''simple docstring'''





						__a     :			Union[str, Any]          =    ids_tensor([self.batch_size, self.seq_length]		,							self.vocab_size				)

						__a     :			Dict          =    None
						if self.use_input_mask:
											__a     :			Optional[int]          =    random_attention_mask([self.batch_size, self.seq_length]				)

						__a     :			Optional[Any]          =    None
						if self.use_token_type_ids:
											__a     :			Optional[int]          =    ids_tensor([self.batch_size, self.seq_length]		,							self.type_vocab_size				)

						__a     :			Optional[Any]          =    None
						__a     :			Optional[int]          =    None
						__a     :			Optional[int]          =    None
						if self.use_labels:
											__a     :			Tuple          =    ids_tensor([self.batch_size]		,							self.type_sequence_label_size				)
											__a     :			Dict          =    ids_tensor([self.batch_size, self.seq_length]		,							self.num_labels				)
											__a     :			Optional[Any]          =    ids_tensor([self.batch_size]		,							self.num_choices				)

						__a     :			List[str]          =    self.get_config()

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




	def  __UpperCAmelCase				(  self				):

						'''simple docstring'''





						return BioGptConfig(
						    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=__a		,							initializer_range=self.initializer_range		,							)




	def  __UpperCAmelCase				(  self		,							__a		,							__a		,							__a		,							__a		,							__a		,							__a		,							__a				):

						'''simple docstring'''





						__a     :			Any          =    BioGptModel(config=__a				)
						model.to(__a				)
						model.eval()
						__a     :			Optional[int]          =    model(__a		,							attention_mask=__a				)
						__a     :			Dict          =    model(__a				)
						self.parent.assertEqual(result.last_hidden_state.shape		,							(self.batch_size, self.seq_length, self.hidden_size)				)




	def  __UpperCAmelCase				(  self		,							__a		,							__a		,							__a		,							__a		,							__a		,							__a		,							__a		,							__a		,							__a		,							):

						'''simple docstring'''





						__a     :			int          =    BioGptForCausalLM(config=__a				)
						model.to(__a				)
						model.eval()
						__a     :			Tuple          =    model(__a		,							attention_mask=__a		,							token_type_ids=__a		,							labels=__a				)
						self.parent.assertEqual(result.logits.shape		,							(self.batch_size, self.seq_length, self.vocab_size)				)




	def  __UpperCAmelCase				(  self		,							__a		,							__a		,							__a		,							__a		,							__a		,							*__a				):

						'''simple docstring'''





						__a     :			Optional[Any]          =    BioGptModel(config=__a				)
						model.to(__a				)
						model.eval()

						# create attention mask
						__a     :			List[Any]          =    torch.ones(input_ids.shape		,							dtype=torch.long		,							device=__a				)
						__a     :			List[Any]          =    self.seq_length // 2
						__a     :			int          =    0

						# first forward pass
						__a       ,							__a     :			Dict          =    model(__a		,							attention_mask=__a				).to_tuple()

						# create hypothetical next token and extent to next_input_ids
						__a     :			int          =    ids_tensor((self.batch_size, 1)		,							config.vocab_size				)

						# change a random masked slice from input_ids
						__a     :			Dict          =    ids_tensor((1,)		,							__a				).item() + 1
						__a     :			int          =    ids_tensor((self.batch_size, 1)		,							config.vocab_size				).squeeze(-1				)
						__a     :			List[Any]          =    random_other_next_tokens

						# append to next input_ids and attn_mask
						__a     :			Tuple          =    torch.cat([input_ids, next_tokens]		,							dim=-1				)
						__a     :			str          =    torch.cat(
						    [attn_mask, torch.ones((attn_mask.shape[0], 1)		,							dtype=torch.long		,							device=__a				)]		,							dim=1		,							)

						# get two different outputs
						__a     :			str          =    model(__a		,							attention_mask=__a				)['last_hidden_state']
						__a     :			Optional[Any]          =    model(__a		,							past_key_values=__a		,							attention_mask=__a				)['last_hidden_state']

						# select random slice
						__a     :			Dict          =    ids_tensor((1,)		,							output_from_past.shape[-1]				).item()
						__a     :			List[str]          =    output_from_no_past[:, -1, random_slice_idx].detach()
						__a     :			Tuple          =    output_from_past[:, 0, random_slice_idx].detach()

						# test that outputs are equal for slice
						self.parent.assertTrue(torch.allclose(__a		,							__a		,							atol=1E-3				)				)




	def  __UpperCAmelCase				(  self		,							__a		,							__a		,							__a		,							__a		,							__a		,							*__a				):

						'''simple docstring'''





						__a     :			Optional[Any]          =    BioGptModel(config=__a				).to(__a				).eval()

						__a     :			List[Any]          =    torch.ones(input_ids.shape		,							dtype=torch.long		,							device=__a				)

						# first forward pass
						__a     :			Optional[int]          =    model(__a		,							attention_mask=__a		,							use_cache=__a				)

						__a       ,							__a     :			Any          =    outputs.to_tuple()

						# create hypothetical multiple next token and extent to next_input_ids
						__a     :			int          =    ids_tensor((self.batch_size, 3)		,							config.vocab_size				)
						__a     :			Optional[int]          =    ids_tensor((self.batch_size, 3)		,							2				)

						# append to next input_ids and
						__a     :			Any          =    torch.cat([input_ids, next_tokens]		,							dim=-1				)
						__a     :			str          =    torch.cat([attention_mask, next_attn_mask]		,							dim=-1				)

						__a     :			List[Any]          =    model(__a		,							attention_mask=__a				)['last_hidden_state']
						__a     :			List[Any]          =    model(__a		,							attention_mask=__a		,							past_key_values=__a				)[
						    'last_hidden_state'
						]

						# select random slice
						__a     :			Optional[int]          =    ids_tensor((1,)		,							output_from_past.shape[-1]				).item()
						__a     :			Union[str, Any]          =    output_from_no_past[:, -3:, random_slice_idx].detach()
						__a     :			List[Any]          =    output_from_past[:, :, random_slice_idx].detach()

						self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]				)

						# test that outputs are equal for slice
						self.parent.assertTrue(torch.allclose(__a		,							__a		,							atol=1E-3				)				)




	def  __UpperCAmelCase				(  self		,							__a		,							__a		,							__a		,							__a		,							__a		,							*__a		,							__a=False				):

						'''simple docstring'''





						__a     :			Optional[int]          =    BioGptForCausalLM(__a				)
						model.to(__a				)
						if gradient_checkpointing:
											model.gradient_checkpointing_enable()

						__a     :			Optional[int]          =    model(__a		,							labels=__a				)
						self.parent.assertEqual(result.loss.shape		,							()				)
						self.parent.assertEqual(result.logits.shape		,							(self.batch_size, self.seq_length, self.vocab_size)				)
						result.loss.backward()




	def  __UpperCAmelCase				(  self		,							__a		,							*__a				):

						'''simple docstring'''





						__a     :			Dict          =    BioGptModel(__a				)
						__a     :			Union[str, Any]          =    model.config.initializer_range / math.sqrt(2 * model.config.num_hidden_layers				)
						for key in model.state_dict().keys():
											if "c_proj" in key and "weight" in key:
																self.parent.assertLessEqual(abs(torch.std(model.state_dict()[key]				) - model_std				)		,							0.001				)
																self.parent.assertLessEqual(abs(torch.mean(model.state_dict()[key]				) - 0.0				)		,							0.01				)




	def  __UpperCAmelCase				(  self		,							__a		,							__a		,							__a		,							__a		,							__a		,							*__a				):

						'''simple docstring'''





						__a     :			Tuple          =    self.num_labels
						__a     :			Union[str, Any]          =    BioGptForTokenClassification(__a				)
						model.to(__a				)
						model.eval()
						__a     :			List[Any]          =    model(__a		,							attention_mask=__a		,							token_type_ids=__a				)
						self.parent.assertEqual(result.logits.shape		,							(self.batch_size, self.seq_length, self.num_labels)				)





	def  __UpperCAmelCase				(  self				):

						'''simple docstring'''





						__a     :			int          =    self.prepare_config_and_inputs()
						(
						    (
						    __a  
						)       ,							(
						    __a  
						)       ,							(
						    __a  
						)       ,							(
						    __a  
						)       ,							(
						    __a  
						)       ,							(
						    __a  
						)       ,							(
						    __a  
						)       ,							  
						)     :			Union[str, Any]          =    config_and_inputs
						__a     :			Union[str, Any]          =    {'input_ids': input_ids, 'attention_mask': input_mask}
						return config, inputs_dict





@require_torch
class 				__UpperCamelCase							(						UpperCamelCase__						, UpperCamelCase__						, UpperCamelCase__						, unittest.TestCase   ):
	A_           =       (
	    (BioGptModel, BioGptForCausalLM, BioGptForSequenceClassification, BioGptForTokenClassification)
	    if is_torch_available()
	    else ()
	)
	A_           =       (BioGptForCausalLM,) if is_torch_available() else ()
	A_           =       (
	    {
	        'feature-extraction': BioGptModel,
	        'text-classification': BioGptForSequenceClassification,
	        'text-generation': BioGptForCausalLM,
	        'token-classification': BioGptForTokenClassification,
	        'zero-shot': BioGptForSequenceClassification,
	    }
	    if is_torch_available()
	    else {}
	)
	A_           =       False




	def  __UpperCAmelCase				(  self				):

						'''simple docstring'''





						__a     :			Optional[int]          =    BioGptModelTester(self				)
						__a     :			Optional[int]          =    ConfigTester(self		,							config_class=__a		,							hidden_size=37				)




	def  __UpperCAmelCase				(  self				):

						'''simple docstring'''





						self.config_tester.run_common_tests()




	def  __UpperCAmelCase				(  self				):

						'''simple docstring'''





						__a     :			Any          =    self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*__a				)




	def  __UpperCAmelCase				(  self				):

						'''simple docstring'''





						__a     :			Tuple          =    self.model_tester.prepare_config_and_inputs()
						for type in ["absolute", "relative_key", "relative_key_query"]:
											__a     :			Tuple          =    type
											self.model_tester.create_and_check_model(*__a				)




	def  __UpperCAmelCase				(  self				):

						'''simple docstring'''





						__a     :			Optional[int]          =    self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_biogpt_model_attention_mask_past(*__a				)




	def  __UpperCAmelCase				(  self				):

						'''simple docstring'''





						__a     :			int          =    self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_forward_and_backwards(*__a		,							gradient_checkpointing=__a				)




	def  __UpperCAmelCase				(  self				):

						'''simple docstring'''





						__a     :			List[str]          =    self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_biogpt_model_past_large_inputs(*__a				)




	def  __UpperCAmelCase				(  self				):

						'''simple docstring'''





						__a     :			Union[str, Any]          =    self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_biogpt_weight_initialization(*__a				)




	def  __UpperCAmelCase				(  self				):

						'''simple docstring'''





						__a     :			Tuple          =    self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_biogpt_for_token_classification(*__a				)




	@slow
	def  __UpperCAmelCase				(  self				):

						'''simple docstring'''





						__a     :			Any          =    BioGptForCausalLM.from_pretrained('microsoft/biogpt'				)
						model.to(__a				)
						__a     :			List[str]          =    BioGptTokenizer.from_pretrained('microsoft/biogpt'				)

						__a     :			Dict          =    'left'

						# Define PAD Token = EOS Token = 50256
						__a     :			Tuple          =    tokenizer.eos_token
						__a     :			Tuple          =    model.config.eos_token_id

						# use different length sentences to test batching
						__a     :			Tuple          =    [
						    'Hello, my dog is a little',
						    'Today, I',
						]

						__a     :			int          =    tokenizer(__a		,							return_tensors='pt'		,							padding=__a				)
						__a     :			int          =    inputs['input_ids'].to(__a				)

						__a     :			List[str]          =    model.generate(
						    input_ids=__a		,							attention_mask=inputs['attention_mask'].to(__a				)		,							)

						__a     :			Union[str, Any]          =    tokenizer(sentences[0]		,							return_tensors='pt'				).input_ids.to(__a				)
						__a     :			str          =    model.generate(input_ids=__a				)

						__a     :			List[Any]          =    inputs_non_padded.shape[-1] - inputs['attention_mask'][-1].long().sum().cpu().item()
						__a     :			Any          =    tokenizer(sentences[1]		,							return_tensors='pt'				).input_ids.to(__a				)
						__a     :			Union[str, Any]          =    model.generate(input_ids=__a		,							max_length=model.config.max_length - num_paddings				)

						__a     :			str          =    tokenizer.batch_decode(__a		,							skip_special_tokens=__a				)
						__a     :			List[str]          =    tokenizer.decode(output_non_padded[0]		,							skip_special_tokens=__a				)
						__a     :			Optional[Any]          =    tokenizer.decode(output_padded[0]		,							skip_special_tokens=__a				)

						__a     :			List[str]          =    [
						    'Hello, my dog is a little bit bigger than a little bit.',
						    'Today, I have a good idea of how to use the information',
						]
						self.assertListEqual(__a		,							__a				)
						self.assertListEqual(__a		,							[non_padded_sentence, padded_sentence]				)




	@slow
	def  __UpperCAmelCase				(  self				):

						'''simple docstring'''





						for model_name in BIOGPT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
											__a     :			int          =    BioGptModel.from_pretrained(__a				)
											self.assertIsNotNone(__a				)




	def  __UpperCAmelCase				(  self				):

						'''simple docstring'''





						__a       ,							__a     :			List[Any]          =    self.model_tester.prepare_config_and_inputs_for_common()
						__a     :			Union[str, Any]          =    3
						__a     :			List[str]          =    input_dict['input_ids']
						__a     :			str          =    input_ids.ne(1				).to(__a				)
						__a     :			Dict          =    ids_tensor([self.model_tester.batch_size]		,							self.model_tester.type_sequence_label_size				)
						__a     :			Union[str, Any]          =    BioGptForSequenceClassification(__a				)
						model.to(__a				)
						model.eval()
						__a     :			Optional[Any]          =    model(__a		,							attention_mask=__a		,							labels=__a				)
						self.assertEqual(result.logits.shape		,							(self.model_tester.batch_size, self.model_tester.num_labels)				)





	def  __UpperCAmelCase				(  self				):

						'''simple docstring'''





						__a       ,							__a     :			Any          =    self.model_tester.prepare_config_and_inputs_for_common()
						__a     :			Tuple          =    3
						__a     :			int          =    'multi_label_classification'
						__a     :			Dict          =    input_dict['input_ids']
						__a     :			Union[str, Any]          =    input_ids.ne(1				).to(__a				)
						__a     :			List[Any]          =    ids_tensor(
						    [self.model_tester.batch_size, config.num_labels]		,							self.model_tester.type_sequence_label_size				).to(torch.float				)
						__a     :			Optional[int]          =    BioGptForSequenceClassification(__a				)
						model.to(__a				)
						model.eval()
						__a     :			Optional[Any]          =    model(__a		,							attention_mask=__a		,							labels=__a				)
						self.assertEqual(result.logits.shape		,							(self.model_tester.batch_size, self.model_tester.num_labels)				)





@require_torch
class 				__UpperCamelCase							(						unittest.TestCase   ):




	@slow
	def  __UpperCAmelCase				(  self				):

						'''simple docstring'''





						__a     :			List[Any]          =    BioGptForCausalLM.from_pretrained('microsoft/biogpt'				)
						__a     :			Any          =    torch.tensor([[2, 4805, 9, 656, 21]]				)
						__a     :			Tuple          =    model(__a				)[0]

						__a     :			str          =    4_2384

						__a     :			str          =    torch.Size((1, 5, vocab_size)				)
						self.assertEqual(output.shape		,							__a				)

						__a     :			str          =    torch.tensor(
						    [[[-9.5236, -9.8918, 10.4557], [-11.0469, -9.6423, 8.1022], [-8.8664, -7.8826, 5.5325]]]				)

						self.assertTrue(torch.allclose(output[:, :3, :3]		,							__a		,							atol=1E-4				)				)





	@slow
	def  __UpperCAmelCase				(  self				):

						'''simple docstring'''





						__a     :			List[Any]          =    BioGptTokenizer.from_pretrained('microsoft/biogpt'				)
						__a     :			Optional[int]          =    BioGptForCausalLM.from_pretrained('microsoft/biogpt'				)
						model.to(__a				)

						torch.manual_seed(0				)
						__a     :			Union[str, Any]          =    tokenizer('COVID-19 is'		,							return_tensors='pt'				).to(__a				)
						__a     :			Optional[int]          =    model.generate(
						    **__a		,							min_length=100		,							max_length=1024		,							num_beams=5		,							early_stopping=__a		,							)
						__a     :			List[Any]          =    tokenizer.decode(output_ids[0]		,							skip_special_tokens=__a				)

						__a     :			int          =    (
						    'COVID-19 is a global pandemic caused by severe acute respiratory syndrome coronavirus 2 (SARS-CoV-2), the'
						    ' causative agent of coronavirus disease 2019 (COVID-19), which has spread to more than 200 countries and'
						    ' territories, including the United States (US), Canada, Australia, New Zealand, the United Kingdom (UK),'
						    ' and the United States of America (USA), as of March 11, 2020, with more than 800,000 confirmed cases and'
						    ' more than 800,000 deaths.'
						)
						self.assertEqual(__a		,							__a				)





'''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_				(      A__  ,      A__						):
	a_      				=						[]
	for part_id in partition_order:
		a_      				=						df.where(F'''SPARK_PARTITION_ID() = {part_id}'''						).collect()
		for row_idx, row in enumerate(A__						):
			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_				(      ):
	a_      				=						pyspark.sql.SparkSession.builder.master("""local[*]"""						).appName("""pyspark"""						).getOrCreate()
	a_      				=						spark.range(1_00						).repartition(1						)
	a_      				=						Spark(A__						)
	# 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_				(      ):
	a_      				=						pyspark.sql.SparkSession.builder.master("""local[*]"""						).appName("""pyspark"""						).getOrCreate()
	a_      				=						spark.range(10						).repartition(2						)
	a_      				=						[1, 0]
	a_      				=						_generate_iterable_examples(A__  ,      A__						)  # Reverse the partitions.
	a_      				=						_get_expected_row_ids_and_row_dicts_for_partition_order(A__  ,      A__						)

	for i, (row_id, row_dict) in enumerate(generate_fn()						):
		a_				,  a_      				=						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_				(      ):
	a_      				=						pyspark.sql.SparkSession.builder.master("""local[*]"""						).appName("""pyspark"""						).getOrCreate()
	a_      				=						spark.range(10						).repartition(1						)
	a_      				=						SparkExamplesIterable(A__						)
	assert it.n_shards == 1
	for i, (row_id, row_dict) in enumerate(A__						):
		assert row_id == F'''0_{i}'''
		assert row_dict == {"id": i}






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

		a_      				=						SparkExamplesIterable(A__						).shuffle_data_sources(A__						)
		assert shuffled_it.n_shards == 3
		for i, (row_id, row_dict) in enumerate(A__						):
			a_				,  a_      				=						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_				(      ):
	a_      				=						pyspark.sql.SparkSession.builder.master("""local[*]"""						).appName("""pyspark"""						).getOrCreate()
	a_      				=						spark.range(20						).repartition(4						)

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

	# Partitions 1 and 3
	a_      				=						SparkExamplesIterable(A__						).shard_data_sources(worker_id=1  ,      num_workers=2						)
	assert shard_it_a.n_shards == 2
	a_      				=						_get_expected_row_ids_and_row_dicts_for_partition_order(A__  ,      [1, 3]						)
	for i, (row_id, row_dict) in enumerate(A__						):
		a_				,  a_      				=						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_				(      ):
	a_      				=						pyspark.sql.SparkSession.builder.master("""local[*]"""						).appName("""pyspark"""						).getOrCreate()
	a_      				=						spark.range(1_00						).repartition(1						)
	a_      				=						Spark(A__						)
	# 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
