


import os
import tempfile
from pathlib import Path
from unittest import TestCase

import pyarrow as pa
import pytest

from datasets.arrow_dataset import Dataset
from datasets.arrow_reader import ArrowReader, BaseReader, FileInstructions, ReadInstruction, make_file_instructions
from datasets.info import DatasetInfo
from datasets.splits import NamedSplit, Split, SplitDict, SplitInfo

from .utils import assert_arrow_memory_doesnt_increase, assert_arrow_memory_increases





class       _a							(       _lowercase):




	def 							UpperCAmelCase__(      self		:	int	,      _SCREAMING_SNAKE_CASE		:	List[Any]	,      _SCREAMING_SNAKE_CASE		:	Optional[int]=False	)-> Dict:
			lowerCAmelCase__  ,			lowerCAmelCase__  ,			lowerCAmelCase__     :    int					  =      (
			    filename_skip_take['''filename'''],
			    filename_skip_take['''skip'''] if '''skip''' in filename_skip_take else None,
			    filename_skip_take['''take'''] if '''take''' in filename_skip_take else None,
			)
			open(os.path.join(_SCREAMING_SNAKE_CASE	)	,      '''wb'''	).close()
			lowerCAmelCase__     :    int					  =      pa.Table.from_pydict({'''filename''': [Path(_SCREAMING_SNAKE_CASE	).name] * 100}	)
			if take == -1:
					lowerCAmelCase__     :    int					  =      len(_SCREAMING_SNAKE_CASE	) - skip
			if skip is not None and take is not None:
					lowerCAmelCase__     :    Dict					  =      pa_table.slice(_SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	)
			return pa_table





