


import unittest

from transformers import ErnieMConfig, is_torch_available
from transformers.testing_utils import require_torch, 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 (
           ErnieMForInformationExtraction,
           ErnieMForMultipleChoice,
           ErnieMForQuestionAnswering,
           ErnieMForSequenceClassification,
           ErnieMForTokenClassification,
           ErnieMModel,
       )
       from transformers.models.ernie_m.modeling_ernie_m import ERNIE_M_PRETRAINED_MODEL_ARCHIVE_LIST


class 	a__						:

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

       def     UpperCAmelCase( self   :  Optional[Any]     ):
            a_     :  List[str]				      =     ids_tensor([self.batch_size, self.seq_length]   ,				self.vocab_size     )

            a_     :  List[Any]				      =     None
            if self.use_input_mask:
                 a_     :  Tuple				      =     random_attention_mask([self.batch_size, self.seq_length]     )

            a_     :  Dict				      =     None
            a_     :  List[Any]				      =     None
            a_     :  List[Any]				      =     None
            if self.use_labels:
                 a_     :  Optional[Any]				      =     ids_tensor([self.batch_size]   ,				self.type_sequence_label_size     )
                 a_     :  str				      =     ids_tensor([self.batch_size, self.seq_length]   ,				self.num_labels     )
                 a_     :  Optional[Any]				      =     ids_tensor([self.batch_size]   ,				self.num_choices     )

            a_     :  Any				      =     self.get_config()

            return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels

       def     UpperCAmelCase( self   :  int     ):
            a_     :  Any				      =     ids_tensor([self.batch_size, self.seq_length]   ,				self.vocab_size     )

            a_     :  List[str]				      =     None
            if self.use_input_mask:
                 a_     :  Tuple				      =     random_attention_mask([self.batch_size, self.seq_length]     )
            a_     :  Optional[int]				      =     self.get_config()

            return config, input_ids, input_mask

       def     UpperCAmelCase( self   :  Optional[int]     ):
            return ErnieMConfig(
                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   ,				initializer_range=self.initializer_range   ,				)

       def     UpperCAmelCase( self   :  Any   ,				lowerCamelCase_   :  List[Any]   ,				lowerCamelCase_   :  Dict   ,				lowerCamelCase_   :  Optional[Any]   ,				lowerCamelCase_   :  Any   ,				lowerCamelCase_   :  Dict   ,				lowerCamelCase_   :  Union[str, Any]     ):
            a_     :  List[Any]				      =     ErnieMModel(config=A_     )
            model.to(A_     )
            model.eval()
            a_     :  Optional[int]				      =     model(A_   ,				return_dict=A_     )
            self.parent.assertEqual(result.last_hidden_state.shape   ,				(self.batch_size, self.seq_length, self.hidden_size)     )

       def     UpperCAmelCase( self   :  Tuple   ,				lowerCamelCase_   :  Optional[Any]   ,				lowerCamelCase_   :  int   ,				lowerCamelCase_   :  Optional[Any]   ,				lowerCamelCase_   :  Union[str, Any]   ,				lowerCamelCase_   :  Union[str, Any]   ,				lowerCamelCase_   :  str     ):
            a_     :  Union[str, Any]				      =     ErnieMForQuestionAnswering(config=A_     )
            model.to(A_     )
            model.eval()
            a_     :  Any				      =     model(
                A_   ,				attention_mask=A_   ,				start_positions=A_   ,				end_positions=A_   ,				)
            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     UpperCAmelCase( self   :  Union[str, Any]   ,				lowerCamelCase_   :  Any   ,				lowerCamelCase_   :  Tuple   ,				lowerCamelCase_   :  List[Any]   ,				lowerCamelCase_   :  Union[str, Any]   ,				lowerCamelCase_   :  int   ,				lowerCamelCase_   :  Any     ):
            a_     :  Optional[Any]				      =     ErnieMForInformationExtraction(config=A_     )
            model.to(A_     )
            model.eval()
            a_     :  Optional[Any]				      =     torch.ones_like(A_   ,				dtype=torch.floataa     )
            a_     :  Any				      =     model(
                A_   ,				attention_mask=A_   ,				start_positions=A_   ,				end_positions=A_   ,				)
            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     UpperCAmelCase( self   :  Optional[int]   ,				lowerCamelCase_   :  List[Any]   ,				lowerCamelCase_   :  Union[str, Any]   ,				lowerCamelCase_   :  Any   ,				lowerCamelCase_   :  Optional[int]   ,				lowerCamelCase_   :  int   ,				lowerCamelCase_   :  int     ):
            a_     :  str				      =     self.num_labels
            a_     :  Optional[int]				      =     ErnieMForSequenceClassification(A_     )
            model.to(A_     )
            model.eval()
            a_     :  List[str]				      =     model(A_   ,				attention_mask=A_   ,				labels=A_     )
            self.parent.assertEqual(result.logits.shape   ,				(self.batch_size, self.num_labels)     )

       def     UpperCAmelCase( self   :  Union[str, Any]   ,				lowerCamelCase_   :  Union[str, Any]   ,				lowerCamelCase_   :  Any   ,				lowerCamelCase_   :  Tuple   ,				lowerCamelCase_   :  Union[str, Any]   ,				lowerCamelCase_   :  int   ,				lowerCamelCase_   :  Dict     ):
            a_     :  Tuple				      =     self.num_labels
            a_     :  Tuple				      =     ErnieMForTokenClassification(config=A_     )
            model.to(A_     )
            model.eval()
            input_ids.to(A_     )
            input_mask.to(A_     )
            token_labels.to(A_     )

            a_     :  Optional[Any]				      =     model(A_   ,				attention_mask=A_   ,				labels=A_     )

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

       def     UpperCAmelCase( self   :  List[str]   ,				lowerCamelCase_   :  Optional[Any]   ,				lowerCamelCase_   :  Optional[Any]   ,				lowerCamelCase_   :  Optional[int]   ,				lowerCamelCase_   :  List[str]   ,				lowerCamelCase_   :  Dict   ,				lowerCamelCase_   :  Any     ):
            a_     :  Union[str, Any]				      =     self.num_choices
            a_     :  str				      =     ErnieMForMultipleChoice(config=A_     )
            model.to(A_     )
            model.eval()
            a_     :  Union[str, Any]				      =     input_ids.unsqueeze(1     ).expand(-1   ,				self.num_choices   ,				-1     ).contiguous()
            a_     :  Dict				      =     input_mask.unsqueeze(1     ).expand(-1   ,				self.num_choices   ,				-1     ).contiguous()
            a_     :  Optional[Any]				      =     model(
                A_   ,				attention_mask=A_   ,				labels=A_   ,				)
            self.parent.assertEqual(result.logits.shape   ,				(self.batch_size, self.num_choices)     )

       def     UpperCAmelCase( self   :  Tuple     ):
            a_     :  List[Any]				      =     self.prepare_config_and_inputs()
            (
                (
                a_    
            )      ,      (
                a_    
            )      ,      (
                a_    
            )      ,      (
                a_    
            )      ,      (
                a_    
            )      ,      (
                a_    
            )      ,          
            )     :  Optional[int]				      =     config_and_inputs
            a_     :  Optional[Any]				      =     {"""input_ids""": input_ids, """attention_mask""": input_mask}
            return config, inputs_dict


