




import math
from typing import Optional

import torch
import torch.utils.checkpoint as checkpoint
from torch import nn

from ....modeling_utils import PreTrainedModel
from ....utils import add_start_docstrings, logging
from ...bert.modeling_bert import BertModel
from .configuration_retribert import RetriBertConfig


__UpperCamelCase      : Optional[Any]					      =   logging.get_logger(__name__)

__UpperCamelCase      : Any					      =   [
    "yjernite/retribert-base-uncased",
    # See all RetriBert models at https://huggingface.co/models?filter=retribert
]

class 	__lowerCAmelCase			(			__magic_name__				):
							UpperCamelCase__        =						RetriBertConfig
							UpperCamelCase__        =						None
							UpperCamelCase__        =						'''retribert'''
							def    lowerCamelCase__							(		self		:str	,   __magic_name__		:Any       ):







								'''simple docstring'''







								if isinstance(__magic_name__	,   nn.Linear       ):
									module.weight.data.normal_(mean=0.0	,   std=self.config.initializer_range       )
									if module.bias is not None:
										module.bias.data.zero_()
								elif isinstance(__magic_name__	,   nn.Embedding       ):
									module.weight.data.normal_(mean=0.0	,   std=self.config.initializer_range       )
									if module.padding_idx is not None:
										module.weight.data[module.padding_idx].zero_()
								elif isinstance(__magic_name__	,   nn.LayerNorm       ):
									module.bias.data.zero_()
									module.weight.data.fill_(1.0       )







__UpperCamelCase      : List[Any]					      =   R"\n\n    This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the\n    library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads\n    etc.)\n\n    This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.\n    Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage\n    and behavior.\n\n    Parameters:\n        config ([`RetriBertConfig`]): Model configuration class with all the parameters of the model.\n            Initializing with a config file does not load the weights associated with the model, only the\n            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.\n"

@add_start_docstrings(
    '''Bert Based model to embed queries or document for document retrieval.'''   , __magic_name__   , )
class 	__lowerCAmelCase			(			__magic_name__				):
							def __init__(		self		:Optional[Any]	,   __magic_name__		:RetriBertConfig       ):







								'''simple docstring'''







								super().__init__(__magic_name__       )
								a		       =							config.projection_dim

								a		       =							BertModel(__magic_name__       )
								a		       =							None if config.share_encoders else BertModel(__magic_name__       )
								a		       =							nn.Dropout(config.hidden_dropout_prob       )
								a		       =							nn.Linear(config.hidden_size	,   config.projection_dim	,   bias=__magic_name__       )
								a		       =							nn.Linear(config.hidden_size	,   config.projection_dim	,   bias=__magic_name__       )

								a		       =							nn.CrossEntropyLoss(reduction="""mean"""       )

								# Initialize weights and apply final processing
								self.post_init()
							def    lowerCamelCase__							(		self		:List[Any]	,   __magic_name__		:List[Any]	,   __magic_name__		:List[str]	,   __magic_name__		:Optional[Any]	,   __magic_name__		:str=-1	,   ):







								'''simple docstring'''







								if checkpoint_batch_size < 0 or input_ids.shape[0] < checkpoint_batch_size:
									return sent_encoder(__magic_name__	,   attention_mask=__magic_name__       )[1]
								else:
									# prepare implicit variables
									a		       =							input_ids.device
									a		       =							input_ids.size()
									a		       =							torch.zeros(__magic_name__	,   dtype=torch.long	,   device=__magic_name__       )
									a		       =							[None] * sent_encoder.config.num_hidden_layers
									a		       =							sent_encoder.get_extended_attention_mask(
									    __magic_name__	,   __magic_name__       )

									# define function for checkpointing
									def partial_encode(*__magic_name__		:Any       ):
										a		       =							sent_encoder.encoder(
										    inputs[0]	,   attention_mask=inputs[1]	,   head_mask=__magic_name__	,   )
										a		       =							encoder_outputs[0]
										a		       =							sent_encoder.pooler(__magic_name__       )
										return pooled_output

									# run embedding layer on everything at once
									a		       =							sent_encoder.embeddings(
									    input_ids=__magic_name__	,   position_ids=__magic_name__	,   token_type_ids=__magic_name__	,   inputs_embeds=__magic_name__       )
									# run encoding and pooling on one mini-batch at a time
									a		       =							[]
									for b in range(math.ceil(input_ids.shape[0] / checkpoint_batch_size       )       ):
										a		       =							embedding_output[b * checkpoint_batch_size : (b + 1) * checkpoint_batch_size]
										a		       =							extended_attention_mask[b * checkpoint_batch_size : (b + 1) * checkpoint_batch_size]
										a		       =							checkpoint.checkpoint(__magic_name__	,   __magic_name__	,   __magic_name__       )
										pooled_output_list.append(__magic_name__       )
									return torch.cat(__magic_name__	,   dim=0       )
							def    lowerCamelCase__							(		self		:Dict	,   __magic_name__		:Dict	,   __magic_name__		:Union[str, Any]=None	,   __magic_name__		:List[str]=-1	,   ):







								'''simple docstring'''







								a		       =							self.embed_sentences_checkpointed(
								    __magic_name__	,   __magic_name__	,   self.bert_query	,   __magic_name__	,   )
								return self.project_query(__magic_name__       )
							def    lowerCamelCase__							(		self		:Tuple	,   __magic_name__		:str	,   __magic_name__		:List[Any]=None	,   __magic_name__		:List[str]=-1	,   ):







								'''simple docstring'''







								a		       =							self.embed_sentences_checkpointed(
								    __magic_name__	,   __magic_name__	,   self.bert_query if self.bert_doc is None else self.bert_doc	,   __magic_name__	,   )
								return self.project_doc(__magic_name__       )





							def    lowerCamelCase__							(		self		:Any	,   __magic_name__		:torch.LongTensor	,   __magic_name__		:Optional[torch.FloatTensor]	,   __magic_name__		:torch.LongTensor	,   __magic_name__		:Optional[torch.FloatTensor]	,   __magic_name__		:int = -1	,   ):







								'''simple docstring'''







								a		       =							input_ids_query.device
								a		       =							self.embed_questions(__magic_name__	,   __magic_name__	,   __magic_name__       )
								a		       =							self.embed_answers(__magic_name__	,   __magic_name__	,   __magic_name__       )
								a		       =							torch.mm(__magic_name__	,   a_reps.t()       )
								a		       =							self.ce_loss(__magic_name__	,   torch.arange(compare_scores.shape[1]       ).to(__magic_name__       )       )
								a		       =							self.ce_loss(compare_scores.t()	,   torch.arange(compare_scores.shape[0]       ).to(__magic_name__       )       )
								a		       =							(loss_qa + loss_aq) / 2
								return loss







