


import math
import unittest

from transformers import BioGptConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
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 (
						    BioGptForCausalLM,
						    BioGptForSequenceClassification,
						    BioGptForTokenClassification,
						    BioGptModel,
						    BioGptTokenizer,
						)
						from transformers.models.biogpt.modeling_biogpt import BIOGPT_PRETRAINED_MODEL_ARCHIVE_LIST







class 				snake_case__:







						'''simple docstring'''




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




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

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

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

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

										lowerCAmelCase_ :      Dict    				=				self.get_config()

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




						def 						lowercase_ (				self							)					->    Optional[Any]:
										return BioGptConfig(
										    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   ,						is_decoder=__lowercase   ,						initializer_range=self.initializer_range   ,						)




						def 						lowercase_ (				self   ,						__lowercase   ,						__lowercase   ,						__lowercase   ,						__lowercase   ,						__lowercase   ,						__lowercase   ,						__lowercase							)					->    int:
										lowerCAmelCase_ :      int    				=				BioGptModel(config=__lowercase							)
										model.to(__lowercase							)
										model.eval()
										lowerCAmelCase_ :      Dict    				=				model(__lowercase   ,						attention_mask=__lowercase							)
										lowerCAmelCase_ :      int    				=				model(__lowercase							)
										self.parent.assertEqual(result.last_hidden_state.shape   ,						(self.batch_size, self.seq_length, self.hidden_size)							)




						def 						lowercase_ (				self   ,						__lowercase   ,						__lowercase   ,						__lowercase   ,						__lowercase   ,						__lowercase   ,						__lowercase   ,						__lowercase   ,						__lowercase   ,						__lowercase   ,						)					->    Union[str, Any]:
										lowerCAmelCase_ :      int    				=				BioGptForCausalLM(config=__lowercase							)
										model.to(__lowercase							)
										model.eval()
										lowerCAmelCase_ :      Dict    				=				model(__lowercase   ,						attention_mask=__lowercase   ,						token_type_ids=__lowercase   ,						labels=__lowercase							)
										self.parent.assertEqual(result.logits.shape   ,						(self.batch_size, self.seq_length, self.vocab_size)							)




						def 						lowercase_ (				self   ,						__lowercase   ,						__lowercase   ,						__lowercase   ,						__lowercase   ,						__lowercase   ,						*__lowercase							)					->    List[str]:
										lowerCAmelCase_ :      Any    				=				BioGptModel(config=__lowercase							)
										model.to(__lowercase							)
										model.eval()

										# create attention mask
										lowerCAmelCase_ :      Dict    				=				torch.ones(input_ids.shape   ,						dtype=torch.long   ,						device=__lowercase							)
										lowerCAmelCase_ :      str    				=				self.seq_length // 2
										lowerCAmelCase_ :      Tuple    				=				0

										# first forward pass
										lowerCAmelCase_				,  lowerCAmelCase_ :      Any    				=				model(__lowercase   ,						attention_mask=__lowercase							).to_tuple()

										# create hypothetical next token and extent to next_input_ids
										lowerCAmelCase_ :      Dict    				=				ids_tensor((self.batch_size, 1)   ,						config.vocab_size							)

										# change a random masked slice from input_ids
										lowerCAmelCase_ :      Any    				=				ids_tensor((1,)   ,						__lowercase							).item() + 1
										lowerCAmelCase_ :      List[str]    				=				ids_tensor((self.batch_size, 1)   ,						config.vocab_size							).squeeze(-1							)
										lowerCAmelCase_ :      Tuple    				=				random_other_next_tokens

										# append to next input_ids and attn_mask
										lowerCAmelCase_ :      Tuple    				=				torch.cat([input_ids, next_tokens]   ,						dim=-1							)
										lowerCAmelCase_ :      Dict    				=				torch.cat(
										    [attn_mask, torch.ones((attn_mask.shape[0], 1)   ,						dtype=torch.long   ,						device=__lowercase							)]   ,						dim=1   ,						)

										# get two different outputs
										lowerCAmelCase_ :      Optional[Any]    				=				model(__lowercase   ,						attention_mask=__lowercase							)['''last_hidden_state''']
										lowerCAmelCase_ :      Union[str, Any]    				=				model(__lowercase   ,						past_key_values=__lowercase   ,						attention_mask=__lowercase							)['''last_hidden_state''']

										# select random slice
										lowerCAmelCase_ :      Optional[Any]    				=				ids_tensor((1,)   ,						output_from_past.shape[-1]							).item()
										lowerCAmelCase_ :      List[str]    				=				output_from_no_past[:, -1, random_slice_idx].detach()
										lowerCAmelCase_ :      List[str]    				=				output_from_past[:, 0, random_slice_idx].detach()

										# test that outputs are equal for slice
										self.parent.assertTrue(torch.allclose(__lowercase   ,						__lowercase   ,						atol=1e-3							)							)




						def 						lowercase_ (				self   ,						__lowercase   ,						__lowercase   ,						__lowercase   ,						__lowercase   ,						__lowercase   ,						*__lowercase							)					->    List[Any]:
										lowerCAmelCase_ :      Union[str, Any]    				=				BioGptModel(config=__lowercase							).to(__lowercase							).eval()

										lowerCAmelCase_ :      Optional[int]    				=				torch.ones(input_ids.shape   ,						dtype=torch.long   ,						device=__lowercase							)

										# first forward pass
										lowerCAmelCase_ :      List[Any]    				=				model(__lowercase   ,						attention_mask=__lowercase   ,						use_cache=__lowercase							)

										lowerCAmelCase_				,  lowerCAmelCase_ :      Optional[int]    				=				outputs.to_tuple()

										# create hypothetical multiple next token and extent to next_input_ids
										lowerCAmelCase_ :      Union[str, Any]    				=				ids_tensor((self.batch_size, 3)   ,						config.vocab_size							)
										lowerCAmelCase_ :      int    				=				ids_tensor((self.batch_size, 3)   ,						2							)

										# append to next input_ids and
										lowerCAmelCase_ :      Dict    				=				torch.cat([input_ids, next_tokens]   ,						dim=-1							)
										lowerCAmelCase_ :      int    				=				torch.cat([attention_mask, next_attn_mask]   ,						dim=-1							)

										lowerCAmelCase_ :      Optional[Any]    				=				model(__lowercase   ,						attention_mask=__lowercase							)['''last_hidden_state''']
										lowerCAmelCase_ :      Optional[Any]    				=				model(__lowercase   ,						attention_mask=__lowercase   ,						past_key_values=__lowercase							)[
										    '''last_hidden_state'''
										]

										# select random slice
										lowerCAmelCase_ :      int    				=				ids_tensor((1,)   ,						output_from_past.shape[-1]							).item()
										lowerCAmelCase_ :      int    				=				output_from_no_past[:, -3:, random_slice_idx].detach()
										lowerCAmelCase_ :      Union[str, Any]    				=				output_from_past[:, :, random_slice_idx].detach()

										self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]							)

										# test that outputs are equal for slice
										self.parent.assertTrue(torch.allclose(__lowercase   ,						__lowercase   ,						atol=1e-3							)							)




						def 						lowercase_ (				self   ,						__lowercase   ,						__lowercase   ,						__lowercase   ,						__lowercase   ,						__lowercase   ,						*__lowercase   ,						__lowercase=False							)					->    List[Any]:
										lowerCAmelCase_ :      List[Any]    				=				BioGptForCausalLM(__lowercase							)
										model.to(__lowercase							)
										if gradient_checkpointing:
														model.gradient_checkpointing_enable()

										lowerCAmelCase_ :      Optional[int]    				=				model(__lowercase   ,						labels=__lowercase							)
										self.parent.assertEqual(result.loss.shape   ,						()							)
										self.parent.assertEqual(result.logits.shape   ,						(self.batch_size, self.seq_length, self.vocab_size)							)
										result.loss.backward()




						def 						lowercase_ (				self   ,						__lowercase   ,						*__lowercase							)					->    Any:
										lowerCAmelCase_ :      List[str]    				=				BioGptModel(__lowercase							)
										lowerCAmelCase_ :      List[str]    				=				model.config.initializer_range / math.sqrt(2 * model.config.num_hidden_layers							)
										for key in model.state_dict().keys():
														if "c_proj" in key and "weight" in key:
																		self.parent.assertLessEqual(abs(torch.std(model.state_dict()[key]							) - model_std							)   ,						0.0_01							)
																		self.parent.assertLessEqual(abs(torch.mean(model.state_dict()[key]							) - 0.0							)   ,						0.01							)




						def 						lowercase_ (				self   ,						__lowercase   ,						__lowercase   ,						__lowercase   ,						__lowercase   ,						__lowercase   ,						*__lowercase							)					->    List[Any]:
										lowerCAmelCase_ :      int    				=				self.num_labels
										lowerCAmelCase_ :      str    				=				BioGptForTokenClassification(__lowercase							)
										model.to(__lowercase							)
										model.eval()
										lowerCAmelCase_ :      List[Any]    				=				model(__lowercase   ,						attention_mask=__lowercase   ,						token_type_ids=__lowercase							)
										self.parent.assertEqual(result.logits.shape   ,						(self.batch_size, self.seq_length, self.num_labels)							)




						def 						lowercase_ (				self							)					->    Dict:
										lowerCAmelCase_ :      List[str]    				=				self.prepare_config_and_inputs()
										(
										    (
										    lowerCAmelCase_
										)				,  (
										    lowerCAmelCase_
										)				,  (
										    lowerCAmelCase_
										)				,  (
										    lowerCAmelCase_
										)				,  (
										    lowerCAmelCase_
										)				,  (
										    lowerCAmelCase_
										)				,  (
										    lowerCAmelCase_
										)				,  
										) :      int    				=				config_and_inputs
										lowerCAmelCase_ :      str    				=				{'''input_ids''': input_ids, '''attention_mask''': input_mask}
										return config, inputs_dict