class       _a							(       _lowercase):




	def 							UpperCAmelCase__(      self		:	Any	)-> str:
			lowerCAmelCase__     :    int					  =      '''my_name'''
			lowerCAmelCase__     :    str					  =      SplitInfo(name='''train'''	,      num_examples=100	)
			lowerCAmelCase__     :    Dict					  =      SplitInfo(name='''test'''	,      num_examples=100	)
			lowerCAmelCase__     :    Optional[Any]					  =      [train_info, test_info]
			lowerCAmelCase__     :    Optional[Any]					  =      SplitDict()
			split_dict.add(_SCREAMING_SNAKE_CASE	)
			split_dict.add(_SCREAMING_SNAKE_CASE	)
			lowerCAmelCase__     :    str					  =      DatasetInfo(splits=_SCREAMING_SNAKE_CASE	)

			with tempfile.TemporaryDirectory() as tmp_dir:
					lowerCAmelCase__     :    Tuple					  =      ReaderTest(_SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	)

					lowerCAmelCase__     :    Dict					  =      '''test[:33%]'''
					lowerCAmelCase__     :    Optional[int]					  =      Dataset(**reader.read(_SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	)	)
					self.assertEqual(dset['''filename'''][0]	,      F'{name}-test'	)
					self.assertEqual(dset.num_rows	,      33	)
					self.assertEqual(dset.num_columns	,      1	)

					lowerCAmelCase__     :    Any					  =      ['''train''', '''test[:33%]''']
					lowerCAmelCase__     :    Dict					  =      [Split.TRAIN, ReadInstruction.from_spec('''test[:33%]'''	)]
					for instructions in [instructionsa, instructionsa]:
							lowerCAmelCase__     :    List[str]					  =      [reader.read(_SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	) for instr in instructions]
							lowerCAmelCase__  ,			lowerCAmelCase__     :    Optional[int]					  =      (Dataset(**_SCREAMING_SNAKE_CASE	) for dataset_kwargs in datasets_kwargs)
							self.assertEqual(train_dset['''filename'''][0]	,      F'{name}-train'	)
							self.assertEqual(train_dset.num_rows	,      100	)
							self.assertEqual(train_dset.num_columns	,      1	)
							self.assertIsInstance(train_dset.split	,      _SCREAMING_SNAKE_CASE	)
							self.assertEqual(str(train_dset.split	)	,      '''train'''	)
							self.assertEqual(test_dset['''filename'''][0]	,      F'{name}-test'	)
							self.assertEqual(test_dset.num_rows	,      33	)
							self.assertEqual(test_dset.num_columns	,      1	)
							self.assertIsInstance(test_dset.split	,      _SCREAMING_SNAKE_CASE	)
							self.assertEqual(str(test_dset.split	)	,      '''test[:33%]'''	)
							del train_dset, test_dset




	def 							UpperCAmelCase__(      self		:	Tuple	)-> List[Any]:
			lowerCAmelCase__     :    Optional[int]					  =      '''my_name'''
			lowerCAmelCase__     :    str					  =      SplitInfo(name='''train'''	,      num_examples=1000	,      shard_lengths=[100] * 10	)
			lowerCAmelCase__     :    List[Any]					  =      [train_info]
			lowerCAmelCase__     :    Tuple					  =      SplitDict()
			split_dict.add(_SCREAMING_SNAKE_CASE	)
			lowerCAmelCase__     :    List[Any]					  =      DatasetInfo(splits=_SCREAMING_SNAKE_CASE	)

			with tempfile.TemporaryDirectory() as tmp_dir:
					lowerCAmelCase__     :    Optional[Any]					  =      ReaderTest(_SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	)

					lowerCAmelCase__     :    Optional[int]					  =      '''train[:33%]'''
					lowerCAmelCase__     :    str					  =      Dataset(**reader.read(_SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	)	)
					self.assertEqual(dset['''filename'''][0]	,      F'{name}-train-00000-of-00010'	)
					self.assertEqual(dset['''filename'''][-1]	,      F'{name}-train-00003-of-00010'	)
					self.assertEqual(dset.num_rows	,      330	)
					self.assertEqual(dset.num_columns	,      1	)




	def 							UpperCAmelCase__(      self		:	Union[str, Any]	)-> Any:
			lowerCAmelCase__     :    str					  =      SplitInfo(name='''train'''	,      num_examples=100	)
			lowerCAmelCase__     :    List[str]					  =      SplitInfo(name='''test'''	,      num_examples=100	)
			lowerCAmelCase__     :    int					  =      SplitDict()
			split_dict.add(_SCREAMING_SNAKE_CASE	)
			split_dict.add(_SCREAMING_SNAKE_CASE	)
			lowerCAmelCase__     :    Optional[int]					  =      DatasetInfo(splits=_SCREAMING_SNAKE_CASE	)

			with tempfile.TemporaryDirectory() as tmp_dir:
					lowerCAmelCase__     :    Union[str, Any]					  =      ReaderTest(_SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	)

					lowerCAmelCase__     :    Tuple					  =      [
					    {'''filename''': os.path.join(_SCREAMING_SNAKE_CASE	,      '''train'''	)},
					    {'''filename''': os.path.join(_SCREAMING_SNAKE_CASE	,      '''test'''	), '''skip''': 10, '''take''': 10},
					]
					lowerCAmelCase__     :    str					  =      Dataset(**reader.read_files(_SCREAMING_SNAKE_CASE	,      original_instructions='''train+test[10:20]'''	)	)
					self.assertEqual(dset.num_rows	,      110	)
					self.assertEqual(dset.num_columns	,      1	)
					del dset






@pytest.mark.parametrize('''in_memory'''			,     [False, True] )
def   lowerCamelCase_					( _a			,     _a			,     _a ):




		"""simple docstring"""






		lowerCAmelCase__     :    Optional[int]					  =      arrow_file
		with assert_arrow_memory_increases() if in_memory else assert_arrow_memory_doesnt_increase():
				lowerCAmelCase__     :    Optional[int]					  =      ArrowReader.read_table(_a			,     in_memory=_a )
		assert table.shape == dataset.data.shape
		assert set(table.column_names ) == set(dataset.data.column_names )
		assert dict(table.to_pydict() ) == dict(dataset.data.to_pydict() )  # to_pydict returns OrderedDict






