





from __future__ import annotations

import unittest

from transformers import XGLMConfig, XGLMTokenizer, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
				import tensorflow as tf

				from transformers.models.xglm.modeling_tf_xglm import (
				    TF_XGLM_PRETRAINED_MODEL_ARCHIVE_LIST,
				    TFXGLMForCausalLM,
				    TFXGLMModel,
				)

@require_tf
class 	A_       :
	UpperCAmelCase__					=			XGLMConfig
	UpperCAmelCase__					=			{}
	UpperCAmelCase__					=			'''gelu'''
	def __init__(    self							,					_A							,					_A=1_4							,					_A=7							,					_A=True							,					_A=True							,					_A=True							,					_A=9_9							,					_A=3_2							,					_A=2							,					_A=4							,					_A=3_7							,					_A="gelu"							,					_A=0.1							,					_A=0.1							,					_A=5_1_2							,					_A=0.02							,					):






				'''simple docstring'''







				UpperCAmelCase   		=				parent
				UpperCAmelCase   		=				batch_size
				UpperCAmelCase   		=				seq_length
				UpperCAmelCase   		=				is_training
				UpperCAmelCase   		=				use_input_mask
				UpperCAmelCase   		=				use_labels
				UpperCAmelCase   		=				vocab_size
				UpperCAmelCase   		=				d_model
				UpperCAmelCase   		=				num_hidden_layers
				UpperCAmelCase   		=				num_attention_heads
				UpperCAmelCase   		=				ffn_dim
				UpperCAmelCase   		=				activation_function
				UpperCAmelCase   		=				activation_dropout
				UpperCAmelCase   		=				attention_dropout
				UpperCAmelCase   		=				max_position_embeddings
				UpperCAmelCase   		=				initializer_range
				UpperCAmelCase   		=				None
				UpperCAmelCase   		=				0
				UpperCAmelCase   		=				2
				UpperCAmelCase   		=				1
	def      _lowercase				(    self		):






				'''simple docstring'''







				return XGLMConfig.from_pretrained('''facebook/xglm-564M'''		)
	def      _lowercase				(    self		):






				'''simple docstring'''







				UpperCAmelCase   		=				tf.clip_by_value(
				    ids_tensor([self.batch_size, self.seq_length]							,					self.vocab_size		)							,					clip_value_min=0							,					clip_value_max=3		)

				UpperCAmelCase   		=				None
				if self.use_input_mask:
							UpperCAmelCase   		=				random_attention_mask([self.batch_size, self.seq_length]		)

				UpperCAmelCase   		=				self.get_config()

				UpperCAmelCase   		=				floats_tensor([self.num_hidden_layers, self.num_attention_heads]							,					2		)

				return (
				    config,
				    input_ids,
				    input_mask,
				    head_mask,
				)
	def      _lowercase				(    self		):






				'''simple docstring'''







				return XGLMConfig(
				    vocab_size=self.vocab_size							,					d_model=self.hidden_size							,					num_layers=self.num_hidden_layers							,					attention_heads=self.num_attention_heads							,					ffn_dim=self.ffn_dim							,					activation_function=self.activation_function							,					activation_dropout=self.activation_dropout							,					attention_dropout=self.attention_dropout							,					max_position_embeddings=self.max_position_embeddings							,					initializer_range=self.initializer_range							,					use_cache=_A							,					bos_token_id=self.bos_token_id							,					eos_token_id=self.eos_token_id							,					pad_token_id=self.pad_token_id							,					return_dict=_A							,					)







	def      _lowercase				(    self		):






				'''simple docstring'''







				UpperCAmelCase   		=				self.prepare_config_and_inputs()

				(
				    (
				    UpperCAmelCase
				)		,			(
				    UpperCAmelCase
				)		,			(
				    UpperCAmelCase
				)		,			(
				    UpperCAmelCase
				)		,			
				)   		=				config_and_inputs

				UpperCAmelCase   		=				{
				    '''input_ids''': input_ids,
				    '''head_mask''': head_mask,
				}

				return config, inputs_dict