import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMParallelScheduler,
    DDPMParallelScheduler,
    StableDiffusionParadigmsPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    require_torch_gpu,
)

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()

class 	__lowerCAmelCase			(			__magic_name__   , __magic_name__   , unittest.TestCase				):
							UpperCamelCase__        =						StableDiffusionParadigmsPipeline
							UpperCamelCase__        =						TEXT_TO_IMAGE_PARAMS
							UpperCamelCase__        =						TEXT_TO_IMAGE_BATCH_PARAMS
							UpperCamelCase__        =						TEXT_TO_IMAGE_IMAGE_PARAMS
							UpperCamelCase__        =						TEXT_TO_IMAGE_IMAGE_PARAMS
							def    lowerCamelCase__							(		self		:Optional[Any]       ):







								'''simple docstring'''







								torch.manual_seed(0       )
								a		       =							UNetaDConditionModel(
								    block_out_channels=(32, 64)	,   layers_per_block=2	,   sample_size=32	,   in_channels=4	,   out_channels=4	,   down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")	,   up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")	,   cross_attention_dim=32	,   attention_head_dim=(2, 4)	,   use_linear_projection=__magic_name__	,   )
								a		       =							DDIMParallelScheduler(
								    beta_start=0.00085	,   beta_end=0.012	,   beta_schedule="""scaled_linear"""	,   clip_sample=__magic_name__	,   set_alpha_to_one=__magic_name__	,   )
								torch.manual_seed(0       )
								a		       =							AutoencoderKL(
								    block_out_channels=[32, 64]	,   in_channels=3	,   out_channels=3	,   down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]	,   up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]	,   latent_channels=4	,   sample_size=128	,   )
								torch.manual_seed(0       )
								a		       =							CLIPTextConfig(
								    bos_token_id=0	,   eos_token_id=2	,   hidden_size=32	,   intermediate_size=37	,   layer_norm_eps=1E-0_5	,   num_attention_heads=4	,   num_hidden_layers=5	,   pad_token_id=1	,   vocab_size=1000	,   hidden_act="""gelu"""	,   projection_dim=512	,   )
								a		       =							CLIPTextModel(__magic_name__       )
								a		       =							CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""       )

								a		       =							{
								    """unet""": unet,
								    """scheduler""": scheduler,
								    """vae""": vae,
								    """text_encoder""": text_encoder,
								    """tokenizer""": tokenizer,
								    """safety_checker""": None,
								    """feature_extractor""": None,
								}
								return components
							def    lowerCamelCase__							(		self		:Any	,   __magic_name__		:List[Any]	,   __magic_name__		:Tuple=0       ):







								'''simple docstring'''







								if str(__magic_name__       ).startswith("""mps"""       ):
									a		       =							torch.manual_seed(__magic_name__       )
								else:
									a		       =							torch.Generator(device=__magic_name__       ).manual_seed(__magic_name__       )
								a		       =							{
								    """prompt""": """a photograph of an astronaut riding a horse""",
								    """generator""": generator,
								    """num_inference_steps""": 10,
								    """guidance_scale""": 6.0,
								    """output_type""": """numpy""",
								    """parallel""": 3,
								    """debug""": True,
								}
								return inputs
							def    lowerCamelCase__							(		self		:Dict       ):







								'''simple docstring'''







								a		       =							"""cpu"""  # ensure determinism for the device-dependent torch.Generator
								a		       =							self.get_dummy_components()
								a		       =							StableDiffusionParadigmsPipeline(**__magic_name__       )
								a		       =							sd_pipe.to(__magic_name__       )
								sd_pipe.set_progress_bar_config(disable=__magic_name__       )

								a		       =							self.get_dummy_inputs(__magic_name__       )
								a		       =							sd_pipe(**__magic_name__       ).images
								a		       =							image[0, -3:, -3:, -1]
								assert image.shape == (1, 64, 64, 3)

								a		       =							np.array([0.4773, 0.5417, 0.4723, 0.4925, 0.5631, 0.4752, 0.5240, 0.4935, 0.5023]       )

								assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1E-2
							def    lowerCamelCase__							(		self		:Any       ):







								'''simple docstring'''







								a		       =							"""cpu"""  # ensure determinism for the device-dependent torch.Generator
								a		       =							self.get_dummy_components()
								torch.manual_seed(0       )
								a		       =							DDPMParallelScheduler()
								torch.manual_seed(0       )
								a		       =							StableDiffusionParadigmsPipeline(**__magic_name__       )
								a		       =							sd_pipe.to(__magic_name__       )
								sd_pipe.set_progress_bar_config(disable=__magic_name__       )

								a		       =							self.get_dummy_inputs(__magic_name__       )
								a		       =							sd_pipe(**__magic_name__       ).images
								a		       =							image[0, -3:, -3:, -1]
								assert image.shape == (1, 64, 64, 3)

								a		       =							np.array([0.3573, 0.4420, 0.4960, 0.4799, 0.3796, 0.3879, 0.4819, 0.4365, 0.4468]       )

								assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1E-2
							def    lowerCamelCase__							(		self		:Optional[int]       ):







								'''simple docstring'''







								super().test_inference_batch_consistent(batch_sizes=[1, 2]       )
							def    lowerCamelCase__							(		self		:str       ):







								'''simple docstring'''







								super().test_inference_batch_single_identical(batch_size=2	,   expected_max_diff=3E-3       )





							def    lowerCamelCase__							(		self		:Optional[Any]       ):







								'''simple docstring'''







								a		       =							"""cpu"""  # ensure determinism for the device-dependent torch.Generator
								a		       =							self.get_dummy_components()
								a		       =							StableDiffusionParadigmsPipeline(**__magic_name__       )
								a		       =							sd_pipe.to(__magic_name__       )
								sd_pipe.set_progress_bar_config(disable=__magic_name__       )

								a		       =							self.get_dummy_inputs(__magic_name__       )
								a		       =							"""french fries"""
								a		       =							sd_pipe(**__magic_name__	,   negative_prompt=__magic_name__       )
								a		       =							output.images
								a		       =							image[0, -3:, -3:, -1]

								assert image.shape == (1, 64, 64, 3)

								a		       =							np.array([0.4771, 0.5420, 0.4683, 0.4918, 0.5636, 0.4725, 0.5230, 0.4923, 0.5015]       )

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







