


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 	__lowercase    :





							"""simple docstring"""







							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										torch.manual_seed(0       )
										SCREAMING_SNAKE_CASE_		:     List[Any]					=					TaEncoderModel.from_pretrained('hf-internal-testing/tiny-random-t5'       )

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

										torch.manual_seed(0       )
										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					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       )
										SCREAMING_SNAKE_CASE_		:     int					=					DDPMScheduler(
										    num_train_timesteps=1_0_0_0			,						beta_schedule='squaredcos_cap_v2'			,						beta_start=0.0_001			,						beta_end=0.02			,						thresholding=lowerCAmelCase__			,						dynamic_thresholding_ratio=0.95			,						sample_max_value=1.0			,						prediction_type='epsilon'			,						variance_type='learned_range'			,						)

										torch.manual_seed(0       )
										SCREAMING_SNAKE_CASE_		:     int					=					IFWatermarker()

										return {
										    "text_encoder": text_encoder,
										    "tokenizer": tokenizer,
										    "unet": unet,
										    "scheduler": scheduler,
										    "watermarker": watermarker,
										    "safety_checker": None,
										    "feature_extractor": None,
										}






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




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

										torch.manual_seed(0       )
										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					AutoTokenizer.from_pretrained('hf-internal-testing/tiny-random-t5'       )

										torch.manual_seed(0       )
										SCREAMING_SNAKE_CASE_		:     List[Any]					=					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.414			,						time_embedding_act_fn='gelu'			,						time_embedding_dim=3_2			,						)
										unet.set_attn_processor(AttnAddedKVProcessor()       )  # For reproducibility tests

										torch.manual_seed(0       )
										SCREAMING_SNAKE_CASE_		:     Any					=					DDPMScheduler(
										    num_train_timesteps=1_0_0_0			,						beta_schedule='squaredcos_cap_v2'			,						beta_start=0.0_001			,						beta_end=0.02			,						thresholding=lowerCAmelCase__			,						dynamic_thresholding_ratio=0.95			,						sample_max_value=1.0			,						prediction_type='epsilon'			,						variance_type='learned_range'			,						)

										torch.manual_seed(0       )
										SCREAMING_SNAKE_CASE_		:     List[str]					=					DDPMScheduler(
										    num_train_timesteps=1_0_0_0			,						beta_schedule='squaredcos_cap_v2'			,						beta_start=0.0_001			,						beta_end=0.02			,						)

										torch.manual_seed(0       )
										SCREAMING_SNAKE_CASE_		:     List[str]					=					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 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     int					=					self.get_dummy_components()
										SCREAMING_SNAKE_CASE_		:     Tuple					=					self.pipeline_class(**lowerCAmelCase__       )
										pipe.to(lowerCAmelCase__       )
										pipe.set_progress_bar_config(disable=lowerCAmelCase__       )

										SCREAMING_SNAKE_CASE_		:     Dict					=					self.get_dummy_inputs(lowerCAmelCase__       )

										SCREAMING_SNAKE_CASE_		:     int					=					inputs['prompt']
										SCREAMING_SNAKE_CASE_		:     Dict					=					inputs['generator']
										SCREAMING_SNAKE_CASE_		:     int					=					inputs['num_inference_steps']
										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					inputs['output_type']

										if "image" in inputs:
													SCREAMING_SNAKE_CASE_		:     List[Any]					=					inputs['image']
										else:
													SCREAMING_SNAKE_CASE_		:     Tuple					=					None

										if "mask_image" in inputs:
													SCREAMING_SNAKE_CASE_		:     Dict					=					inputs['mask_image']
										else:
													SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					None

										if "original_image" in inputs:
													SCREAMING_SNAKE_CASE_		:     Optional[int]					=					inputs['original_image']
										else:
													SCREAMING_SNAKE_CASE_		:     List[Any]					=					None

										SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		:     Dict					=					pipe.encode_prompt(lowerCAmelCase__       )

										# inputs with prompt converted to embeddings
										SCREAMING_SNAKE_CASE_		:     int					=					{
										    '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:
													SCREAMING_SNAKE_CASE_		:     int					=					image

										if mask_image is not None:
													SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					mask_image

										if original_image is not None:
													SCREAMING_SNAKE_CASE_		:     str					=					original_image

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

										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					pipe(**lowerCAmelCase__       )[0]

										with tempfile.TemporaryDirectory() as tmpdir:
													pipe.save_pretrained(lowerCAmelCase__       )
													SCREAMING_SNAKE_CASE_		:     str					=					self.pipeline_class.from_pretrained(lowerCAmelCase__       )
													pipe_loaded.to(lowerCAmelCase__       )
													pipe_loaded.set_progress_bar_config(disable=lowerCAmelCase__       )

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

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

										SCREAMING_SNAKE_CASE_		:     int					=					self.get_dummy_inputs(lowerCAmelCase__       )

										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					inputs['generator']
										SCREAMING_SNAKE_CASE_		:     Tuple					=					inputs['num_inference_steps']
										SCREAMING_SNAKE_CASE_		:     List[str]					=					inputs['output_type']

										# inputs with prompt converted to embeddings
										SCREAMING_SNAKE_CASE_		:     Any					=					{
										    '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:
													SCREAMING_SNAKE_CASE_		:     List[Any]					=					image

										if mask_image is not None:
													SCREAMING_SNAKE_CASE_		:     List[str]					=					mask_image

										if original_image is not None:
													SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					original_image

										SCREAMING_SNAKE_CASE_		:     int					=					pipe_loaded(**lowerCAmelCase__       )[0]

										SCREAMING_SNAKE_CASE_		:     Dict					=					np.abs(to_np(lowerCAmelCase__       ) - to_np(lowerCAmelCase__       )       ).max()
										self.assertLess(lowerCAmelCase__			,						1E-4       )






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					self.get_dummy_components()
										SCREAMING_SNAKE_CASE_		:     Any					=					self.pipeline_class(**lowerCAmelCase__       )
										pipe.to(lowerCAmelCase__       )
										pipe.set_progress_bar_config(disable=lowerCAmelCase__       )

										SCREAMING_SNAKE_CASE_		:     List[str]					=					self.get_dummy_inputs(lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     List[str]					=					pipe(**lowerCAmelCase__       )[0]

										with tempfile.TemporaryDirectory() as tmpdir:
													pipe.save_pretrained(lowerCAmelCase__       )
													SCREAMING_SNAKE_CASE_		:     List[Any]					=					self.pipeline_class.from_pretrained(lowerCAmelCase__       )
													pipe_loaded.to(lowerCAmelCase__       )
													pipe_loaded.set_progress_bar_config(disable=lowerCAmelCase__       )

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

										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					self.get_dummy_inputs(lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     List[str]					=					pipe_loaded(**lowerCAmelCase__       )[0]

										SCREAMING_SNAKE_CASE_		:     Dict					=					np.abs(to_np(lowerCAmelCase__       ) - to_np(lowerCAmelCase__       )       ).max()
										self.assertLess(lowerCAmelCase__			,						1E-4       )









import os
import shutil
import tempfile
from unittest import TestCase
from unittest.mock import patch

import numpy as np
from datasets import Dataset

from transformers.models.realm.configuration_realm import RealmConfig
from transformers.models.realm.retrieval_realm import _REALM_BLOCK_RECORDS_FILENAME, RealmRetriever
from transformers.models.realm.tokenization_realm import VOCAB_FILES_NAMES, RealmTokenizer







class 	__lowercase    (__SCREAMING_SNAKE_CASE  ):





							"""simple docstring"""







							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Tuple					=					tempfile.mkdtemp()
										SCREAMING_SNAKE_CASE_		:     Dict					=					5

										# Realm tok
										SCREAMING_SNAKE_CASE_		:     Dict					=					[
										    '[UNK]',
										    '[CLS]',
										    '[SEP]',
										    '[PAD]',
										    '[MASK]',
										    'test',
										    'question',
										    'this',
										    'is',
										    'the',
										    'first',
										    'second',
										    'third',
										    'fourth',
										    'fifth',
										    'record',
										    'want',
										    '##want',
										    '##ed',
										    'wa',
										    'un',
										    'runn',
										    '##ing',
										    ',',
										    'low',
										    'lowest',
										]
										SCREAMING_SNAKE_CASE_		:     Tuple					=					os.path.join(self.tmpdirname			,						'realm_tokenizer'       )
										os.makedirs(lowerCAmelCase__			,						exist_ok=lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     Tuple					=					os.path.join(lowerCAmelCase__			,						VOCAB_FILES_NAMES['vocab_file']       )
										with open(self.vocab_file			,						'w'			,						encoding='utf-8'       ) as vocab_writer:
													vocab_writer.write(''.join([x + '\n' for x in vocab_tokens]       )       )

										SCREAMING_SNAKE_CASE_		:     str					=					os.path.join(self.tmpdirname			,						'realm_block_records'       )
										os.makedirs(lowerCAmelCase__			,						exist_ok=lowerCAmelCase__       )






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										return RealmTokenizer.from_pretrained(os.path.join(self.tmpdirname			,						'realm_tokenizer'       )       )






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										shutil.rmtree(self.tmpdirname       )






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					RealmConfig(num_block_records=self.num_block_records       )
										return config






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					Dataset.from_dict(
										    {
										        'id': ['0', '1'],
										        'question': ['foo', 'bar'],
										        'answers': [['Foo', 'Bar'], ['Bar']],
										    }       )
										return dataset






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Dict					=					np.array(
										    [
										        b'This is the first record',
										        b'This is the second record',
										        b'This is the third record',
										        b'This is the fourth record',
										        b'This is the fifth record',
										        b'This is a longer longer longer record',
										    ]			,						dtype=lowerCAmelCase__			,						)
										return block_records






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Any					=					RealmRetriever(
										    block_records=self.get_dummy_block_records()			,						tokenizer=self.get_tokenizer()			,						)
										return retriever






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Any					=					self.get_config()
										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					self.get_dummy_retriever()
										SCREAMING_SNAKE_CASE_		:     List[str]					=					retriever.tokenizer

										SCREAMING_SNAKE_CASE_		:     Any					=					np.array([0, 3]			,						dtype='long'       )
										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					tokenizer(['Test question']       ).input_ids
										SCREAMING_SNAKE_CASE_		:     int					=					tokenizer(
										    ['the fourth']			,						add_special_tokens=lowerCAmelCase__			,						return_token_type_ids=lowerCAmelCase__			,						return_attention_mask=lowerCAmelCase__			,						).input_ids
										SCREAMING_SNAKE_CASE_		:     Any					=					config.reader_seq_len

										SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		:     str					=					retriever(
										    lowerCAmelCase__			,						lowerCAmelCase__			,						answer_ids=lowerCAmelCase__			,						max_length=lowerCAmelCase__			,						return_tensors='np'       )

										self.assertEqual(len(lowerCAmelCase__       )			,						2       )
										self.assertEqual(len(lowerCAmelCase__       )			,						2       )
										self.assertEqual(len(lowerCAmelCase__       )			,						2       )
										self.assertEqual(concat_inputs.input_ids.shape			,						(2, 1_0)       )
										self.assertEqual(concat_inputs.attention_mask.shape			,						(2, 1_0)       )
										self.assertEqual(concat_inputs.token_type_ids.shape			,						(2, 1_0)       )
										self.assertEqual(concat_inputs.special_tokens_mask.shape			,						(2, 1_0)       )
										self.assertEqual(
										    tokenizer.convert_ids_to_tokens(concat_inputs.input_ids[0]       )			,						['[CLS]', 'test', 'question', '[SEP]', 'this', 'is', 'the', 'first', 'record', '[SEP]']			,						)
										self.assertEqual(
										    tokenizer.convert_ids_to_tokens(concat_inputs.input_ids[1]       )			,						['[CLS]', 'test', 'question', '[SEP]', 'this', 'is', 'the', 'fourth', 'record', '[SEP]']			,						)






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Dict					=					self.get_config()
										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					self.get_dummy_retriever()
										SCREAMING_SNAKE_CASE_		:     Tuple					=					retriever.tokenizer

										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					np.array([0, 3, 5]			,						dtype='long'       )
										SCREAMING_SNAKE_CASE_		:     List[Any]					=					tokenizer(['Test question']       ).input_ids
										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					tokenizer(
										    ['the fourth', 'longer longer']			,						add_special_tokens=lowerCAmelCase__			,						return_token_type_ids=lowerCAmelCase__			,						return_attention_mask=lowerCAmelCase__			,						).input_ids
										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					config.reader_seq_len

										SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		:     Tuple					=					retriever(
										    lowerCAmelCase__			,						lowerCAmelCase__			,						answer_ids=lowerCAmelCase__			,						max_length=lowerCAmelCase__			,						return_tensors='np'       )

										self.assertEqual([False, True, True]			,						lowerCAmelCase__       )
										self.assertEqual([[-1, -1, -1], [6, -1, -1], [6, 7, 8]]			,						lowerCAmelCase__       )
										self.assertEqual([[-1, -1, -1], [7, -1, -1], [7, 8, 9]]			,						lowerCAmelCase__       )






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Dict					=					self.get_dummy_retriever()
										retriever.save_pretrained(os.path.join(self.tmpdirname			,						'realm_block_records'       )       )

										# Test local path
										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					retriever.from_pretrained(os.path.join(self.tmpdirname			,						'realm_block_records'       )       )
										self.assertEqual(retriever.block_records[0]			,						b'This is the first record'       )

										# Test mocked remote path
										with patch('transformers.models.realm.retrieval_realm.hf_hub_download'       ) as mock_hf_hub_download:
													SCREAMING_SNAKE_CASE_		:     Optional[int]					=					os.path.join(
													    os.path.join(self.tmpdirname			,						'realm_block_records'       )			,						_REALM_BLOCK_RECORDS_FILENAME       )
													SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					RealmRetriever.from_pretrained('google/realm-cc-news-pretrained-openqa'       )

										self.assertEqual(retriever.block_records[0]			,						b'This is the first record'       )