@require_torch
class 	a__						(    _SCREAMING_SNAKE_CASE    ,							_SCREAMING_SNAKE_CASE    ,							unittest.TestCase     ):
       lowerCamelCase__:  int				     =  (
           (
               ErnieMModel,
               ErnieMForMultipleChoice,
               ErnieMForQuestionAnswering,
               ErnieMForSequenceClassification,
               ErnieMForTokenClassification,
           )
           if is_torch_available()
           else ()
       )
       lowerCamelCase__:  Any				     =  ()
       lowerCamelCase__:  List[Any]				     =  (
           {
               "feature-extraction": ErnieMModel,
               "question-answering": ErnieMForQuestionAnswering,
               "text-classification": ErnieMForSequenceClassification,
               "token-classification": ErnieMForTokenClassification,
               "zero-shot": ErnieMForSequenceClassification,
           }
           if is_torch_available()
           else {}
       )
       lowerCamelCase__:  str				     =  False

       def     UpperCAmelCase( self   :  List[Any]   ,				lowerCamelCase_   :  Tuple   ,				lowerCamelCase_   :  int   ,				lowerCamelCase_   :  Optional[int]   ,				lowerCamelCase_   :  int   ,				lowerCamelCase_   :  Optional[int]     ):
            if pipeline_test_casse_name == "QAPipelineTests":
                 return True

            return False

       def     UpperCAmelCase( self   :  Any     ):
            a_     :  Optional[int]				      =     ErnieMModelTester(self     )
            a_     :  Union[str, Any]				      =     ConfigTester(self   ,				config_class=A_   ,				hidden_size=3_7     )

       def     UpperCAmelCase( self   :  Optional[Any]     ):
            self.config_tester.run_common_tests()

       def     UpperCAmelCase( self   :  Union[str, Any]     ):
            a_     :  Dict				      =     self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_model(*A_     )

       def     UpperCAmelCase( self   :  Dict     ):
            a_     :  int				      =     self.model_tester.prepare_config_and_inputs()
            for type in ["absolute", "relative_key", "relative_key_query"]:
                 a_     :  Dict				      =     type
                 self.model_tester.create_and_check_model(*A_     )

       def     UpperCAmelCase( self   :  Dict     ):
            a_     :  Optional[int]				      =     self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_for_multiple_choice(*A_     )

       def     UpperCAmelCase( self   :  Optional[int]     ):
            a_     :  Tuple				      =     self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_for_question_answering(*A_     )

       def     UpperCAmelCase( self   :  Union[str, Any]     ):
            a_     :  Optional[Any]				      =     self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_for_sequence_classification(*A_     )

       def     UpperCAmelCase( self   :  int     ):
            a_     :  Dict				      =     self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_for_information_extraction(*A_     )

       def     UpperCAmelCase( self   :  str     ):
            a_     :  Optional[int]				      =     self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_for_token_classification(*A_     )

       @slow
       def     UpperCAmelCase( self   :  Tuple     ):
            for model_name in ERNIE_M_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 a_     :  Optional[int]				      =     ErnieMModel.from_pretrained(A_     )
                 self.assertIsNotNone(A_     )


