




from __future__ import annotations

import unittest

from transformers import AlbertConfig, is_tf_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_tf, 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 tensorflow as tf

   from transformers import TF_MODEL_FOR_PRETRAINING_MAPPING
   from transformers.models.albert.modeling_tf_albert import (
       TF_ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
       TFAlbertForMaskedLM,
       TFAlbertForMultipleChoice,
       TFAlbertForPreTraining,
       TFAlbertForQuestionAnswering,
       TFAlbertForSequenceClassification,
       TFAlbertForTokenClassification,
       TFAlbertModel,
   )

class      UpperCAmelCase						:






   def __init__(	self   :       List[Any]						,					UpperCamelCase_   :       Any						,					UpperCamelCase_   :       str=13						,					UpperCamelCase_   :       List[Any]=7						,					UpperCamelCase_   :       Optional[int]=True						,					UpperCamelCase_   :       Dict=True						,					UpperCamelCase_   :       List[Any]=True						,					UpperCamelCase_   :       Dict=True						,					UpperCamelCase_   :       Optional[int]=99						,					UpperCamelCase_   :       List[str]=16						,					UpperCamelCase_   :       int=32						,					UpperCamelCase_   :       Tuple=2						,					UpperCamelCase_   :       List[str]=4						,					UpperCamelCase_   :       Tuple=37						,					UpperCamelCase_   :       Optional[Any]="gelu"						,					UpperCamelCase_   :       Optional[Any]=0.1						,					UpperCamelCase_   :       Optional[Any]=0.1						,					UpperCamelCase_   :       Optional[Any]=5_12						,					UpperCamelCase_   :       Tuple=16						,					UpperCamelCase_   :       Tuple=2						,					UpperCamelCase_   :       Union[str, Any]=0.02						,					UpperCamelCase_   :       Optional[Any]=3						,					UpperCamelCase_   :       Any=4						,					UpperCamelCase_   :       List[str]=None						,					)->			str:
     _lowercase: Optional[int]						     =    parent
     _lowercase: Any						     =    13
     _lowercase: Tuple						     =    7
     _lowercase: str						     =    True
     _lowercase: int						     =    True
     _lowercase: Optional[Any]						     =    True
     _lowercase: Optional[Any]						     =    True
     _lowercase: Any						     =    99
     _lowercase: List[str]						     =    16
     _lowercase: str						     =    32
     _lowercase: Optional[int]						     =    2
     _lowercase: Optional[int]						     =    4
     _lowercase: Any						     =    37
     _lowercase: List[Any]						     =    'gelu'
     _lowercase: Optional[int]						     =    0.1
     _lowercase: List[str]						     =    0.1
     _lowercase: Union[str, Any]						     =    5_12
     _lowercase: List[str]						     =    16
     _lowercase: List[Any]						     =    2
     _lowercase: int						     =    0.02
     _lowercase: List[str]						     =    3
     _lowercase: Dict						     =    4
     _lowercase: Union[str, Any]						     =    None






   def  _a						(	self   :       Optional[int]     )->			Union[str, Any]:
     _lowercase: Dict						     =    ids_tensor([self.batch_size, self.seq_length]						,					self.vocab_size     )

     _lowercase: Optional[int]						     =    None
     if self.use_input_mask:
       _lowercase: Union[str, Any]						     =    random_attention_mask([self.batch_size, self.seq_length]     )

     _lowercase: Union[str, Any]						     =    None
     if self.use_token_type_ids:
       _lowercase: Tuple						     =    ids_tensor([self.batch_size, self.seq_length]						,					self.type_vocab_size     )

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

     _lowercase: Optional[int]						     =    AlbertConfig(
         vocab_size=self.vocab_size						,					hidden_size=self.hidden_size						,					embedding_size=self.embedding_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						,					)

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






   def  _a						(	self   :       str						,					UpperCamelCase_   :       List[str]						,					UpperCamelCase_   :       Optional[int]						,					UpperCamelCase_   :       List[str]						,					UpperCamelCase_   :       int						,					UpperCamelCase_   :       Optional[int]						,					UpperCamelCase_   :       int						,					UpperCamelCase_   :       Tuple     )->			int:
     _lowercase: Tuple						     =    TFAlbertModel(config=UpperCamelCase_     )
     # inputs = {'input_ids': input_ids,
     #           'attention_mask': input_mask,
     #           'token_type_ids': token_type_ids}
     # sequence_output, pooled_output = model(**inputs)
     _lowercase: str						     =    {'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
     _lowercase: Dict						     =    model(UpperCamelCase_     )

     _lowercase: Tuple						     =    [input_ids, input_mask]
     _lowercase: List[str]						     =    model(UpperCamelCase_     )

     _lowercase: Any						     =    model(UpperCamelCase_     )

     self.parent.assertEqual(result.last_hidden_state.shape						,					(self.batch_size, self.seq_length, self.hidden_size)     )
     self.parent.assertEqual(result.pooler_output.shape						,					(self.batch_size, self.hidden_size)     )






   def  _a						(	self   :       Dict						,					UpperCamelCase_   :       Tuple						,					UpperCamelCase_   :       Optional[int]						,					UpperCamelCase_   :       int						,					UpperCamelCase_   :       int						,					UpperCamelCase_   :       Optional[Any]						,					UpperCamelCase_   :       Tuple						,					UpperCamelCase_   :       Union[str, Any]     )->			List[str]:
     _lowercase: str						     =    self.num_labels
     _lowercase: Union[str, Any]						     =    TFAlbertForPreTraining(config=UpperCamelCase_     )
     _lowercase: List[str]						     =    {'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
     _lowercase: List[str]						     =    model(UpperCamelCase_     )
     self.parent.assertEqual(result.prediction_logits.shape						,					(self.batch_size, self.seq_length, self.vocab_size)     )
     self.parent.assertEqual(result.sop_logits.shape						,					(self.batch_size, self.num_labels)     )






   def  _a						(	self   :       Dict						,					UpperCamelCase_   :       str						,					UpperCamelCase_   :       Dict						,					UpperCamelCase_   :       Optional[Any]						,					UpperCamelCase_   :       List[str]						,					UpperCamelCase_   :       Any						,					UpperCamelCase_   :       List[str]						,					UpperCamelCase_   :       Dict     )->			List[Any]:
     _lowercase: Optional[Any]						     =    TFAlbertForMaskedLM(config=UpperCamelCase_     )
     _lowercase: List[str]						     =    {'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
     _lowercase: Optional[int]						     =    model(UpperCamelCase_     )
     self.parent.assertEqual(result.logits.shape						,					(self.batch_size, self.seq_length, self.vocab_size)     )






   def  _a						(	self   :       List[Any]						,					UpperCamelCase_   :       List[str]						,					UpperCamelCase_   :       str						,					UpperCamelCase_   :       Any						,					UpperCamelCase_   :       int						,					UpperCamelCase_   :       Any						,					UpperCamelCase_   :       int						,					UpperCamelCase_   :       Union[str, Any]     )->			List[Any]:
     _lowercase: int						     =    self.num_labels
     _lowercase: Union[str, Any]						     =    TFAlbertForSequenceClassification(config=UpperCamelCase_     )
     _lowercase: Tuple						     =    {'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
     _lowercase: Dict						     =    model(UpperCamelCase_     )
     self.parent.assertEqual(result.logits.shape						,					(self.batch_size, self.num_labels)     )






   def  _a						(	self   :       Dict						,					UpperCamelCase_   :       List[Any]						,					UpperCamelCase_   :       Tuple						,					UpperCamelCase_   :       Any						,					UpperCamelCase_   :       List[str]						,					UpperCamelCase_   :       List[Any]						,					UpperCamelCase_   :       Tuple						,					UpperCamelCase_   :       Union[str, Any]     )->			Any:
     _lowercase: Dict						     =    TFAlbertForQuestionAnswering(config=UpperCamelCase_     )
     _lowercase: Tuple						     =    {'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
     _lowercase: Optional[Any]						     =    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   :       Union[str, Any]						,					UpperCamelCase_   :       List[Any]						,					UpperCamelCase_   :       List[str]						,					UpperCamelCase_   :       List[str]						,					UpperCamelCase_   :       Union[str, Any]						,					UpperCamelCase_   :       Any						,					UpperCamelCase_   :       int						,					UpperCamelCase_   :       List[Any]     )->			List[str]:
     _lowercase: Optional[Any]						     =    self.num_choices
     _lowercase: Optional[int]						     =    TFAlbertForMultipleChoice(config=UpperCamelCase_     )
     _lowercase: str						     =    tf.tile(tf.expand_dims(UpperCamelCase_						,					1     )						,					(1, self.num_choices, 1)     )
     _lowercase: int						     =    tf.tile(tf.expand_dims(UpperCamelCase_						,					1     )						,					(1, self.num_choices, 1)     )
     _lowercase: List[str]						     =    tf.tile(tf.expand_dims(UpperCamelCase_						,					1     )						,					(1, self.num_choices, 1)     )
     _lowercase: Tuple						     =    {
         'input_ids': multiple_choice_inputs_ids,
         'attention_mask': multiple_choice_input_mask,
         'token_type_ids': multiple_choice_token_type_ids,
     }
     _lowercase: List[Any]						     =    model(UpperCamelCase_     )
     self.parent.assertListEqual(list(result['logits'].shape     )						,					[self.batch_size, self.num_choices]     )






   def  _a						(	self   :       List[Any]						,					UpperCamelCase_   :       List[Any]						,					UpperCamelCase_   :       Any						,					UpperCamelCase_   :       str						,					UpperCamelCase_   :       Union[str, Any]						,					UpperCamelCase_   :       Any						,					UpperCamelCase_   :       List[str]						,					UpperCamelCase_   :       Tuple     )->			str:
     _lowercase: int						     =    self.num_labels
     _lowercase: List[str]						     =    TFAlbertForTokenClassification(config=UpperCamelCase_     )
     _lowercase: List[Any]						     =    {
         'input_ids': input_ids,
         'attention_mask': input_mask,
         'token_type_ids': token_type_ids,
     }
     _lowercase: Tuple						     =    model(UpperCamelCase_     )
     self.parent.assertListEqual(list(result['logits'].shape     )						,					[self.batch_size, self.seq_length, self.num_labels]     )






   def  _a						(	self   :       int     )->			List[Any]:
     _lowercase: Dict						     =    self.prepare_config_and_inputs()
     (
         (
         _lowercase
     )     ,					(
         _lowercase
     )     ,					(
         _lowercase
     )     ,					(
         _lowercase
     )     ,					(
         _lowercase
     )     ,					(
         _lowercase
     )     ,					(
         _lowercase
     )     ,					
     ): Union[str, Any]						     =    config_and_inputs
     _lowercase: List[Any]						     =    {'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
     return config, inputs_dict

@require_tf
class      UpperCAmelCase						(			lowercase_   ,					lowercase_   ,					unittest.TestCase   ):
   __SCREAMING_SNAKE_CASE		:							Dict							=					(
       (
           TFAlbertModel,
           TFAlbertForPreTraining,
           TFAlbertForMaskedLM,
           TFAlbertForSequenceClassification,
           TFAlbertForQuestionAnswering,
           TFAlbertForTokenClassification,
           TFAlbertForMultipleChoice,
       )
       if is_tf_available()
       else ()
   )
   __SCREAMING_SNAKE_CASE		:							Union[str, Any]							=					(
       {
           """feature-extraction""": TFAlbertModel,
           """fill-mask""": TFAlbertForMaskedLM,
           """question-answering""": TFAlbertForQuestionAnswering,
           """text-classification""": TFAlbertForSequenceClassification,
           """token-classification""": TFAlbertForTokenClassification,
           """zero-shot""": TFAlbertForSequenceClassification,
       }
       if is_tf_available()
       else {}
   )
   __SCREAMING_SNAKE_CASE		:							str							=					False
   __SCREAMING_SNAKE_CASE		:							Optional[Any]							=					False






   def  _a						(	self   :       Any						,					UpperCamelCase_   :       Optional[int]						,					UpperCamelCase_   :       int						,					UpperCamelCase_   :       Union[str, Any]=False     )->			Optional[Any]:
     _lowercase: Tuple						     =    super()._prepare_for_class(UpperCamelCase_						,					UpperCamelCase_						,					return_labels=UpperCamelCase_     )

     if return_labels:
       if model_class in get_values(UpperCamelCase_     ):
         _lowercase: Optional[Any]						     =    tf.zeros(self.model_tester.batch_size						,					dtype=tf.intaa     )

     return inputs_dict






   def  _a						(	self   :       List[Any]     )->			List[Any]:
     _lowercase: Dict						     =    TFAlbertModelTester(self     )
     _lowercase: List[Any]						     =    ConfigTester(self						,					config_class=UpperCamelCase_						,					hidden_size=37     )






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






   def  _a						(	self   :       Optional[int]     )->			str:
     _lowercase: str						     =    self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_albert_model(*UpperCamelCase_     )






   def  _a						(	self   :       int     )->			Optional[Any]:
     _lowercase: Tuple						     =    self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_albert_for_pretraining(*UpperCamelCase_     )






   def  _a						(	self   :       List[str]     )->			Dict:
     _lowercase: str						     =    self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_albert_for_masked_lm(*UpperCamelCase_     )






   def  _a						(	self   :       List[str]     )->			Optional[Any]:
     _lowercase: List[str]						     =    self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_albert_for_multiple_choice(*UpperCamelCase_     )






   def  _a						(	self   :       Optional[int]     )->			List[str]:
     _lowercase: List[Any]						     =    self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_albert_for_sequence_classification(*UpperCamelCase_     )






   def  _a						(	self   :       int     )->			Optional[int]:
     _lowercase: Optional[Any]						     =    self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_albert_for_question_answering(*UpperCamelCase_     )






   @slow
   def  _a						(	self   :       List[str]     )->			int:







     for model_name in TF_ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
       _lowercase: Union[str, Any]						     =    TFAlbertModel.from_pretrained(UpperCamelCase_     )
       self.assertIsNotNone(UpperCamelCase_     )

@require_tf
class      UpperCAmelCase						(			unittest.TestCase   ):






   @slow
   def  _a						(	self   :       List[str]     )->			Tuple:
     _lowercase: Optional[Any]						     =    TFAlbertForPreTraining.from_pretrained('albert-base-v2'     )
     _lowercase: str						     =    tf.constant([[0, 1, 2, 3, 4, 5]]     )
     _lowercase: Optional[Any]						     =    model(UpperCamelCase_     )[0]

     _lowercase: Optional[int]						     =    [1, 6, 3_00_00]
     self.assertEqual(output.shape						,					UpperCamelCase_     )

     _lowercase: Optional[int]						     =    tf.constant(
         [
             [
                 [4.595_668, 0.74_462_754, -1.818_147],
                 [4.5_954_347, 0.7_454_184, -1.8_188_258],
                 [4.5_954_905, 0.7_448_235, -1.8_182_316],
             ]
         ]     )
     tf.debugging.assert_near(output[:, :3, :3]						,					UpperCamelCase_						,					atol=1E-4     )






# limitations under the License.


from typing import Optional, Tuple, Union

import torch

from diffusers.pipeline_utils import DiffusionPipeline, ImagePipelineOutput

class      UpperCAmelCase						(			A   ):






   def __init__(	self   :       Optional[Any]						,					UpperCamelCase_   :       Optional[Any]						,					UpperCamelCase_   :       Tuple     )->			Any:
     super().__init__()
     self.register_modules(unet=UpperCamelCase_						,					scheduler=UpperCamelCase_     )






   @torch.no_grad()
   def __call__(	self   :       Optional[int]						,					UpperCamelCase_   :       int = 1						,					UpperCamelCase_   :       Optional[torch.Generator] = None						,					UpperCamelCase_   :       int = 50						,					UpperCamelCase_   :       Optional[str] = "pil"						,					UpperCamelCase_   :       bool = True						,					**UpperCamelCase_   :       Dict						,					)->			Union[ImagePipelineOutput, Tuple]:
     _lowercase: List[str]						     =    torch.randn(
         (batch_size, self.unet.config.in_channels, self.unet.config.sample_size, self.unet.config.sample_size)						,					generator=UpperCamelCase_						,					)
     _lowercase: str						     =    image.to(self.device     )

     # set step values
     self.scheduler.set_timesteps(UpperCamelCase_     )

     for t in self.progress_bar(self.scheduler.timesteps     ):
       # 1. predict noise model_output
       _lowercase: int						     =    self.unet(UpperCamelCase_						,					UpperCamelCase_     ).sample

       # 2. predict previous mean of image x_t-1 and add variance depending on eta
       # eta corresponds to η in paper and should be between [0, 1]
       # do x_t -> x_t-1
       _lowercase: List[Any]						     =    self.scheduler.step(UpperCamelCase_						,					UpperCamelCase_						,					UpperCamelCase_     ).prev_sample

     _lowercase: str						     =    (image / 2 + 0.5).clamp(0						,					1     )
     _lowercase: List[str]						     =    image.cpu().permute(0						,					2						,					3						,					1     ).numpy()
     if output_type == "pil":
       _lowercase: Any						     =    self.numpy_to_pil(UpperCamelCase_     )

     if not return_dict:
       return (image,), "This is a local test"

     return ImagePipelineOutput(images=UpperCamelCase_     ), "This is a local test"