@require_torch
class 				snake_case__(							UpperCAmelCase__,						UpperCAmelCase__,						UpperCAmelCase__,						unittest.TestCase							):







						'''simple docstring'''




						SCREAMING_SNAKE_CASE__      :  Optional[Any]					=							(
						    (BioGptModel, BioGptForCausalLM, BioGptForSequenceClassification, BioGptForTokenClassification)
						    if is_torch_available()
						    else ()
						)
						SCREAMING_SNAKE_CASE__      :  Any					=							(BioGptForCausalLM,) if is_torch_available() else ()
						SCREAMING_SNAKE_CASE__      :  List[Any]					=							(
						    {
						        """feature-extraction""": BioGptModel,
						        """text-classification""": BioGptForSequenceClassification,
						        """text-generation""": BioGptForCausalLM,
						        """token-classification""": BioGptForTokenClassification,
						        """zero-shot""": BioGptForSequenceClassification,
						    }
						    if is_torch_available()
						    else {}
						)
						SCREAMING_SNAKE_CASE__      :  Any					=							False




						def 						lowercase_ (				self							)					->    Dict:
										lowerCAmelCase_ :      str    				=				BioGptModelTester(self							)
										lowerCAmelCase_ :      Any    				=				ConfigTester(self   ,						config_class=__lowercase   ,						hidden_size=3_7							)




						def 						lowercase_ (				self							)					->    Dict:
										self.config_tester.run_common_tests()




						def 						lowercase_ (				self							)					->    Optional[Any]:
										lowerCAmelCase_ :      List[Any]    				=				self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*__lowercase							)




						def 						lowercase_ (				self							)					->    Dict:
										lowerCAmelCase_ :      List[str]    				=				self.model_tester.prepare_config_and_inputs()
										for type in ["absolute", "relative_key", "relative_key_query"]:
														lowerCAmelCase_ :      List[Any]    				=				type
														self.model_tester.create_and_check_model(*__lowercase							)




						def 						lowercase_ (				self							)					->    List[Any]:
										lowerCAmelCase_ :      List[Any]    				=				self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_biogpt_model_attention_mask_past(*__lowercase							)




						def 						lowercase_ (				self							)					->    str:
										lowerCAmelCase_ :      List[str]    				=				self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_forward_and_backwards(*__lowercase   ,						gradient_checkpointing=__lowercase							)




						def 						lowercase_ (				self							)					->    List[Any]:
										lowerCAmelCase_ :      List[str]    				=				self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_biogpt_model_past_large_inputs(*__lowercase							)




						def 						lowercase_ (				self							)					->    Dict:
										lowerCAmelCase_ :      List[Any]    				=				self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_biogpt_weight_initialization(*__lowercase							)




						def 						lowercase_ (				self							)					->    List[str]:
										lowerCAmelCase_ :      List[Any]    				=				self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_biogpt_for_token_classification(*__lowercase							)




						@slow
						def 						lowercase_ (				self							)					->    Optional[Any]:
										lowerCAmelCase_ :      Optional[Any]    				=				BioGptForCausalLM.from_pretrained('''microsoft/biogpt'''							)
										model.to(__lowercase							)
										lowerCAmelCase_ :      Optional[Any]    				=				BioGptTokenizer.from_pretrained('''microsoft/biogpt'''							)

										lowerCAmelCase_ :      int    				=				'''left'''

										# Define PAD Token = EOS Token = 50256
										lowerCAmelCase_ :      int    				=				tokenizer.eos_token
										lowerCAmelCase_ :      List[str]    				=				model.config.eos_token_id

										# use different length sentences to test batching
										lowerCAmelCase_ :      Optional[Any]    				=				[
										    '''Hello, my dog is a little''',
										    '''Today, I''',
										]

										lowerCAmelCase_ :      int    				=				tokenizer(__lowercase   ,						return_tensors='''pt'''   ,						padding=__lowercase							)
										lowerCAmelCase_ :      Tuple    				=				inputs['''input_ids'''].to(__lowercase							)

										lowerCAmelCase_ :      Union[str, Any]    				=				model.generate(
										    input_ids=__lowercase   ,						attention_mask=inputs['''attention_mask'''].to(__lowercase							)   ,						)

										lowerCAmelCase_ :      Tuple    				=				tokenizer(sentences[0]   ,						return_tensors='''pt'''							).input_ids.to(__lowercase							)
										lowerCAmelCase_ :      Optional[Any]    				=				model.generate(input_ids=__lowercase							)

										lowerCAmelCase_ :      Any    				=				inputs_non_padded.shape[-1] - inputs['''attention_mask'''][-1].long().sum().cpu().item()
										lowerCAmelCase_ :      int    				=				tokenizer(sentences[1]   ,						return_tensors='''pt'''							).input_ids.to(__lowercase							)
										lowerCAmelCase_ :      str    				=				model.generate(input_ids=__lowercase   ,						max_length=model.config.max_length - num_paddings							)

										lowerCAmelCase_ :      Tuple    				=				tokenizer.batch_decode(__lowercase   ,						skip_special_tokens=__lowercase							)
										lowerCAmelCase_ :      Optional[int]    				=				tokenizer.decode(output_non_padded[0]   ,						skip_special_tokens=__lowercase							)
										lowerCAmelCase_ :      List[Any]    				=				tokenizer.decode(output_padded[0]   ,						skip_special_tokens=__lowercase							)

										lowerCAmelCase_ :      List[str]    				=				[
										    '''Hello, my dog is a little bit bigger than a little bit.''',
										    '''Today, I have a good idea of how to use the information''',
										]
										self.assertListEqual(__lowercase   ,						__lowercase							)
										self.assertListEqual(__lowercase   ,						[non_padded_sentence, padded_sentence]							)




						@slow
						def 						lowercase_ (				self							)					->    Optional[int]:
										for model_name in BIOGPT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
														lowerCAmelCase_ :      str    				=				BioGptModel.from_pretrained(__lowercase							)
														self.assertIsNotNone(__lowercase							)




						def 						lowercase_ (				self							)					->    str:
										lowerCAmelCase_				,  lowerCAmelCase_ :      int    				=				self.model_tester.prepare_config_and_inputs_for_common()
										lowerCAmelCase_ :      Tuple    				=				3
										lowerCAmelCase_ :      Tuple    				=				input_dict['''input_ids''']
										lowerCAmelCase_ :      Optional[Any]    				=				input_ids.ne(1							).to(__lowercase							)
										lowerCAmelCase_ :      Union[str, Any]    				=				ids_tensor([self.model_tester.batch_size]   ,						self.model_tester.type_sequence_label_size							)
										lowerCAmelCase_ :      Optional[Any]    				=				BioGptForSequenceClassification(__lowercase							)
										model.to(__lowercase							)
										model.eval()
										lowerCAmelCase_ :      int    				=				model(__lowercase   ,						attention_mask=__lowercase   ,						labels=__lowercase							)
										self.assertEqual(result.logits.shape   ,						(self.model_tester.batch_size, self.model_tester.num_labels)							)




						def 						lowercase_ (				self							)					->    Tuple:
										lowerCAmelCase_				,  lowerCAmelCase_ :      Tuple    				=				self.model_tester.prepare_config_and_inputs_for_common()
										lowerCAmelCase_ :      List[Any]    				=				3
										lowerCAmelCase_ :      List[str]    				=				'''multi_label_classification'''
										lowerCAmelCase_ :      Union[str, Any]    				=				input_dict['''input_ids''']
										lowerCAmelCase_ :      int    				=				input_ids.ne(1							).to(__lowercase							)
										lowerCAmelCase_ :      Dict    				=				ids_tensor(
										    [self.model_tester.batch_size, config.num_labels]   ,						self.model_tester.type_sequence_label_size							).to(torch.float							)
										lowerCAmelCase_ :      List[str]    				=				BioGptForSequenceClassification(__lowercase							)
										model.to(__lowercase							)
										model.eval()
										lowerCAmelCase_ :      Dict    				=				model(__lowercase   ,						attention_mask=__lowercase   ,						labels=__lowercase							)
										self.assertEqual(result.logits.shape   ,						(self.model_tester.batch_size, self.model_tester.num_labels)							)







