from __future__ import annotations

import unittest

from transformers import AutoTokenizer, GPTJConfig, is_tf_available
from transformers.testing_utils import require_tf, slow, tooslow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin
from ...utils.test_modeling_tf_core import TFCoreModelTesterMixin


if is_tf_available():
				import tensorflow as tf

				from transformers.models.gptj.modeling_tf_gptj import (
				    TFGPTJForCausalLM,
				    TFGPTJForQuestionAnswering,
				    TFGPTJForSequenceClassification,
				    TFGPTJModel,
				    shape_list,
				)





class 					__snake_case    :




					def __init__(  self       ,			__UpperCamelCase						)	->    int:



									'''simple docstring'''







									snake_case__   :      List[str]   				= parent
									snake_case__   :      List[Any]   				= 13
									snake_case__   :      Dict   				= 7
									snake_case__   :      List[Any]   				= True
									snake_case__   :      List[Any]   				= True
									snake_case__   :      List[Any]   				= True
									snake_case__   :      List[str]   				= True
									snake_case__   :      Dict   				= True
									snake_case__   :      str   				= 99
									snake_case__   :      int   				= 32
									snake_case__   :      Union[str, Any]   				= 4
									snake_case__   :      Any   				= 2
									snake_case__   :      Any   				= 4
									snake_case__   :      Any   				= 37
									snake_case__   :      Optional[Any]   				= 'gelu'
									snake_case__   :      Any   				= 0.1
									snake_case__   :      Any   				= 0.1
									snake_case__   :      int   				= 512
									snake_case__   :      int   				= 16
									snake_case__   :      Tuple   				= 2
									snake_case__   :      Tuple   				= 0.0_2
									snake_case__   :      Dict   				= 3
									snake_case__   :      Optional[Any]   				= 4
									snake_case__   :      List[str]   				= None
									snake_case__   :      str   				= self.vocab_size - 1
									snake_case__   :      Union[str, Any]   				= self.vocab_size - 1
									snake_case__   :      Any   				= self.vocab_size - 1




					def      __a     (  self						)	->    List[str]:



									'''simple docstring'''







									snake_case__   :      Any   				= ids_tensor([self.batch_size, self.seq_length]       ,			self.vocab_size						)

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

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

									snake_case__   :      Union[str, Any]   				= None
									if self.use_mc_token_ids:
													snake_case__   :      Optional[Any]   				= ids_tensor([self.batch_size, self.num_choices]       ,			self.seq_length						)

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

									snake_case__   :      List[Any]   				= GPTJConfig(
									    vocab_size=self.vocab_size       ,			n_embd=self.hidden_size       ,			n_layer=self.num_hidden_layers       ,			n_head=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       ,			n_positions=self.max_position_embeddings       ,			type_vocab_size=self.type_vocab_size       ,			initializer_range=self.initializer_range       ,			bos_token_id=self.bos_token_id       ,			eos_token_id=self.eos_token_id       ,			pad_token_id=self.pad_token_id       ,			rotary_dim=self.rotary_dim       ,			return_dict=__UpperCamelCase       ,			)

									snake_case__   :      int   				= ids_tensor([self.num_hidden_layers, self.num_attention_heads]       ,			2						)

									return (
									    config,
									    input_ids,
									    input_mask,
									    head_mask,
									    token_type_ids,
									    mc_token_ids,
									    sequence_labels,
									    token_labels,
									    choice_labels,
									)




					def      __a     (  self       ,			__UpperCamelCase       ,			__UpperCamelCase       ,			__UpperCamelCase       ,			__UpperCamelCase       ,			__UpperCamelCase       ,			*__UpperCamelCase						)	->    List[str]:



									'''simple docstring'''







									snake_case__   :      Union[str, Any]   				= TFGPTJModel(config=__UpperCamelCase						)
									snake_case__   :      Optional[int]   				= {
									    'input_ids': input_ids,
									    'attention_mask': input_mask,
									    'token_type_ids': token_type_ids,
									}
									snake_case__   :      str   				= model(__UpperCamelCase						)

									snake_case__   :      int   				= [input_ids, None, input_mask]  # None is the input for 'past'
									snake_case__   :      List[str]   				= model(__UpperCamelCase						)

									snake_case__   :      int   				= model(__UpperCamelCase						)

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




					def      __a     (  self       ,			__UpperCamelCase       ,			__UpperCamelCase       ,			__UpperCamelCase       ,			__UpperCamelCase       ,			__UpperCamelCase       ,			*__UpperCamelCase						)	->    List[str]:



									'''simple docstring'''







									snake_case__   :      List[Any]   				= TFGPTJModel(config=__UpperCamelCase						)

									# first forward pass
									snake_case__   :      Optional[Any]   				= model(__UpperCamelCase       ,			token_type_ids=__UpperCamelCase       ,			use_cache=__UpperCamelCase						)
									snake_case__   :      Optional[int]   				= model(__UpperCamelCase       ,			token_type_ids=__UpperCamelCase						)
									snake_case__   :      Tuple   				= model(__UpperCamelCase       ,			token_type_ids=__UpperCamelCase       ,			use_cache=__UpperCamelCase						)

									self.parent.assertTrue(len(__UpperCamelCase						) == len(__UpperCamelCase						)						)
									self.parent.assertTrue(len(__UpperCamelCase						) == len(__UpperCamelCase						) + 1						)

									snake_case__				,    snake_case__   :      Union[str, Any]   				= outputs.to_tuple()

									# create hypothetical next token and extent to next_input_ids
									snake_case__   :      Dict   				= ids_tensor((self.batch_size, 1)       ,			config.vocab_size						)
									snake_case__   :      Tuple   				= ids_tensor([self.batch_size, 1]       ,			self.type_vocab_size						)

									# append to next input_ids and token_type_ids
									snake_case__   :      Dict   				= tf.concat([input_ids, next_tokens]       ,			axis=-1						)
									snake_case__   :      int   				= tf.concat([token_type_ids, next_token_types]       ,			axis=-1						)

									snake_case__   :      Dict   				= model(__UpperCamelCase       ,			token_type_ids=__UpperCamelCase						)['last_hidden_state']
									snake_case__   :      str   				= model(__UpperCamelCase       ,			token_type_ids=__UpperCamelCase       ,			past_key_values=__UpperCamelCase						)[
									    'last_hidden_state'
									]

									# select random slice
									snake_case__   :      List[Any]   				= int(ids_tensor((1,)       ,			shape_list(__UpperCamelCase						)[-1]						)						)
									snake_case__   :      int   				= output_from_no_past[:, -1, random_slice_idx]
									snake_case__   :      int   				= output_from_past[:, 0, random_slice_idx]

									# test that outputs are equal for slice
									tf.debugging.assert_near(__UpperCamelCase       ,			__UpperCamelCase       ,			rtol=1E-6						)




					def      __a     (  self       ,			__UpperCamelCase       ,			__UpperCamelCase       ,			__UpperCamelCase       ,			__UpperCamelCase       ,			__UpperCamelCase       ,			*__UpperCamelCase						)	->    Any:



									'''simple docstring'''







									snake_case__   :      Optional[int]   				= TFGPTJModel(config=__UpperCamelCase						)

									# create attention mask
									snake_case__   :      Union[str, Any]   				= self.seq_length // 2
									snake_case__   :      Union[str, Any]   				= tf.ones((self.batch_size, half_seq_length)       ,			dtype=tf.intaa						)
									snake_case__   :      Optional[int]   				= tf.zeros((self.batch_size, self.seq_length - half_seq_length)       ,			dtype=tf.intaa						)
									snake_case__   :      str   				= tf.concat([attn_mask_begin, attn_mask_end]       ,			axis=1						)

									# first forward pass
									snake_case__				,    snake_case__   :      Optional[int]   				= model(__UpperCamelCase       ,			attention_mask=__UpperCamelCase						).to_tuple()

									# create hypothetical next token and extent to next_input_ids
									snake_case__   :      Optional[int]   				= ids_tensor((self.batch_size, 1)       ,			config.vocab_size						)

									# change a random masked slice from input_ids
									snake_case__   :      Tuple   				= ids_tensor((1,)       ,			__UpperCamelCase						).numpy() + 1
									snake_case__   :      str   				= ids_tensor((self.batch_size, self.seq_length)       ,			config.vocab_size						)
									snake_case__   :      Tuple   				= tf.range(self.seq_length						) == (self.seq_length - random_seq_idx_to_change)
									snake_case__   :      Union[str, Any]   				= tf.transpose(
									    tf.broadcast_to(tf.expand_dims(__UpperCamelCase       ,			-1						)       ,			(self.seq_length, self.batch_size)						)						)
									snake_case__   :      Tuple   				= tf.where(__UpperCamelCase       ,			__UpperCamelCase       ,			__UpperCamelCase						)

									# append to next input_ids and attn_mask
									snake_case__   :      List[Any]   				= tf.concat([input_ids, next_tokens]       ,			axis=-1						)
									snake_case__   :      Any   				= tf.concat([attn_mask, tf.ones((shape_list(__UpperCamelCase						)[0], 1)       ,			dtype=tf.intaa						)]       ,			axis=1						)

									# get two different outputs
									snake_case__   :      Dict   				= model(__UpperCamelCase       ,			attention_mask=__UpperCamelCase						)['last_hidden_state']
									snake_case__   :      str   				= model(__UpperCamelCase       ,			past_key_values=__UpperCamelCase       ,			attention_mask=__UpperCamelCase						)[
									    'last_hidden_state'
									]

									# select random slice
									snake_case__   :      Dict   				= int(ids_tensor((1,)       ,			shape_list(__UpperCamelCase						)[-1]						)						)
									snake_case__   :      Dict   				= output_from_no_past[:, -1, random_slice_idx]
									snake_case__   :      Optional[int]   				= output_from_past[:, 0, random_slice_idx]

									# test that outputs are equal for slice
									tf.debugging.assert_near(__UpperCamelCase       ,			__UpperCamelCase       ,			rtol=1E-12						)




					def      __a     (  self       ,			__UpperCamelCase       ,			__UpperCamelCase       ,			__UpperCamelCase       ,			__UpperCamelCase       ,			__UpperCamelCase       ,			*__UpperCamelCase						)	->    List[Any]:



									'''simple docstring'''







									snake_case__   :      int   				= TFGPTJModel(config=__UpperCamelCase						)

									snake_case__   :      List[Any]   				= input_ids[:1, :]
									snake_case__   :      List[str]   				= input_mask[:1, :]
									snake_case__   :      List[str]   				= token_type_ids[:1, :]
									snake_case__   :      Tuple   				= 1

									# first forward pass
									snake_case__   :      Any   				= model(__UpperCamelCase       ,			attention_mask=__UpperCamelCase       ,			token_type_ids=__UpperCamelCase       ,			use_cache=__UpperCamelCase						)

									snake_case__				,    snake_case__   :      Any   				= outputs.to_tuple()

									# create hypothetical next token and extent to next_input_ids
									snake_case__   :      Optional[int]   				= ids_tensor((self.batch_size, 3)       ,			config.vocab_size						)
									snake_case__   :      List[str]   				= ids_tensor((self.batch_size, 3)       ,			2						)
									snake_case__   :      Any   				= ids_tensor((self.batch_size, 3)       ,			self.type_vocab_size						)

									# append to next input_ids and token_type_ids
									snake_case__   :      Tuple   				= tf.concat([input_ids, next_tokens]       ,			axis=-1						)
									snake_case__   :      List[str]   				= tf.concat([input_mask, next_attn_mask]       ,			axis=-1						)
									snake_case__   :      int   				= tf.concat([token_type_ids, next_token_types]       ,			axis=-1						)

									snake_case__   :      List[Any]   				= model(
									    __UpperCamelCase       ,			token_type_ids=__UpperCamelCase       ,			attention_mask=__UpperCamelCase						)['last_hidden_state']
									snake_case__   :      Union[str, Any]   				= model(
									    __UpperCamelCase       ,			token_type_ids=__UpperCamelCase       ,			attention_mask=__UpperCamelCase       ,			past_key_values=__UpperCamelCase       ,			)['last_hidden_state']
									self.parent.assertTrue(output_from_past.shape[1] == next_tokens.shape[1]						)

									# select random slice
									snake_case__   :      List[str]   				= int(ids_tensor((1,)       ,			shape_list(__UpperCamelCase						)[-1]						)						)
									snake_case__   :      Optional[Any]   				= output_from_no_past[:, -3:, random_slice_idx]
									snake_case__   :      Optional[Any]   				= output_from_past[:, :, random_slice_idx]

									# test that outputs are equal for slice
									tf.debugging.assert_near(__UpperCamelCase       ,			__UpperCamelCase       ,			rtol=1E-3						)




					def      __a     (  self       ,			__UpperCamelCase       ,			__UpperCamelCase       ,			__UpperCamelCase       ,			__UpperCamelCase       ,			__UpperCamelCase       ,			*__UpperCamelCase						)	->    Optional[int]:



									'''simple docstring'''







									snake_case__   :      int   				= TFGPTJForCausalLM(config=__UpperCamelCase						)
									snake_case__   :      Dict   				= {
									    'input_ids': input_ids,
									    'attention_mask': input_mask,
									    'token_type_ids': token_type_ids,
									}
									snake_case__   :      Any   				= model(__UpperCamelCase						)
									self.parent.assertEqual(result.logits.shape       ,			(self.batch_size, self.seq_length, self.vocab_size)						)




					def      __a     (  self						)	->    Union[str, Any]:



									'''simple docstring'''







									snake_case__   :      Optional[Any]   				= self.prepare_config_and_inputs()

									(
									    (
									    snake_case__
									)				,    (
									    snake_case__
									)				,    (
									    snake_case__
									)				,    (
									    snake_case__
									)				,    (
									    snake_case__
									)				,    (
									    snake_case__
									)				,    (
									    snake_case__
									)				,    (
									    snake_case__
									)				,    (
									    snake_case__
									)				,    
									)   :      Optional[Any]   				= config_and_inputs

									snake_case__   :      Union[str, Any]   				= {
									    'input_ids': input_ids,
									    'token_type_ids': token_type_ids,
									    'attention_mask': input_mask,
									}
									return config, inputs_dict





