

"""simple docstring"""

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


a__ :	List[str]						= logging.get_logger(__name__)

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


class 			_snake_case     (						lowerCAmelCase__       ):
      __snake_case 			=     RetriBertConfig
      __snake_case 			=     None
      __snake_case 			=     'retribert'







      def        _UpperCamelCase    (self     ,							UpperCAmelCase_							)			->						Tuple:

           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							)



a__ :	Optional[Any]						= R"""

    This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
    library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
    etc.)

    This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
    Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
    and behavior.

    Parameters:
        config ([`RetriBertConfig`]): Model configuration class with all the parameters of the model.
            Initializing with a config file does not load the weights associated with the model, only the
            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
"""


@add_start_docstrings(
    'Bert Based model to embed queries or document for document retrieval.'						,			lowerCAmelCase__						,			)
class 			_snake_case     (						lowerCAmelCase__       ):







      def __init__(self     ,							UpperCAmelCase_							)			->						Union[str, Any]:
           super().__init__(_lowerCamelCase							)
           __lowercase												=			config.projection_dim

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

           __lowercase												=			nn.CrossEntropyLoss(reduction='''mean'''							)

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







      def        _UpperCamelCase    (self     ,							UpperCAmelCase_     ,							UpperCAmelCase_     ,							UpperCAmelCase_     ,							UpperCAmelCase_=-1     ,							)			->						List[str]:

           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
                __lowercase												=			input_ids.device
                __lowercase												=			input_ids.size()
                __lowercase												=			torch.zeros(_lowerCamelCase     ,							dtype=torch.long     ,							device=_lowerCamelCase							)
                __lowercase												=			[None] * sent_encoder.config.num_hidden_layers
                __lowercase												=			sent_encoder.get_extended_attention_mask(
                    _lowerCamelCase     ,							_lowerCamelCase							)

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

                # run embedding layer on everything at once
                __lowercase												=			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
                __lowercase												=			[]
                for b in range(math.ceil(input_ids.shape[0] / checkpoint_batch_size							)							):
                     __lowercase												=			embedding_output[b * checkpoint_batch_size : (b + 1) * checkpoint_batch_size]
                     __lowercase												=			extended_attention_mask[b * checkpoint_batch_size : (b + 1) * checkpoint_batch_size]
                     __lowercase												=			checkpoint.checkpoint(_lowerCamelCase     ,							_lowerCamelCase     ,							_lowerCamelCase							)
                     pooled_output_list.append(_lowerCamelCase							)
                return torch.cat(_lowerCamelCase     ,							dim=0							)







      def        _UpperCamelCase    (self     ,							UpperCAmelCase_     ,							UpperCAmelCase_=None     ,							UpperCAmelCase_=-1     ,							)			->						List[Any]:
           __lowercase												=			self.embed_sentences_checkpointed(
               _lowerCamelCase     ,							_lowerCamelCase     ,							self.bert_query     ,							_lowerCamelCase     ,							)
           return self.project_query(_lowerCamelCase							)







      def        _UpperCamelCase    (self     ,							UpperCAmelCase_     ,							UpperCAmelCase_=None     ,							UpperCAmelCase_=-1     ,							)			->						Optional[int]:
           __lowercase												=			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     ,							UpperCAmelCase_     ,							UpperCAmelCase_     ,							UpperCAmelCase_     ,							UpperCAmelCase_     ,							UpperCAmelCase_ = -1     ,							)			->						List[Any]:
           __lowercase												=			input_ids_query.device
           __lowercase												=			self.embed_questions(_lowerCamelCase     ,							_lowerCamelCase     ,							_lowerCamelCase							)
           __lowercase												=			self.embed_answers(_lowerCamelCase     ,							_lowerCamelCase     ,							_lowerCamelCase							)
           __lowercase												=			torch.mm(_lowerCamelCase     ,							a_reps.t()							)
           __lowercase												=			self.ce_loss(_lowerCamelCase     ,							torch.arange(compare_scores.shape[1]							).to(_lowerCamelCase							)							)
           __lowercase												=			self.ce_loss(compare_scores.t()     ,							torch.arange(compare_scores.shape[0]							).to(_lowerCamelCase							)							)
           __lowercase												=			(loss_qa + loss_aq) / 2
           return loss