@require_torch
class 				snake_case__(							unittest.TestCase							):







						'''simple docstring'''




						@slow
						def 						lowercase_ (				self							)					->    Union[str, Any]:
										lowerCAmelCase_ :      Optional[int]    				=				BioGptForCausalLM.from_pretrained('''microsoft/biogpt'''							)
										lowerCAmelCase_ :      Dict    				=				torch.tensor([[2, 4_8_0_5, 9, 6_5_6, 2_1]]							)
										lowerCAmelCase_ :      str    				=				model(__lowercase							)[0]

										lowerCAmelCase_ :      Tuple    				=				4_2_3_8_4

										lowerCAmelCase_ :      str    				=				torch.Size((1, 5, vocab_size)							)
										self.assertEqual(output.shape   ,						__lowercase							)

										lowerCAmelCase_ :      Tuple    				=				torch.tensor(
										    [[[-9.52_36, -9.89_18, 10.45_57], [-11.04_69, -9.64_23, 8.10_22], [-8.86_64, -7.88_26, 5.53_25]]]							)

										self.assertTrue(torch.allclose(output[:, :3, :3]   ,						__lowercase   ,						atol=1e-4							)							)




						@slow
						def 						lowercase_ (				self							)					->    List[str]:
										lowerCAmelCase_ :      List[str]    				=				BioGptTokenizer.from_pretrained('''microsoft/biogpt'''							)
										lowerCAmelCase_ :      Optional[Any]    				=				BioGptForCausalLM.from_pretrained('''microsoft/biogpt'''							)
										model.to(__lowercase							)

										torch.manual_seed(0							)
										lowerCAmelCase_ :      Dict    				=				tokenizer('''COVID-19 is'''   ,						return_tensors='''pt'''							).to(__lowercase							)
										lowerCAmelCase_ :      Optional[int]    				=				model.generate(
										    **__lowercase   ,						min_length=1_0_0   ,						max_length=1_0_2_4   ,						num_beams=5   ,						early_stopping=__lowercase   ,						)
										lowerCAmelCase_ :      int    				=				tokenizer.decode(output_ids[0]   ,						skip_special_tokens=__lowercase							)

										lowerCAmelCase_ :      List[Any]    				=				(
										    '''COVID-19 is a global pandemic caused by severe acute respiratory syndrome coronavirus 2 (SARS-CoV-2), the'''
										    ''' causative agent of coronavirus disease 2019 (COVID-19), which has spread to more than 200 countries and'''
										    ''' territories, including the United States (US), Canada, Australia, New Zealand, the United Kingdom (UK),'''
										    ''' and the United States of America (USA), as of March 11, 2020, with more than 800,000 confirmed cases and'''
										    ''' more than 800,000 deaths.'''
										)
										self.assertEqual(__lowercase   ,						__lowercase							)