@require_tf
class 					__snake_case    (     _lowerCamelCase       ,_lowerCamelCase       ,_lowerCamelCase       ,unittest.TestCase      ):
					__lowerCamelCase			   =  (
					    (TFGPTJForCausalLM, TFGPTJForSequenceClassification, TFGPTJForQuestionAnswering, TFGPTJModel)
					    if is_tf_available()
					    else ()
					)

					__lowerCamelCase			   =  (TFGPTJForCausalLM,) if is_tf_available() else ()
					__lowerCamelCase			   =  (
					    {
					        """feature-extraction""": TFGPTJModel,
					        """question-answering""": TFGPTJForQuestionAnswering,
					        """text-classification""": TFGPTJForSequenceClassification,
					        """text-generation""": TFGPTJForCausalLM,
					        """zero-shot""": TFGPTJForSequenceClassification,
					    }
					    if is_tf_available()
					    else {}
					)
					__lowerCamelCase			   =  False
					__lowerCamelCase			   =  False
					__lowerCamelCase			   =  False
					__lowerCamelCase			   =  False




					def      __a     (  self       ,			__UpperCamelCase       ,			__UpperCamelCase       ,			__UpperCamelCase       ,			__UpperCamelCase       ,			__UpperCamelCase						)	->    List[Any]:



									'''simple docstring'''







									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      __a     (  self						)	->    Union[str, Any]:



									'''simple docstring'''







									snake_case__   :      Optional[Any]   				= TFGPTJModelTester(self						)
									snake_case__   :      Tuple   				= ConfigTester(self       ,			config_class=__UpperCamelCase       ,			n_embd=37						)




					def      __a     (  self						)	->    Dict:



									'''simple docstring'''







									self.config_tester.run_common_tests()




					def      __a     (  self						)	->    List[Any]:



									'''simple docstring'''







									snake_case__   :      List[str]   				= self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_gptj_model(*__UpperCamelCase						)




					def      __a     (  self						)	->    str:



									'''simple docstring'''







									snake_case__   :      int   				= self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_gptj_model_past(*__UpperCamelCase						)




					def      __a     (  self						)	->    int:



									'''simple docstring'''







									snake_case__   :      Any   				= self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_gptj_model_attention_mask_past(*__UpperCamelCase						)




					def      __a     (  self						)	->    Dict:



									'''simple docstring'''







									snake_case__   :      Optional[int]   				= self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_gptj_model_past_large_inputs(*__UpperCamelCase						)




					def      __a     (  self						)	->    Optional[int]:



									'''simple docstring'''







									snake_case__   :      List[Any]   				= self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_gptj_lm_head_model(*__UpperCamelCase						)




					@slow
					@unittest.skipIf(
					    not is_tf_available() or len(tf.config.list_physical_devices('GPU'						)						) > 0       ,			'skip testing on GPU for now to avoid GPU OOM.'       ,			)
					def      __a     (  self						)	->    int:



									'''simple docstring'''







									snake_case__   :      Union[str, Any]   				= TFGPTJModel.from_pretrained('EleutherAI/gpt-j-6B'       ,			from_pt=__UpperCamelCase						)
									self.assertIsNotNone(__UpperCamelCase						)




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



									'''simple docstring'''







									super().test_resize_token_embeddings()





