
'''simple docstring'''






import os
import tempfile
import unittest
from pathlib import Path

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


if is_tf_available():
 import tensorflow as tf

 from transformers import TensorFlowBenchmark, TensorFlowBenchmarkArguments

@require_tf
class  UpperCAmelCase__						(       unittest.TestCase ):



    def   lowerCamelCase_  (	self      :       Optional[int],__A      :       Optional[Any]     ):

          for model_result in results.values():
                for batch_size, sequence_length in zip(model_result["bs"],model_result["ss"]     ):
                      _lowerCamelCase				:   List[str]								=   model_result["""result"""][batch_size][sequence_length]
                      self.assertIsNotNone(_lowercase     )



    def   lowerCamelCase_  (	self      :       str     ):
          _lowerCamelCase				:   Dict								=   """sshleifer/tiny-gpt2"""
          _lowerCamelCase				:   List[Any]								=   TensorFlowBenchmarkArguments(
              models=[MODEL_ID],training=_lowercase,inference=_lowercase,sequence_lengths=[8],batch_sizes=[1],eager_mode=_lowercase,multi_process=_lowercase,)
          _lowerCamelCase				:   Dict								=   TensorFlowBenchmark(_lowercase     )
          _lowerCamelCase				:   str								=   benchmark.run()
          self.check_results_dict_not_empty(results.time_inference_result     )
          self.check_results_dict_not_empty(results.memory_inference_result     )



    def   lowerCamelCase_  (	self      :       Union[str, Any]     ):
          _lowerCamelCase				:   Optional[Any]								=   """sgugger/tiny-distilbert-classification"""
          _lowerCamelCase				:   List[Any]								=   TensorFlowBenchmarkArguments(
              models=[MODEL_ID],training=_lowercase,inference=_lowercase,sequence_lengths=[8],batch_sizes=[1],multi_process=_lowercase,only_pretrain_model=_lowercase,)
          _lowerCamelCase				:   Optional[Any]								=   TensorFlowBenchmark(_lowercase     )
          _lowerCamelCase				:   str								=   benchmark.run()
          self.check_results_dict_not_empty(results.time_inference_result     )
          self.check_results_dict_not_empty(results.memory_inference_result     )



    def   lowerCamelCase_  (	self      :       List[str]     ):
          _lowerCamelCase				:   Optional[Any]								=   """sshleifer/tiny-gpt2"""
          _lowerCamelCase				:   Tuple								=   TensorFlowBenchmarkArguments(
              models=[MODEL_ID],training=_lowercase,inference=_lowercase,sequence_lengths=[8],batch_sizes=[1],multi_process=_lowercase,)
          _lowerCamelCase				:   int								=   TensorFlowBenchmark(_lowercase     )
          _lowerCamelCase				:   Tuple								=   benchmark.run()
          self.check_results_dict_not_empty(results.time_inference_result     )
          self.check_results_dict_not_empty(results.memory_inference_result     )



    def   lowerCamelCase_  (	self      :       Dict     ):
          _lowerCamelCase				:   Union[str, Any]								=   """sshleifer/tiny-gpt2"""
          _lowerCamelCase				:   str								=   AutoConfig.from_pretrained(_lowercase     )
          _lowerCamelCase				:   List[Any]								=   TensorFlowBenchmarkArguments(
              models=[MODEL_ID],training=_lowercase,inference=_lowercase,sequence_lengths=[8],batch_sizes=[1],eager_mode=_lowercase,multi_process=_lowercase,)
          _lowerCamelCase				:   List[Any]								=   TensorFlowBenchmark(_lowercase,[config]     )
          _lowerCamelCase				:   Union[str, Any]								=   benchmark.run()
          self.check_results_dict_not_empty(results.time_inference_result     )
          self.check_results_dict_not_empty(results.memory_inference_result     )



    def   lowerCamelCase_  (	self      :       List[str]     ):
          _lowerCamelCase				:   Optional[Any]								=   """sshleifer/tiny-gpt2"""
          _lowerCamelCase				:   Optional[int]								=   AutoConfig.from_pretrained(_lowercase     )
          _lowerCamelCase				:   Optional[int]								=   TensorFlowBenchmarkArguments(
              models=[MODEL_ID],training=_lowercase,inference=_lowercase,sequence_lengths=[8],batch_sizes=[1],multi_process=_lowercase,)
          _lowerCamelCase				:   List[str]								=   TensorFlowBenchmark(_lowercase,[config]     )
          _lowerCamelCase				:   Union[str, Any]								=   benchmark.run()
          self.check_results_dict_not_empty(results.time_inference_result     )
          self.check_results_dict_not_empty(results.memory_inference_result     )



    def   lowerCamelCase_  (	self      :       List[Any]     ):
          _lowerCamelCase				:   Any								=   """sshleifer/tiny-gpt2"""
          _lowerCamelCase				:   Optional[Any]								=   TensorFlowBenchmarkArguments(
              models=[MODEL_ID],training=_lowercase,inference=_lowercase,sequence_lengths=[8],batch_sizes=[1],multi_process=_lowercase,)
          _lowerCamelCase				:   List[Any]								=   TensorFlowBenchmark(_lowercase     )
          _lowerCamelCase				:   List[str]								=   benchmark.run()
          self.check_results_dict_not_empty(results.time_train_result     )
          self.check_results_dict_not_empty(results.memory_train_result     )



    def   lowerCamelCase_  (	self      :       Optional[int]     ):
          _lowerCamelCase				:   Optional[int]								=   """sshleifer/tiny-gpt2"""
          _lowerCamelCase				:   List[Any]								=   AutoConfig.from_pretrained(_lowercase     )
          _lowerCamelCase				:   Any								=   TensorFlowBenchmarkArguments(
              models=[MODEL_ID],training=_lowercase,inference=_lowercase,sequence_lengths=[8],batch_sizes=[1],multi_process=_lowercase,)
          _lowerCamelCase				:   Optional[Any]								=   TensorFlowBenchmark(_lowercase,[config]     )
          _lowerCamelCase				:   Union[str, Any]								=   benchmark.run()
          self.check_results_dict_not_empty(results.time_train_result     )
          self.check_results_dict_not_empty(results.memory_train_result     )



    def   lowerCamelCase_  (	self      :       Dict     ):
          _lowerCamelCase				:   Dict								=   """patrickvonplaten/t5-tiny-random"""
          _lowerCamelCase				:   int								=   AutoConfig.from_pretrained(_lowercase     )
          _lowerCamelCase				:   Tuple								=   TensorFlowBenchmarkArguments(
              models=[MODEL_ID],training=_lowercase,inference=_lowercase,sequence_lengths=[8],batch_sizes=[1],multi_process=_lowercase,)
          _lowerCamelCase				:   Tuple								=   TensorFlowBenchmark(_lowercase,configs=[config]     )
          _lowerCamelCase				:   int								=   benchmark.run()
          self.check_results_dict_not_empty(results.time_inference_result     )
          self.check_results_dict_not_empty(results.memory_inference_result     )



    @unittest.skipIf(is_tf_available() and len(tf.config.list_physical_devices("GPU"     )     ) == 0,"Cannot do xla on CPU."     )
    def   lowerCamelCase_  (	self      :       int     ):
          _lowerCamelCase				:   Dict								=   """sshleifer/tiny-gpt2"""
          _lowerCamelCase				:   List[str]								=   TensorFlowBenchmarkArguments(
              models=[MODEL_ID],training=_lowercase,inference=_lowercase,sequence_lengths=[8],batch_sizes=[1],use_xla=_lowercase,multi_process=_lowercase,)
          _lowerCamelCase				:   int								=   TensorFlowBenchmark(_lowercase     )
          _lowerCamelCase				:   List[Any]								=   benchmark.run()
          self.check_results_dict_not_empty(results.time_inference_result     )
          self.check_results_dict_not_empty(results.memory_inference_result     )



    def   lowerCamelCase_  (	self      :       Union[str, Any]     ):
          _lowerCamelCase				:   str								=   """sshleifer/tiny-gpt2"""
          with tempfile.TemporaryDirectory() as tmp_dir:
                _lowerCamelCase				:   str								=   TensorFlowBenchmarkArguments(
                    models=[MODEL_ID],inference=_lowercase,save_to_csv=_lowercase,sequence_lengths=[8],batch_sizes=[1],inference_time_csv_file=os.path.join(_lowercase,"inf_time.csv"     ),inference_memory_csv_file=os.path.join(_lowercase,"inf_mem.csv"     ),env_info_csv_file=os.path.join(_lowercase,"env.csv"     ),multi_process=_lowercase,)
                _lowerCamelCase				:   Union[str, Any]								=   TensorFlowBenchmark(_lowercase     )
                benchmark.run()
                self.assertTrue(Path(os.path.join(_lowercase,"inf_time.csv"     )     ).exists()     )
                self.assertTrue(Path(os.path.join(_lowercase,"inf_mem.csv"     )     ).exists()     )
                self.assertTrue(Path(os.path.join(_lowercase,"env.csv"     )     ).exists()     )



    def   lowerCamelCase_  (	self      :       Tuple     ):
          _lowerCamelCase				:   Dict								=   """sshleifer/tiny-gpt2"""

          def _check_summary_is_not_empty(__A      :       List[str]     ):
                self.assertTrue(hasattr(_lowercase,"sequential"     )     )
                self.assertTrue(hasattr(_lowercase,"cumulative"     )     )
                self.assertTrue(hasattr(_lowercase,"current"     )     )
                self.assertTrue(hasattr(_lowercase,"total"     )     )

          with tempfile.TemporaryDirectory() as tmp_dir:
                _lowerCamelCase				:   Dict								=   TensorFlowBenchmarkArguments(
                    models=[MODEL_ID],inference=_lowercase,sequence_lengths=[8],batch_sizes=[1],log_filename=os.path.join(_lowercase,"log.txt"     ),log_print=_lowercase,trace_memory_line_by_line=_lowercase,eager_mode=_lowercase,multi_process=_lowercase,)
                _lowerCamelCase				:   Union[str, Any]								=   TensorFlowBenchmark(_lowercase     )
                _lowerCamelCase				:   List[str]								=   benchmark.run()
                _check_summary_is_not_empty(result.inference_summary     )
                self.assertTrue(Path(os.path.join(_lowercase,"log.txt"     )     ).exists()     )

