


'''simple docstring'''

from __future__ import annotations

import unittest

from transformers import is_tf_available
from transformers.testing_utils import require_sentencepiece, require_tf, require_tokenizers, 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 numpy as np
   import tensorflow as tf

   from transformers import (
       TF_FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
       FlaubertConfig,
       TFFlaubertForMultipleChoice,
       TFFlaubertForQuestionAnsweringSimple,
       TFFlaubertForSequenceClassification,
       TFFlaubertForTokenClassification,
       TFFlaubertModel,
       TFFlaubertWithLMHeadModel,
   )





class 						__SCREAMING_SNAKE_CASE					:






     def __init__(			self     ,							__UpperCamelCase     ,							)       ->    Any:
          _a						=		parent
          _a						=		13
          _a						=		7
          _a						=		True
          _a						=		True
          _a						=		True
          _a						=		True
          _a						=		True
          _a						=		False
          _a						=		False
          _a						=		False
          _a						=		2
          _a						=		99
          _a						=		0
          _a						=		32
          _a						=		2
          _a						=		4
          _a						=		0.1
          _a						=		0.1
          _a						=		512
          _a						=		16
          _a						=		2
          _a						=		0.02
          _a						=		3
          _a						=		4
          _a						=		"last"
          _a						=		True
          _a						=		None
          _a						=		0






     def  a_					(			self		)       ->    Optional[Any]:
          _a						=		ids_tensor([self.batch_size, self.seq_length]     ,							self.vocab_size		)
          _a						=		random_attention_mask([self.batch_size, self.seq_length]     ,							dtype=tf.floataa		)

          _a						=		None
          if self.use_input_lengths:
               _a						=		(
                   ids_tensor([self.batch_size]     ,							vocab_size=2		) + self.seq_length - 2
               )  # small variation of seq_length

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

          _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]     ,							2     ,							dtype=tf.floataa		)
               _a						=		ids_tensor([self.batch_size]     ,							self.num_choices		)

          _a						=		FlaubertConfig(
              vocab_size=self.vocab_size     ,							n_special=self.n_special     ,							emb_dim=self.hidden_size     ,							n_layers=self.num_hidden_layers     ,							n_heads=self.num_attention_heads     ,							dropout=self.hidden_dropout_prob     ,							attention_dropout=self.attention_probs_dropout_prob     ,							gelu_activation=self.gelu_activation     ,							sinusoidal_embeddings=self.sinusoidal_embeddings     ,							asm=self.asm     ,							causal=self.causal     ,							n_langs=self.n_langs     ,							max_position_embeddings=self.max_position_embeddings     ,							initializer_range=self.initializer_range     ,							summary_type=self.summary_type     ,							use_proj=self.use_proj     ,							bos_token_id=self.bos_token_id     ,							)

          return (
              config,
              input_ids,
              token_type_ids,
              input_lengths,
              sequence_labels,
              token_labels,
              is_impossible_labels,
              choice_labels,
              input_mask,
          )






     def  a_					(			self     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							)       ->    str:
          _a						=		TFFlaubertModel(config=__UpperCamelCase		)
          _a						=		{"input_ids": input_ids, "lengths": input_lengths, "langs": token_type_ids}
          _a						=		model(__UpperCamelCase		)

          _a						=		[input_ids, input_mask]
          _a						=		model(__UpperCamelCase		)
          self.parent.assertEqual(result.last_hidden_state.shape     ,							(self.batch_size, self.seq_length, self.hidden_size)		)






     def  a_					(			self     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							)       ->    Tuple:
          _a						=		TFFlaubertWithLMHeadModel(__UpperCamelCase		)

          _a						=		{"input_ids": input_ids, "lengths": input_lengths, "langs": token_type_ids}
          _a						=		model(__UpperCamelCase		)

          self.parent.assertEqual(result.logits.shape     ,							(self.batch_size, self.seq_length, self.vocab_size)		)






     def  a_					(			self     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							)       ->    Optional[Any]:
          _a						=		TFFlaubertForQuestionAnsweringSimple(__UpperCamelCase		)

          _a						=		{"input_ids": input_ids, "lengths": input_lengths}

          _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  a_					(			self     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							)       ->    str:
          _a						=		TFFlaubertForSequenceClassification(__UpperCamelCase		)

          _a						=		{"input_ids": input_ids, "lengths": input_lengths}

          _a						=		model(__UpperCamelCase		)

          self.parent.assertEqual(result.logits.shape     ,							(self.batch_size, self.type_sequence_label_size)		)






     def  a_					(			self     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							)       ->    Optional[Any]:
          _a						=		self.num_labels
          _a						=		TFFlaubertForTokenClassification(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.num_labels)		)






     def  a_					(			self     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							)       ->    int:
          _a						=		self.num_choices
          _a						=		TFFlaubertForMultipleChoice(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.assertEqual(result.logits.shape     ,							(self.batch_size, self.num_choices)		)






     def  a_					(			self		)       ->    List[str]:
          _a						=		self.prepare_config_and_inputs()
          (
              (
              _a
          )			,			(
              _a
          )			,			(
              _a
          )			,			(
              _a
          )			,			(
              _a
          )			,			(
              _a
          )			,			(
              _a
          )			,			(
              _a
          )			,			(
              _a
          )			,			
          )						=		config_and_inputs
          _a						=		{
              "input_ids": input_ids,
              "token_type_ids": token_type_ids,
              "langs": token_type_ids,
              "lengths": input_lengths,
          }
          return config, inputs_dict







@require_tf
class 						__SCREAMING_SNAKE_CASE					(      lowerCamelCase__ ,			lowerCamelCase__ ,			unittest.TestCase ):
     UpperCAmelCase 					=				(
         (
             TFFlaubertModel,
             TFFlaubertWithLMHeadModel,
             TFFlaubertForSequenceClassification,
             TFFlaubertForQuestionAnsweringSimple,
             TFFlaubertForTokenClassification,
             TFFlaubertForMultipleChoice,
         )
         if is_tf_available()
         else ()
     )
     UpperCAmelCase 					=				(
         (TFFlaubertWithLMHeadModel,) if is_tf_available() else ()
     )  # TODO (PVP): Check other models whether language generation is also applicable
     UpperCAmelCase 					=				(
         {
             '''feature-extraction''': TFFlaubertModel,
             '''fill-mask''': TFFlaubertWithLMHeadModel,
             '''question-answering''': TFFlaubertForQuestionAnsweringSimple,
             '''text-classification''': TFFlaubertForSequenceClassification,
             '''token-classification''': TFFlaubertForTokenClassification,
             '''zero-shot''': TFFlaubertForSequenceClassification,
         }
         if is_tf_available()
         else {}
     )
     UpperCAmelCase 					=				False
     UpperCAmelCase 					=				False






     def  a_					(			self     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase		)       ->    int:
          if (
              pipeline_test_casse_name == "QAPipelineTests"
              and tokenizer_name is not None
              and not tokenizer_name.endswith("Fast"		)
          ):
               # `QAPipelineTests` fails for a few models when the slower tokenizer are used.
               # (The slower tokenizers were never used for pipeline tests before the pipeline testing rework)
               # TODO: check (and possibly fix) the `QAPipelineTests` with slower tokenizer
               return True

          return False






     def  a_					(			self		)       ->    Union[str, Any]:
          _a						=		TFFlaubertModelTester(self		)
          _a						=		ConfigTester(self     ,							config_class=__UpperCamelCase     ,							emb_dim=37		)






     def  a_					(			self		)       ->    Optional[int]:
          self.config_tester.run_common_tests()






     def  a_					(			self		)       ->    Dict:
          _a						=		self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_flaubert_model(*__UpperCamelCase		)






     def  a_					(			self		)       ->    Optional[int]:
          _a						=		self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_flaubert_lm_head(*__UpperCamelCase		)






     def  a_					(			self		)       ->    List[str]:
          _a						=		self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_flaubert_qa(*__UpperCamelCase		)






     def  a_					(			self		)       ->    List[Any]:
          _a						=		self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_flaubert_sequence_classif(*__UpperCamelCase		)






     def  a_					(			self		)       ->    str:
          _a						=		self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_flaubert_for_token_classification(*__UpperCamelCase		)






     def  a_					(			self		)       ->    str:
          _a						=		self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_flaubert_for_multiple_choice(*__UpperCamelCase		)






     @slow
     def  a_					(			self		)       ->    Optional[int]:
          for model_name in TF_FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
               _a						=		TFFlaubertModel.from_pretrained(__UpperCamelCase		)
               self.assertIsNotNone(__UpperCamelCase		)







@require_tf
@require_sentencepiece
@require_tokenizers
class 						__SCREAMING_SNAKE_CASE					(      unittest.TestCase ):






     @slow
     def  a_					(			self		)       ->    Dict:
          _a						=		TFFlaubertModel.from_pretrained("jplu/tf-flaubert-small-cased"		)

          _a						=		tf.convert_to_tensor(
              [[0, 158, 735, 2_592, 1_424, 6_727, 82, 1]]     ,							dtype=tf.intaa     ,							)  # "J'aime flaubert !"

          _a						=		model(__UpperCamelCase		)[0]
          _a						=		tf.TensorShape((1, 8, 512)		)
          self.assertEqual(output.shape     ,							__UpperCamelCase		)
          # compare the actual values for a slice.
          _a						=		tf.convert_to_tensor(
              [
                  [
                      [-1.8_76_87_73, -1.56_65_55, 0.27_07_24_18],
                      [-1.6_92_00_38, -0.5_87_35_05, 1.9_32_95_99],
                      [-2.9_56_39_85, -1.6_99_38_35, 1.7_97_20_52],
                  ]
              ]     ,							dtype=tf.floataa     ,							)

          self.assertTrue(np.allclose(output[:, :3, :3].numpy()     ,							expected_slice.numpy()     ,							atol=1e-4		)		)





'''simple docstring'''

import gc
import tempfile
import unittest

import numpy as np
import torch

from diffusers import VersatileDiffusionTextToImagePipeline
from diffusers.utils.testing_utils import nightly, require_torch_gpu, torch_device


lowercase__				     =		False





class 						__SCREAMING_SNAKE_CASE					(      unittest.TestCase ):
     pass







@nightly
@require_torch_gpu
class 						__SCREAMING_SNAKE_CASE					(      unittest.TestCase ):






     def  a_					(			self		)       ->    Dict:
          # clean up the VRAM after each test
          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()






     def  a_					(			self		)       ->    List[str]:
          _a						=		VersatileDiffusionTextToImagePipeline.from_pretrained("shi-labs/versatile-diffusion"		)
          # remove text_unet
          pipe.remove_unused_weights()
          pipe.to(__UpperCamelCase		)
          pipe.set_progress_bar_config(disable=__UpperCamelCase		)

          _a						=		"A painting of a squirrel eating a burger "
          _a						=		torch.manual_seed(0		)
          _a						=		pipe(
              prompt=__UpperCamelCase     ,							generator=__UpperCamelCase     ,							guidance_scale=7.5     ,							num_inference_steps=2     ,							output_type="numpy"		).images

          with tempfile.TemporaryDirectory() as tmpdirname:
               pipe.save_pretrained(__UpperCamelCase		)
               _a						=		VersatileDiffusionTextToImagePipeline.from_pretrained(__UpperCamelCase		)
          pipe.to(__UpperCamelCase		)
          pipe.set_progress_bar_config(disable=__UpperCamelCase		)

          _a						=		generator.manual_seed(0		)
          _a						=		pipe(
              prompt=__UpperCamelCase     ,							generator=__UpperCamelCase     ,							guidance_scale=7.5     ,							num_inference_steps=2     ,							output_type="numpy"		).images

          assert np.abs(image - new_image		).sum() < 1e-5, "Models don't have the same forward pass"






     def  a_					(			self		)       ->    List[Any]:
          _a						=		VersatileDiffusionTextToImagePipeline.from_pretrained(
              "shi-labs/versatile-diffusion"     ,							torch_dtype=torch.floataa		)
          pipe.to(__UpperCamelCase		)
          pipe.set_progress_bar_config(disable=__UpperCamelCase		)

          _a						=		"A painting of a squirrel eating a burger "
          _a						=		torch.manual_seed(0		)
          _a						=		pipe(
              prompt=__UpperCamelCase     ,							generator=__UpperCamelCase     ,							guidance_scale=7.5     ,							num_inference_steps=50     ,							output_type="numpy"		).images

          _a						=		image[0, 253:256, 253:256, -1]

          assert image.shape == (1, 512, 512, 3)
          _a						=		np.array([0.33_67, 0.31_69, 0.26_56, 0.38_70, 0.47_90, 0.37_96, 0.40_09, 0.48_78, 0.47_78]		)

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