import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import TransformeraDModel, VQDiffusionPipeline, VQDiffusionScheduler, VQModel
from diffusers.pipelines.vq_diffusion.pipeline_vq_diffusion import LearnedClassifierFreeSamplingEmbeddings
from diffusers.utils import load_numpy, slow, torch_device
from diffusers.utils.testing_utils import require_torch_gpu


_UpperCAmelCase			:  Any     =False







class 				snake_case__(							unittest.TestCase							):







						'''simple docstring'''




						def 						lowercase_ (				self							)					->    Optional[int]:
										# clean up the VRAM after each test
										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()




						@property
						def 						lowercase_ (				self							)					->    Union[str, Any]:
										return 1_2




						@property
						def 						lowercase_ (				self							)					->    Any:
										return 1_2




						@property
						def 						lowercase_ (				self							)					->    Optional[Any]:
										return 3_2




						@property
						def 						lowercase_ (				self							)					->    int:
										torch.manual_seed(0							)
										lowerCAmelCase_ :      Any    				=				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   ,						num_vq_embeddings=self.num_embed   ,						vq_embed_dim=3   ,						)
										return model




						@property
						def 						lowercase_ (				self							)					->    Dict:
										lowerCAmelCase_ :      Optional[Any]    				=				CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''							)
										return tokenizer




						@property
						def 						lowercase_ (				self							)					->    int:
										torch.manual_seed(0							)
										lowerCAmelCase_ :      Union[str, Any]    				=				CLIPTextConfig(
										    bos_token_id=0   ,						eos_token_id=2   ,						hidden_size=self.text_embedder_hidden_size   ,						intermediate_size=3_7   ,						layer_norm_eps=1e-05   ,						num_attention_heads=4   ,						num_hidden_layers=5   ,						pad_token_id=1   ,						vocab_size=1_0_0_0   ,						)
										return CLIPTextModel(__lowercase							)




						@property
						def 						lowercase_ (				self							)					->    List[str]:
										torch.manual_seed(0							)

										lowerCAmelCase_ :      Union[str, Any]    				=				1_2
										lowerCAmelCase_ :      int    				=				1_2

										lowerCAmelCase_ :      Union[str, Any]    				=				{
										    '''attention_bias''': True,
										    '''cross_attention_dim''': 3_2,
										    '''attention_head_dim''': height * width,
										    '''num_attention_heads''': 1,
										    '''num_vector_embeds''': self.num_embed,
										    '''num_embeds_ada_norm''': self.num_embeds_ada_norm,
										    '''norm_num_groups''': 3_2,
										    '''sample_size''': width,
										    '''activation_fn''': '''geglu-approximate''',
										}

										lowerCAmelCase_ :      List[str]    				=				TransformeraDModel(**__lowercase							)
										return model




						def 						lowercase_ (				self							)					->    str:
										lowerCAmelCase_ :      List[Any]    				=				'''cpu'''

										lowerCAmelCase_ :      Any    				=				self.dummy_vqvae
										lowerCAmelCase_ :      str    				=				self.dummy_text_encoder
										lowerCAmelCase_ :      Union[str, Any]    				=				self.dummy_tokenizer
										lowerCAmelCase_ :      int    				=				self.dummy_transformer
										lowerCAmelCase_ :      List[str]    				=				VQDiffusionScheduler(self.num_embed							)
										lowerCAmelCase_ :      Union[str, Any]    				=				LearnedClassifierFreeSamplingEmbeddings(learnable=__lowercase							)

										lowerCAmelCase_ :      Dict    				=				VQDiffusionPipeline(
										    vqvae=__lowercase   ,						text_encoder=__lowercase   ,						tokenizer=__lowercase   ,						transformer=__lowercase   ,						scheduler=__lowercase   ,						learned_classifier_free_sampling_embeddings=__lowercase   ,						)
										lowerCAmelCase_ :      int    				=				pipe.to(__lowercase							)
										pipe.set_progress_bar_config(disable=__lowercase							)

										lowerCAmelCase_ :      Any    				=				'''teddy bear playing in the pool'''

										lowerCAmelCase_ :      int    				=				torch.Generator(device=__lowercase							).manual_seed(0							)
										lowerCAmelCase_ :      Tuple    				=				pipe([prompt]   ,						generator=__lowercase   ,						num_inference_steps=2   ,						output_type='''np'''							)
										lowerCAmelCase_ :      Union[str, Any]    				=				output.images

										lowerCAmelCase_ :      List[Any]    				=				torch.Generator(device=__lowercase							).manual_seed(0							)
										lowerCAmelCase_ :      List[Any]    				=				pipe(
										    [prompt]   ,						generator=__lowercase   ,						output_type='''np'''   ,						return_dict=__lowercase   ,						num_inference_steps=2							)[0]

										lowerCAmelCase_ :      List[str]    				=				image[0, -3:, -3:, -1]
										lowerCAmelCase_ :      str    				=				image_from_tuple[0, -3:, -3:, -1]

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

										lowerCAmelCase_ :      Optional[int]    				=				np.array([0.65_51, 0.61_68, 0.50_08, 0.56_76, 0.56_59, 0.42_95, 0.60_73, 0.55_99, 0.49_92]							)

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




						def 						lowercase_ (				self							)					->    List[str]:
										lowerCAmelCase_ :      Optional[Any]    				=				'''cpu'''

										lowerCAmelCase_ :      str    				=				self.dummy_vqvae
										lowerCAmelCase_ :      Dict    				=				self.dummy_text_encoder
										lowerCAmelCase_ :      List[Any]    				=				self.dummy_tokenizer
										lowerCAmelCase_ :      Union[str, Any]    				=				self.dummy_transformer
										lowerCAmelCase_ :      Tuple    				=				VQDiffusionScheduler(self.num_embed							)
										lowerCAmelCase_ :      str    				=				LearnedClassifierFreeSamplingEmbeddings(
										    learnable=__lowercase   ,						hidden_size=self.text_embedder_hidden_size   ,						length=tokenizer.model_max_length							)

										lowerCAmelCase_ :      List[str]    				=				VQDiffusionPipeline(
										    vqvae=__lowercase   ,						text_encoder=__lowercase   ,						tokenizer=__lowercase   ,						transformer=__lowercase   ,						scheduler=__lowercase   ,						learned_classifier_free_sampling_embeddings=__lowercase   ,						)
										lowerCAmelCase_ :      Union[str, Any]    				=				pipe.to(__lowercase							)
										pipe.set_progress_bar_config(disable=__lowercase							)

										lowerCAmelCase_ :      Any    				=				'''teddy bear playing in the pool'''

										lowerCAmelCase_ :      List[str]    				=				torch.Generator(device=__lowercase							).manual_seed(0							)
										lowerCAmelCase_ :      Dict    				=				pipe([prompt]   ,						generator=__lowercase   ,						num_inference_steps=2   ,						output_type='''np'''							)
										lowerCAmelCase_ :      str    				=				output.images

										lowerCAmelCase_ :      List[Any]    				=				torch.Generator(device=__lowercase							).manual_seed(0							)
										lowerCAmelCase_ :      Union[str, Any]    				=				pipe(
										    [prompt]   ,						generator=__lowercase   ,						output_type='''np'''   ,						return_dict=__lowercase   ,						num_inference_steps=2							)[0]

										lowerCAmelCase_ :      List[str]    				=				image[0, -3:, -3:, -1]
										lowerCAmelCase_ :      str    				=				image_from_tuple[0, -3:, -3:, -1]

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

										lowerCAmelCase_ :      Union[str, Any]    				=				np.array([0.66_93, 0.60_75, 0.49_59, 0.57_01, 0.55_83, 0.43_33, 0.61_71, 0.56_84, 0.49_88]							)

										assert np.abs(image_slice.flatten() - expected_slice							).max() < 2.0
										assert np.abs(image_from_tuple_slice.flatten() - expected_slice							).max() < 1e-2







