






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  __lowercase :







     '''simple docstring'''




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


         """simple docstring"""



         __A             =		parent
         __A             =		13
         __A             =		7
         __A             =		True
         __A             =		True
         __A             =		True
         __A             =		True
         __A             =		99
         __A             =		16
         __A             =		32
         __A             =		2
         __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




     def 						lowerCAmelCase_       (			self    :     Dict			):


         """simple docstring"""



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

         __A             =		None
         if self.use_input_mask:
             __A             =		random_attention_mask([self.batch_size, self.seq_length]			)

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

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

         __A             =		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 						lowerCAmelCase_       (			self    :     Tuple	,		UpperCamelCase_    :     Union[str, Any]	,		UpperCamelCase_    :     Union[str, Any]	,		UpperCamelCase_    :     Dict	,		UpperCamelCase_    :     Dict	,		UpperCamelCase_    :     List[str]	,		UpperCamelCase_    :     Tuple	,		UpperCamelCase_    :     Optional[int]			):


         """simple docstring"""



         __A             =		TFAlbertModel(config=UpperCamelCase_			)
         # inputs = {'input_ids': input_ids,
         #           'attention_mask': input_mask,
         #           'token_type_ids': token_type_ids}
         # sequence_output, pooled_output = model(**inputs)
         __A             =		{"""input_ids""": input_ids, """attention_mask""": input_mask, """token_type_ids""": token_type_ids}
         __A             =		model(UpperCamelCase_			)

         __A             =		[input_ids, input_mask]
         __A             =		model(UpperCamelCase_			)

         __A             =		model(UpperCamelCase_			)

         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 						lowerCAmelCase_       (			self    :     int	,		UpperCamelCase_    :     List[Any]	,		UpperCamelCase_    :     Union[str, Any]	,		UpperCamelCase_    :     Tuple	,		UpperCamelCase_    :     Union[str, Any]	,		UpperCamelCase_    :     Optional[int]	,		UpperCamelCase_    :     Optional[int]	,		UpperCamelCase_    :     List[str]			):


         """simple docstring"""



         __A             =		self.num_labels
         __A             =		TFAlbertForPreTraining(config=UpperCamelCase_			)
         __A             =		{"""input_ids""": input_ids, """attention_mask""": input_mask, """token_type_ids""": token_type_ids}
         __A             =		model(UpperCamelCase_			)
         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 						lowerCAmelCase_       (			self    :     Optional[int]	,		UpperCamelCase_    :     str	,		UpperCamelCase_    :     Tuple	,		UpperCamelCase_    :     List[Any]	,		UpperCamelCase_    :     Optional[Any]	,		UpperCamelCase_    :     Union[str, Any]	,		UpperCamelCase_    :     List[Any]	,		UpperCamelCase_    :     Union[str, Any]			):


         """simple docstring"""



         __A             =		TFAlbertForMaskedLM(config=UpperCamelCase_			)
         __A             =		{"""input_ids""": input_ids, """attention_mask""": input_mask, """token_type_ids""": token_type_ids}
         __A             =		model(UpperCamelCase_			)
         self.parent.assertEqual(result.logits.shape	,		(self.batch_size, self.seq_length, self.vocab_size)			)




     def 						lowerCAmelCase_       (			self    :     Tuple	,		UpperCamelCase_    :     Any	,		UpperCamelCase_    :     Optional[Any]	,		UpperCamelCase_    :     List[Any]	,		UpperCamelCase_    :     int	,		UpperCamelCase_    :     Any	,		UpperCamelCase_    :     Optional[int]	,		UpperCamelCase_    :     Tuple			):


         """simple docstring"""



         __A             =		self.num_labels
         __A             =		TFAlbertForSequenceClassification(config=UpperCamelCase_			)
         __A             =		{"""input_ids""": input_ids, """attention_mask""": input_mask, """token_type_ids""": token_type_ids}
         __A             =		model(UpperCamelCase_			)
         self.parent.assertEqual(result.logits.shape	,		(self.batch_size, self.num_labels)			)




     def 						lowerCAmelCase_       (			self    :     List[Any]	,		UpperCamelCase_    :     Optional[Any]	,		UpperCamelCase_    :     Union[str, Any]	,		UpperCamelCase_    :     Optional[int]	,		UpperCamelCase_    :     Optional[Any]	,		UpperCamelCase_    :     Dict	,		UpperCamelCase_    :     List[Any]	,		UpperCamelCase_    :     int			):


         """simple docstring"""



         __A             =		TFAlbertForQuestionAnswering(config=UpperCamelCase_			)
         __A             =		{"""input_ids""": input_ids, """attention_mask""": input_mask, """token_type_ids""": token_type_ids}
         __A             =		model(UpperCamelCase_			)
         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 						lowerCAmelCase_       (			self    :     int	,		UpperCamelCase_    :     Optional[Any]	,		UpperCamelCase_    :     List[Any]	,		UpperCamelCase_    :     Optional[Any]	,		UpperCamelCase_    :     List[str]	,		UpperCamelCase_    :     Optional[int]	,		UpperCamelCase_    :     Optional[int]	,		UpperCamelCase_    :     Optional[Any]			):


         """simple docstring"""



         __A             =		self.num_choices
         __A             =		TFAlbertForMultipleChoice(config=UpperCamelCase_			)
         __A             =		tf.tile(tf.expand_dims(UpperCamelCase_	,		1			)	,		(1, self.num_choices, 1)			)
         __A             =		tf.tile(tf.expand_dims(UpperCamelCase_	,		1			)	,		(1, self.num_choices, 1)			)
         __A             =		tf.tile(tf.expand_dims(UpperCamelCase_	,		1			)	,		(1, self.num_choices, 1)			)
         __A             =		{
             """input_ids""": multiple_choice_inputs_ids,
             """attention_mask""": multiple_choice_input_mask,
             """token_type_ids""": multiple_choice_token_type_ids,
         }
         __A             =		model(UpperCamelCase_			)
         self.parent.assertListEqual(list(result["""logits"""].shape			)	,		[self.batch_size, self.num_choices]			)




     def 						lowerCAmelCase_       (			self    :     int	,		UpperCamelCase_    :     str	,		UpperCamelCase_    :     Dict	,		UpperCamelCase_    :     Union[str, Any]	,		UpperCamelCase_    :     Union[str, Any]	,		UpperCamelCase_    :     Optional[int]	,		UpperCamelCase_    :     List[Any]	,		UpperCamelCase_    :     Optional[Any]			):


         """simple docstring"""



         __A             =		self.num_labels
         __A             =		TFAlbertForTokenClassification(config=UpperCamelCase_			)
         __A             =		{
             """input_ids""": input_ids,
             """attention_mask""": input_mask,
             """token_type_ids""": token_type_ids,
         }
         __A             =		model(UpperCamelCase_			)
         self.parent.assertListEqual(list(result["""logits"""].shape			)	,		[self.batch_size, self.seq_length, self.num_labels]			)





     def 						lowerCAmelCase_       (			self    :     List[Any]			):


         """simple docstring"""



         __A             =		self.prepare_config_and_inputs()
         (
             (
             __A
         )      ,		(
             __A
         )      ,		(
             __A
         )      ,		(
             __A
         )      ,		(
             __A
         )      ,		(
             __A
         )      ,		(
             __A
         )      ,		
         )             =		config_and_inputs
         __A             =		{"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
         return config, inputs_dict


@require_tf
class  __lowercase (		lowercase_     , lowercase_     , unittest.TestCase					):







     '''simple docstring'''
     SCREAMING_SNAKE_CASE				  =			(
         (
             TFAlbertModel,
             TFAlbertForPreTraining,
             TFAlbertForMaskedLM,
             TFAlbertForSequenceClassification,
             TFAlbertForQuestionAnswering,
             TFAlbertForTokenClassification,
             TFAlbertForMultipleChoice,
         )
         if is_tf_available()
         else ()
     )
     SCREAMING_SNAKE_CASE				  =			(
         {
             "feature-extraction": TFAlbertModel,
             "fill-mask": TFAlbertForMaskedLM,
             "question-answering": TFAlbertForQuestionAnswering,
             "text-classification": TFAlbertForSequenceClassification,
             "token-classification": TFAlbertForTokenClassification,
             "zero-shot": TFAlbertForSequenceClassification,
         }
         if is_tf_available()
         else {}
     )
     SCREAMING_SNAKE_CASE				  =			False
     SCREAMING_SNAKE_CASE				  =			False




     def 						lowerCAmelCase_       (			self    :     int	,		UpperCamelCase_    :     List[str]	,		UpperCamelCase_    :     Union[str, Any]	,		UpperCamelCase_    :     Optional[Any]=False			):


         """simple docstring"""



         __A             =		super()._prepare_for_class(UpperCamelCase_	,		UpperCamelCase_	,		return_labels=UpperCamelCase_			)

         if return_labels:
             if model_class in get_values(UpperCamelCase_			):
                 __A             =		tf.zeros(self.model_tester.batch_size	,		dtype=tf.intaa			)

         return inputs_dict




     def 						lowerCAmelCase_       (			self    :     Optional[Any]			):


         """simple docstring"""



         __A             =		TFAlbertModelTester(self			)
         __A             =		ConfigTester(self	,		config_class=UpperCamelCase_	,		hidden_size=37			)




     def 						lowerCAmelCase_       (			self    :     List[str]			):


         """simple docstring"""



         self.config_tester.run_common_tests()




     def 						lowerCAmelCase_       (			self    :     Optional[Any]			):


         """simple docstring"""



         __A             =		self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_albert_model(*UpperCamelCase_			)




     def 						lowerCAmelCase_       (			self    :     Tuple			):


         """simple docstring"""



         __A             =		self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_albert_for_pretraining(*UpperCamelCase_			)




     def 						lowerCAmelCase_       (			self    :     List[Any]			):


         """simple docstring"""



         __A             =		self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_albert_for_masked_lm(*UpperCamelCase_			)




     def 						lowerCAmelCase_       (			self    :     Tuple			):


         """simple docstring"""



         __A             =		self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_albert_for_multiple_choice(*UpperCamelCase_			)




     def 						lowerCAmelCase_       (			self    :     Any			):


         """simple docstring"""



         __A             =		self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_albert_for_sequence_classification(*UpperCamelCase_			)




     def 						lowerCAmelCase_       (			self    :     Any			):


         """simple docstring"""



         __A             =		self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_albert_for_question_answering(*UpperCamelCase_			)





     @slow
     def 						lowerCAmelCase_       (			self    :     List[Any]			):


         """simple docstring"""



         for model_name in TF_ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             __A             =		TFAlbertModel.from_pretrained(UpperCamelCase_			)
             self.assertIsNotNone(UpperCamelCase_			)


@require_tf
class  __lowercase (		unittest.TestCase					):







     '''simple docstring'''




     @slow
     def 						lowerCAmelCase_       (			self    :     Optional[int]			):


         """simple docstring"""



         __A             =		TFAlbertForPreTraining.from_pretrained("""albert-base-v2"""			)
         __A             =		tf.constant([[0, 1, 2, 3, 4, 5]]			)
         __A             =		model(UpperCamelCase_			)[0]

         __A             =		[1, 6, 30_000]
         self.assertEqual(output.shape	,		UpperCamelCase_			)

         __A             =		tf.constant(
             [
                 [
                     [4.595668, 0.74462754, -1.818147],
                     [4.5954347, 0.7454184, -1.8188258],
                     [4.5954905, 0.7448235, -1.8182316],
                 ]
             ]			)
         tf.debugging.assert_near(output[:, :3, :3]	,		UpperCamelCase_	,		atol=1e-4			)












import importlib
import inspect
from functools import wraps
from typing import TYPE_CHECKING, Optional, Union

from .download.streaming_download_manager import (
    xbasename,
    xdirname,
    xet_parse,
    xexists,
    xgetsize,
    xglob,
    xgzip_open,
    xisdir,
    xisfile,
    xjoin,
    xlistdir,
    xnumpy_load,
    xopen,
    xpandas_read_csv,
    xpandas_read_excel,
    xPath,
    xrelpath,
    xsio_loadmat,
    xsplit,
    xsplitext,
    xwalk,
    xxml_dom_minidom_parse,
)
from .utils.logging import get_logger
from .utils.patching import patch_submodule
from .utils.py_utils import get_imports


__a :    List[str]								=				get_logger(__name__)


if TYPE_CHECKING:
  from .builder import DatasetBuilder



def      _SCREAMING_SNAKE_CASE  (	__lowercase				:     Dict     ,			__lowercase				:     Optional[Union[str, bool]] = None      )  -> List[str]:





    """simple docstring"""

    __A             =		importlib.import_module(__lowercase      )

    # TODO(QL): always update the module to add subsequent new authentication
    if hasattr(__lowercase     ,			"""_patched_for_streaming"""      ) and module._patched_for_streaming:
        return

    def wrap_auth(__lowercase				:     int      ):
        @wraps(__lowercase      )
        def wrapper(*__lowercase				:     Tuple     ,			**__lowercase				:     List[str]      ):
            return function(*__lowercase     ,			use_auth_token=__lowercase     ,			**__lowercase      )

        __A             =		"""wrap_auth"""
        return wrapper

    # open files in a streaming fashion
    patch_submodule(__lowercase     ,			"""open"""     ,			wrap_auth(__lowercase      )      ).start()
    patch_submodule(__lowercase     ,			"""os.listdir"""     ,			wrap_auth(__lowercase      )      ).start()
    patch_submodule(__lowercase     ,			"""os.walk"""     ,			wrap_auth(__lowercase      )      ).start()
    patch_submodule(__lowercase     ,			"""glob.glob"""     ,			wrap_auth(__lowercase      )      ).start()
    # allow to navigate in remote zip files
    patch_submodule(__lowercase     ,			"""os.path.join"""     ,			__lowercase      ).start()
    patch_submodule(__lowercase     ,			"""os.path.dirname"""     ,			__lowercase      ).start()
    patch_submodule(__lowercase     ,			"""os.path.basename"""     ,			__lowercase      ).start()
    patch_submodule(__lowercase     ,			"""os.path.relpath"""     ,			__lowercase      ).start()
    patch_submodule(__lowercase     ,			"""os.path.split"""     ,			__lowercase      ).start()
    patch_submodule(__lowercase     ,			"""os.path.splitext"""     ,			__lowercase      ).start()
    # allow checks on paths
    patch_submodule(__lowercase     ,			"""os.path.exists"""     ,			wrap_auth(__lowercase      )      ).start()
    patch_submodule(__lowercase     ,			"""os.path.isdir"""     ,			wrap_auth(__lowercase      )      ).start()
    patch_submodule(__lowercase     ,			"""os.path.isfile"""     ,			wrap_auth(__lowercase      )      ).start()
    patch_submodule(__lowercase     ,			"""os.path.getsize"""     ,			wrap_auth(__lowercase      )      ).start()
    patch_submodule(__lowercase     ,			"""pathlib.Path"""     ,			__lowercase      ).start()
    # file readers
    patch_submodule(__lowercase     ,			"""gzip.open"""     ,			wrap_auth(__lowercase      )      ).start()
    patch_submodule(__lowercase     ,			"""numpy.load"""     ,			wrap_auth(__lowercase      )      ).start()
    patch_submodule(__lowercase     ,			"""pandas.read_csv"""     ,			wrap_auth(__lowercase      )     ,			attrs=["""__version__"""]      ).start()
    patch_submodule(__lowercase     ,			"""pandas.read_excel"""     ,			wrap_auth(__lowercase      )     ,			attrs=["""__version__"""]      ).start()
    patch_submodule(__lowercase     ,			"""scipy.io.loadmat"""     ,			wrap_auth(__lowercase      )     ,			attrs=["""__version__"""]      ).start()
    patch_submodule(__lowercase     ,			"""xml.etree.ElementTree.parse"""     ,			wrap_auth(__lowercase      )      ).start()
    patch_submodule(__lowercase     ,			"""xml.dom.minidom.parse"""     ,			wrap_auth(__lowercase      )      ).start()
    __A             =		True






def      _SCREAMING_SNAKE_CASE  (	__lowercase				:     "DatasetBuilder"      )  -> Union[str, Any]:





    """simple docstring"""

    extend_module_for_streaming(builder.__module__     ,			use_auth_token=builder.use_auth_token      )
    # if needed, we also have to extend additional internal imports (like wmt14 -> wmt_utils)
    if not builder.__module__.startswith("""datasets."""      ):  # check that it's not a packaged builder like csv
        for imports in get_imports(inspect.getfile(builder.__class__      )      ):
            if imports[0] == "internal":
                __A             =		imports[1]
                __A             =		""".""".join(builder.__module__.split("""."""      )[:-1] + [internal_import_name]      )
                extend_module_for_streaming(__lowercase     ,			use_auth_token=builder.use_auth_token      )

    # builders can inherit from other builders that might use streaming functionality
    # (for example, ImageFolder and AudioFolder inherit from FolderBuilder which implements examples generation)
    # but these parents builders are not patched automatically as they are not instantiated, so we patch them here
    from .builder import DatasetBuilder

    __A             =		[
        cls.__module__
        for cls in type(__lowercase      ).__mro__[1:]  # make sure it's not the same module we've already patched
        if issubclass(cls     ,			__lowercase      ) and cls.__module__ != DatasetBuilder.__module__
    ]  # check it's not a standard builder from datasets.builder
    for module in parent_builder_modules:
        extend_module_for_streaming(__lowercase     ,			use_auth_token=builder.use_auth_token      )




