





"""simple docstring"""





import os
import tempfile
import unittest

from transformers import FlaubertConfig, is_torch_available
from transformers.testing_utils import require_torch, require_torch_gpu, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
       import torch

       from transformers import (
           FlaubertForMultipleChoice,
           FlaubertForQuestionAnswering,
           FlaubertForQuestionAnsweringSimple,
           FlaubertForSequenceClassification,
           FlaubertForTokenClassification,
           FlaubertModel,
           FlaubertWithLMHeadModel,
       )
       from transformers.models.flaubert.modeling_flaubert import FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST

class 		lowercase					(					_UpperCAmelCase					):





      def __init__(       self      ,						lowercase      ,						lowercase=13      ,						lowercase=7      ,						lowercase=True      ,						lowercase=True      ,						lowercase=True      ,						lowercase=True      ,						lowercase=True      ,						lowercase=False      ,						lowercase=False      ,						lowercase=False      ,						lowercase=2      ,						lowercase=99      ,						lowercase=0      ,						lowercase=32      ,						lowercase=5      ,						lowercase=4      ,						lowercase=0.1      ,						lowercase=0.1      ,						lowercase=512      ,						lowercase=12      ,						lowercase=2      ,						lowercase=0.02      ,						lowercase=3      ,						lowercase=4      ,						lowercase="last"      ,						lowercase=None      ,						lowercase=None      ,						)		->     int:
            lowerCAmelCase  						=       parent
            lowerCAmelCase  						=       batch_size
            lowerCAmelCase  						=       seq_length
            lowerCAmelCase  						=       is_training
            lowerCAmelCase  						=       use_input_lengths
            lowerCAmelCase  						=       use_token_type_ids
            lowerCAmelCase  						=       use_labels
            lowerCAmelCase  						=       gelu_activation
            lowerCAmelCase  						=       sinusoidal_embeddings
            lowerCAmelCase  						=       causal
            lowerCAmelCase  						=       asm
            lowerCAmelCase  						=       n_langs
            lowerCAmelCase  						=       vocab_size
            lowerCAmelCase  						=       n_special
            lowerCAmelCase  						=       hidden_size
            lowerCAmelCase  						=       num_hidden_layers
            lowerCAmelCase  						=       num_attention_heads
            lowerCAmelCase  						=       hidden_dropout_prob
            lowerCAmelCase  						=       attention_probs_dropout_prob
            lowerCAmelCase  						=       max_position_embeddings
            lowerCAmelCase  						=       type_vocab_size
            lowerCAmelCase  						=       type_sequence_label_size
            lowerCAmelCase  						=       initializer_range
            lowerCAmelCase  						=       num_labels
            lowerCAmelCase  						=       num_choices
            lowerCAmelCase  						=       summary_type
            lowerCAmelCase  						=       use_proj
            lowerCAmelCase  						=       scope





      def     _snake_case			(       self )		->     int:
            lowerCAmelCase  						=       ids_tensor([self.batch_size, self.seq_length]      ,						self.vocab_size )
            lowerCAmelCase  						=       random_attention_mask([self.batch_size, self.seq_length] )

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

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

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

            lowerCAmelCase  						=       self.get_config()

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





      def     _snake_case			(       self )		->     List[Any]:
            return 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      ,						)





      def     _snake_case			(       self      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						)		->     Any:
            lowerCAmelCase  						=       FlaubertModel(config=lowercase )
            model.to(lowercase )
            model.eval()
            lowerCAmelCase  						=       model(lowercase      ,						lengths=lowercase      ,						langs=lowercase )
            lowerCAmelCase  						=       model(lowercase      ,						langs=lowercase )
            lowerCAmelCase  						=       model(lowercase )
            self.parent.assertEqual(result.last_hidden_state.shape      ,						(self.batch_size, self.seq_length, self.hidden_size) )





      def     _snake_case			(       self      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						)		->     Tuple:
            lowerCAmelCase  						=       FlaubertWithLMHeadModel(lowercase )
            model.to(lowercase )
            model.eval()

            lowerCAmelCase  						=       model(lowercase      ,						token_type_ids=lowercase      ,						labels=lowercase )
            self.parent.assertEqual(result.loss.shape      ,						() )
            self.parent.assertEqual(result.logits.shape      ,						(self.batch_size, self.seq_length, self.vocab_size) )





      def     _snake_case			(       self      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						)		->     str:
            lowerCAmelCase  						=       FlaubertForQuestionAnsweringSimple(lowercase )
            model.to(lowercase )
            model.eval()

            lowerCAmelCase  						=       model(lowercase )

            lowerCAmelCase  						=       model(lowercase      ,						start_positions=lowercase      ,						end_positions=lowercase )
            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     _snake_case			(       self      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						)		->     Dict:
            lowerCAmelCase  						=       FlaubertForQuestionAnswering(lowercase )
            model.to(lowercase )
            model.eval()

            lowerCAmelCase  						=       model(lowercase )

            lowerCAmelCase  						=       model(
                lowercase      ,						start_positions=lowercase      ,						end_positions=lowercase      ,						cls_index=lowercase      ,						is_impossible=lowercase      ,						p_mask=lowercase      ,						)

            lowerCAmelCase  						=       model(
                lowercase      ,						start_positions=lowercase      ,						end_positions=lowercase      ,						cls_index=lowercase      ,						is_impossible=lowercase      ,						)

            ((lowerCAmelCase)						,		)  						=       result_with_labels.to_tuple()

            lowerCAmelCase  						=       model(lowercase      ,						start_positions=lowercase      ,						end_positions=lowercase )

            ((lowerCAmelCase)						,		)  						=       result_with_labels.to_tuple()

            self.parent.assertEqual(result_with_labels.loss.shape      ,						() )
            self.parent.assertEqual(result.start_top_log_probs.shape      ,						(self.batch_size, model.config.start_n_top) )
            self.parent.assertEqual(result.start_top_index.shape      ,						(self.batch_size, model.config.start_n_top) )
            self.parent.assertEqual(
                result.end_top_log_probs.shape      ,						(self.batch_size, model.config.start_n_top * model.config.end_n_top) )
            self.parent.assertEqual(
                result.end_top_index.shape      ,						(self.batch_size, model.config.start_n_top * model.config.end_n_top) )
            self.parent.assertEqual(result.cls_logits.shape      ,						(self.batch_size,) )





      def     _snake_case			(       self      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						)		->     int:
            lowerCAmelCase  						=       FlaubertForSequenceClassification(lowercase )
            model.to(lowercase )
            model.eval()

            lowerCAmelCase  						=       model(lowercase )
            lowerCAmelCase  						=       model(lowercase      ,						labels=lowercase )

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





      def     _snake_case			(       self      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						)		->     int:
            lowerCAmelCase  						=       self.num_labels
            lowerCAmelCase  						=       FlaubertForTokenClassification(lowercase )
            model.to(lowercase )
            model.eval()

            lowerCAmelCase  						=       model(lowercase      ,						attention_mask=lowercase      ,						labels=lowercase )
            self.parent.assertEqual(result.logits.shape      ,						(self.batch_size, self.seq_length, self.num_labels) )





      def     _snake_case			(       self      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						)		->     Tuple:
            lowerCAmelCase  						=       self.num_choices
            lowerCAmelCase  						=       FlaubertForMultipleChoice(config=lowercase )
            model.to(lowercase )
            model.eval()
            lowerCAmelCase  						=       input_ids.unsqueeze(1 ).expand(-1      ,						self.num_choices      ,						-1 ).contiguous()
            lowerCAmelCase  						=       token_type_ids.unsqueeze(1 ).expand(-1      ,						self.num_choices      ,						-1 ).contiguous()
            lowerCAmelCase  						=       input_mask.unsqueeze(1 ).expand(-1      ,						self.num_choices      ,						-1 ).contiguous()
            lowerCAmelCase  						=       model(
                lowercase      ,						attention_mask=lowercase      ,						token_type_ids=lowercase      ,						labels=lowercase      ,						)
            self.parent.assertEqual(result.logits.shape      ,						(self.batch_size, self.num_choices) )





      def     _snake_case			(       self )		->     List[Any]:
            lowerCAmelCase  						=       self.prepare_config_and_inputs()
            (
                (
                lowerCAmelCase
            )						,		(
                lowerCAmelCase
            )						,		(
                lowerCAmelCase
            )						,		(
                lowerCAmelCase
            )						,		(
                lowerCAmelCase
            )						,		(
                lowerCAmelCase
            )						,		(
                lowerCAmelCase
            )						,		(
                lowerCAmelCase
            )						,		(
                lowerCAmelCase
            )						,		
            )  						=       config_and_inputs
            lowerCAmelCase  						=       {
                """input_ids""": input_ids,
                """token_type_ids""": token_type_ids,
                """lengths""": input_lengths,
                """attention_mask""": input_mask,
            }
            return config, inputs_dict







