"""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


UpperCamelCase_							:							Dict				 =		logging.get_logger(__name__)

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


class       _lowercase		( __lowerCamelCase	):
				lowerCamelCase__							=							RetriBertConfig
				lowerCamelCase__							=							None
				lowerCamelCase__							=							'retribert'

				def 							A			(  self							:     List[str]							,     _lowercase							:     Union[str, Any]  ):


										'''simple docstring'''
										if isinstance(_lowercase							,     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(_lowercase							,     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(_lowercase							,     nn.LayerNorm  ):
																module.bias.data.zero_()
																module.weight.data.fill_(1.0  )
UpperCamelCase_							:							List[str]				 =		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       _lowercase		( __lowerCamelCase	):

				def __init__(  self							:     Optional[int]							,     _lowercase							:     RetriBertConfig  ):


										'''simple docstring'''
										super().__init__(_lowercase  )
										UpperCAmelCase__       						=		config.projection_dim

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

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

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

				def 							A			(  self							:     Dict							,     _lowercase							:     Tuple							,     _lowercase							:     List[str]							,     _lowercase							:     Any							,     _lowercase							:     str=-1							,     ):


										'''simple docstring'''
										# reproduces BERT forward pass with checkpointing
										if checkpoint_batch_size < 0 or input_ids.shape[0] < checkpoint_batch_size:
																return sent_encoder(_lowercase							,     attention_mask=_lowercase  )[1]
										else:
																# prepare implicit variables
																UpperCAmelCase__       						=		input_ids.device
																UpperCAmelCase__       						=		input_ids.size()
																UpperCAmelCase__       						=		torch.zeros(_lowercase							,     dtype=torch.long							,     device=_lowercase  )
																UpperCAmelCase__       						=		[None] * sent_encoder.config.num_hidden_layers
																UpperCAmelCase__       						=		sent_encoder.get_extended_attention_mask(
																    _lowercase							,     _lowercase  )

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

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

				def 							A			(  self							:     List[str]							,     _lowercase							:     str							,     _lowercase							:     Optional[Any]=None							,     _lowercase							:     List[str]=-1							,     ):


										'''simple docstring'''
										UpperCAmelCase__       						=		self.embed_sentences_checkpointed(
										    _lowercase							,     _lowercase							,     self.bert_query							,     _lowercase							,     )
										return self.project_query(_lowercase  )

				def 							A			(  self							:     str							,     _lowercase							:     List[str]							,     _lowercase							:     int=None							,     _lowercase							:     Optional[int]=-1							,     ):


										'''simple docstring'''
										UpperCAmelCase__       						=		self.embed_sentences_checkpointed(
										    _lowercase							,     _lowercase							,     self.bert_query if self.bert_doc is None else self.bert_doc							,     _lowercase							,     )
										return self.project_doc(_lowercase  )

				def 							A			(  self							:     List[Any]							,     _lowercase							:     torch.LongTensor							,     _lowercase							:     Optional[torch.FloatTensor]							,     _lowercase							:     torch.LongTensor							,     _lowercase							:     Optional[torch.FloatTensor]							,     _lowercase							:     int = -1							,     ):


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


"""simple docstring"""







import os
import warnings
from pathlib import Path
from typing import Optional, Union

from .dynamic_module_utils import custom_object_save
from .tokenization_utils_base import PreTrainedTokenizerBase
from .utils import PushToHubMixin, copy_func, direct_transformers_import, logging


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

# Dynamically import the Transformers module to grab the attribute classes of the processor form their names.
UpperCamelCase_							:							Optional[Any]				 =		direct_transformers_import(Path(__file__).parent)


UpperCamelCase_							:							Optional[int]				 =		{
    """AutoTokenizer""": """PreTrainedTokenizerBase""",
    """AutoFeatureExtractor""": """FeatureExtractionMixin""",
    """AutoImageProcessor""": """ImageProcessingMixin""",
}


class       _lowercase		( __lowerCamelCase	):
				lowerCamelCase__							=							['feature_extractor', 'tokenizer']
				# Names need to be attr_class for attr in attributes
				lowerCamelCase__							=							None
				lowerCamelCase__							=							None
				lowerCamelCase__							=							None

				def __init__(  self							:     Any							,     *_lowercase							:     Optional[int]							,     **_lowercase							:     int  ):


										'''simple docstring'''
										# Sanitize args and kwargs
										for key in kwargs:
																if key not in self.attributes:
																						raise TypeError(f'Unexpected keyword argument {key}.'  )
										for arg, attribute_name in zip(_lowercase							,     self.attributes  ):
																if attribute_name in kwargs:
																						raise TypeError(f'Got multiple values for argument {attribute_name}.'  )
																else:
																						UpperCAmelCase__       						=		arg

										if len(_lowercase  ) != len(self.attributes  ):
																raise ValueError(
																    f'This processor requires {len(self.attributes  )} arguments: {", ".join(self.attributes  )}. Got '
																    f'{len(_lowercase  )} arguments instead.'  )

										# Check each arg is of the proper class (this will also catch a user initializing in the wrong order)
										for attribute_name, arg in kwargs.items():
																UpperCAmelCase__       						=		getattr(self							,     f'{attribute_name}_class'  )
																# Nothing is ever going to be an instance of "AutoXxx", in that case we check the base class.
																UpperCAmelCase__       						=		AUTO_TO_BASE_CLASS_MAPPING.get(_lowercase							,     _lowercase  )
																if isinstance(_lowercase							,     _lowercase  ):
																						UpperCAmelCase__       						=		tuple(getattr(_lowercase							,     _lowercase  ) for n in class_name if n is not None  )
																else:
																						UpperCAmelCase__       						=		getattr(_lowercase							,     _lowercase  )

																if not isinstance(_lowercase							,     _lowercase  ):
																						raise ValueError(
																						    f'Received a {type(_lowercase  ).__name__} for argument {attribute_name}, but a {class_name} was expected.'  )

																setattr(self							,     _lowercase							,     _lowercase  )

				def __repr__(  self							:     Tuple  ):


										'''simple docstring'''
										UpperCAmelCase__       						=		[f'- {name}: {repr(getattr(self							,     _lowercase  )  )}' for name in self.attributes]
										UpperCAmelCase__       						=		"""\n""".join(_lowercase  )
										return f'{self.__class__.__name__}:\n{attributes_repr}'

				def 							A			(  self							:     Tuple							,     _lowercase							:     Dict							,     _lowercase							:     bool = False							,     **_lowercase							:     List[str]  ):


										'''simple docstring'''
										os.makedirs(_lowercase							,     exist_ok=_lowercase  )

										if push_to_hub:
																UpperCAmelCase__       						=		kwargs.pop("""commit_message"""							,     _lowercase  )
																UpperCAmelCase__       						=		kwargs.pop("""repo_id"""							,     save_directory.split(os.path.sep  )[-1]  )
																UpperCAmelCase__       						=		self._create_repo(_lowercase							,     **_lowercase  )
																UpperCAmelCase__       						=		self._get_files_timestamps(_lowercase  )
										# If we have a custom config, we copy the file defining it in the folder and set the attributes so it can be
										# loaded from the Hub.
										if self._auto_class is not None:
																UpperCAmelCase__       						=		[getattr(self							,     _lowercase  ) for attribute_name in self.attributes]
																UpperCAmelCase__       						=		[(a.init_kwargs if isinstance(_lowercase							,     _lowercase  ) else a) for a in attrs]
																custom_object_save(self							,     _lowercase							,     config=_lowercase  )

										for attribute_name in self.attributes:
																UpperCAmelCase__       						=		getattr(self							,     _lowercase  )
																# Include the processor class in the attribute config so this processor can then be reloaded with the
																# `AutoProcessor` API.
																if hasattr(_lowercase							,     """_set_processor_class"""  ):
																						attribute._set_processor_class(self.__class__.__name__  )
																attribute.save_pretrained(_lowercase  )

										if self._auto_class is not None:
																# We added an attribute to the init_kwargs of the tokenizers, which needs to be cleaned up.
																for attribute_name in self.attributes:
																						UpperCAmelCase__       						=		getattr(self							,     _lowercase  )
																						if isinstance(_lowercase							,     _lowercase  ):
																												del attribute.init_kwargs["auto_map"]

										if push_to_hub:
																self._upload_modified_files(
																    _lowercase							,     _lowercase							,     _lowercase							,     commit_message=_lowercase							,     token=kwargs.get("""use_auth_token"""  )							,     )

				@classmethod
				def 							A			(  cls							:     Dict							,     _lowercase							:     Union[str, os.PathLike]							,     _lowercase							:     Optional[Union[str, os.PathLike]] = None							,     _lowercase							:     bool = False							,     _lowercase							:     bool = False							,     _lowercase							:     Optional[Union[str, bool]] = None							,     _lowercase							:     str = "main"							,     **_lowercase							:     Dict							,     ):


										'''simple docstring'''
										UpperCAmelCase__       						=		cache_dir
										UpperCAmelCase__       						=		force_download
										UpperCAmelCase__       						=		local_files_only
										UpperCAmelCase__       						=		revision

										UpperCAmelCase__       						=		kwargs.pop("""use_auth_token"""							,     _lowercase  )
										if use_auth_token is not None:
																warnings.warn(
																    """The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers."""							,     _lowercase  )
																if token is not None:
																						raise ValueError(
																						    """`token` and `use_auth_token` are both specified. Please set only the argument `token`."""  )
																UpperCAmelCase__       						=		use_auth_token

										if token is not None:
																# change to `token` in a follow-up PR
																UpperCAmelCase__       						=		token

										UpperCAmelCase__       						=		cls._get_arguments_from_pretrained(_lowercase							,     **_lowercase  )
										return cls(*_lowercase  )

				@classmethod
				def 							A			(  cls							:     List[Any]							,     _lowercase							:     Dict="AutoProcessor"  ):


										'''simple docstring'''
										if not isinstance(_lowercase							,     _lowercase  ):
																UpperCAmelCase__       						=		auto_class.__name__

										import transformers.models.auto as auto_module

										if not hasattr(_lowercase							,     _lowercase  ):
																raise ValueError(f'{auto_class} is not a valid auto class.'  )

										UpperCAmelCase__       						=		auto_class

				@classmethod
				def 							A			(  cls							:     Any							,     _lowercase							:     Tuple							,     **_lowercase							:     List[Any]  ):


										'''simple docstring'''
										UpperCAmelCase__       						=		[]
										for attribute_name in cls.attributes:
																UpperCAmelCase__       						=		getattr(cls							,     f'{attribute_name}_class'  )
																if isinstance(_lowercase							,     _lowercase  ):
																						UpperCAmelCase__       						=		tuple(getattr(_lowercase							,     _lowercase  ) if n is not None else None for n in class_name  )
																						UpperCAmelCase__       						=		kwargs.get("""use_fast"""							,     _lowercase  )
																						if use_fast and classes[1] is not None:
																												UpperCAmelCase__       						=		classes[1]
																						else:
																												UpperCAmelCase__       						=		classes[0]
																else:
																						UpperCAmelCase__       						=		getattr(_lowercase							,     _lowercase  )

																args.append(attribute_class.from_pretrained(_lowercase							,     **_lowercase  )  )
										return args

				@property
				def 							A			(  self							:     int  ):


										'''simple docstring'''
										UpperCAmelCase__       						=		getattr(self							,     self.attributes[0]  )
										return getattr(_lowercase							,     """model_input_names"""							,     _lowercase  )
UpperCamelCase_							:							List[Any]				 =		copy_func(ProcessorMixin.push_to_hub)
if ProcessorMixin.push_to_hub.__doc__ is not None:
						UpperCamelCase_							:							Dict				 =		ProcessorMixin.push_to_hub.__doc__.format(
						    object="""processor""", object_class="""AutoProcessor""", object_files="""processor files"""
						)

