

"""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     :    Optional[int]    =		logging.get_logger(__name__)

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





class 			__UpperCamelCase (   _A						):
		SCREAMING_SNAKE_CASE								=    RetriBertConfig
		SCREAMING_SNAKE_CASE								=    None
		SCREAMING_SNAKE_CASE								=    "retribert"


		def 					SCREAMING_SNAKE_CASE__   (self :		Any     ,     __SCREAMING_SNAKE_CASE :		Optional[int]):
							if isinstance(__SCREAMING_SNAKE_CASE     ,     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(__SCREAMING_SNAKE_CASE     ,     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(__SCREAMING_SNAKE_CASE     ,     nn.LayerNorm):
												module.bias.data.zero_()
												module.weight.data.fill_(1.0)

__A     :    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."     ,		_A     ,		)
class 			__UpperCamelCase (   _A						):


		def __init__(self :		Optional[Any]     ,     __SCREAMING_SNAKE_CASE :		RetriBertConfig):
							super().__init__(__SCREAMING_SNAKE_CASE)
							A			=					config.projection_dim

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

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

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


		def 					SCREAMING_SNAKE_CASE__   (self :		Any     ,     __SCREAMING_SNAKE_CASE :		List[str]     ,     __SCREAMING_SNAKE_CASE :		Optional[int]     ,     __SCREAMING_SNAKE_CASE :		Any     ,     __SCREAMING_SNAKE_CASE :		int=-1     ,     ):
							# reproduces BERT forward pass with checkpointing
							if checkpoint_batch_size < 0 or input_ids.shape[0] < checkpoint_batch_size:
												return sent_encoder(__SCREAMING_SNAKE_CASE     ,     attention_mask=__SCREAMING_SNAKE_CASE)[1]
							else:
												# prepare implicit variables
												A			=					input_ids.device
												A			=					input_ids.size()
												A			=					torch.zeros(__SCREAMING_SNAKE_CASE     ,     dtype=torch.long     ,     device=__SCREAMING_SNAKE_CASE)
												A			=					[None] * sent_encoder.config.num_hidden_layers
												A			=					sent_encoder.get_extended_attention_mask(
												    __SCREAMING_SNAKE_CASE     ,     __SCREAMING_SNAKE_CASE)

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

												# run embedding layer on everything at once
												A			=					sent_encoder.embeddings(
												    input_ids=__SCREAMING_SNAKE_CASE     ,     position_ids=__SCREAMING_SNAKE_CASE     ,     token_type_ids=__SCREAMING_SNAKE_CASE     ,     inputs_embeds=__SCREAMING_SNAKE_CASE)
												# 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(__SCREAMING_SNAKE_CASE     ,     __SCREAMING_SNAKE_CASE     ,     __SCREAMING_SNAKE_CASE)
																	pooled_output_list.append(__SCREAMING_SNAKE_CASE)
												return torch.cat(__SCREAMING_SNAKE_CASE     ,     dim=0)


		def 					SCREAMING_SNAKE_CASE__   (self :		Optional[int]     ,     __SCREAMING_SNAKE_CASE :		List[str]     ,     __SCREAMING_SNAKE_CASE :		Optional[Any]=None     ,     __SCREAMING_SNAKE_CASE :		Any=-1     ,     ):
							A			=					self.embed_sentences_checkpointed(
							    __SCREAMING_SNAKE_CASE     ,     __SCREAMING_SNAKE_CASE     ,     self.bert_query     ,     __SCREAMING_SNAKE_CASE     ,     )
							return self.project_query(__SCREAMING_SNAKE_CASE)


		def 					SCREAMING_SNAKE_CASE__   (self :		List[Any]     ,     __SCREAMING_SNAKE_CASE :		int     ,     __SCREAMING_SNAKE_CASE :		Union[str, Any]=None     ,     __SCREAMING_SNAKE_CASE :		Tuple=-1     ,     ):
							A			=					self.embed_sentences_checkpointed(
							    __SCREAMING_SNAKE_CASE     ,     __SCREAMING_SNAKE_CASE     ,     self.bert_query if self.bert_doc is None else self.bert_doc     ,     __SCREAMING_SNAKE_CASE     ,     )
							return self.project_doc(__SCREAMING_SNAKE_CASE)


		def 					SCREAMING_SNAKE_CASE__   (self :		Union[str, Any]     ,     __SCREAMING_SNAKE_CASE :		torch.LongTensor     ,     __SCREAMING_SNAKE_CASE :		Optional[torch.FloatTensor]     ,     __SCREAMING_SNAKE_CASE :		torch.LongTensor     ,     __SCREAMING_SNAKE_CASE :		Optional[torch.FloatTensor]     ,     __SCREAMING_SNAKE_CASE :		int = -1     ,     ):
							A			=					input_ids_query.device
							A			=					self.embed_questions(__SCREAMING_SNAKE_CASE     ,     __SCREAMING_SNAKE_CASE     ,     __SCREAMING_SNAKE_CASE)
							A			=					self.embed_answers(__SCREAMING_SNAKE_CASE     ,     __SCREAMING_SNAKE_CASE     ,     __SCREAMING_SNAKE_CASE)
							A			=					torch.mm(__SCREAMING_SNAKE_CASE     ,     a_reps.t())
							A			=					self.ce_loss(__SCREAMING_SNAKE_CASE     ,     torch.arange(compare_scores.shape[1]).to(__SCREAMING_SNAKE_CASE))
							A			=					self.ce_loss(compare_scores.t()     ,     torch.arange(compare_scores.shape[0]).to(__SCREAMING_SNAKE_CASE))
							A			=					(loss_qa + loss_aq) / 2
							return loss








"""simple docstring"""




import argparse
import collections
import importlib.util
import os
import re


# All paths are set with the intent you should run this script from the root of the repo with the command
# python utils/check_table.py
__A     :    List[str]    =		'src/diffusers'
__A     :    Optional[int]    =		'docs/source/en'
__A     :    Union[str, Any]    =		'.'
def    __SCREAMING_SNAKE_CASE						(						lowercase__			,							lowercase__			,							lowercase__      ):
					"""simple docstring"""
					with open(lowercase__			,							"r"			,							encoding="utf-8"			,							newline="\n"      ) as f:
										A			=					f.readlines()
					# Find the start prompt.
					A			=					0
					while not lines[start_index].startswith(lowercase__      ):
										start_index += 1
					start_index += 1

					A			=					start_index
					while not lines[end_index].startswith(lowercase__      ):
										end_index += 1
					end_index -= 1

					while len(lines[start_index]      ) <= 1:
										start_index += 1
					while len(lines[end_index]      ) <= 1:
										end_index -= 1
					end_index += 1
					return "".join(lines[start_index:end_index]      ), start_index, end_index, lines


# Add here suffixes that are used to identify models, separated by |
__A     :    Any    =		'Model|Encoder|Decoder|ForConditionalGeneration'
# Regexes that match TF/Flax/PT model names.
__A     :    Union[str, Any]    =		re.compile(R'TF(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration)')
__A     :    int    =		re.compile(R'Flax(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration)')
# Will match any TF or Flax model too so need to be in an else branch afterthe two previous regexes.
__A     :    Any    =		re.compile(R'(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration)')


# This is to make sure the diffusers module imported is the one in the repo.
__A     :    Optional[int]    =		importlib.util.spec_from_file_location(
    'diffusers',
    os.path.join(TRANSFORMERS_PATH, '__init__.py'),
    submodule_search_locations=[TRANSFORMERS_PATH],
)
__A     :    Tuple    =		spec.loader.load_module()
def    __SCREAMING_SNAKE_CASE						(						lowercase__      ):
					"""simple docstring"""
					A			=					re.finditer(".+?(?:(?<=[a-z])(?=[A-Z])|(?<=[A-Z])(?=[A-Z][a-z])|$)"			,							lowercase__      )
					return [m.group(0      ) for m in matches]
def    __SCREAMING_SNAKE_CASE						(						lowercase__			,							lowercase__      ):
					"""simple docstring"""
					A			=					2 if text == "✅" or text == "❌" else len(lowercase__      )
					A			=					(width - text_length) // 2
					A			=					width - text_length - left_indent
					return " " * left_indent + text + " " * right_indent
def    __SCREAMING_SNAKE_CASE						(						):
					"""simple docstring"""
					A			=					diffusers_module.models.auto.configuration_auto.CONFIG_MAPPING_NAMES
					A			=					{
					    name: config_mapping_names[code]
					    for code, name in diffusers_module.MODEL_NAMES_MAPPING.items()
					    if code in config_mapping_names
					}
					A			=					{name: config.replace("ConfigMixin"			,							""      ) for name, config in model_name_to_config.items()}

					# Dictionaries flagging if each model prefix has a slow/fast tokenizer, backend in PT/TF/Flax.
					A			=					collections.defaultdict(lowercase__      )
					A			=					collections.defaultdict(lowercase__      )
					A			=					collections.defaultdict(lowercase__      )
					A			=					collections.defaultdict(lowercase__      )
					A			=					collections.defaultdict(lowercase__      )

					# Let's lookup through all diffusers object (once).
					for attr_name in dir(lowercase__      ):
										A			=					None
										if attr_name.endswith("Tokenizer"      ):
															A			=					slow_tokenizers
															A			=					attr_name[:-9]
										elif attr_name.endswith("TokenizerFast"      ):
															A			=					fast_tokenizers
															A			=					attr_name[:-13]
										elif _re_tf_models.match(lowercase__      ) is not None:
															A			=					tf_models
															A			=					_re_tf_models.match(lowercase__      ).groups()[0]
										elif _re_flax_models.match(lowercase__      ) is not None:
															A			=					flax_models
															A			=					_re_flax_models.match(lowercase__      ).groups()[0]
										elif _re_pt_models.match(lowercase__      ) is not None:
															A			=					pt_models
															A			=					_re_pt_models.match(lowercase__      ).groups()[0]

										if lookup_dict is not None:
															while len(lowercase__      ) > 0:
																				if attr_name in model_name_to_prefix.values():
																									A			=					True
																									break
																				# Try again after removing the last word in the name
																				A			=					"".join(camel_case_split(lowercase__      )[:-1]      )

    # Let's build that table!
					A			=					list(model_name_to_config.keys()      )
					model_names.sort(key=str.lower      )
					A			=					["Model", "Tokenizer slow", "Tokenizer fast", "PyTorch support", "TensorFlow support", "Flax Support"]
					# We'll need widths to properly display everything in the center (+2 is to leave one extra space on each side).
					A			=					[len(lowercase__      ) + 2 for c in columns]
					A			=					max([len(lowercase__      ) for name in model_names]      ) + 2

					# Build the table per se
					A			=					"|" + "|".join([_center_text(lowercase__			,							lowercase__      ) for c, w in zip(lowercase__			,							lowercase__      )]      ) + "|\n"
					# Use ":-----:" format to center-aligned table cell texts
					table += "|" + "|".join([":" + "-" * (w - 2) + ":" for w in widths]      ) + "|\n"

					A			=					{True: "✅", False: "❌"}
					for name in model_names:
										A			=					model_name_to_prefix[name]
										A			=					[
										    name,
										    check[slow_tokenizers[prefix]],
										    check[fast_tokenizers[prefix]],
										    check[pt_models[prefix]],
										    check[tf_models[prefix]],
										    check[flax_models[prefix]],
										]
										table += "|" + "|".join([_center_text(lowercase__			,							lowercase__      ) for l, w in zip(lowercase__			,							lowercase__      )]      ) + "|\n"
					return table
def    __SCREAMING_SNAKE_CASE						(						lowercase__=False      ):
					"""simple docstring"""
					A			,    A			,    A			,    A			=					_find_text_in_file(
					    filename=os.path.join(lowercase__			,							"index.mdx"      )			,							start_prompt="<!--This table is updated automatically from the auto modules"			,							end_prompt="<!-- End table-->"			,							)
					A			=					get_model_table_from_auto_modules()

					if current_table != new_table:
										if overwrite:
															with open(os.path.join(lowercase__			,							"index.mdx"      )			,							"w"			,							encoding="utf-8"			,							newline="\n"      ) as f:
																				f.writelines(lines[:start_index] + [new_table] + lines[end_index:]      )
										else:
															raise ValueError(
															    "The model table in the `index.mdx` has not been updated. Run `make fix-copies` to fix this."      )


if __name__ == "__main__":
	__A     :    Any    =		argparse.ArgumentParser()
	parser.add_argument('--fix_and_overwrite', action='store_true', help='Whether to fix inconsistencies.')
	__A     :    str    =		parser.parse_args()

	check_model_table(args.fix_and_overwrite)





