



from __future__ import annotations

import unittest

from transformers import is_tf_available
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_XLM_PRETRAINED_MODEL_ARCHIVE_LIST,
		    TFXLMForMultipleChoice,
		    TFXLMForQuestionAnsweringSimple,
		    TFXLMForSequenceClassification,
		    TFXLMForTokenClassification,
		    TFXLMModel,
		    TFXLMWithLMHeadModel,
		    XLMConfig,
		)




class 				lowerCAmelCase__	:







					def __init__(							self   :  List[Any]      ,		_A   :  List[Any]      ,		):
											A__						:				str         =			parent
											A__						:				Tuple         =			13
											A__						:				str         =			7
											A__						:				List[str]         =			True
											A__						:				Optional[Any]         =			True
											A__						:				Tuple         =			True
											A__						:				Tuple         =			True
											A__						:				int         =			True
											A__						:				int         =			False
											A__						:				Optional[Any]         =			False
											A__						:				Any         =			False
											A__						:				Any         =			2
											A__						:				str         =			99
											A__						:				int         =			0
											A__						:				Any         =			32
											A__						:				Tuple         =			2
											A__						:				Optional[int]         =			4
											A__						:				List[str]         =			0.1
											A__						:				List[Any]         =			0.1
											A__						:				List[str]         =			512
											A__						:				int         =			16
											A__						:				Any         =			2
											A__						:				Dict         =			0.02
											A__						:				Any         =			3
											A__						:				str         =			4
											A__						:				Any         =			'last'
											A__						:				int         =			True
											A__						:				List[str]         =			None
											A__						:				Union[str, Any]         =			0







					def 			_lowercase      (							self   :  Optional[Any]):
											A__						:				List[str]         =			ids_tensor([self.batch_size, self.seq_length]      ,		self.vocab_size)
											A__						:				str         =			random_attention_mask([self.batch_size, self.seq_length]      ,		dtype=tf.floataa)

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

											A__						:				int         =			None
											if self.use_token_type_ids:
																	A__						:				Union[str, Any]         =			ids_tensor([self.batch_size, self.seq_length]      ,		self.n_langs)

											A__						:				Optional[Any]         =			None
											A__						:				Optional[Any]         =			None
											A__						:				Tuple         =			None
											if self.use_labels:
																	A__						:				Tuple         =			ids_tensor([self.batch_size]      ,		self.type_sequence_label_size)
																	A__						:				Dict         =			ids_tensor([self.batch_size, self.seq_length]      ,		self.num_labels)
																	A__						:				Optional[int]         =			ids_tensor([self.batch_size]      ,		2      ,		dtype=tf.floataa)
																	A__						:				int         =			ids_tensor([self.batch_size]      ,		self.num_choices)

											A__						:				int         =			XLMConfig(
											    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 			_lowercase      (							self   :  List[Any]      ,		_A   :  str      ,		_A   :  List[str]      ,		_A   :  Optional[Any]      ,		_A   :  Dict      ,		_A   :  Tuple      ,		_A   :  Optional[Any]      ,		_A   :  List[Any]      ,		_A   :  Union[str, Any]      ,		_A   :  int      ,		):
											A__						:				str         =			TFXLMModel(config=_lowerCamelCase)
											A__						:				Optional[Any]         =			{'input_ids': input_ids, 'lengths': input_lengths, 'langs': token_type_ids}
											A__						:				str         =			model(_lowerCamelCase)

											A__						:				Dict         =			[input_ids, input_mask]
											A__						:				Tuple         =			model(_lowerCamelCase)
											self.parent.assertEqual(result.last_hidden_state.shape      ,		(self.batch_size, self.seq_length, self.hidden_size))







					def 			_lowercase      (							self   :  Dict      ,		_A   :  List[str]      ,		_A   :  int      ,		_A   :  Optional[int]      ,		_A   :  Any      ,		_A   :  Optional[int]      ,		_A   :  Optional[Any]      ,		_A   :  Tuple      ,		_A   :  List[Any]      ,		_A   :  Any      ,		):
											A__						:				int         =			TFXLMWithLMHeadModel(_lowerCamelCase)

											A__						:				List[str]         =			{'input_ids': input_ids, 'lengths': input_lengths, 'langs': token_type_ids}
											A__						:				List[str]         =			model(_lowerCamelCase)

											A__						:				Tuple         =			outputs

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







					def 			_lowercase      (							self   :  int      ,		_A   :  Union[str, Any]      ,		_A   :  str      ,		_A   :  Dict      ,		_A   :  int      ,		_A   :  Optional[Any]      ,		_A   :  Union[str, Any]      ,		_A   :  str      ,		_A   :  Optional[int]      ,		_A   :  List[str]      ,		):
											A__						:				Tuple         =			TFXLMForQuestionAnsweringSimple(_lowerCamelCase)

											A__						:				Dict         =			{'input_ids': input_ids, 'lengths': input_lengths}

											A__						:				Optional[Any]         =			model(_lowerCamelCase)

											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 			_lowercase      (							self   :  Optional[int]      ,		_A   :  Optional[Any]      ,		_A   :  Union[str, Any]      ,		_A   :  int      ,		_A   :  List[str]      ,		_A   :  List[Any]      ,		_A   :  int      ,		_A   :  Union[str, Any]      ,		_A   :  Any      ,		_A   :  str      ,		):
											A__						:				Optional[Any]         =			TFXLMForSequenceClassification(_lowerCamelCase)

											A__						:				Union[str, Any]         =			{'input_ids': input_ids, 'lengths': input_lengths}

											A__						:				Dict         =			model(_lowerCamelCase)

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







					def 			_lowercase      (							self   :  Any      ,		_A   :  Optional[int]      ,		_A   :  Union[str, Any]      ,		_A   :  Optional[int]      ,		_A   :  Optional[Any]      ,		_A   :  List[Any]      ,		_A   :  str      ,		_A   :  Any      ,		_A   :  List[Any]      ,		_A   :  Dict      ,		):
											A__						:				int         =			self.num_labels
											A__						:				int         =			TFXLMForTokenClassification(config=_lowerCamelCase)
											A__						:				List[str]         =			{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
											A__						:				List[str]         =			model(_lowerCamelCase)
											self.parent.assertEqual(result.logits.shape      ,		(self.batch_size, self.seq_length, self.num_labels))







					def 			_lowercase      (							self   :  str      ,		_A   :  Dict      ,		_A   :  Optional[Any]      ,		_A   :  List[str]      ,		_A   :  Optional[Any]      ,		_A   :  Tuple      ,		_A   :  Union[str, Any]      ,		_A   :  List[str]      ,		_A   :  Tuple      ,		_A   :  Any      ,		):
											A__						:				Any         =			self.num_choices
											A__						:				List[Any]         =			TFXLMForMultipleChoice(config=_lowerCamelCase)
											A__						:				Tuple         =			tf.tile(tf.expand_dims(_lowerCamelCase      ,		1)      ,		(1, self.num_choices, 1))
											A__						:				Optional[int]         =			tf.tile(tf.expand_dims(_lowerCamelCase      ,		1)      ,		(1, self.num_choices, 1))
											A__						:				Union[str, Any]         =			tf.tile(tf.expand_dims(_lowerCamelCase      ,		1)      ,		(1, self.num_choices, 1))
											A__						:				str         =			{
											    'input_ids': multiple_choice_inputs_ids,
											    'attention_mask': multiple_choice_input_mask,
											    'token_type_ids': multiple_choice_token_type_ids,
											}
											A__						:				Dict         =			model(_lowerCamelCase)
											self.parent.assertEqual(result.logits.shape      ,		(self.batch_size, self.num_choices))







					def 			_lowercase      (							self   :  Dict):
											A__						:				List[str]         =			self.prepare_config_and_inputs()
											(
											    A__       
											)						:				Any         =			config_and_inputs
											A__						:				Union[str, Any]         =			{
											    'input_ids': input_ids,
											    'token_type_ids': token_type_ids,
											    'langs': token_type_ids,
											    'lengths': input_lengths,
											}
											return config, inputs_dict




@require_tf
class 				lowerCAmelCase__	(						_lowerCamelCase       ,				_lowerCamelCase       ,				unittest.TestCase    ):
					__A  :     List[Any]     =    (
					    (
					        TFXLMModel,
					        TFXLMWithLMHeadModel,
					        TFXLMForSequenceClassification,
					        TFXLMForQuestionAnsweringSimple,
					        TFXLMForTokenClassification,
					        TFXLMForMultipleChoice,
					    )
					    if is_tf_available()
					    else ()
					)
					__A  :     Tuple     =    (
					    (TFXLMWithLMHeadModel,) if is_tf_available() else ()
					)  # TODO (PVP): Check other models whether language generation is also applicable
					__A  :     List[str]     =    (
					    {
					        'feature-extraction': TFXLMModel,
					        'fill-mask': TFXLMWithLMHeadModel,
					        'question-answering': TFXLMForQuestionAnsweringSimple,
					        'text-classification': TFXLMForSequenceClassification,
					        'text-generation': TFXLMWithLMHeadModel,
					        'token-classification': TFXLMForTokenClassification,
					        'zero-shot': TFXLMForSequenceClassification,
					    }
					    if is_tf_available()
					    else {}
					)
					__A  :     str     =    False
					__A  :     Dict     =    False







					def 			_lowercase      (							self   :  int      ,		_A   :  List[str]      ,		_A   :  Tuple      ,		_A   :  str      ,		_A   :  int      ,		_A   :  str):
											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 			_lowercase      (							self   :  int):
											A__						:				Optional[Any]         =			TFXLMModelTester(self)
											A__						:				Union[str, Any]         =			ConfigTester(self      ,		config_class=_lowerCamelCase      ,		emb_dim=37)







					def 			_lowercase      (							self   :  int):
											self.config_tester.run_common_tests()







					def 			_lowercase      (							self   :  Any):
											A__						:				Optional[int]         =			self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_xlm_model(*_lowerCamelCase)







					def 			_lowercase      (							self   :  Any):
											A__						:				Optional[int]         =			self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_xlm_lm_head(*_lowerCamelCase)







					def 			_lowercase      (							self   :  List[Any]):
											A__						:				Optional[Any]         =			self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_xlm_qa(*_lowerCamelCase)







					def 			_lowercase      (							self   :  Union[str, Any]):
											A__						:				Optional[Any]         =			self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_xlm_sequence_classif(*_lowerCamelCase)







					def 			_lowercase      (							self   :  List[str]):
											A__						:				Any         =			self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_xlm_for_token_classification(*_lowerCamelCase)







					def 			_lowercase      (							self   :  int):
											A__						:				Dict         =			self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_xlm_for_multiple_choice(*_lowerCamelCase)







					@slow
					def 			_lowercase      (							self   :  str):
											for model_name in TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																	A__						:				List[Any]         =			TFXLMModel.from_pretrained(_lowerCamelCase)
																	self.assertIsNotNone(_lowerCamelCase)




@require_tf
class 				lowerCAmelCase__	(						unittest.TestCase    ):







					@slow
					def 			_lowercase      (							self   :  Optional[int]):
											A__						:				Union[str, Any]         =			TFXLMWithLMHeadModel.from_pretrained("xlm-mlm-en-2048")
											A__						:				Union[str, Any]         =			tf.convert_to_tensor([[14, 447]]      ,		dtype=tf.intaa)  # the president
											A__						:				List[Any]         =			[
											    14,
											    447,
											    14,
											    447,
											    14,
											    447,
											    14,
											    447,
											    14,
											    447,
											    14,
											    447,
											    14,
											    447,
											    14,
											    447,
											    14,
											    447,
											    14,
											    447,
											]  # the president the president the president the president the president the president the president the president the president the president
											# TODO(PVP): this and other input_ids I tried for generation give pretty bad results. Not sure why. Model might just not be made for auto-regressive inference
											A__						:				str         =			model.generate(_lowerCamelCase      ,		do_sample=_lowerCamelCase)
											self.assertListEqual(output_ids[0].numpy().tolist()      ,		_lowerCamelCase)
'''simple docstring'''
import random
import unittest

import numpy as np
import torch

from diffusers import DDIMScheduler, LDMSuperResolutionPipeline, UNetaDModel, VQModel
from diffusers.utils import PIL_INTERPOLATION, floats_tensor, load_image, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch


enable_full_determinism()






class       __lowercase		(    unittest.TestCase    ):







     """simple docstring"""






     @property
     def lowerCAmelCase  (		self     ):
         __UpperCamelCase						:						Dict		     =	1
         __UpperCamelCase						:						Union[str, Any]		     =	3
         __UpperCamelCase						:						Tuple		     =	(3_2, 3_2)

         __UpperCamelCase						:						List[Any]		     =	floats_tensor((batch_size, num_channels) + sizes  ,      rng=random.Random(0     )     ).to(_lowerCamelCase     )
         return image







     @property
     def lowerCAmelCase  (		self     ):
         torch.manual_seed(0     )
         __UpperCamelCase						:						Dict		     =	UNetaDModel(
             block_out_channels=(3_2, 6_4)  ,      layers_per_block=2  ,      sample_size=3_2  ,      in_channels=6  ,      out_channels=3  ,      down_block_types=('DownBlock2D', 'AttnDownBlock2D')  ,      up_block_types=('AttnUpBlock2D', 'UpBlock2D')  ,      )
         return model







     @property
     def lowerCAmelCase  (		self     ):
         torch.manual_seed(0     )
         __UpperCamelCase						:						List[str]		     =	VQModel(
             block_out_channels=[3_2, 6_4]  ,      in_channels=3  ,      out_channels=3  ,      down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']  ,      up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']  ,      latent_channels=3  ,      )
         return model







     def lowerCAmelCase  (		self     ):
         __UpperCamelCase						:						Tuple		     =	'cpu'
         __UpperCamelCase						:						Union[str, Any]		     =	self.dummy_uncond_unet
         __UpperCamelCase						:						List[str]		     =	DDIMScheduler()
         __UpperCamelCase						:						Any		     =	self.dummy_vq_model

         __UpperCamelCase						:						Tuple		     =	LDMSuperResolutionPipeline(unet=_lowerCamelCase  ,      vqvae=_lowerCamelCase  ,      scheduler=_lowerCamelCase     )
         ldm.to(_lowerCamelCase     )
         ldm.set_progress_bar_config(disable=_lowerCamelCase     )

         __UpperCamelCase						:						int		     =	self.dummy_image.to(_lowerCamelCase     )

         __UpperCamelCase						:						Union[str, Any]		     =	torch.Generator(device=_lowerCamelCase     ).manual_seed(0     )
         __UpperCamelCase						:						Union[str, Any]		     =	ldm(image=_lowerCamelCase  ,      generator=_lowerCamelCase  ,      num_inference_steps=2  ,      output_type='numpy'     ).images

         __UpperCamelCase						:						str		     =	image[0, -3:, -3:, -1]

         assert image.shape == (1, 6_4, 6_4, 3)
         __UpperCamelCase						:						str		     =	np.array([0.8_6_7_8, 0.8_2_4_5, 0.6_3_8_1, 0.6_8_3_0, 0.4_3_8_5, 0.5_5_9_9, 0.4_6_4_1, 0.6_2_0_1, 0.5_1_5_0]     )

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







     @unittest.skipIf(torch_device != 'cuda'  ,      'This test requires a GPU'     )
     def lowerCAmelCase  (		self     ):
         __UpperCamelCase						:						Optional[Any]		     =	self.dummy_uncond_unet
         __UpperCamelCase						:						Optional[int]		     =	DDIMScheduler()
         __UpperCamelCase						:						Any		     =	self.dummy_vq_model

         # put models in fp16
         __UpperCamelCase						:						str		     =	unet.half()
         __UpperCamelCase						:						Optional[Any]		     =	vqvae.half()

         __UpperCamelCase						:						Optional[int]		     =	LDMSuperResolutionPipeline(unet=_lowerCamelCase  ,      vqvae=_lowerCamelCase  ,      scheduler=_lowerCamelCase     )
         ldm.to(_lowerCamelCase     )
         ldm.set_progress_bar_config(disable=_lowerCamelCase     )

         __UpperCamelCase						:						Dict		     =	self.dummy_image.to(_lowerCamelCase     )

         __UpperCamelCase						:						Optional[Any]		     =	ldm(_lowerCamelCase  ,      num_inference_steps=2  ,      output_type='numpy'     ).images

         assert image.shape == (1, 6_4, 6_4, 3)






@slow
@require_torch
class       __lowercase		(    unittest.TestCase    ):







     """simple docstring"""






     def lowerCAmelCase  (		self     ):
         __UpperCamelCase						:						List[Any]		     =	load_image(
             'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
             '/vq_diffusion/teddy_bear_pool.png'     )
         __UpperCamelCase						:						List[Any]		     =	init_image.resize((6_4, 6_4)  ,      resample=PIL_INTERPOLATION['lanczos']     )

         __UpperCamelCase						:						List[str]		     =	LDMSuperResolutionPipeline.from_pretrained('duongna/ldm-super-resolution'  ,      device_map='auto'     )
         ldm.set_progress_bar_config(disable=_lowerCamelCase     )

         __UpperCamelCase						:						Union[str, Any]		     =	torch.manual_seed(0     )
         __UpperCamelCase						:						Dict		     =	ldm(image=_lowerCamelCase  ,      generator=_lowerCamelCase  ,      num_inference_steps=2_0  ,      output_type='numpy'     ).images

         __UpperCamelCase						:						List[Any]		     =	image[0, -3:, -3:, -1]

         assert image.shape == (1, 2_5_6, 2_5_6, 3)
         __UpperCamelCase						:						str		     =	np.array([0.7_6_4_4, 0.7_6_7_9, 0.7_6_4_2, 0.7_6_3_3, 0.7_6_6_6, 0.7_5_6_0, 0.7_4_2_5, 0.7_2_5_7, 0.6_9_0_7]     )

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