@require_tf
@tooslow
# Marked as @tooslow due to GPU OOM -- but still useful to run locally. Requires ~39GB of RAM.
class 					__snake_case    (     unittest.TestCase      ):




					def      __a     (  self						)	->    Tuple:



									'''simple docstring'''







									snake_case__   :      Any   				= TFGPTJForCausalLM.from_pretrained('EleutherAI/gpt-j-6B'       ,			from_pt=__UpperCamelCase						)
									snake_case__   :      Any   				= tf.convert_to_tensor([[464, 3290]]       ,			dtype=tf.intaa						)  # The dog
									# fmt: off
									# The dog is a man's best friend. It is a loyal companion, and it is a friend
									snake_case__   :      Tuple   				= [464, 3290, 318, 257, 582, 338, 1266, 1545, 13, 632, 318, 257, 9112, 15185, 11, 290, 340, 318, 257, 1545]
									# fmt: on
									snake_case__   :      Tuple   				= model.generate(__UpperCamelCase       ,			do_sample=__UpperCamelCase						)
									self.assertListEqual(output_ids[0].numpy().tolist()       ,			__UpperCamelCase						)




					def      __a     (  self						)	->    Dict:



									'''simple docstring'''







									snake_case__   :      Optional[int]   				= AutoTokenizer.from_pretrained('EleutherAI/gpt-j-6B'       ,			revision='float16'						)
									snake_case__   :      List[Any]   				= TFGPTJForCausalLM.from_pretrained('EleutherAI/gpt-j-6B'       ,			revision='float16'       ,			from_pt=__UpperCamelCase						)

									snake_case__   :      Optional[int]   				= tokenizer('Today is a nice day and'       ,			return_tensors='tf'						)
									# forces the generation to happen on CPU, to avoid GPU-related quirks
									with tf.device(':/CPU:0'						):
													snake_case__   :      List[Any]   				= model.generate(**__UpperCamelCase       ,			do_sample=__UpperCamelCase       ,			seed=[42, 0]						)
									snake_case__   :      Tuple   				= tokenizer.decode(output_ids[0]       ,			skip_special_tokens=__UpperCamelCase						)

									snake_case__   :      Dict   				= 'Today is a nice day and I’m going to go for a walk. I’'
									self.assertEqual(__UpperCamelCase       ,			__UpperCamelCase						)




					def      __a     (  self						)	->    Any:



									'''simple docstring'''







									snake_case__   :      Any   				= TFGPTJForCausalLM.from_pretrained('EleutherAI/gpt-j-6B'       ,			revision='float16'       ,			from_pt=__UpperCamelCase						)
									snake_case__   :      Any   				= AutoTokenizer.from_pretrained('EleutherAI/gpt-j-6B'       ,			revision='float16'						)

									snake_case__   :      Optional[Any]   				= 'left'

									# Define PAD Token = EOS Token = 50256
									snake_case__   :      str   				= tokenizer.eos_token
									snake_case__   :      Tuple   				= model.config.eos_token_id

									# use different length sentences to test batching
									snake_case__   :      Tuple   				= [
									    'Hello, my dog is a little',
									    'Today, I',
									]
									snake_case__   :      Optional[int]   				= [
									    'Hello, my dog is a little over a year old and has been diagnosed with hip dysplasia',
									    'Today, I’m going to be talking about a topic that’',
									]
									return model, tokenizer, sentences, expected_output_sentences




					def      __a     (  self						)	->    Optional[int]:



									'''simple docstring'''







									snake_case__				,    snake_case__				,    snake_case__				,    snake_case__   :      int   				= self._get_beam_search_test_objects()

									snake_case__   :      Dict   				= tokenizer(__UpperCamelCase       ,			return_tensors='tf'       ,			padding=__UpperCamelCase						)
									snake_case__   :      Tuple   				= model.generate(**__UpperCamelCase       ,			do_sample=__UpperCamelCase       ,			num_beams=2						)
									snake_case__   :      int   				= tokenizer.batch_decode(__UpperCamelCase       ,			skip_special_tokens=__UpperCamelCase						)
									self.assertListEqual(__UpperCamelCase       ,			__UpperCamelCase						)




					def      __a     (  self						)	->    Tuple:



									'''simple docstring'''







									snake_case__				,    snake_case__				,    snake_case__				,    snake_case__   :      List[str]   				= self._get_beam_search_test_objects()

									snake_case__   :      Tuple   				= tokenizer(__UpperCamelCase       ,			return_tensors='tf'       ,			padding=__UpperCamelCase						)
									snake_case__   :      Union[str, Any]   				= tokenizer(sentences[0]       ,			return_tensors='tf'						)
									snake_case__   :      int   				= model.generate(**__UpperCamelCase       ,			do_sample=__UpperCamelCase       ,			num_beams=2						)
									snake_case__   :      Optional[Any]   				= (
									    shape_list(inputs_non_padded['input_ids']						)[-1]
									    - tf.reduce_sum(tf.cast(inputs['attention_mask'][-1]       ,			tf.intaa						)						).numpy()
									)
									snake_case__   :      Tuple   				= tokenizer(sentences[1]       ,			return_tensors='tf'						)
									snake_case__   :      Optional[int]   				= model.generate(
									    **__UpperCamelCase       ,			do_sample=__UpperCamelCase       ,			num_beams=2       ,			max_length=model.config.max_length - num_paddings						)
									snake_case__   :      Optional[int]   				= tokenizer.decode(output_non_padded[0]       ,			skip_special_tokens=__UpperCamelCase						)
									snake_case__   :      str   				= tokenizer.decode(output_padded[0]       ,			skip_special_tokens=__UpperCamelCase						)
									self.assertListEqual(__UpperCamelCase       ,			[non_padded_sentence, padded_sentence]						)




					def      __a     (  self						)	->    int:



									'''simple docstring'''







									snake_case__				,    snake_case__				,    snake_case__				,    snake_case__   :      Any   				= self._get_beam_search_test_objects()

									snake_case__   :      List[Any]   				= tokenizer(__UpperCamelCase       ,			return_tensors='tf'       ,			padding=__UpperCamelCase						)
									snake_case__   :      Optional[Any]   				= tf.function(model.generate       ,			jit_compile=__UpperCamelCase						)
									snake_case__   :      Tuple   				= xla_generate(**__UpperCamelCase       ,			do_sample=__UpperCamelCase       ,			num_beams=2						)
									snake_case__   :      int   				= tokenizer.batch_decode(__UpperCamelCase       ,			skip_special_tokens=__UpperCamelCase						)
									self.assertListEqual(__UpperCamelCase       ,			__UpperCamelCase						)