@require_torch
class 	a__						(    unittest.TestCase     ):

       @slow
       def     UpperCAmelCase( self   :  Union[str, Any]     ):
            a_     :  str				      =     ErnieMModel.from_pretrained("""susnato/ernie-m-base_pytorch"""     )
            model.eval()
            a_     :  Any				      =     torch.tensor([[0, 1, 2, 3, 4, 5]]     )
            a_     :  Union[str, Any]				      =     model(A_     )[0]

            # TODO Replace vocab size
            a_     :  List[str]				      =     7_6_8

            a_     :  List[str]				      =     torch.Size((1, 6, hidden_size)     )
            self.assertEqual(output.shape   ,				A_     )

            a_     :  List[Any]				      =     torch.tensor(
                [[[-0.0_0_1_2, 0.1_2_4_5, -0.0_2_1_4], [-0.0_7_4_2, 0.0_2_4_4, -0.0_7_7_1], [-0.0_3_3_3, 0.1_1_6_4, -0.1_5_5_4]]]     )

            self.assertTrue(torch.allclose(output[:, :3, :3]   ,				A_   ,				atol=1E-3     )     )





from dataclasses import dataclass
from typing import List, Optional, Union

import numpy as np
import PIL

from ...utils import BaseOutput, is_invisible_watermark_available, is_torch_available, is_transformers_available


@dataclass
class 							lowercase   (						_SCREAMING_SNAKE_CASE  ):
		__lowercase			:			Union[List[PIL.Image.Image], np.ndarray]


if is_transformers_available() and is_torch_available() and is_invisible_watermark_available():
					from .pipeline_stable_diffusion_xl import StableDiffusionXLPipeline
					from .pipeline_stable_diffusion_xl_imgaimg import StableDiffusionXLImgaImgPipeline
