'''simple docstring'''
from __future__ import annotations

import copy
import tempfile
import unittest

from transformers import CONFIG_MAPPING, AutoConfig, BertConfig, GPTaConfig, TaConfig, TapasConfig, is_tf_available
from transformers.testing_utils import (
    DUMMY_UNKNOWN_IDENTIFIER,
    SMALL_MODEL_IDENTIFIER,
    RequestCounter,
    require_tensorflow_probability,
    require_tf,
    slow,
)

from ..bert.test_modeling_bert import BertModelTester


if is_tf_available():
      from transformers import (
          TFAutoModel,
          TFAutoModelForCausalLM,
          TFAutoModelForMaskedLM,
          TFAutoModelForPreTraining,
          TFAutoModelForQuestionAnswering,
          TFAutoModelForSeqaSeqLM,
          TFAutoModelForSequenceClassification,
          TFAutoModelForTableQuestionAnswering,
          TFAutoModelForTokenClassification,
          TFAutoModelWithLMHead,
          TFBertForMaskedLM,
          TFBertForPreTraining,
          TFBertForQuestionAnswering,
          TFBertForSequenceClassification,
          TFBertModel,
          TFFunnelBaseModel,
          TFFunnelModel,
          TFGPTaLMHeadModel,
          TFRobertaForMaskedLM,
          TFTaForConditionalGeneration,
          TFTapasForQuestionAnswering,
      )
      from transformers.models.auto.modeling_tf_auto import (
          TF_MODEL_FOR_CAUSAL_LM_MAPPING,
          TF_MODEL_FOR_MASKED_LM_MAPPING,
          TF_MODEL_FOR_PRETRAINING_MAPPING,
          TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING,
          TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
          TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
          TF_MODEL_MAPPING,
      )
      from transformers.models.bert.modeling_tf_bert import TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST
      from transformers.models.gpta.modeling_tf_gpta import TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST
      from transformers.models.ta.modeling_tf_ta import TF_T5_PRETRAINED_MODEL_ARCHIVE_LIST
      from transformers.models.tapas.modeling_tf_tapas import TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST







class lowerCamelCase							(	lowerCamelCase					):




  '''simple docstring'''



  lowerCAmelCase__										=	'''new-model'''


if is_tf_available():







      class lowerCamelCase							(	lowerCamelCase					):




        '''simple docstring'''



        lowerCAmelCase__										=	NewModelConfig