import gc
import tempfile
import unittest

import numpy as np
import torch

from diffusers import VersatileDiffusionDualGuidedPipeline
from diffusers.utils.testing_utils import load_image, nightly, require_torch_gpu, torch_device


lowerCAmelCase__	:  Union[str, Any]						=  False





@nightly
@require_torch_gpu
class 					__snake_case    (     unittest.TestCase      ):




					def      __a     (  self						)	->    str:



									'''simple docstring'''







									super().tearDown()
									gc.collect()
									torch.cuda.empty_cache()




					def      __a     (  self						)	->    Optional[int]:



									'''simple docstring'''







									snake_case__   :      str   				= VersatileDiffusionDualGuidedPipeline.from_pretrained('shi-labs/versatile-diffusion'						)
									# remove text_unet
									pipe.remove_unused_weights()
									pipe.to(__UpperCamelCase						)
									pipe.set_progress_bar_config(disable=__UpperCamelCase						)

									snake_case__   :      Union[str, Any]   				= load_image(
									    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/versatile_diffusion/benz.jpg'						)

									snake_case__   :      Dict   				= torch.manual_seed(0						)
									snake_case__   :      Tuple   				= pipe(
									    prompt='first prompt'       ,			image=__UpperCamelCase       ,			text_to_image_strength=0.7_5       ,			generator=__UpperCamelCase       ,			guidance_scale=7.5       ,			num_inference_steps=2       ,			output_type='numpy'       ,			).images

									with tempfile.TemporaryDirectory() as tmpdirname:
													pipe.save_pretrained(__UpperCamelCase						)
													snake_case__   :      List[Any]   				= VersatileDiffusionDualGuidedPipeline.from_pretrained(__UpperCamelCase						)

									pipe.to(__UpperCamelCase						)
									pipe.set_progress_bar_config(disable=__UpperCamelCase						)

									snake_case__   :      str   				= generator.manual_seed(0						)
									snake_case__   :      List[str]   				= pipe(
									    prompt='first prompt'       ,			image=__UpperCamelCase       ,			text_to_image_strength=0.7_5       ,			generator=__UpperCamelCase       ,			guidance_scale=7.5       ,			num_inference_steps=2       ,			output_type='numpy'       ,			).images

									assert np.abs(image - new_image						).sum() < 1E-5, "Models don't have the same forward pass"




					def      __a     (  self						)	->    Optional[Any]:



									'''simple docstring'''







									snake_case__   :      Union[str, Any]   				= VersatileDiffusionDualGuidedPipeline.from_pretrained('shi-labs/versatile-diffusion'						)
									pipe.remove_unused_weights()
									pipe.to(__UpperCamelCase						)
									pipe.set_progress_bar_config(disable=__UpperCamelCase						)

									snake_case__   :      str   				= 'cyberpunk 2077'
									snake_case__   :      Dict   				= load_image(
									    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/versatile_diffusion/benz.jpg'						)
									snake_case__   :      Union[str, Any]   				= torch.manual_seed(0						)
									snake_case__   :      Union[str, Any]   				= pipe(
									    prompt=__UpperCamelCase       ,			image=__UpperCamelCase       ,			text_to_image_strength=0.7_5       ,			generator=__UpperCamelCase       ,			guidance_scale=7.5       ,			num_inference_steps=50       ,			output_type='numpy'       ,			).images

									snake_case__   :      int   				= image[0, 253:256, 253:256, -1]

									assert image.shape == (1, 512, 512, 3)
									snake_case__   :      List[str]   				= np.array([0.0_7_8_7, 0.0_8_4_9, 0.0_8_2_6, 0.0_8_1_2, 0.0_8_0_7, 0.0_7_9_5, 0.0_8_1_8, 0.0_7_9_8, 0.0_7_7_9]						)

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