@pytest.mark.parametrize('''in_memory'''			,     [False, True] )
def   lowerCamelCase_					( _a			,     _a			,     _a ):




		"""simple docstring"""






		lowerCAmelCase__     :    Any					  =      arrow_file
		lowerCAmelCase__     :    List[Any]					  =      ArrowReader(''''''			,     _a )
		with assert_arrow_memory_increases() if in_memory else assert_arrow_memory_doesnt_increase():
				lowerCAmelCase__     :    Tuple					  =      reader.read_files([{'''filename''': filename}]			,     in_memory=_a )
		assert dataset_kwargs.keys() == {"arrow_table", "info", "split"}
		lowerCAmelCase__     :    Union[str, Any]					  =      dataset_kwargs['''arrow_table''']
		assert table.shape == dataset.data.shape
		assert set(table.column_names ) == set(dataset.data.column_names )
		assert dict(table.to_pydict() ) == dict(dataset.data.to_pydict() )  # to_pydict returns OrderedDict






def   lowerCamelCase_					( ):




		"""simple docstring"""






		assert ReadInstruction('''train'''			,     to=10			,     unit='''abs''' ).to_spec() == "train[:10]"
		assert ReadInstruction('''train'''			,     from_=-80			,     to=10			,     unit='''%''' ).to_spec() == "train[-80%:10%]"

		lowerCAmelCase__     :    List[str]					  =      '''train+test'''
		assert ReadInstruction.from_spec(_a ).to_spec() == spec_train_test

		lowerCAmelCase__     :    Optional[int]					  =      '''train[2:10]'''
		assert ReadInstruction.from_spec(_a ).to_spec() == spec_train_abs

		lowerCAmelCase__     :    Union[str, Any]					  =      '''train[15%:-20%]'''
		assert ReadInstruction.from_spec(_a ).to_spec() == spec_train_pct

		lowerCAmelCase__     :    Tuple					  =      '''train[:10%](closest)'''
		assert ReadInstruction.from_spec(_a ).to_spec() == "train[:10%]"

		lowerCAmelCase__     :    List[str]					  =      '''train[:10%](pct1_dropremainder)'''
		assert ReadInstruction.from_spec(_a ).to_spec() == spec_train_pct_rounding

		lowerCAmelCase__     :    Optional[Any]					  =      '''train[:10%](pct1_dropremainder)+test[-10%:](pct1_dropremainder)'''
		assert ReadInstruction.from_spec(_a ).to_spec() == spec_train_test_pct_rounding






def   lowerCamelCase_					( ):




		"""simple docstring"""






		lowerCAmelCase__     :    Dict					  =      '''dummy'''
		lowerCAmelCase__     :    int					  =      [SplitInfo(name='''train'''			,     num_examples=100 )]
		lowerCAmelCase__     :    Tuple					  =      '''train[:33%]'''
		lowerCAmelCase__     :    Dict					  =      '''arrow'''
		lowerCAmelCase__     :    List[Any]					  =      '''prefix'''

		lowerCAmelCase__     :    Optional[int]					  =      make_file_instructions(_a			,     _a			,     _a			,     _a			,     _a )
		assert isinstance(_a			,     _a )
		assert file_instructions.num_examples == 33
		assert file_instructions.file_instructions == [
		    {"filename": os.path.join(_a			,     f'{name}-train.arrow' ), "skip": 0, "take": 33}
		]

		lowerCAmelCase__     :    Optional[Any]					  =      [SplitInfo(name='''train'''			,     num_examples=100			,     shard_lengths=[10] * 10 )]
		lowerCAmelCase__     :    Tuple					  =      make_file_instructions(_a			,     _a			,     _a			,     _a			,     _a )
		assert isinstance(_a			,     _a )
		assert file_instructions.num_examples == 33
		assert file_instructions.file_instructions == [
		    {"filename": os.path.join(_a			,     f'{name}-train-00000-of-00010.arrow' ), "skip": 0, "take": -1},
		    {"filename": os.path.join(_a			,     f'{name}-train-00001-of-00010.arrow' ), "skip": 0, "take": -1},
		    {"filename": os.path.join(_a			,     f'{name}-train-00002-of-00010.arrow' ), "skip": 0, "take": -1},
		    {"filename": os.path.join(_a			,     f'{name}-train-00003-of-00010.arrow' ), "skip": 0, "take": 3},
		]