@require_tf
class lowerCamelCase							(	unittest.TestCase					):




  '''simple docstring'''



  @slow
  def  lowerCAmelCase__ (      self			:       int  )				->int:
        UpperCAmelCase_            =  '''bert-base-cased'''
        UpperCAmelCase_            =  AutoConfig.from_pretrained(UpperCAmelCase__  )
        self.assertIsNotNone(UpperCAmelCase__  )
        self.assertIsInstance(UpperCAmelCase__  ,							UpperCAmelCase__  )

        UpperCAmelCase_            =  TFAutoModel.from_pretrained(UpperCAmelCase__  )
        self.assertIsNotNone(UpperCAmelCase__  )
        self.assertIsInstance(UpperCAmelCase__  ,							UpperCAmelCase__  )






  @slow
  def  lowerCAmelCase__ (      self			:       int  )				->Any:
        UpperCAmelCase_            =  '''bert-base-cased'''
        UpperCAmelCase_            =  AutoConfig.from_pretrained(UpperCAmelCase__  )
        self.assertIsNotNone(UpperCAmelCase__  )
        self.assertIsInstance(UpperCAmelCase__  ,							UpperCAmelCase__  )

        UpperCAmelCase_            =  TFAutoModelForPreTraining.from_pretrained(UpperCAmelCase__  )
        self.assertIsNotNone(UpperCAmelCase__  )
        self.assertIsInstance(UpperCAmelCase__  ,							UpperCAmelCase__  )






  @slow
  def  lowerCAmelCase__ (      self			:       Tuple  )				->str:
        for model_name in TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
              UpperCAmelCase_            =  AutoConfig.from_pretrained(UpperCAmelCase__  )
              self.assertIsNotNone(UpperCAmelCase__  )
              self.assertIsInstance(UpperCAmelCase__  ,							UpperCAmelCase__  )

              UpperCAmelCase_            =  TFAutoModelForCausalLM.from_pretrained(UpperCAmelCase__  )
              UpperCAmelCase_      ,		UpperCAmelCase_            =  TFAutoModelForCausalLM.from_pretrained(UpperCAmelCase__  ,							output_loading_info=UpperCAmelCase__  )
              self.assertIsNotNone(UpperCAmelCase__  )
              self.assertIsInstance(UpperCAmelCase__  ,							UpperCAmelCase__  )






  @slow
  def  lowerCAmelCase__ (      self			:       Tuple  )				->str:
        for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
              UpperCAmelCase_            =  AutoConfig.from_pretrained(UpperCAmelCase__  )
              self.assertIsNotNone(UpperCAmelCase__  )
              self.assertIsInstance(UpperCAmelCase__  ,							UpperCAmelCase__  )

              UpperCAmelCase_            =  TFAutoModelWithLMHead.from_pretrained(UpperCAmelCase__  )
              self.assertIsNotNone(UpperCAmelCase__  )
              self.assertIsInstance(UpperCAmelCase__  ,							UpperCAmelCase__  )






  @slow
  def  lowerCAmelCase__ (      self			:       Optional[Any]  )				->List[str]:
        for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
              UpperCAmelCase_            =  AutoConfig.from_pretrained(UpperCAmelCase__  )
              self.assertIsNotNone(UpperCAmelCase__  )
              self.assertIsInstance(UpperCAmelCase__  ,							UpperCAmelCase__  )

              UpperCAmelCase_            =  TFAutoModelForMaskedLM.from_pretrained(UpperCAmelCase__  )
              UpperCAmelCase_      ,		UpperCAmelCase_            =  TFAutoModelForMaskedLM.from_pretrained(UpperCAmelCase__  ,							output_loading_info=UpperCAmelCase__  )
              self.assertIsNotNone(UpperCAmelCase__  )
              self.assertIsInstance(UpperCAmelCase__  ,							UpperCAmelCase__  )






  @slow
  def  lowerCAmelCase__ (      self			:       Tuple  )				->int:
        for model_name in TF_T5_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
              UpperCAmelCase_            =  AutoConfig.from_pretrained(UpperCAmelCase__  )
              self.assertIsNotNone(UpperCAmelCase__  )
              self.assertIsInstance(UpperCAmelCase__  ,							UpperCAmelCase__  )

              UpperCAmelCase_            =  TFAutoModelForSeqaSeqLM.from_pretrained(UpperCAmelCase__  )
              UpperCAmelCase_      ,		UpperCAmelCase_            =  TFAutoModelForSeqaSeqLM.from_pretrained(UpperCAmelCase__  ,							output_loading_info=UpperCAmelCase__  )
              self.assertIsNotNone(UpperCAmelCase__  )
              self.assertIsInstance(UpperCAmelCase__  ,							UpperCAmelCase__  )






  @slow
  def  lowerCAmelCase__ (      self			:       List[str]  )				->Any:
        # for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        for model_name in ["bert-base-uncased"]:
              UpperCAmelCase_            =  AutoConfig.from_pretrained(UpperCAmelCase__  )
              self.assertIsNotNone(UpperCAmelCase__  )
              self.assertIsInstance(UpperCAmelCase__  ,							UpperCAmelCase__  )

              UpperCAmelCase_            =  TFAutoModelForSequenceClassification.from_pretrained(UpperCAmelCase__  )
              self.assertIsNotNone(UpperCAmelCase__  )
              self.assertIsInstance(UpperCAmelCase__  ,							UpperCAmelCase__  )






  @slow
  def  lowerCAmelCase__ (      self			:       List[str]  )				->Any:
        # for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        for model_name in ["bert-base-uncased"]:
              UpperCAmelCase_            =  AutoConfig.from_pretrained(UpperCAmelCase__  )
              self.assertIsNotNone(UpperCAmelCase__  )
              self.assertIsInstance(UpperCAmelCase__  ,							UpperCAmelCase__  )

              UpperCAmelCase_            =  TFAutoModelForQuestionAnswering.from_pretrained(UpperCAmelCase__  )
              self.assertIsNotNone(UpperCAmelCase__  )
              self.assertIsInstance(UpperCAmelCase__  ,							UpperCAmelCase__  )






  @slow
  @require_tensorflow_probability
  def  lowerCAmelCase__ (      self			:       Optional[Any]  )				->int:
        for model_name in TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST[5:6]:
              UpperCAmelCase_            =  AutoConfig.from_pretrained(UpperCAmelCase__  )
              self.assertIsNotNone(UpperCAmelCase__  )
              self.assertIsInstance(UpperCAmelCase__  ,							UpperCAmelCase__  )

              UpperCAmelCase_            =  TFAutoModelForTableQuestionAnswering.from_pretrained(UpperCAmelCase__  )
              UpperCAmelCase_      ,		UpperCAmelCase_            =  TFAutoModelForTableQuestionAnswering.from_pretrained(
                  UpperCAmelCase__  ,							output_loading_info=UpperCAmelCase__  )
              self.assertIsNotNone(UpperCAmelCase__  )
              self.assertIsInstance(UpperCAmelCase__  ,							UpperCAmelCase__  )






  def  lowerCAmelCase__ (      self			:       int  )				->Optional[int]:
        UpperCAmelCase_            =  TFAutoModelWithLMHead.from_pretrained(UpperCAmelCase__  )
        self.assertIsInstance(UpperCAmelCase__  ,							UpperCAmelCase__  )
        self.assertEqual(model.num_parameters()  ,							1_4410  )
        self.assertEqual(model.num_parameters(only_trainable=UpperCAmelCase__  )  ,							1_4410  )






  def  lowerCAmelCase__ (      self			:       int  )				->int:
        UpperCAmelCase_            =  TFAutoModelWithLMHead.from_pretrained(UpperCAmelCase__  )
        self.assertIsInstance(UpperCAmelCase__  ,							UpperCAmelCase__  )
        self.assertEqual(model.num_parameters()  ,							1_4410  )
        self.assertEqual(model.num_parameters(only_trainable=UpperCAmelCase__  )  ,							1_4410  )






  def  lowerCAmelCase__ (      self			:       List[str]  )				->List[Any]:
        # For the auto model mapping, FunnelConfig has two models: FunnelModel and FunnelBaseModel
        UpperCAmelCase_            =  TFAutoModel.from_pretrained('''sgugger/funnel-random-tiny'''  )
        self.assertIsInstance(UpperCAmelCase__  ,							UpperCAmelCase__  )

        UpperCAmelCase_            =  copy.deepcopy(model.config  )
        UpperCAmelCase_            =  ['''FunnelBaseModel''']
        UpperCAmelCase_            =  TFAutoModel.from_config(UpperCAmelCase__  )
        self.assertIsInstance(UpperCAmelCase__  ,							UpperCAmelCase__  )

        with tempfile.TemporaryDirectory() as tmp_dir:
              model.save_pretrained(UpperCAmelCase__  )
              UpperCAmelCase_            =  TFAutoModel.from_pretrained(UpperCAmelCase__  )
              self.assertIsInstance(UpperCAmelCase__  ,							UpperCAmelCase__  )






  def  lowerCAmelCase__ (      self			:       str  )				->Optional[int]:
        try:
              AutoConfig.register('''new-model'''  ,							UpperCAmelCase__  )

              UpperCAmelCase_            =  [
                  TFAutoModel,
                  TFAutoModelForCausalLM,
                  TFAutoModelForMaskedLM,
                  TFAutoModelForPreTraining,
                  TFAutoModelForQuestionAnswering,
                  TFAutoModelForSequenceClassification,
                  TFAutoModelForTokenClassification,
              ]

              for auto_class in auto_classes:
                    with self.subTest(auto_class.__name__  ):
                          # Wrong config class will raise an error
                          with self.assertRaises(UpperCAmelCase__  ):
                                auto_class.register(UpperCAmelCase__  ,							UpperCAmelCase__  )
                          auto_class.register(UpperCAmelCase__  ,							UpperCAmelCase__  )
                          # Trying to register something existing in the Transformers library will raise an error
                          with self.assertRaises(UpperCAmelCase__  ):
                                auto_class.register(UpperCAmelCase__  ,							UpperCAmelCase__  )

                          # Now that the config is registered, it can be used as any other config with the auto-API
                          UpperCAmelCase_            =  BertModelTester(self  ).get_config()
                          UpperCAmelCase_            =  NewModelConfig(**tiny_config.to_dict()  )
                          UpperCAmelCase_            =  auto_class.from_config(UpperCAmelCase__  )
                          self.assertIsInstance(UpperCAmelCase__  ,							UpperCAmelCase__  )

                          with tempfile.TemporaryDirectory() as tmp_dir:
                                model.save_pretrained(UpperCAmelCase__  )
                                UpperCAmelCase_            =  auto_class.from_pretrained(UpperCAmelCase__  )
                                self.assertIsInstance(UpperCAmelCase__  ,							UpperCAmelCase__  )

        finally:
              if "new-model" in CONFIG_MAPPING._extra_content:
                    del CONFIG_MAPPING._extra_content["new-model"]
              for mapping in (
                  TF_MODEL_MAPPING,
                  TF_MODEL_FOR_PRETRAINING_MAPPING,
                  TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING,
                  TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
                  TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
                  TF_MODEL_FOR_CAUSAL_LM_MAPPING,
                  TF_MODEL_FOR_MASKED_LM_MAPPING,
              ):
                    if NewModelConfig in mapping._extra_content:
                          del mapping._extra_content[NewModelConfig]






  def  lowerCAmelCase__ (      self			:       Optional[int]  )				->Tuple:
        with self.assertRaisesRegex(
            UpperCAmelCase__  ,							'''bert-base is not a local folder and is not a valid model identifier'''  ):
              UpperCAmelCase_            =  TFAutoModel.from_pretrained('''bert-base'''  )






  def  lowerCAmelCase__ (      self			:       Union[str, Any]  )				->str:
        with self.assertRaisesRegex(
            UpperCAmelCase__  ,							r'''aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)'''  ):
              UpperCAmelCase_            =  TFAutoModel.from_pretrained(UpperCAmelCase__  ,							revision='''aaaaaa'''  )






  def  lowerCAmelCase__ (      self			:       Any  )				->Union[str, Any]:
        with self.assertRaisesRegex(
            UpperCAmelCase__  ,							'''hf-internal-testing/config-no-model does not appear to have a file named pytorch_model.bin'''  ,							):
              UpperCAmelCase_            =  TFAutoModel.from_pretrained('''hf-internal-testing/config-no-model'''  )






  def  lowerCAmelCase__ (      self			:       Union[str, Any]  )				->str:
        with self.assertRaisesRegex(UpperCAmelCase__  ,							'''Use `from_pt=True` to load this model'''  ):
              UpperCAmelCase_            =  TFAutoModel.from_pretrained('''hf-internal-testing/tiny-bert-pt-only'''  )






  def  lowerCAmelCase__ (      self			:       List[str]  )				->Dict:
        # Make sure we have cached the model.
        UpperCAmelCase_            =  TFAutoModel.from_pretrained('''hf-internal-testing/tiny-random-bert'''  )
        with RequestCounter() as counter:
              UpperCAmelCase_            =  TFAutoModel.from_pretrained('''hf-internal-testing/tiny-random-bert'''  )
              self.assertEqual(counter.get_request_count  ,							0  )
              self.assertEqual(counter.head_request_count  ,							1  )
              self.assertEqual(counter.other_request_count  ,							0  )

        # With a sharded checkpoint
        UpperCAmelCase_            =  TFAutoModel.from_pretrained('''ArthurZ/tiny-random-bert-sharded'''  )
        with RequestCounter() as counter:
              UpperCAmelCase_            =  TFAutoModel.from_pretrained('''ArthurZ/tiny-random-bert-sharded'''  )
              self.assertEqual(counter.get_request_count  ,							0  )
              self.assertEqual(counter.head_request_count  ,							1  )
              self.assertEqual(counter.other_request_count  ,							0  )