"""simple docstring"""

import glob
import os
from typing import Dict, List, Union

import torch

from diffusers.utils import is_safetensors_available


if is_safetensors_available():
							import safetensors.torch

from huggingface_hub import snapshot_download

from diffusers import DiffusionPipeline, __version__
from diffusers.schedulers.scheduling_utils import SCHEDULER_CONFIG_NAME
from diffusers.utils import CONFIG_NAME, DIFFUSERS_CACHE, ONNX_WEIGHTS_NAME, WEIGHTS_NAME


class 			_snake_case     (						UpperCamelCase       ):







						def __init__(self							)			->						Optional[Any]:
											self.register_to_config()
											super().__init__()







						def        _UpperCamelCase    (self     ,							UpperCAmelCase_     ,							UpperCAmelCase_							)			->						List[Any]:
											if dicta == dicta:
																return True
											else:
																__lowercase						, __lowercase												=			self._remove_meta_keys(UpperCAmelCase_							)
																__lowercase						, __lowercase												=			self._remove_meta_keys(UpperCAmelCase_							)
																if configa == configa:
																					print(f'''Warning !: Mismatch in keys {meta_keysa} and {meta_keysa}.'''							)
																					return True
											return False







						def        _UpperCamelCase    (self     ,							UpperCAmelCase_							)			->						str:
											__lowercase												=			[]
											__lowercase												=			config_dict.copy()
											for key in config_dict.keys():
																if key.startswith('''_'''							):
																					temp_dict.pop(UpperCAmelCase_							)
																					meta_keys.append(UpperCAmelCase_							)
											return (temp_dict, meta_keys)







						@torch.no_grad()
						def        _UpperCamelCase    (self     ,							UpperCAmelCase_     ,							**UpperCAmelCase_							)			->						Optional[Any]:
											__lowercase												=			kwargs.pop('''cache_dir'''     ,							UpperCAmelCase_							)
											__lowercase												=			kwargs.pop('''resume_download'''     ,							UpperCAmelCase_							)
											__lowercase												=			kwargs.pop('''force_download'''     ,							UpperCAmelCase_							)
											__lowercase												=			kwargs.pop('''proxies'''     ,							UpperCAmelCase_							)
											__lowercase												=			kwargs.pop('''local_files_only'''     ,							UpperCAmelCase_							)
											__lowercase												=			kwargs.pop('''use_auth_token'''     ,							UpperCAmelCase_							)
											__lowercase												=			kwargs.pop('''revision'''     ,							UpperCAmelCase_							)
											__lowercase												=			kwargs.pop('''torch_dtype'''     ,							UpperCAmelCase_							)
											__lowercase												=			kwargs.pop('''device_map'''     ,							UpperCAmelCase_							)

											__lowercase												=			kwargs.pop('''alpha'''     ,							0.5							)
											__lowercase												=			kwargs.pop('''interp'''     ,							UpperCAmelCase_							)

											print('''Received list'''     ,							UpperCAmelCase_							)
											print(f'''Combining with alpha={alpha}, interpolation mode={interp}'''							)

											__lowercase												=			len(UpperCAmelCase_							)
											# Ignore result from model_index_json comparision of the two checkpoints
											__lowercase												=			kwargs.pop('''force'''     ,							UpperCAmelCase_							)

											# If less than 2 checkpoints, nothing to merge. If more than 3, not supported for now.
											if checkpoint_count > 3 or checkpoint_count < 2:
																raise ValueError(
																    '''Received incorrect number of checkpoints to merge. Ensure that either 2 or 3 checkpoints are being'''
																    ''' passed.'''							)

											print('''Received the right number of checkpoints'''							)
											# chkpt0, chkpt1 = pretrained_model_name_or_path_list[0:2]
											# chkpt2 = pretrained_model_name_or_path_list[2] if checkpoint_count == 3 else None

											# Validate that the checkpoints can be merged
											# Step 1: Load the model config and compare the checkpoints. We'll compare the model_index.json first while ignoring the keys starting with '_'
											__lowercase												=			[]
											for pretrained_model_name_or_path in pretrained_model_name_or_path_list:
																__lowercase												=			DiffusionPipeline.load_config(
																    UpperCAmelCase_     ,							cache_dir=UpperCAmelCase_     ,							resume_download=UpperCAmelCase_     ,							force_download=UpperCAmelCase_     ,							proxies=UpperCAmelCase_     ,							local_files_only=UpperCAmelCase_     ,							use_auth_token=UpperCAmelCase_     ,							revision=UpperCAmelCase_     ,							)
																config_dicts.append(UpperCAmelCase_							)

											__lowercase												=			True
											for idx in range(1     ,							len(UpperCAmelCase_							)							):
																comparison_result &= self._compare_model_configs(config_dicts[idx - 1]     ,							config_dicts[idx]							)
																if not force and comparison_result is False:
																					raise ValueError('''Incompatible checkpoints. Please check model_index.json for the models.'''							)
																					print(config_dicts[0]     ,							config_dicts[1]							)
											print('''Compatible model_index.json files found'''							)
											# Step 2: Basic Validation has succeeded. Let's download the models and save them into our local files.
											__lowercase												=			[]
											for pretrained_model_name_or_path, config_dict in zip(UpperCAmelCase_     ,							UpperCAmelCase_							):
																__lowercase												=			[k for k in config_dict.keys() if not k.startswith('''_'''							)]
																__lowercase												=			[os.path.join(UpperCAmelCase_     ,							'''*'''							) for k in folder_names]
																allow_patterns += [
																    WEIGHTS_NAME,
																    SCHEDULER_CONFIG_NAME,
																    CONFIG_NAME,
																    ONNX_WEIGHTS_NAME,
																    DiffusionPipeline.config_name,
																]
																__lowercase												=			config_dict.get('''_class_name'''							)
																__lowercase												=			{'''diffusers''': __version__, '''pipeline_class''': requested_pipeline_class}

																__lowercase												=			(
																    pretrained_model_name_or_path
																    if os.path.isdir(UpperCAmelCase_							)
																    else snapshot_download(
																        UpperCAmelCase_     ,							cache_dir=UpperCAmelCase_     ,							resume_download=UpperCAmelCase_     ,							proxies=UpperCAmelCase_     ,							local_files_only=UpperCAmelCase_     ,							use_auth_token=UpperCAmelCase_     ,							revision=UpperCAmelCase_     ,							allow_patterns=UpperCAmelCase_     ,							user_agent=UpperCAmelCase_     ,							)
																)
																print('''Cached Folder'''     ,							UpperCAmelCase_							)
																cached_folders.append(UpperCAmelCase_							)

											# Step 3:-
											# Load the first checkpoint as a diffusion pipeline and modify its module state_dict in place
											__lowercase												=			DiffusionPipeline.from_pretrained(
											    cached_folders[0]     ,							torch_dtype=UpperCAmelCase_     ,							device_map=UpperCAmelCase_							)
											final_pipe.to(self.device							)

											__lowercase												=			None
											if len(UpperCAmelCase_							) > 2:
																__lowercase												=			os.path.join(cached_folders[2]							)

											if interp == "sigmoid":
																__lowercase												=			CheckpointMergerPipeline.sigmoid
											elif interp == "inv_sigmoid":
																__lowercase												=			CheckpointMergerPipeline.inv_sigmoid
											elif interp == "add_diff":
																__lowercase												=			CheckpointMergerPipeline.add_difference
											else:
																__lowercase												=			CheckpointMergerPipeline.weighted_sum

											# Find each module's state dict.
											for attr in final_pipe.config.keys():
																if not attr.startswith('''_'''							):
																					__lowercase												=			os.path.join(cached_folders[1]     ,							UpperCAmelCase_							)
																					if os.path.exists(UpperCAmelCase_							):
																										__lowercase												=			[
																										    *glob.glob(os.path.join(UpperCAmelCase_     ,							'''*.safetensors'''							)							),
																										    *glob.glob(os.path.join(UpperCAmelCase_     ,							'''*.bin'''							)							),
																										]
																										__lowercase												=			files[0] if len(UpperCAmelCase_							) > 0 else None
																					if len(UpperCAmelCase_							) < 3:
																										__lowercase												=			None
																					else:
																										__lowercase												=			os.path.join(cached_folders[2]     ,							UpperCAmelCase_							)
																										if os.path.exists(UpperCAmelCase_							):
																															__lowercase												=			[
																															    *glob.glob(os.path.join(UpperCAmelCase_     ,							'''*.safetensors'''							)							),
																															    *glob.glob(os.path.join(UpperCAmelCase_     ,							'''*.bin'''							)							),
																															]
																															__lowercase												=			files[0] if len(UpperCAmelCase_							) > 0 else None
                # For an attr if both checkpoint_path_1 and 2 are None, ignore.
                # If atleast one is present, deal with it according to interp method, of course only if the state_dict keys match.
																					if checkpoint_path_a is None and checkpoint_path_a is None:
																										print(f'''Skipping {attr}: not present in 2nd or 3d model'''							)
																										continue
																					try:
																										__lowercase												=			getattr(UpperCAmelCase_     ,							UpperCAmelCase_							)
																										if isinstance(UpperCAmelCase_     ,							UpperCAmelCase_							):  # ignore requires_safety_checker boolean
																															continue
																										__lowercase												=			getattr(UpperCAmelCase_     ,							'''state_dict'''							)
																										__lowercase												=			theta_a()

																										__lowercase												=			getattr(UpperCAmelCase_     ,							'''load_state_dict'''							)
																										__lowercase												=			(
																										    safetensors.torch.load_file(UpperCAmelCase_							)
																										    if (is_safetensors_available() and checkpoint_path_a.endswith('''.safetensors'''							))
																										    else torch.load(UpperCAmelCase_     ,							map_location='''cpu'''							)
																										)
																										__lowercase												=			None
																										if checkpoint_path_a:
																															__lowercase												=			(
																															    safetensors.torch.load_file(UpperCAmelCase_							)
																															    if (is_safetensors_available() and checkpoint_path_a.endswith('''.safetensors'''							))
																															    else torch.load(UpperCAmelCase_     ,							map_location='''cpu'''							)
																															)

																										if not theta_a.keys() == theta_a.keys():
																															print(f'''Skipping {attr}: key mismatch'''							)
																															continue
																										if theta_a and not theta_a.keys() == theta_a.keys():
																															print(f'''Skipping {attr}:y mismatch'''							)
																					except Exception as e:
																										print(f'''Skipping {attr} do to an unexpected error: {str(UpperCAmelCase_							)}'''							)
																										continue
																					print(f'''MERGING {attr}'''							)

																					for key in theta_a.keys():
																										if theta_a:
																															__lowercase												=			theta_func(theta_a[key]     ,							theta_a[key]     ,							theta_a[key]     ,							UpperCAmelCase_							)
																										else:
																															__lowercase												=			theta_func(theta_a[key]     ,							theta_a[key]     ,							UpperCAmelCase_     ,							UpperCAmelCase_							)

																					del theta_a
																					del theta_a
																					update_theta_a(UpperCAmelCase_							)

																					del theta_a
											return final_pipe







						@staticmethod
						def        _UpperCamelCase    (UpperCAmelCase_     ,							UpperCAmelCase_     ,							UpperCAmelCase_     ,							UpperCAmelCase_							)			->						List[Any]:
											return ((1 - alpha) * thetaa) + (alpha * thetaa)







						@staticmethod
						def        _UpperCamelCase    (UpperCAmelCase_     ,							UpperCAmelCase_     ,							UpperCAmelCase_     ,							UpperCAmelCase_							)			->						List[Any]:
											__lowercase												=			alpha * alpha * (3 - (2 * alpha))
											return thetaa + ((thetaa - thetaa) * alpha)







						@staticmethod
						def        _UpperCamelCase    (UpperCAmelCase_     ,							UpperCAmelCase_     ,							UpperCAmelCase_     ,							UpperCAmelCase_							)			->						List[Any]:
											import math

											__lowercase												=			0.5 - math.sin(math.asin(1.0 - 2.0 * alpha							) / 3.0							)
											return thetaa + ((thetaa - thetaa) * alpha)







						@staticmethod
						def        _UpperCamelCase    (UpperCAmelCase_     ,							UpperCAmelCase_     ,							UpperCAmelCase_     ,							UpperCAmelCase_							)			->						List[str]:
											return thetaa + (thetaa - thetaa) * (1.0 - alpha)