'''simple docstring'''






import unittest

import numpy as np

from transformers import AlbertConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
 import jax.numpy as jnp

 from transformers.models.albert.modeling_flax_albert import (
     FlaxAlbertForMaskedLM,
     FlaxAlbertForMultipleChoice,
     FlaxAlbertForPreTraining,
     FlaxAlbertForQuestionAnswering,
     FlaxAlbertForSequenceClassification,
     FlaxAlbertForTokenClassification,
     FlaxAlbertModel,
 )

class  UpperCAmelCase__						(       unittest.TestCase ):



    def __init__(	self      :       Union[str, Any],__A      :       Union[str, Any],__A      :       Optional[int]=1_3,__A      :       Optional[Any]=7,__A      :       int=True,__A      :       Any=True,__A      :       Dict=True,__A      :       List[Any]=True,__A      :       Optional[Any]=9_9,__A      :       List[str]=3_2,__A      :       int=5,__A      :       str=4,__A      :       int=3_7,__A      :       Optional[Any]="gelu",__A      :       Union[str, Any]=0.1,__A      :       str=0.1,__A      :       Any=5_1_2,__A      :       Any=1_6,__A      :       Optional[int]=2,__A      :       int=0.02,__A      :       Optional[int]=4,):
          _lowerCamelCase				:   Union[str, Any]								=   parent
          _lowerCamelCase				:   Tuple								=   batch_size
          _lowerCamelCase				:   List[str]								=   seq_length
          _lowerCamelCase				:   Optional[int]								=   is_training
          _lowerCamelCase				:   str								=   use_attention_mask
          _lowerCamelCase				:   str								=   use_token_type_ids
          _lowerCamelCase				:   List[str]								=   use_labels
          _lowerCamelCase				:   Tuple								=   vocab_size
          _lowerCamelCase				:   Dict								=   hidden_size
          _lowerCamelCase				:   str								=   num_hidden_layers
          _lowerCamelCase				:   List[str]								=   num_attention_heads
          _lowerCamelCase				:   Optional[Any]								=   intermediate_size
          _lowerCamelCase				:   Optional[int]								=   hidden_act
          _lowerCamelCase				:   Any								=   hidden_dropout_prob
          _lowerCamelCase				:   List[Any]								=   attention_probs_dropout_prob
          _lowerCamelCase				:   Tuple								=   max_position_embeddings
          _lowerCamelCase				:   str								=   type_vocab_size
          _lowerCamelCase				:   str								=   type_sequence_label_size
          _lowerCamelCase				:   int								=   initializer_range
          _lowerCamelCase				:   List[Any]								=   num_choices



    def   lowerCamelCase_  (	self      :       Optional[int]     ):
          _lowerCamelCase				:   List[Any]								=   ids_tensor([self.batch_size, self.seq_length],self.vocab_size     )

          _lowerCamelCase				:   Union[str, Any]								=   None
          if self.use_attention_mask:
                _lowerCamelCase				:   Tuple								=   random_attention_mask([self.batch_size, self.seq_length]     )

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

          _lowerCamelCase				:   int								=   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,is_decoder=__A,initializer_range=self.initializer_range,)

          return config, input_ids, token_type_ids, attention_mask



    def   lowerCamelCase_  (	self      :       Dict     ):
          _lowerCamelCase				:   str								=   self.prepare_config_and_inputs()
          _lowerCamelCase						,		_lowerCamelCase						,		_lowerCamelCase						,		_lowerCamelCase				:   Union[str, Any]								=   config_and_inputs
          _lowerCamelCase				:   Any								=   {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": attention_mask}
          return config, inputs_dict


@require_flax
class  UpperCAmelCase__						(       A						,	unittest.TestCase ):
    lowerCAmelCase_				    =		(
        (
            FlaxAlbertModel,
            FlaxAlbertForPreTraining,
            FlaxAlbertForMaskedLM,
            FlaxAlbertForMultipleChoice,
            FlaxAlbertForQuestionAnswering,
            FlaxAlbertForSequenceClassification,
            FlaxAlbertForTokenClassification,
            FlaxAlbertForQuestionAnswering,
        )
        if is_flax_available()
        else ()
    )



    def   lowerCamelCase_  (	self      :       Any     ):
          _lowerCamelCase				:   List[str]								=   FlaxAlbertModelTester(self     )



    @slow
    def   lowerCamelCase_  (	self      :       Union[str, Any]     ):
          for model_class_name in self.all_model_classes:
                _lowerCamelCase				:   str								=   model_class_name.from_pretrained("albert-base-v2"     )
                _lowerCamelCase				:   Any								=   model(np.ones((1, 1)     )     )
                self.assertIsNotNone(__A     )


@require_flax
class  UpperCAmelCase__						(       unittest.TestCase ):



    @slow
    def   lowerCamelCase_  (	self      :       Tuple     ):
          _lowerCamelCase				:   Optional[int]								=   FlaxAlbertModel.from_pretrained("albert-base-v2"     )
          _lowerCamelCase				:   int								=   np.array([[0, 3_4_5, 2_3_2, 3_2_8, 7_4_0, 1_4_0, 1_6_9_5, 6_9, 6_0_7_8, 1_5_8_8, 2]]     )
          _lowerCamelCase				:   Union[str, Any]								=   np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]     )
          _lowerCamelCase				:   Optional[Any]								=   model(__A,attention_mask=__A     )[0]
          _lowerCamelCase				:   Dict								=   (1, 1_1, 7_6_8)
          self.assertEqual(output.shape,__A     )
          _lowerCamelCase				:   Tuple								=   np.array(
              [[[-0.6513, 1.5035, -0.2766], [-0.6515, 1.5046, -0.2780], [-0.6512, 1.5049, -0.2784]]]     )

          self.assertTrue(jnp.allclose(output[:, 1:4, 1:4],__A,atol=1e-4     )     )