


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


lowerCAmelCase__       :	Optional[int]     				=logging.get_logger(__name__)

lowerCAmelCase__       :	Tuple     				=[
    'yjernite/retribert-base-uncased',
    # See all RetriBert models at https://huggingface.co/models?filter=retribert
]







class      __lowercase  (__SCREAMING_SNAKE_CASE   ):





							"""simple docstring"""







							_UpperCAmelCase     =		RetriBertConfig
							_UpperCAmelCase     =		None
							_UpperCAmelCase     =		"""retribert"""






							def  UpperCamelCase__			(	self       ,		lowerCAmelCase__    ):





										"""simple docstring"""




										if isinstance(lowerCAmelCase__       ,		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(lowerCAmelCase__       ,		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(lowerCAmelCase__       ,		nn.LayerNorm    ):
													module.bias.data.zero_()
													module.weight.data.fill_(1.0    )







lowerCAmelCase__       :	int     				=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."""     ,    __SCREAMING_SNAKE_CASE     ,    )
class      __lowercase  (__SCREAMING_SNAKE_CASE   ):





							"""simple docstring"""







							def __init__(	self       ,		lowerCAmelCase__    ):





										"""simple docstring"""




										super().__init__(lowerCAmelCase__    )
										SCREAMING_SNAKE_CASE_						:			Tuple       						= config.projection_dim

										SCREAMING_SNAKE_CASE_						:			Optional[Any]       						= BertModel(lowerCAmelCase__    )
										SCREAMING_SNAKE_CASE_						:			List[str]       						= None if config.share_encoders else BertModel(lowerCAmelCase__    )
										SCREAMING_SNAKE_CASE_						:			Any       						= nn.Dropout(config.hidden_dropout_prob    )
										SCREAMING_SNAKE_CASE_						:			str       						= nn.Linear(config.hidden_size       ,		config.projection_dim       ,		bias=lowerCAmelCase__    )
										SCREAMING_SNAKE_CASE_						:			Tuple       						= nn.Linear(config.hidden_size       ,		config.projection_dim       ,		bias=lowerCAmelCase__    )

										SCREAMING_SNAKE_CASE_						:			int       						= nn.CrossEntropyLoss(reduction='mean'    )

										# Initialize weights and apply final processing
										self.post_init()






							def  UpperCamelCase__			(	self       ,		lowerCAmelCase__       ,		lowerCAmelCase__       ,		lowerCAmelCase__       ,		lowerCAmelCase__=-1       ,		):





										"""simple docstring"""




										if checkpoint_batch_size < 0 or input_ids.shape[0] < checkpoint_batch_size:
													return sent_encoder(lowerCAmelCase__       ,		attention_mask=lowerCAmelCase__    )[1]
										else:
													# prepare implicit variables
													SCREAMING_SNAKE_CASE_						:			List[str]       						= input_ids.device
													SCREAMING_SNAKE_CASE_						:			Optional[Any]       						= input_ids.size()
													SCREAMING_SNAKE_CASE_						:			str       						= torch.zeros(lowerCAmelCase__       ,		dtype=torch.long       ,		device=lowerCAmelCase__    )
													SCREAMING_SNAKE_CASE_						:			Optional[Any]       						= [None] * sent_encoder.config.num_hidden_layers
													SCREAMING_SNAKE_CASE_						:			torch.Tensor       						= sent_encoder.get_extended_attention_mask(
													    lowerCAmelCase__       ,		lowerCAmelCase__    )

													# define function for checkpointing
													def partial_encode(*lowerCAmelCase__    ):
																SCREAMING_SNAKE_CASE_						:			List[str]       						= sent_encoder.encoder(
																    inputs[0]       ,		attention_mask=inputs[1]       ,		head_mask=lowerCAmelCase__       ,		)
																SCREAMING_SNAKE_CASE_						:			Tuple       						= encoder_outputs[0]
																SCREAMING_SNAKE_CASE_						:			List[str]       						= sent_encoder.pooler(lowerCAmelCase__    )
																return pooled_output

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






							def  UpperCamelCase__			(	self       ,		lowerCAmelCase__       ,		lowerCAmelCase__=None       ,		lowerCAmelCase__=-1       ,		):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_						:			Any       						= self.embed_sentences_checkpointed(
										    lowerCAmelCase__       ,		lowerCAmelCase__       ,		self.bert_query       ,		lowerCAmelCase__       ,		)
										return self.project_query(lowerCAmelCase__    )






							def  UpperCamelCase__			(	self       ,		lowerCAmelCase__       ,		lowerCAmelCase__=None       ,		lowerCAmelCase__=-1       ,		):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_						:			Union[str, Any]       						= self.embed_sentences_checkpointed(
										    lowerCAmelCase__       ,		lowerCAmelCase__       ,		self.bert_query if self.bert_doc is None else self.bert_doc       ,		lowerCAmelCase__       ,		)
										return self.project_doc(lowerCAmelCase__    )






							def  UpperCamelCase__			(	self       ,		lowerCAmelCase__       ,		lowerCAmelCase__       ,		lowerCAmelCase__       ,		lowerCAmelCase__       ,		lowerCAmelCase__ = -1       ,		):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_						:			Any       						= input_ids_query.device
										SCREAMING_SNAKE_CASE_						:			str       						= self.embed_questions(lowerCAmelCase__       ,		lowerCAmelCase__       ,		lowerCAmelCase__    )
										SCREAMING_SNAKE_CASE_						:			Tuple       						= self.embed_answers(lowerCAmelCase__       ,		lowerCAmelCase__       ,		lowerCAmelCase__    )
										SCREAMING_SNAKE_CASE_						:			Any       						= torch.mm(lowerCAmelCase__       ,		a_reps.t()    )
										SCREAMING_SNAKE_CASE_						:			List[Any]       						= self.ce_loss(lowerCAmelCase__       ,		torch.arange(compare_scores.shape[1]    ).to(lowerCAmelCase__    )    )
										SCREAMING_SNAKE_CASE_						:			Any       						= self.ce_loss(compare_scores.t()       ,		torch.arange(compare_scores.shape[0]    ).to(lowerCAmelCase__    )    )
										SCREAMING_SNAKE_CASE_						:			Tuple       						= (loss_qa + loss_aq) / 2
										return loss







from typing import Callable, List, Optional, Union

import torch

from ...models import UNetaDModel
from ...schedulers import CMStochasticIterativeScheduler
from ...utils import (
    is_accelerate_available,
    is_accelerate_version,
    logging,
    randn_tensor,
    replace_example_docstring,
)
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput


_A  					= logging.get_logger(__name__)  # pylint: disable=invalid-name


_A  					= "\n    Examples:\n        ```py\n        >>> import torch\n\n        >>> from diffusers import ConsistencyModelPipeline\n\n        >>> device = \"cuda\"\n        >>> # Load the cd_imagenet64_l2 checkpoint.\n        >>> model_id_or_path = \"openai/diffusers-cd_imagenet64_l2\"\n        >>> pipe = ConsistencyModelPipeline.from_pretrained(model_id_or_path, torch_dtype=torch.float16)\n        >>> pipe.to(device)\n\n        >>> # Onestep Sampling\n        >>> image = pipe(num_inference_steps=1).images[0]\n        >>> image.save(\"cd_imagenet64_l2_onestep_sample.png\")\n\n        >>> # Onestep sampling, class-conditional image generation\n        >>> # ImageNet-64 class label 145 corresponds to king penguins\n        >>> image = pipe(num_inference_steps=1, class_labels=145).images[0]\n        >>> image.save(\"cd_imagenet64_l2_onestep_sample_penguin.png\")\n\n        >>> # Multistep sampling, class-conditional image generation\n        >>> # Timesteps can be explicitly specified; the particular timesteps below are from the original Github repo:\n        >>> # https://github.com/openai/consistency_models/blob/main/scripts/launch.sh#L77\n        >>> image = pipe(num_inference_steps=None, timesteps=[22, 0], class_labels=145).images[0]\n        >>> image.save(\"cd_imagenet64_l2_multistep_sample_penguin.png\")\n        ```\n"







class      lowerCamelCase		(  A_  ):





     def __init__(self      :    Tuple				,    _A      :    UNetaDModel				,    _A      :    CMStochasticIterativeScheduler   )						->						None:
         super().__init__()

         self.register_modules(
             unet=_A				,    scheduler=_A				,    )

         snake_case      =					None





     def 		UpperCAmelCase(self      :    Dict				,    _A      :    List[str]=0   )						->						Any:
         if is_accelerate_available() and is_accelerate_version(">="				,    "0.14.0"   ):
             from accelerate import cpu_offload
         else:
             raise ImportError("`enable_sequential_cpu_offload` requires `accelerate v0.14.0` or higher"   )

         snake_case      =					torch.device(f'cuda:{gpu_id}'   )

         if self.device.type != "cpu":
             self.to("cpu"				,    silence_dtype_warnings=_A   )
             torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)

         for cpu_offloaded_model in [self.unet]:
             cpu_offload(_A				,    _A   )

         if self.safety_checker is not None:
             cpu_offload(self.safety_checker				,    execution_device=_A				,    offload_buffers=_A   )





     def 		UpperCAmelCase(self      :    str				,    _A      :    Optional[Any]=0   )						->						Union[str, Any]:
         if is_accelerate_available() and is_accelerate_version(">="				,    "0.17.0.dev0"   ):
             from accelerate import cpu_offload_with_hook
         else:
             raise ImportError("`enable_model_cpu_offload` requires `accelerate v0.17.0` or higher."   )

         snake_case      =					torch.device(f'cuda:{gpu_id}'   )

         if self.device.type != "cpu":
             self.to("cpu"				,    silence_dtype_warnings=_A   )
             torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)

         snake_case      =					None
         for cpu_offloaded_model in [self.unet]:
             snake_case     ,							snake_case      =					cpu_offload_with_hook(_A				,    _A				,    prev_module_hook=_A   )

         if self.safety_checker is not None:
             snake_case     ,							snake_case      =					cpu_offload_with_hook(self.safety_checker				,    _A				,    prev_module_hook=_A   )

         # We'll offload the last model manually.
         snake_case      =					hook





     @property
     # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
     def 		UpperCAmelCase(self      :    Dict   )						->						Any:
         if not hasattr(self.unet				,    "_hf_hook"   ):
             return self.device
         for module in self.unet.modules():
             if (
                 hasattr(_A				,    "_hf_hook"   )
                 and hasattr(module._hf_hook				,    "execution_device"   )
                 and module._hf_hook.execution_device is not None
             ):
                 return torch.device(module._hf_hook.execution_device   )
         return self.device





     def 		UpperCAmelCase(self      :    Dict				,    _A      :    List[str]				,    _A      :    Tuple				,    _A      :    List[str]				,    _A      :    List[Any]				,    _A      :    Dict				,    _A      :    List[str]				,    _A      :    List[str]				,    _A      :    Optional[Any]=None   )						->						Tuple:
         snake_case      =					(batch_size, num_channels, height, width)
         if isinstance(_A				,    _A   ) and len(_A   ) != batch_size:
             raise ValueError(
                 f'You have passed a list of generators of length {len(_A   )}, but requested an effective batch'
                 f' size of {batch_size}. Make sure the batch size matches the length of the generators.'   )

         if latents is None:
             snake_case      =					randn_tensor(_A				,    generator=_A				,    device=_A				,    dtype=_A   )
         else:
             snake_case      =					latents.to(device=_A				,    dtype=_A   )

         # scale the initial noise by the standard deviation required by the scheduler
         snake_case      =					latents * self.scheduler.init_noise_sigma
         return latents





     def 		UpperCAmelCase(self      :    List[Any]				,    _A      :    torch.FloatTensor				,    _A      :    str = "pil"   )						->						str:
         if output_type not in ["pt", "np", "pil"]:
             raise ValueError(
                 f'output_type={output_type} is not supported. Make sure to choose one of [\'pt\', \'np\', or \'pil\']'   )

         # Equivalent to diffusers.VaeImageProcessor.denormalize
         snake_case      =					(sample / 2 + 0.5).clamp(0				,    1   )
         if output_type == "pt":
             return sample

         # Equivalent to diffusers.VaeImageProcessor.pt_to_numpy
         snake_case      =					sample.cpu().permute(0				,    2				,    3				,    1   ).numpy()
         if output_type == "np":
             return sample

         # Output_type must be 'pil'
         snake_case      =					self.numpy_to_pil(_A   )
         return sample





     def 		UpperCAmelCase(self      :    str				,    _A      :    List[Any]				,    _A      :    Optional[Any]				,    _A      :    List[Any]=None   )						->						List[str]:
         if self.unet.config.num_class_embeds is not None:
             if isinstance(_A				,    _A   ):
                 snake_case      =					torch.tensor(_A				,    dtype=torch.int   )
             elif isinstance(_A				,    _A   ):
                 assert batch_size == 1, "Batch size must be 1 if classes is an int"
                 snake_case      =					torch.tensor([class_labels]				,    dtype=torch.int   )
             elif class_labels is None:
                 # Randomly generate batch_size class labels
                 # TODO: should use generator here? int analogue of randn_tensor is not exposed in ...utils
                 snake_case      =					torch.randint(0				,    self.unet.config.num_class_embeds				,    size=(batch_size,)   )
             snake_case      =					class_labels.to(_A   )
         else:
             snake_case      =					None
         return class_labels





     def 		UpperCAmelCase(self      :    Optional[Any]				,    _A      :    Tuple				,    _A      :    Union[str, Any]				,    _A      :    Dict				,    _A      :    Union[str, Any]				,    _A      :    Any				,    _A      :    str   )						->						Optional[Any]:
         if num_inference_steps is None and timesteps is None:
             raise ValueError("Exactly one of `num_inference_steps` or `timesteps` must be supplied."   )

         if num_inference_steps is not None and timesteps is not None:
             logger.warning(
                 f'Both `num_inference_steps`: {num_inference_steps} and `timesteps`: {timesteps} are supplied;'
                 " `timesteps` will be used over `num_inference_steps`."   )

         if latents is not None:
             snake_case      =					(batch_size, 3, img_size, img_size)
             if latents.shape != expected_shape:
                 raise ValueError(f'The shape of latents is {latents.shape} but is expected to be {expected_shape}.'   )

         if (callback_steps is None) or (
             callback_steps is not None and (not isinstance(_A				,    _A   ) or callback_steps <= 0)
         ):
             raise ValueError(
                 f'`callback_steps` has to be a positive integer but is {callback_steps} of type'
                 f' {type(_A   )}.'   )





     @torch.no_grad()
     @replace_example_docstring(_A   )
     def __call__(self      :    Union[str, Any]				,    _A      :    int = 1				,    _A      :    Optional[Union[torch.Tensor, List[int], int]] = None				,    _A      :    int = 1				,    _A      :    List[int] = None				,    _A      :    Optional[Union[torch.Generator, List[torch.Generator]]] = None				,    _A      :    Optional[torch.FloatTensor] = None				,    _A      :    Optional[str] = "pil"				,    _A      :    bool = True				,    _A      :    Optional[Callable[[int, int, torch.FloatTensor], None]] = None				,    _A      :    int = 1				,    )						->						List[str]:
         snake_case      =					self.unet.config.sample_size
         snake_case      =					self._execution_device

         # 1. Check inputs
         self.check_inputs(_A				,    _A				,    _A				,    _A				,    _A				,    _A   )

         # 2. Prepare image latents
         # Sample image latents x_0 ~ N(0, sigma_0^2 * I)
         snake_case      =					self.prepare_latents(
             batch_size=_A				,    num_channels=self.unet.config.in_channels				,    height=_A				,    width=_A				,    dtype=self.unet.dtype				,    device=_A				,    generator=_A				,    latents=_A				,    )

         # 3. Handle class_labels for class-conditional models
         snake_case      =					self.prepare_class_labels(_A				,    _A				,    class_labels=_A   )

         # 4. Prepare timesteps
         if timesteps is not None:
             self.scheduler.set_timesteps(timesteps=_A				,    device=_A   )
             snake_case      =					self.scheduler.timesteps
             snake_case      =					len(_A   )
         else:
             self.scheduler.set_timesteps(_A   )
             snake_case      =					self.scheduler.timesteps

         # 5. Denoising loop
         # Multistep sampling: implements Algorithm 1 in the paper
         with self.progress_bar(total=_A   ) as progress_bar:
             for i, t in enumerate(_A   ):
                 snake_case      =					self.scheduler.scale_model_input(_A				,    _A   )
                 snake_case      =					self.unet(_A				,    _A				,    class_labels=_A				,    return_dict=_A   )[0]

                 snake_case      =					self.scheduler.step(_A				,    _A				,    _A				,    generator=_A   )[0]

                 # call the callback, if provided
                 progress_bar.update()
                 if callback is not None and i % callback_steps == 0:
                     callback(_A				,    _A				,    _A   )

        # 6. Post-process image sample
         snake_case      =					self.postprocess_image(_A				,    output_type=_A   )

         # Offload last model to CPU
         if hasattr(self				,    "final_offload_hook"   ) and self.final_offload_hook is not None:
             self.final_offload_hook.offload()

         if not return_dict:
             return (image,)

         return ImagePipelineOutput(images=_A   )