'''simple docstring'''
import gc
import unittest

import numpy as np
import torch

from diffusers import StableDiffusionKDiffusionPipeline
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()







@slow
@require_torch_gpu
class lowerCamelCase							(	unittest.TestCase					):




  '''simple docstring'''



  def  lowerCAmelCase__ (      self			:       Any  )				->Optional[int]:
        # clean up the VRAM after each test
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()






  def  lowerCAmelCase__ (      self			:       str  )				->List[str]:
        UpperCAmelCase_            =  StableDiffusionKDiffusionPipeline.from_pretrained('''CompVis/stable-diffusion-v1-4'''  )
        UpperCAmelCase_            =  sd_pipe.to(UpperCAmelCase__  )
        sd_pipe.set_progress_bar_config(disable=UpperCAmelCase__  )

        sd_pipe.set_scheduler('''sample_euler'''  )

        UpperCAmelCase_            =  '''A painting of a squirrel eating a burger'''
        UpperCAmelCase_            =  torch.manual_seed(0  )
        UpperCAmelCase_            =  sd_pipe([prompt]  ,							generator=UpperCAmelCase__  ,							guidance_scale=9.0  ,							num_inference_steps=20  ,							output_type='''np'''  )

        UpperCAmelCase_            =  output.images

        UpperCAmelCase_            =  image[0, -3:, -3:, -1]

        assert image.shape == (1, 512, 512, 3)
        UpperCAmelCase_            =  np.array([0.0447, 0.0492, 0.0468, 0.0408, 0.0383, 0.0408, 0.0354, 0.0380, 0.0339]  )

        assert np.abs(image_slice.flatten() - expected_slice  ).max() < 1e-2






  def  lowerCAmelCase__ (      self			:       List[str]  )				->int:
        UpperCAmelCase_            =  StableDiffusionKDiffusionPipeline.from_pretrained('''stabilityai/stable-diffusion-2-1-base'''  )
        UpperCAmelCase_            =  sd_pipe.to(UpperCAmelCase__  )
        sd_pipe.set_progress_bar_config(disable=UpperCAmelCase__  )

        sd_pipe.set_scheduler('''sample_euler'''  )

        UpperCAmelCase_            =  '''A painting of a squirrel eating a burger'''
        UpperCAmelCase_            =  torch.manual_seed(0  )
        UpperCAmelCase_            =  sd_pipe([prompt]  ,							generator=UpperCAmelCase__  ,							guidance_scale=9.0  ,							num_inference_steps=20  ,							output_type='''np'''  )

        UpperCAmelCase_            =  output.images

        UpperCAmelCase_            =  image[0, -3:, -3:, -1]

        assert image.shape == (1, 512, 512, 3)
        UpperCAmelCase_            =  np.array([0.1237, 0.1320, 0.1438, 0.1359, 0.1390, 0.1132, 0.1277, 0.1175, 0.1112]  )

        assert np.abs(image_slice.flatten() - expected_slice  ).max() < 5e-1






  def  lowerCAmelCase__ (      self			:       List[Any]  )				->Optional[int]:
        UpperCAmelCase_            =  StableDiffusionKDiffusionPipeline.from_pretrained('''stabilityai/stable-diffusion-2-1-base'''  )
        UpperCAmelCase_            =  sd_pipe.to(UpperCAmelCase__  )
        sd_pipe.set_progress_bar_config(disable=UpperCAmelCase__  )

        sd_pipe.set_scheduler('''sample_dpmpp_2m'''  )

        UpperCAmelCase_            =  '''A painting of a squirrel eating a burger'''
        UpperCAmelCase_            =  torch.manual_seed(0  )
        UpperCAmelCase_            =  sd_pipe(
            [prompt]  ,							generator=UpperCAmelCase__  ,							guidance_scale=7.5  ,							num_inference_steps=15  ,							output_type='''np'''  ,							use_karras_sigmas=UpperCAmelCase__  ,							)

        UpperCAmelCase_            =  output.images

        UpperCAmelCase_            =  image[0, -3:, -3:, -1]

        assert image.shape == (1, 512, 512, 3)
        UpperCAmelCase_            =  np.array(
            [0.1138_1689, 0.1211_2921, 0.138_9457, 0.1254_9606, 0.124_4964, 0.1083_1517, 0.1156_2866, 0.1086_7816, 0.1049_9048]  )

        assert np.abs(image_slice.flatten() - expected_slice  ).max() < 1e-2