@slow
@require_torch_gpu
class 	__lowerCAmelCase			(			unittest.TestCase				):
							def    lowerCamelCase__							(		self		:Optional[int]       ):







								'''simple docstring'''







								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()
							def    lowerCamelCase__							(		self		:str	,   __magic_name__		:Union[str, Any]=0       ):







								'''simple docstring'''







								a		       =							torch.Generator(device=__magic_name__       ).manual_seed(__magic_name__       )
								a		       =							{
								    """prompt""": """a photograph of an astronaut riding a horse""",
								    """generator""": generator,
								    """num_inference_steps""": 10,
								    """guidance_scale""": 7.5,
								    """output_type""": """numpy""",
								    """parallel""": 3,
								    """debug""": True,
								}
								return inputs





							def    lowerCamelCase__							(		self		:Dict       ):







								'''simple docstring'''







								a		       =							"""stabilityai/stable-diffusion-2-base"""
								a		       =							DDIMParallelScheduler.from_pretrained(__magic_name__	,   subfolder="""scheduler"""       )
								a		       =							StableDiffusionParadigmsPipeline.from_pretrained(__magic_name__	,   scheduler=__magic_name__	,   safety_checker=__magic_name__       )
								pipe.to(__magic_name__       )
								pipe.set_progress_bar_config(disable=__magic_name__       )
								pipe.enable_attention_slicing()

								a		       =							self.get_inputs()
								a		       =							pipe(**__magic_name__       ).images
								a		       =							image[0, -3:, -3:, -1].flatten()

								assert image.shape == (1, 512, 512, 3)

								a		       =							np.array([0.9622, 0.9602, 0.9748, 0.9591, 0.9630, 0.9691, 0.9661, 0.9631, 0.9741]       )

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

