






"""simple docstring"""




import unittest

from transformers import AlbertConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, slow, torch_device

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 (
     MODEL_FOR_PRETRAINING_MAPPING,
     AlbertForMaskedLM,
     AlbertForMultipleChoice,
     AlbertForPreTraining,
     AlbertForQuestionAnswering,
     AlbertForSequenceClassification,
     AlbertForTokenClassification,
     AlbertModel,
 )
 from transformers.models.albert.modeling_albert import ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST






class        UpperCamelCase			:


 def __init__(     self		:						int							,						UpperCAmelCase__		:						Optional[int]							,						UpperCAmelCase__		:						List[Any]=13							,						UpperCAmelCase__		:						Tuple=7							,						UpperCAmelCase__		:						List[Any]=True							,						UpperCAmelCase__		:						Union[str, Any]=True							,						UpperCAmelCase__		:						int=True							,						UpperCAmelCase__		:						List[Any]=True							,						UpperCAmelCase__		:						Tuple=99							,						UpperCAmelCase__		:						List[str]=16							,						UpperCAmelCase__		:						Optional[int]=36							,						UpperCAmelCase__		:						Optional[int]=6							,						UpperCAmelCase__		:						Optional[int]=6							,						UpperCAmelCase__		:						Tuple=6							,						UpperCAmelCase__		:						Any=37							,						UpperCAmelCase__		:						Dict="gelu"							,						UpperCAmelCase__		:						List[str]=0.1							,						UpperCAmelCase__		:						int=0.1							,						UpperCAmelCase__		:						Dict=512							,						UpperCAmelCase__		:						Tuple=16							,						UpperCAmelCase__		:						Tuple=2							,						UpperCAmelCase__		:						List[Any]=0.0_2							,						UpperCAmelCase__		:						Tuple=3							,						UpperCAmelCase__		:						Optional[Any]=4							,						UpperCAmelCase__		:						List[str]=None							,						)					->       List[str]:
   _a  :       int             =	parent
   _a  :       List[Any]             =	batch_size
   _a  :       int             =	seq_length
   _a  :       str             =	is_training
   _a  :       Optional[int]             =	use_input_mask
   _a  :       Optional[Any]             =	use_token_type_ids
   _a  :       Tuple             =	use_labels
   _a  :       Dict             =	vocab_size
   _a  :       List[Any]             =	embedding_size
   _a  :       int             =	hidden_size
   _a  :       Union[str, Any]             =	num_hidden_layers
   _a  :       Union[str, Any]             =	num_hidden_groups
   _a  :       int             =	num_attention_heads
   _a  :       List[Any]             =	intermediate_size
   _a  :       List[str]             =	hidden_act
   _a  :       List[Any]             =	hidden_dropout_prob
   _a  :       Union[str, Any]             =	attention_probs_dropout_prob
   _a  :       Any             =	max_position_embeddings
   _a  :       List[str]             =	type_vocab_size
   _a  :       Dict             =	type_sequence_label_size
   _a  :       List[str]             =	initializer_range
   _a  :       Union[str, Any]             =	num_labels
   _a  :       Union[str, Any]             =	num_choices
   _a  :       Tuple             =	scope


 def _lowercase	(     self		:						List[str]							)					->       Optional[int]:
   _a  :       Tuple             =	ids_tensor([self.batch_size, self.seq_length]							,						self.vocab_size							)

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

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

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

   _a  :       Dict             =	self.get_config()

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


 def _lowercase	(     self		:						List[str]							)					->       int:
   return AlbertConfig(
       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							,						initializer_range=self.initializer_range							,						num_hidden_groups=self.num_hidden_groups							,						)


 def _lowercase	(     self		:						List[str]							,						UpperCAmelCase__		:						str							,						UpperCAmelCase__		:						Any							,						UpperCAmelCase__		:						Union[str, Any]							,						UpperCAmelCase__		:						List[str]							,						UpperCAmelCase__		:						int							,						UpperCAmelCase__		:						Tuple							,						UpperCAmelCase__		:						Any							)					->       Union[str, Any]:
   _a  :       Dict             =	AlbertModel(config=__lowercase							)
   model.to(__lowercase							)
   model.eval()
   _a  :       List[str]             =	model(__lowercase							,						attention_mask=__lowercase							,						token_type_ids=__lowercase							)
   _a  :       int             =	model(__lowercase							,						token_type_ids=__lowercase							)
   _a  :       Dict             =	model(__lowercase							)
   self.parent.assertEqual(result.last_hidden_state.shape							,						(self.batch_size, self.seq_length, self.hidden_size)							)
   self.parent.assertEqual(result.pooler_output.shape							,						(self.batch_size, self.hidden_size)							)


 def _lowercase	(     self		:						List[Any]							,						UpperCAmelCase__		:						List[Any]							,						UpperCAmelCase__		:						int							,						UpperCAmelCase__		:						Dict							,						UpperCAmelCase__		:						Optional[Any]							,						UpperCAmelCase__		:						str							,						UpperCAmelCase__		:						List[str]							,						UpperCAmelCase__		:						List[str]							)					->       Tuple:
   _a  :       Optional[int]             =	AlbertForPreTraining(config=__lowercase							)
   model.to(__lowercase							)
   model.eval()
   _a  :       List[Any]             =	model(
       __lowercase							,						attention_mask=__lowercase							,						token_type_ids=__lowercase							,						labels=__lowercase							,						sentence_order_label=__lowercase							,						)
   self.parent.assertEqual(result.prediction_logits.shape							,						(self.batch_size, self.seq_length, self.vocab_size)							)
   self.parent.assertEqual(result.sop_logits.shape							,						(self.batch_size, config.num_labels)							)


 def _lowercase	(     self		:						List[Any]							,						UpperCAmelCase__		:						str							,						UpperCAmelCase__		:						Union[str, Any]							,						UpperCAmelCase__		:						Any							,						UpperCAmelCase__		:						List[str]							,						UpperCAmelCase__		:						Tuple							,						UpperCAmelCase__		:						Optional[Any]							,						UpperCAmelCase__		:						Any							)					->       Tuple:
   _a  :       Dict             =	AlbertForMaskedLM(config=__lowercase							)
   model.to(__lowercase							)
   model.eval()
   _a  :       Dict             =	model(__lowercase							,						attention_mask=__lowercase							,						token_type_ids=__lowercase							,						labels=__lowercase							)
   self.parent.assertEqual(result.logits.shape							,						(self.batch_size, self.seq_length, self.vocab_size)							)


 def _lowercase	(     self		:						int							,						UpperCAmelCase__		:						Tuple							,						UpperCAmelCase__		:						Tuple							,						UpperCAmelCase__		:						Optional[int]							,						UpperCAmelCase__		:						Dict							,						UpperCAmelCase__		:						int							,						UpperCAmelCase__		:						List[str]							,						UpperCAmelCase__		:						List[Any]							)					->       Union[str, Any]:
   _a  :       List[str]             =	AlbertForQuestionAnswering(config=__lowercase							)
   model.to(__lowercase							)
   model.eval()
   _a  :       Optional[int]             =	model(
       __lowercase							,						attention_mask=__lowercase							,						token_type_ids=__lowercase							,						start_positions=__lowercase							,						end_positions=__lowercase							,						)
   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)							)


 def _lowercase	(     self		:						Optional[Any]							,						UpperCAmelCase__		:						List[Any]							,						UpperCAmelCase__		:						str							,						UpperCAmelCase__		:						str							,						UpperCAmelCase__		:						Optional[Any]							,						UpperCAmelCase__		:						int							,						UpperCAmelCase__		:						Dict							,						UpperCAmelCase__		:						Any							)					->       Any:
   _a  :       Optional[int]             =	self.num_labels
   _a  :       Optional[Any]             =	AlbertForSequenceClassification(__lowercase							)
   model.to(__lowercase							)
   model.eval()
   _a  :       Tuple             =	model(__lowercase							,						attention_mask=__lowercase							,						token_type_ids=__lowercase							,						labels=__lowercase							)
   self.parent.assertEqual(result.logits.shape							,						(self.batch_size, self.num_labels)							)


 def _lowercase	(     self		:						Tuple							,						UpperCAmelCase__		:						Tuple							,						UpperCAmelCase__		:						Any							,						UpperCAmelCase__		:						Optional[Any]							,						UpperCAmelCase__		:						Tuple							,						UpperCAmelCase__		:						Optional[Any]							,						UpperCAmelCase__		:						Any							,						UpperCAmelCase__		:						int							)					->       List[Any]:
   _a  :       Any             =	self.num_labels
   _a  :       Any             =	AlbertForTokenClassification(config=__lowercase							)
   model.to(__lowercase							)
   model.eval()
   _a  :       Any             =	model(__lowercase							,						attention_mask=__lowercase							,						token_type_ids=__lowercase							,						labels=__lowercase							)
   self.parent.assertEqual(result.logits.shape							,						(self.batch_size, self.seq_length, self.num_labels)							)


 def _lowercase	(     self		:						Optional[Any]							,						UpperCAmelCase__		:						Optional[Any]							,						UpperCAmelCase__		:						List[Any]							,						UpperCAmelCase__		:						List[Any]							,						UpperCAmelCase__		:						Any							,						UpperCAmelCase__		:						Optional[int]							,						UpperCAmelCase__		:						List[Any]							,						UpperCAmelCase__		:						int							)					->       List[Any]:
   _a  :       Any             =	self.num_choices
   _a  :       Tuple             =	AlbertForMultipleChoice(config=__lowercase							)
   model.to(__lowercase							)
   model.eval()
   _a  :       Dict             =	input_ids.unsqueeze(1							).expand(-1							,						self.num_choices							,						-1							).contiguous()
   _a  :       Tuple             =	token_type_ids.unsqueeze(1							).expand(-1							,						self.num_choices							,						-1							).contiguous()
   _a  :       Optional[Any]             =	input_mask.unsqueeze(1							).expand(-1							,						self.num_choices							,						-1							).contiguous()
   _a  :       Optional[int]             =	model(
       __lowercase							,						attention_mask=__lowercase							,						token_type_ids=__lowercase							,						labels=__lowercase							,						)
   self.parent.assertEqual(result.logits.shape							,						(self.batch_size, self.num_choices)							)







 def _lowercase	(     self		:						Tuple							)					->       List[Any]:
   _a  :       Optional[Any]             =	self.prepare_config_and_inputs()
   (
       _a 
   )  :       Dict             =	config_and_inputs
   _a  :       List[Any]             =	{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
   return config, inputs_dict







@require_torch
class        UpperCamelCase			(	__lowerCamelCase		,     __lowerCamelCase		,     unittest.TestCase      ):
 UpperCamelCase					:       str			=					(
     (
         AlbertModel,
         AlbertForPreTraining,
         AlbertForMaskedLM,
         AlbertForMultipleChoice,
         AlbertForSequenceClassification,
         AlbertForTokenClassification,
         AlbertForQuestionAnswering,
     )
     if is_torch_available()
     else ()
 )
 UpperCamelCase					:       Tuple			=					(
     {
         """feature-extraction""": AlbertModel,
         """fill-mask""": AlbertForMaskedLM,
         """question-answering""": AlbertForQuestionAnswering,
         """text-classification""": AlbertForSequenceClassification,
         """token-classification""": AlbertForTokenClassification,
         """zero-shot""": AlbertForSequenceClassification,
     }
     if is_torch_available()
     else {}
 )
 UpperCamelCase					:       Dict			=					True


 def _lowercase	(     self		:						Tuple							,						UpperCAmelCase__		:						List[Any]							,						UpperCAmelCase__		:						Any							,						UpperCAmelCase__		:						int=False							)					->       List[Any]:
   _a  :       str             =	super()._prepare_for_class(__lowercase							,						__lowercase							,						return_labels=__lowercase							)

   if return_labels:
     if model_class in get_values(__lowercase							):
       _a  :       Dict             =	torch.zeros(
           (self.model_tester.batch_size, self.model_tester.seq_length)							,						dtype=torch.long							,						device=__lowercase							)
       _a  :       List[Any]             =	torch.zeros(
           self.model_tester.batch_size							,						dtype=torch.long							,						device=__lowercase							)
   return inputs_dict


 def _lowercase	(     self		:						Tuple							)					->       Dict:
   _a  :       str             =	AlbertModelTester(self							)
   _a  :       List[str]             =	ConfigTester(self							,						config_class=__lowercase							,						hidden_size=37							)


 def _lowercase	(     self		:						List[Any]							)					->       Tuple:
   self.config_tester.run_common_tests()


 def _lowercase	(     self		:						Optional[Any]							)					->       int:
   _a  :       List[str]             =	self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_model(*__lowercase							)


 def _lowercase	(     self		:						Union[str, Any]							)					->       List[Any]:
   _a  :       Dict             =	self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_pretraining(*__lowercase							)


 def _lowercase	(     self		:						int							)					->       List[str]:
   _a  :       List[Any]             =	self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_masked_lm(*__lowercase							)


 def _lowercase	(     self		:						List[Any]							)					->       List[str]:
   _a  :       List[str]             =	self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_multiple_choice(*__lowercase							)


 def _lowercase	(     self		:						Union[str, Any]							)					->       Optional[Any]:
   _a  :       List[Any]             =	self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_question_answering(*__lowercase							)


 def _lowercase	(     self		:						Tuple							)					->       Any:
   _a  :       List[str]             =	self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_sequence_classification(*__lowercase							)


 def _lowercase	(     self		:						Tuple							)					->       str:
   _a  :       List[Any]             =	self.model_tester.prepare_config_and_inputs()
   for type in ["absolute", "relative_key", "relative_key_query"]:
     _a  :       int             =	type
     self.model_tester.create_and_check_model(*__lowercase							)







 @slow
 def _lowercase	(     self		:						int							)					->       List[Any]:
   for model_name in ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
     _a  :       Optional[Any]             =	AlbertModel.from_pretrained(__lowercase							)
     self.assertIsNotNone(__lowercase							)







@require_torch
class        UpperCamelCase			(	unittest.TestCase      ):


 @slow
 def _lowercase	(     self		:						List[Any]							)					->       Union[str, Any]:
   _a  :       Tuple             =	AlbertModel.from_pretrained("""albert-base-v2"""							)
   _a  :       Tuple             =	torch.tensor([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]							)
   _a  :       List[Any]             =	torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]							)
   with torch.no_grad():
     _a  :       Optional[Any]             =	model(__lowercase							,						attention_mask=__lowercase							)[0]
   _a  :       Union[str, Any]             =	torch.Size((1, 11, 768)							)
   self.assertEqual(output.shape							,						__lowercase							)
   _a  :       Optional[Any]             =	torch.tensor(
       [[[-0.6_5_1_3, 1.5_0_3_5, -0.2_7_6_6], [-0.6_5_1_5, 1.5_0_4_6, -0.2_7_8_0], [-0.6_5_1_2, 1.5_0_4_9, -0.2_7_8_4]]]							)

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






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