@require_torch
class 		lowercase					(					_UpperCAmelCase     ,  _UpperCAmelCase     ,  unittest.TestCase					):
      _SCREAMING_SNAKE_CASE					=      (
          (
              FlaubertModel,
              FlaubertWithLMHeadModel,
              FlaubertForQuestionAnswering,
              FlaubertForQuestionAnsweringSimple,
              FlaubertForSequenceClassification,
              FlaubertForTokenClassification,
              FlaubertForMultipleChoice,
          )
          if is_torch_available()
          else ()
      )
      _SCREAMING_SNAKE_CASE					=      (
          {
              'feature-extraction': FlaubertModel,
              'fill-mask': FlaubertWithLMHeadModel,
              'question-answering': FlaubertForQuestionAnsweringSimple,
              'text-classification': FlaubertForSequenceClassification,
              'token-classification': FlaubertForTokenClassification,
              'zero-shot': FlaubertForSequenceClassification,
          }
          if is_torch_available()
          else {}
      )





      def     _snake_case			(       self      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase )		->     Optional[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     _snake_case			(       self      ,						lowercase      ,						lowercase      ,						lowercase=False )		->     Optional[Any]:
            lowerCAmelCase  						=       super()._prepare_for_class(lowercase      ,						lowercase      ,						return_labels=lowercase )

            if return_labels:
                  if model_class.__name__ == "FlaubertForQuestionAnswering":
                        lowerCAmelCase  						=       torch.zeros(
                            self.model_tester.batch_size      ,						dtype=torch.long      ,						device=lowercase )
                        lowerCAmelCase  						=       torch.zeros(
                            self.model_tester.batch_size      ,						dtype=torch.long      ,						device=lowercase )

            return inputs_dict





      def     _snake_case			(       self )		->     List[str]:
            lowerCAmelCase  						=       FlaubertModelTester(self )
            lowerCAmelCase  						=       ConfigTester(self      ,						config_class=lowercase      ,						emb_dim=37 )





      def     _snake_case			(       self )		->     Tuple:
            self.config_tester.run_common_tests()





      def     _snake_case			(       self )		->     Tuple:
            lowerCAmelCase  						=       self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_flaubert_model(*lowercase )





      def     _snake_case			(       self )		->     Tuple:
            lowerCAmelCase  						=       self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_flaubert_lm_head(*lowercase )





      def     _snake_case			(       self )		->     Optional[int]:
            lowerCAmelCase  						=       self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_flaubert_simple_qa(*lowercase )





      def     _snake_case			(       self )		->     Union[str, Any]:
            lowerCAmelCase  						=       self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_flaubert_qa(*lowercase )





      def     _snake_case			(       self )		->     Any:
            lowerCAmelCase  						=       self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_flaubert_sequence_classif(*lowercase )





      def     _snake_case			(       self )		->     Any:
            lowerCAmelCase  						=       self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_flaubert_token_classif(*lowercase )





      def     _snake_case			(       self )		->     Any:
            lowerCAmelCase  						=       self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_flaubert_multiple_choice(*lowercase )





      @slow
      def     _snake_case			(       self )		->     Tuple:
            for model_name in FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                  lowerCAmelCase  						=       FlaubertModel.from_pretrained(lowercase )
                  self.assertIsNotNone(lowercase )





      @slow
      @require_torch_gpu
      def     _snake_case			(       self )		->     List[Any]:
            lowerCAmelCase						,		lowerCAmelCase  						=       self.model_tester.prepare_config_and_inputs_for_common()
            for model_class in self.all_model_classes:
                  # FlauBertForMultipleChoice behaves incorrectly in JIT environments.
                  if model_class == FlaubertForMultipleChoice:
                        return

                  lowerCAmelCase  						=       True
                  lowerCAmelCase  						=       model_class(config=lowercase )

                  lowerCAmelCase  						=       self._prepare_for_class(lowercase      ,						lowercase )
                  lowerCAmelCase  						=       torch.jit.trace(
                      lowercase      ,						(inputs_dict["""input_ids"""].to("""cpu""" ), inputs_dict["""attention_mask"""].to("""cpu""" )) )

                  with tempfile.TemporaryDirectory() as tmp:
                        torch.jit.save(lowercase      ,						os.path.join(lowercase      ,						"""traced_model.pt""" ) )
                        lowerCAmelCase  						=       torch.jit.load(os.path.join(lowercase      ,						"""traced_model.pt""" )      ,						map_location=lowercase )
                        loaded(inputs_dict["""input_ids"""].to(lowercase )      ,						inputs_dict["""attention_mask"""].to(lowercase ) )







@require_torch
class 		lowercase					(					unittest.TestCase					):





      @slow
      def     _snake_case			(       self )		->     Optional[int]:
            lowerCAmelCase  						=       FlaubertModel.from_pretrained("""flaubert/flaubert_base_cased""" )
            lowerCAmelCase  						=       torch.tensor([[0, 345, 232, 328, 740, 140, 1_695, 69, 6_078, 1_588, 2]] )
            with torch.no_grad():
                  lowerCAmelCase  						=       model(lowercase )[0]
            lowerCAmelCase  						=       torch.Size((1, 11, 768) )
            self.assertEqual(output.shape      ,						lowercase )
            lowerCAmelCase  						=       torch.tensor(
                [[[-2.6_251, -1.4_298, -0.0_227], [-2.8_510, -1.6_387, 0.2_258], [-2.8_114, -1.1_832, -0.3_066]]] )

            self.assertTrue(torch.allclose(output[:, :3, :3]      ,						lowercase      ,						atol=1e-4 ) )











import gc
import unittest

from diffusers import FlaxStableDiffusionInpaintPipeline
from diffusers.utils import is_flax_available, load_image, slow
from diffusers.utils.testing_utils import require_flax


if is_flax_available():
						import jax
						import jax.numpy as jnp
						from flax.jax_utils import replicate
						from flax.training.common_utils import shard



@slow
@require_flax
class 				__snake_case						(     unittest.TestCase						):




							def    UpperCAmelCase__			(       self     :       Tuple):
									# clean up the VRAM after each test
									super().tearDown()
									gc.collect()







							def    UpperCAmelCase__			(       self     :       List[str]):
									lowerCAmelCase_			:     Any									= load_image(
									    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
									    '''/sd2-inpaint/init_image.png''')
									lowerCAmelCase_			:     str									= load_image(
									    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint/mask.png''')

									lowerCAmelCase_			:     List[Any]									= '''xvjiarui/stable-diffusion-2-inpainting'''
									lowerCAmelCase_					,  lowerCAmelCase_			:     Optional[Any]									= FlaxStableDiffusionInpaintPipeline.from_pretrained(A_				,				safety_checker=A_)

									lowerCAmelCase_			:     List[str]									= '''Face of a yellow cat, high resolution, sitting on a park bench'''

									lowerCAmelCase_			:     List[Any]									= jax.random.PRNGKey(0)
									lowerCAmelCase_			:     str									= 5_0

									lowerCAmelCase_			:     List[Any]									= jax.device_count()
									lowerCAmelCase_			:     Union[str, Any]									= num_samples * [prompt]
									lowerCAmelCase_			:     str									= num_samples * [init_image]
									lowerCAmelCase_			:     Union[str, Any]									= num_samples * [mask_image]
									lowerCAmelCase_					,  lowerCAmelCase_					,  lowerCAmelCase_			:     Optional[Any]									= pipeline.prepare_inputs(A_				,				A_				,				A_)

									# shard inputs and rng
									lowerCAmelCase_			:     str									= replicate(A_)
									lowerCAmelCase_			:     str									= jax.random.split(A_				,				jax.device_count())
									lowerCAmelCase_			:     List[Any]									= shard(A_)
									lowerCAmelCase_			:     str									= shard(A_)
									lowerCAmelCase_			:     Tuple									= shard(A_)

									lowerCAmelCase_			:     int									= pipeline(
									    A_				,				A_				,				A_				,				A_				,				A_				,				A_				,				jit=A_)

									lowerCAmelCase_			:     Optional[int]									= output.images.reshape(A_				,				5_1_2				,				5_1_2				,				3)

									lowerCAmelCase_			:     List[str]									= images[0, 2_5_3:2_5_6, 2_5_3:2_5_6, -1]

									lowerCAmelCase_			:     int									= jnp.asarray(jax.device_get(image_slice.flatten()))
									lowerCAmelCase_			:     str									= jnp.array(
									    [0.361_1307, 0.3764_9736, 0.375_7408, 0.3821_3953, 0.3929_5167, 0.384_1631, 0.4155_4978, 0.413_7475, 0.421_7084])
									print(F"""output_slice: {output_slice}""")
									assert jnp.abs(output_slice - expected_slice).max() < 1e-2