@slow
@require_torch_gpu
class 				snake_case__(							unittest.TestCase							):







						'''simple docstring'''




						def 						lowercase_ (				self							)					->    Optional[int]:
										# clean up the VRAM after each test
										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()




						def 						lowercase_ (				self							)					->    int:
										lowerCAmelCase_ :      Tuple    				=				load_numpy(
										    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
										    '''/vq_diffusion/teddy_bear_pool_classifier_free_sampling.npy'''							)

										lowerCAmelCase_ :      str    				=				VQDiffusionPipeline.from_pretrained('''microsoft/vq-diffusion-ithq'''							)
										lowerCAmelCase_ :      List[Any]    				=				pipeline.to(__lowercase							)
										pipeline.set_progress_bar_config(disable=__lowercase							)

										# requires GPU generator for gumbel softmax
										# don't use GPU generator in tests though
										lowerCAmelCase_ :      List[Any]    				=				torch.Generator(device=__lowercase							).manual_seed(0							)
										lowerCAmelCase_ :      Optional[int]    				=				pipeline(
										    '''teddy bear playing in the pool'''   ,						num_images_per_prompt=1   ,						generator=__lowercase   ,						output_type='''np'''   ,						)

										lowerCAmelCase_ :      Union[str, Any]    				=				output.images[0]

										assert image.shape == (2_5_6, 2_5_6, 3)
										assert np.abs(expected_image - image							).max() < 2.0