@pytest.mark.parametrize('''name, expected_exception'''			,     [(None, TypeError), ('''''', ValueError)] )
def   lowerCamelCase_					( _a			,     _a ):




		"""simple docstring"""






		lowerCAmelCase__     :    Any					  =      [SplitInfo(name='''train'''			,     num_examples=100 )]
		lowerCAmelCase__     :    Any					  =      '''train'''
		lowerCAmelCase__     :    Any					  =      '''arrow'''
		lowerCAmelCase__     :    str					  =      '''prefix_path'''
		with pytest.raises(_a ):
				lowerCAmelCase__     :    str					  =      make_file_instructions(_a			,     _a			,     _a			,     _a			,     _a )








from __future__ import annotations

import unittest

from transformers import AlbertConfig, is_tf_available
from transformers.models.auto import get_values
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 tensorflow as tf

		from transformers import TF_MODEL_FOR_PRETRAINING_MAPPING
		from transformers.models.albert.modeling_tf_albert import (
		    TF_ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
		    TFAlbertForMaskedLM,
		    TFAlbertForMultipleChoice,
		    TFAlbertForPreTraining,
		    TFAlbertForQuestionAnswering,
		    TFAlbertForSequenceClassification,
		    TFAlbertForTokenClassification,
		    TFAlbertModel,
		)