@require_tf
class 	A_       (a_     ,							a_     ,							unittest.TestCase			):
	UpperCAmelCase__					=			(TFXGLMModel, TFXGLMForCausalLM) if is_tf_available() else ()
	UpperCAmelCase__					=			(TFXGLMForCausalLM,) if is_tf_available() else ()
	UpperCAmelCase__					=			(
	    {'''feature-extraction''': TFXGLMModel, '''text-generation''': TFXGLMForCausalLM} if is_tf_available() else {}
	)
	UpperCAmelCase__					=			False
	UpperCAmelCase__					=			False
	UpperCAmelCase__					=			False
	def      _lowercase				(    self		):






				'''simple docstring'''







				UpperCAmelCase   		=				TFXGLMModelTester(self		)
				UpperCAmelCase   		=				ConfigTester(self							,					config_class=_A							,					n_embd=3_7		)
	def      _lowercase				(    self		):






				'''simple docstring'''







				self.config_tester.run_common_tests()
	@slow
	def      _lowercase				(    self		):






				'''simple docstring'''







				for model_name in TF_XGLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
							UpperCAmelCase   		=				TFXGLMModel.from_pretrained(_A		)
							self.assertIsNotNone(_A		)







	@unittest.skip(reason='''Currently, model embeddings are going to undergo a major refactor.'''		)
	def      _lowercase				(    self		):






				'''simple docstring'''







				super().test_resize_token_embeddings()

@require_tf
class 	A_       (unittest.TestCase			):
	@slow
	def      _lowercase				(    self							,					_A=True		):






				'''simple docstring'''







				UpperCAmelCase   		=				TFXGLMForCausalLM.from_pretrained('''facebook/xglm-564M'''		)
				UpperCAmelCase   		=				tf.convert_to_tensor([[2, 2_6_8, 9_8_6_5]]							,					dtype=tf.intaa		)  # The dog
				# </s> The dog is a very friendly dog. He is very affectionate and loves to play with other
				# fmt: off
				UpperCAmelCase   		=				[2, 2_6_8, 9_8_6_5, 6_7, 1_1, 1_9_8_8, 5_7_2_5_2, 9_8_6_5, 5, 9_8_4, 6_7, 1_9_8_8, 2_1_3_8_3_8, 1_6_5_8, 5_3, 7_0_4_4_6, 3_3, 6_6_5_7, 2_7_8, 1_5_8_1]
				# fmt: on
				UpperCAmelCase   		=				model.generate(_A							,					do_sample=_A							,					num_beams=1		)
				if verify_outputs:
							self.assertListEqual(output_ids[0].numpy().tolist()							,					_A		)
	@slow
	def      _lowercase				(    self		):






				'''simple docstring'''







				UpperCAmelCase   		=				XGLMTokenizer.from_pretrained('''facebook/xglm-564M'''		)
				UpperCAmelCase   		=				TFXGLMForCausalLM.from_pretrained('''facebook/xglm-564M'''		)

				tf.random.set_seed(0		)
				UpperCAmelCase   		=				tokenizer('''Today is a nice day and'''							,					return_tensors='''tf'''		)
				UpperCAmelCase   		=				tokenized.input_ids
				# forces the generation to happen on CPU, to avoid GPU-related quirks (and assure same output regardless of the available devices)
				with tf.device(''':/CPU:0'''		):
							UpperCAmelCase   		=				model.generate(_A							,					do_sample=_A							,					seed=[7, 0]		)
				UpperCAmelCase   		=				tokenizer.decode(output_ids[0]							,					skip_special_tokens=_A		)

				UpperCAmelCase   		=				(
				    '''Today is a nice day and warm evening here over Southern Alberta!! Today when they closed schools due'''
				)
				self.assertEqual(_A							,					_A		)







	@slow
	def      _lowercase				(    self		):






				'''simple docstring'''







				UpperCAmelCase   		=				TFXGLMForCausalLM.from_pretrained('''facebook/xglm-564M'''		)
				UpperCAmelCase   		=				XGLMTokenizer.from_pretrained('''facebook/xglm-564M'''		)

				UpperCAmelCase   		=				'''left'''

				# use different length sentences to test batching
				UpperCAmelCase   		=				[
				    '''This is an extremelly long sentence that only exists to test the ability of the model to cope with '''
				    '''left-padding, such as in batched generation. The output for the sequence below should be the same '''
				    '''regardless of whether left padding is applied or not. When''',
				    '''Hello, my dog is a little''',
				]

				UpperCAmelCase   		=				tokenizer(_A							,					return_tensors='''tf'''							,					padding=_A		)
				UpperCAmelCase   		=				inputs['''input_ids''']

				UpperCAmelCase   		=				model.generate(input_ids=_A							,					attention_mask=inputs['''attention_mask''']							,					max_new_tokens=1_2		)

				UpperCAmelCase   		=				tokenizer(sentences[0]							,					return_tensors='''tf'''		).input_ids
				UpperCAmelCase   		=				model.generate(input_ids=_A							,					max_new_tokens=1_2		)

				UpperCAmelCase   		=				tokenizer(sentences[1]							,					return_tensors='''tf'''		).input_ids
				UpperCAmelCase   		=				model.generate(input_ids=_A							,					max_new_tokens=1_2		)

				UpperCAmelCase   		=				tokenizer.batch_decode(_A							,					skip_special_tokens=_A		)
				UpperCAmelCase   		=				tokenizer.decode(output_non_padded[0]							,					skip_special_tokens=_A		)
				UpperCAmelCase   		=				tokenizer.decode(output_padded[0]							,					skip_special_tokens=_A		)

				UpperCAmelCase   		=				[
				    '''This is an extremelly long sentence that only exists to test the ability of the model to cope with '''
				    '''left-padding, such as in batched generation. The output for the sequence below should be the same '''
				    '''regardless of whether left padding is applied or not. When left padding is applied, the sequence will be '''
				    '''a single''',
				    '''Hello, my dog is a little bit of a shy one, but he is very friendly''',
				]
				self.assertListEqual(_A							,					_A		)
				self.assertListEqual(_A							,					[non_padded_sentence, padded_sentence]		)