class       _a							:




	def __init__(      self		:	Tuple	,      _SCREAMING_SNAKE_CASE		:	List[Any]	,      _SCREAMING_SNAKE_CASE		:	Tuple=13	,      _SCREAMING_SNAKE_CASE		:	Dict=7	,      _SCREAMING_SNAKE_CASE		:	Union[str, Any]=True	,      _SCREAMING_SNAKE_CASE		:	Optional[Any]=True	,      _SCREAMING_SNAKE_CASE		:	Optional[Any]=True	,      _SCREAMING_SNAKE_CASE		:	int=True	,      _SCREAMING_SNAKE_CASE		:	Union[str, Any]=99	,      _SCREAMING_SNAKE_CASE		:	int=16	,      _SCREAMING_SNAKE_CASE		:	Union[str, Any]=32	,      _SCREAMING_SNAKE_CASE		:	Tuple=2	,      _SCREAMING_SNAKE_CASE		:	int=4	,      _SCREAMING_SNAKE_CASE		:	List[str]=37	,      _SCREAMING_SNAKE_CASE		:	Optional[Any]="gelu"	,      _SCREAMING_SNAKE_CASE		:	List[Any]=0.1	,      _SCREAMING_SNAKE_CASE		:	List[Any]=0.1	,      _SCREAMING_SNAKE_CASE		:	Union[str, Any]=512	,      _SCREAMING_SNAKE_CASE		:	Union[str, Any]=16	,      _SCREAMING_SNAKE_CASE		:	str=2	,      _SCREAMING_SNAKE_CASE		:	Optional[int]=0.02	,      _SCREAMING_SNAKE_CASE		:	Any=3	,      _SCREAMING_SNAKE_CASE		:	Dict=4	,      _SCREAMING_SNAKE_CASE		:	int=None	,      )-> str:
			lowerCAmelCase__     :    Any					  =      parent
			lowerCAmelCase__     :    List[Any]					  =      13
			lowerCAmelCase__     :    int					  =      7
			lowerCAmelCase__     :    Optional[int]					  =      True
			lowerCAmelCase__     :    Optional[Any]					  =      True
			lowerCAmelCase__     :    List[str]					  =      True
			lowerCAmelCase__     :    Optional[Any]					  =      True
			lowerCAmelCase__     :    Tuple					  =      99
			lowerCAmelCase__     :    Optional[Any]					  =      16
			lowerCAmelCase__     :    str					  =      32
			lowerCAmelCase__     :    List[str]					  =      2
			lowerCAmelCase__     :    List[str]					  =      4
			lowerCAmelCase__     :    Union[str, Any]					  =      37
			lowerCAmelCase__     :    Tuple					  =      '''gelu'''
			lowerCAmelCase__     :    Any					  =      0.1
			lowerCAmelCase__     :    Optional[Any]					  =      0.1
			lowerCAmelCase__     :    Tuple					  =      512
			lowerCAmelCase__     :    int					  =      16
			lowerCAmelCase__     :    Tuple					  =      2
			lowerCAmelCase__     :    List[str]					  =      0.02
			lowerCAmelCase__     :    Tuple					  =      3
			lowerCAmelCase__     :    str					  =      4
			lowerCAmelCase__     :    Union[str, Any]					  =      None




	def 							UpperCAmelCase__(      self		:	List[str]	)-> Any:
			lowerCAmelCase__     :    str					  =      ids_tensor([self.batch_size, self.seq_length]	,      self.vocab_size	)

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

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

			lowerCAmelCase__     :    str					  =      None
			lowerCAmelCase__     :    Union[str, Any]					  =      None
			lowerCAmelCase__     :    Optional[Any]					  =      None
			if self.use_labels:
					lowerCAmelCase__     :    Union[str, Any]					  =      ids_tensor([self.batch_size]	,      self.type_sequence_label_size	)
					lowerCAmelCase__     :    str					  =      ids_tensor([self.batch_size, self.seq_length]	,      self.num_labels	)
					lowerCAmelCase__     :    str					  =      ids_tensor([self.batch_size]	,      self.num_choices	)

			lowerCAmelCase__     :    List[Any]					  =      AlbertConfig(
			    vocab_size=self.vocab_size	,      hidden_size=self.hidden_size	,      embedding_size=self.embedding_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	,      )

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




	def 							UpperCAmelCase__(      self		:	int	,      _SCREAMING_SNAKE_CASE		:	Optional[Any]	,      _SCREAMING_SNAKE_CASE		:	Union[str, Any]	,      _SCREAMING_SNAKE_CASE		:	str	,      _SCREAMING_SNAKE_CASE		:	Union[str, Any]	,      _SCREAMING_SNAKE_CASE		:	Any	,      _SCREAMING_SNAKE_CASE		:	List[Any]	,      _SCREAMING_SNAKE_CASE		:	List[Any]	)-> List[Any]:
			lowerCAmelCase__     :    int					  =      TFAlbertModel(config=_SCREAMING_SNAKE_CASE	)
			# inputs = {'input_ids': input_ids,
			#           'attention_mask': input_mask,
			#           'token_type_ids': token_type_ids}
			# sequence_output, pooled_output = model(**inputs)
			lowerCAmelCase__     :    str					  =      {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
			lowerCAmelCase__     :    Any					  =      model(_SCREAMING_SNAKE_CASE	)

			lowerCAmelCase__     :    str					  =      [input_ids, input_mask]
			lowerCAmelCase__     :    Optional[Any]					  =      model(_SCREAMING_SNAKE_CASE	)

			lowerCAmelCase__     :    str					  =      model(_SCREAMING_SNAKE_CASE	)

			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 							UpperCAmelCase__(      self		:	Optional[int]	,      _SCREAMING_SNAKE_CASE		:	Tuple	,      _SCREAMING_SNAKE_CASE		:	Optional[int]	,      _SCREAMING_SNAKE_CASE		:	int	,      _SCREAMING_SNAKE_CASE		:	int	,      _SCREAMING_SNAKE_CASE		:	Dict	,      _SCREAMING_SNAKE_CASE		:	List[Any]	,      _SCREAMING_SNAKE_CASE		:	Tuple	)-> Optional[Any]:
			lowerCAmelCase__     :    Any					  =      self.num_labels
			lowerCAmelCase__     :    Dict					  =      TFAlbertForPreTraining(config=_SCREAMING_SNAKE_CASE	)
			lowerCAmelCase__     :    Tuple					  =      {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
			lowerCAmelCase__     :    Any					  =      model(_SCREAMING_SNAKE_CASE	)
			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, self.num_labels)	)




	def 							UpperCAmelCase__(      self		:	Optional[Any]	,      _SCREAMING_SNAKE_CASE		:	int	,      _SCREAMING_SNAKE_CASE		:	List[str]	,      _SCREAMING_SNAKE_CASE		:	Any	,      _SCREAMING_SNAKE_CASE		:	List[str]	,      _SCREAMING_SNAKE_CASE		:	Optional[Any]	,      _SCREAMING_SNAKE_CASE		:	Optional[int]	,      _SCREAMING_SNAKE_CASE		:	Any	)-> Optional[int]:
			lowerCAmelCase__     :    Any					  =      TFAlbertForMaskedLM(config=_SCREAMING_SNAKE_CASE	)
			lowerCAmelCase__     :    Any					  =      {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
			lowerCAmelCase__     :    Dict					  =      model(_SCREAMING_SNAKE_CASE	)
			self.parent.assertEqual(result.logits.shape	,      (self.batch_size, self.seq_length, self.vocab_size)	)




	def 							UpperCAmelCase__(      self		:	Optional[Any]	,      _SCREAMING_SNAKE_CASE		:	Tuple	,      _SCREAMING_SNAKE_CASE		:	Optional[Any]	,      _SCREAMING_SNAKE_CASE		:	Any	,      _SCREAMING_SNAKE_CASE		:	str	,      _SCREAMING_SNAKE_CASE		:	int	,      _SCREAMING_SNAKE_CASE		:	Optional[int]	,      _SCREAMING_SNAKE_CASE		:	Any	)-> List[str]:
			lowerCAmelCase__     :    Union[str, Any]					  =      self.num_labels
			lowerCAmelCase__     :    List[str]					  =      TFAlbertForSequenceClassification(config=_SCREAMING_SNAKE_CASE	)
			lowerCAmelCase__     :    str					  =      {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
			lowerCAmelCase__     :    int					  =      model(_SCREAMING_SNAKE_CASE	)
			self.parent.assertEqual(result.logits.shape	,      (self.batch_size, self.num_labels)	)




	def 							UpperCAmelCase__(      self		:	Dict	,      _SCREAMING_SNAKE_CASE		:	Optional[int]	,      _SCREAMING_SNAKE_CASE		:	List[Any]	,      _SCREAMING_SNAKE_CASE		:	str	,      _SCREAMING_SNAKE_CASE		:	List[Any]	,      _SCREAMING_SNAKE_CASE		:	int	,      _SCREAMING_SNAKE_CASE		:	List[str]	,      _SCREAMING_SNAKE_CASE		:	Optional[int]	)-> int:
			lowerCAmelCase__     :    List[str]					  =      TFAlbertForQuestionAnswering(config=_SCREAMING_SNAKE_CASE	)
			lowerCAmelCase__     :    List[str]					  =      {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
			lowerCAmelCase__     :    Optional[Any]					  =      model(_SCREAMING_SNAKE_CASE	)
			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 							UpperCAmelCase__(      self		:	Tuple	,      _SCREAMING_SNAKE_CASE		:	List[Any]	,      _SCREAMING_SNAKE_CASE		:	int	,      _SCREAMING_SNAKE_CASE		:	Optional[Any]	,      _SCREAMING_SNAKE_CASE		:	str	,      _SCREAMING_SNAKE_CASE		:	Any	,      _SCREAMING_SNAKE_CASE		:	Union[str, Any]	,      _SCREAMING_SNAKE_CASE		:	int	)-> str:
			lowerCAmelCase__     :    str					  =      self.num_choices
			lowerCAmelCase__     :    Tuple					  =      TFAlbertForMultipleChoice(config=_SCREAMING_SNAKE_CASE	)
			lowerCAmelCase__     :    Union[str, Any]					  =      tf.tile(tf.expand_dims(_SCREAMING_SNAKE_CASE	,      1	)	,      (1, self.num_choices, 1)	)
			lowerCAmelCase__     :    Dict					  =      tf.tile(tf.expand_dims(_SCREAMING_SNAKE_CASE	,      1	)	,      (1, self.num_choices, 1)	)
			lowerCAmelCase__     :    Tuple					  =      tf.tile(tf.expand_dims(_SCREAMING_SNAKE_CASE	,      1	)	,      (1, self.num_choices, 1)	)
			lowerCAmelCase__     :    Tuple					  =      {
			    '''input_ids''': multiple_choice_inputs_ids,
			    '''attention_mask''': multiple_choice_input_mask,
			    '''token_type_ids''': multiple_choice_token_type_ids,
			}
			lowerCAmelCase__     :    Optional[Any]					  =      model(_SCREAMING_SNAKE_CASE	)
			self.parent.assertListEqual(list(result['''logits'''].shape	)	,      [self.batch_size, self.num_choices]	)




	def 							UpperCAmelCase__(      self		:	List[Any]	,      _SCREAMING_SNAKE_CASE		:	Tuple	,      _SCREAMING_SNAKE_CASE		:	Optional[int]	,      _SCREAMING_SNAKE_CASE		:	List[str]	,      _SCREAMING_SNAKE_CASE		:	Optional[Any]	,      _SCREAMING_SNAKE_CASE		:	Union[str, Any]	,      _SCREAMING_SNAKE_CASE		:	int	,      _SCREAMING_SNAKE_CASE		:	List[str]	)-> Tuple:
			lowerCAmelCase__     :    Union[str, Any]					  =      self.num_labels
			lowerCAmelCase__     :    Optional[int]					  =      TFAlbertForTokenClassification(config=_SCREAMING_SNAKE_CASE	)
			lowerCAmelCase__     :    Any					  =      {
			    '''input_ids''': input_ids,
			    '''attention_mask''': input_mask,
			    '''token_type_ids''': token_type_ids,
			}
			lowerCAmelCase__     :    Optional[Any]					  =      model(_SCREAMING_SNAKE_CASE	)
			self.parent.assertListEqual(list(result['''logits'''].shape	)	,      [self.batch_size, self.seq_length, self.num_labels]	)




	def 							UpperCAmelCase__(      self		:	Union[str, Any]	)-> Dict:
			lowerCAmelCase__     :    int					  =      self.prepare_config_and_inputs()
			(
			    (
			    lowerCAmelCase__
			)  ,			(
			    lowerCAmelCase__
			)  ,			(
			    lowerCAmelCase__
			)  ,			(
			    lowerCAmelCase__
			)  ,			(
			    lowerCAmelCase__
			)  ,			(
			    lowerCAmelCase__
			)  ,			(
			    lowerCAmelCase__
			)  ,			
			)     :    Union[str, Any]					  =      config_and_inputs
			lowerCAmelCase__     :    Tuple					  =      {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
			return config, inputs_dict





@require_tf
class       _a							(       _lowercase  ,		_lowercase  ,		unittest.TestCase):
	_a						:						Tuple   					=							(
	    (
	        TFAlbertModel,
	        TFAlbertForPreTraining,
	        TFAlbertForMaskedLM,
	        TFAlbertForSequenceClassification,
	        TFAlbertForQuestionAnswering,
	        TFAlbertForTokenClassification,
	        TFAlbertForMultipleChoice,
	    )
	    if is_tf_available()
	    else ()
	)
	_a						:						Optional[Any]   					=							(
	    {
	        '''feature-extraction''': TFAlbertModel,
	        '''fill-mask''': TFAlbertForMaskedLM,
	        '''question-answering''': TFAlbertForQuestionAnswering,
	        '''text-classification''': TFAlbertForSequenceClassification,
	        '''token-classification''': TFAlbertForTokenClassification,
	        '''zero-shot''': TFAlbertForSequenceClassification,
	    }
	    if is_tf_available()
	    else {}
	)
	_a						:						Union[str, Any]   					=							False
	_a						:						int   					=							False




	def 							UpperCAmelCase__(      self		:	Any	,      _SCREAMING_SNAKE_CASE		:	Optional[Any]	,      _SCREAMING_SNAKE_CASE		:	Any	,      _SCREAMING_SNAKE_CASE		:	List[str]=False	)-> Union[str, Any]:
			lowerCAmelCase__     :    Any					  =      super()._prepare_for_class(_SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	,      return_labels=_SCREAMING_SNAKE_CASE	)

			if return_labels:
					if model_class in get_values(_SCREAMING_SNAKE_CASE	):
							lowerCAmelCase__     :    Optional[Any]					  =      tf.zeros(self.model_tester.batch_size	,      dtype=tf.intaa	)

			return inputs_dict




	def 							UpperCAmelCase__(      self		:	List[str]	)-> Tuple:
			lowerCAmelCase__     :    List[Any]					  =      TFAlbertModelTester(self	)
			lowerCAmelCase__     :    List[Any]					  =      ConfigTester(self	,      config_class=_SCREAMING_SNAKE_CASE	,      hidden_size=37	)




	def 							UpperCAmelCase__(      self		:	str	)-> Any:
			self.config_tester.run_common_tests()




	def 							UpperCAmelCase__(      self		:	Optional[Any]	)-> Dict:
			lowerCAmelCase__     :    Any					  =      self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_albert_model(*_SCREAMING_SNAKE_CASE	)




	def 							UpperCAmelCase__(      self		:	Any	)-> str:
			lowerCAmelCase__     :    Any					  =      self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_albert_for_pretraining(*_SCREAMING_SNAKE_CASE	)




	def 							UpperCAmelCase__(      self		:	Optional[Any]	)-> str:
			lowerCAmelCase__     :    List[str]					  =      self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_albert_for_masked_lm(*_SCREAMING_SNAKE_CASE	)




	def 							UpperCAmelCase__(      self		:	Union[str, Any]	)-> Union[str, Any]:
			lowerCAmelCase__     :    int					  =      self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_albert_for_multiple_choice(*_SCREAMING_SNAKE_CASE	)




	def 							UpperCAmelCase__(      self		:	List[Any]	)-> List[str]:
			lowerCAmelCase__     :    Union[str, Any]					  =      self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_albert_for_sequence_classification(*_SCREAMING_SNAKE_CASE	)




	def 							UpperCAmelCase__(      self		:	str	)-> List[str]:
			lowerCAmelCase__     :    Dict					  =      self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_albert_for_question_answering(*_SCREAMING_SNAKE_CASE	)




	@slow
	def 							UpperCAmelCase__(      self		:	str	)-> Optional[int]:
			for model_name in TF_ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
					lowerCAmelCase__     :    Optional[int]					  =      TFAlbertModel.from_pretrained(_SCREAMING_SNAKE_CASE	)
					self.assertIsNotNone(_SCREAMING_SNAKE_CASE	)





@require_tf
class       _a							(       unittest.TestCase):




	@slow
	def 							UpperCAmelCase__(      self		:	str	)-> str:
			lowerCAmelCase__     :    Tuple					  =      TFAlbertForPreTraining.from_pretrained('''albert-base-v2'''	)
			lowerCAmelCase__     :    Tuple					  =      tf.constant([[0, 1, 2, 3, 4, 5]]	)
			lowerCAmelCase__     :    List[str]					  =      model(_SCREAMING_SNAKE_CASE	)[0]

			lowerCAmelCase__     :    Union[str, Any]					  =      [1, 6, 3_0000]
			self.assertEqual(output.shape	,      _SCREAMING_SNAKE_CASE	)

			lowerCAmelCase__     :    int					  =      tf.constant(
			    [
			        [
			            [4.59_5668, 0.7446_2754, -1.81_8147],
			            [4.595_4347, 0.745_4184, -1.818_8258],
			            [4.595_4905, 0.744_8235, -1.818_2316],
			        ]
			    ]	)
			tf.debugging.assert_near(output[:, :3, :3]	,      _SCREAMING_SNAKE_CASE	,      atol=1E-4	)