import tempfile

import numpy as np
import torch
from transformers import AutoTokenizer, TaEncoderModel

from diffusers import DDPMScheduler, UNetaDConditionModel
from diffusers.models.attention_processor import AttnAddedKVProcessor
from diffusers.pipelines.deepfloyd_if import IFWatermarker
from diffusers.utils.testing_utils import torch_device

from ..test_pipelines_common import to_np

class 	A_       :
	def      _lowercase				(    self		):






				'''simple docstring'''







				torch.manual_seed(0		)
				UpperCAmelCase   		=				TaEncoderModel.from_pretrained('''hf-internal-testing/tiny-random-t5'''		)

				torch.manual_seed(0		)
				UpperCAmelCase   		=				AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-random-t5'''		)

				torch.manual_seed(0		)
				UpperCAmelCase   		=				UNetaDConditionModel(
				    sample_size=3_2							,					layers_per_block=1							,					block_out_channels=[3_2, 6_4]							,					down_block_types=[
				        '''ResnetDownsampleBlock2D''',
				        '''SimpleCrossAttnDownBlock2D''',
				    ]							,					mid_block_type='''UNetMidBlock2DSimpleCrossAttn'''							,					up_block_types=['''SimpleCrossAttnUpBlock2D''', '''ResnetUpsampleBlock2D''']							,					in_channels=3							,					out_channels=6							,					cross_attention_dim=3_2							,					encoder_hid_dim=3_2							,					attention_head_dim=8							,					addition_embed_type='''text'''							,					addition_embed_type_num_heads=2							,					cross_attention_norm='''group_norm'''							,					resnet_time_scale_shift='''scale_shift'''							,					act_fn='''gelu'''							,					)
				unet.set_attn_processor(AttnAddedKVProcessor()		)  # For reproducibility tests

				torch.manual_seed(0		)
				UpperCAmelCase   		=				DDPMScheduler(
				    num_train_timesteps=1_0_0_0							,					beta_schedule='''squaredcos_cap_v2'''							,					beta_start=0.00_01							,					beta_end=0.02							,					thresholding=_A							,					dynamic_thresholding_ratio=0.95							,					sample_max_value=1.0							,					prediction_type='''epsilon'''							,					variance_type='''learned_range'''							,					)

				torch.manual_seed(0		)
				UpperCAmelCase   		=				IFWatermarker()

				return {
				    "text_encoder": text_encoder,
				    "tokenizer": tokenizer,
				    "unet": unet,
				    "scheduler": scheduler,
				    "watermarker": watermarker,
				    "safety_checker": None,
				    "feature_extractor": None,
				}
	def      _lowercase				(    self		):






				'''simple docstring'''







				torch.manual_seed(0		)
				UpperCAmelCase   		=				TaEncoderModel.from_pretrained('''hf-internal-testing/tiny-random-t5'''		)

				torch.manual_seed(0		)
				UpperCAmelCase   		=				AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-random-t5'''		)

				torch.manual_seed(0		)
				UpperCAmelCase   		=				UNetaDConditionModel(
				    sample_size=3_2							,					layers_per_block=[1, 2]							,					block_out_channels=[3_2, 6_4]							,					down_block_types=[
				        '''ResnetDownsampleBlock2D''',
				        '''SimpleCrossAttnDownBlock2D''',
				    ]							,					mid_block_type='''UNetMidBlock2DSimpleCrossAttn'''							,					up_block_types=['''SimpleCrossAttnUpBlock2D''', '''ResnetUpsampleBlock2D''']							,					in_channels=6							,					out_channels=6							,					cross_attention_dim=3_2							,					encoder_hid_dim=3_2							,					attention_head_dim=8							,					addition_embed_type='''text'''							,					addition_embed_type_num_heads=2							,					cross_attention_norm='''group_norm'''							,					resnet_time_scale_shift='''scale_shift'''							,					act_fn='''gelu'''							,					class_embed_type='''timestep'''							,					mid_block_scale_factor=1.4_14							,					time_embedding_act_fn='''gelu'''							,					time_embedding_dim=3_2							,					)
				unet.set_attn_processor(AttnAddedKVProcessor()		)  # For reproducibility tests

				torch.manual_seed(0		)
				UpperCAmelCase   		=				DDPMScheduler(
				    num_train_timesteps=1_0_0_0							,					beta_schedule='''squaredcos_cap_v2'''							,					beta_start=0.00_01							,					beta_end=0.02							,					thresholding=_A							,					dynamic_thresholding_ratio=0.95							,					sample_max_value=1.0							,					prediction_type='''epsilon'''							,					variance_type='''learned_range'''							,					)

				torch.manual_seed(0		)
				UpperCAmelCase   		=				DDPMScheduler(
				    num_train_timesteps=1_0_0_0							,					beta_schedule='''squaredcos_cap_v2'''							,					beta_start=0.00_01							,					beta_end=0.02							,					)

				torch.manual_seed(0		)
				UpperCAmelCase   		=				IFWatermarker()

				return {
				    "text_encoder": text_encoder,
				    "tokenizer": tokenizer,
				    "unet": unet,
				    "scheduler": scheduler,
				    "image_noising_scheduler": image_noising_scheduler,
				    "watermarker": watermarker,
				    "safety_checker": None,
				    "feature_extractor": None,
				}
	def      _lowercase				(    self		):






				'''simple docstring'''







				UpperCAmelCase   		=				self.get_dummy_components()
				UpperCAmelCase   		=				self.pipeline_class(**_A		)
				pipe.to(_A		)
				pipe.set_progress_bar_config(disable=_A		)

				UpperCAmelCase   		=				self.get_dummy_inputs(_A		)

				UpperCAmelCase   		=				inputs['''prompt''']
				UpperCAmelCase   		=				inputs['''generator''']
				UpperCAmelCase   		=				inputs['''num_inference_steps''']
				UpperCAmelCase   		=				inputs['''output_type''']

				if "image" in inputs:
							UpperCAmelCase   		=				inputs['''image''']
				else:
							UpperCAmelCase   		=				None

				if "mask_image" in inputs:
							UpperCAmelCase   		=				inputs['''mask_image''']
				else:
							UpperCAmelCase   		=				None

				if "original_image" in inputs:
							UpperCAmelCase   		=				inputs['''original_image''']
				else:
							UpperCAmelCase   		=				None

				UpperCAmelCase		,			UpperCAmelCase   		=				pipe.encode_prompt(_A		)

				# inputs with prompt converted to embeddings
				UpperCAmelCase   		=				{
				    '''prompt_embeds''': prompt_embeds,
				    '''negative_prompt_embeds''': negative_prompt_embeds,
				    '''generator''': generator,
				    '''num_inference_steps''': num_inference_steps,
				    '''output_type''': output_type,
				}

				if image is not None:
							UpperCAmelCase   		=				image

				if mask_image is not None:
							UpperCAmelCase   		=				mask_image

				if original_image is not None:
							UpperCAmelCase   		=				original_image

				# set all optional components to None
				for optional_component in pipe._optional_components:
							setattr(_A							,					_A							,					_A		)

				UpperCAmelCase   		=				pipe(**_A		)[0]

				with tempfile.TemporaryDirectory() as tmpdir:
							pipe.save_pretrained(_A		)
							UpperCAmelCase   		=				self.pipeline_class.from_pretrained(_A		)
							pipe_loaded.to(_A		)
							pipe_loaded.set_progress_bar_config(disable=_A		)

				pipe_loaded.unet.set_attn_processor(AttnAddedKVProcessor()		)  # For reproducibility tests

				for optional_component in pipe._optional_components:
							self.assertTrue(
							    getattr(_A							,					_A		) is None							,					F"""`{optional_component}` did not stay set to None after loading."""							,					)

				UpperCAmelCase   		=				self.get_dummy_inputs(_A		)

				UpperCAmelCase   		=				inputs['''generator''']
				UpperCAmelCase   		=				inputs['''num_inference_steps''']
				UpperCAmelCase   		=				inputs['''output_type''']

				# inputs with prompt converted to embeddings
				UpperCAmelCase   		=				{
				    '''prompt_embeds''': prompt_embeds,
				    '''negative_prompt_embeds''': negative_prompt_embeds,
				    '''generator''': generator,
				    '''num_inference_steps''': num_inference_steps,
				    '''output_type''': output_type,
				}

				if image is not None:
							UpperCAmelCase   		=				image

				if mask_image is not None:
							UpperCAmelCase   		=				mask_image

				if original_image is not None:
							UpperCAmelCase   		=				original_image

				UpperCAmelCase   		=				pipe_loaded(**_A		)[0]

				UpperCAmelCase   		=				np.abs(to_np(_A		) - to_np(_A		)		).max()
				self.assertLess(_A							,					1E-4		)







	def      _lowercase				(    self		):






				'''simple docstring'''







				UpperCAmelCase   		=				self.get_dummy_components()
				UpperCAmelCase   		=				self.pipeline_class(**_A		)
				pipe.to(_A		)
				pipe.set_progress_bar_config(disable=_A		)

				UpperCAmelCase   		=				self.get_dummy_inputs(_A		)
				UpperCAmelCase   		=				pipe(**_A		)[0]

				with tempfile.TemporaryDirectory() as tmpdir:
							pipe.save_pretrained(_A		)
							UpperCAmelCase   		=				self.pipeline_class.from_pretrained(_A		)
							pipe_loaded.to(_A		)
							pipe_loaded.set_progress_bar_config(disable=_A		)

				pipe_loaded.unet.set_attn_processor(AttnAddedKVProcessor()		)  # For reproducibility tests

				UpperCAmelCase   		=				self.get_dummy_inputs(_A		)
				UpperCAmelCase   		=				pipe_loaded(**_A		)[0]

				UpperCAmelCase   		=				np.abs(to_np(_A		) - to_np(_A		)		).max()
				self.assertLess(_A							,					1E-4		)



