
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_  :       List[str]					      =     [
    """yjernite/retribert-base-uncased""",
    # See all RetriBert models at https://huggingface.co/models?filter=retribert
]

class 				_UpperCamelCase (    a__     ):






			'''simple docstring'''

			__UpperCamelCase					:  Optional[int]    					=       RetriBertConfig
			__UpperCamelCase					:  Optional[Any]    					=       None
			__UpperCamelCase					:  Union[str, Any]    					=       "retribert"

			def 					lowerCAmelCase__		(				self       :				Tuple  ,		snake_case_       :				List[str]					):
							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_  :       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."""   ,				a__   ,				)
class 				_UpperCamelCase (    a__     ):






			'''simple docstring'''

			def __init__(				self       :				Union[str, Any]  ,		snake_case_       :				int					):
							super().__init__(lowerCamelCase_					)
							UpperCamelCase_:							List[Any]		=  config.projection_dim

							UpperCamelCase_:							Optional[int]		=  BertModel(lowerCamelCase_					)
							UpperCamelCase_:							Optional[Any]		=  None if config.share_encoders else BertModel(lowerCamelCase_					)
							UpperCamelCase_:							Dict		=  nn.Dropout(config.hidden_dropout_prob					)
							UpperCamelCase_:							Any		=  nn.Linear(config.hidden_size  ,		config.projection_dim  ,		bias=lowerCamelCase_					)
							UpperCamelCase_:							List[str]		=  nn.Linear(config.hidden_size  ,		config.projection_dim  ,		bias=lowerCamelCase_					)

							UpperCamelCase_:							List[Any]		=  nn.CrossEntropyLoss(reduction="""mean"""					)

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

			def 					lowerCAmelCase__		(				self       :				Optional[Any]  ,		snake_case_       :				int  ,		snake_case_       :				Optional[Any]  ,		snake_case_       :				Union[str, Any]  ,		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(lowerCamelCase_  ,		attention_mask=lowerCamelCase_					)[1]
							else:
											# prepare implicit variables
											UpperCamelCase_:							Tuple		=  input_ids.device
											UpperCamelCase_:							int		=  input_ids.size()
											UpperCamelCase_:							int		=  torch.zeros(lowerCamelCase_  ,		dtype=torch.long  ,		device=lowerCamelCase_					)
											UpperCamelCase_:							Any		=  [None] * sent_encoder.config.num_hidden_layers
											UpperCamelCase_:							Optional[Any]		=  sent_encoder.get_extended_attention_mask(
											    lowerCamelCase_  ,		lowerCamelCase_					)

											# define function for checkpointing
											def partial_encode(*snake_case_       :				Union[str, Any]					):
															UpperCamelCase_:							Optional[int]		=  sent_encoder.encoder(
															    inputs[0]  ,		attention_mask=inputs[1]  ,		head_mask=lowerCamelCase_  ,		)
															UpperCamelCase_:							Optional[int]		=  encoder_outputs[0]
															UpperCamelCase_:							Optional[int]		=  sent_encoder.pooler(lowerCamelCase_					)
															return pooled_output

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

			def 					lowerCAmelCase__		(				self       :				Union[str, Any]  ,		snake_case_       :				Any  ,		snake_case_       :				List[Any]=None  ,		snake_case_       :				Dict=-1  ,		):
							UpperCamelCase_:							Optional[int]		=  self.embed_sentences_checkpointed(
							    lowerCamelCase_  ,		lowerCamelCase_  ,		self.bert_query  ,		lowerCamelCase_  ,		)
							return self.project_query(lowerCamelCase_					)

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










'''simple docstring'''






import inspect
import unittest

from transformers import is_torch_available, is_vision_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
		import torch

		from transformers import MODEL_MAPPING, PoolFormerConfig, PoolFormerForImageClassification, PoolFormerModel
		from transformers.models.poolformer.modeling_poolformer import POOLFORMER_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
		from PIL import Image

		from transformers import PoolFormerImageProcessor




class 		a__				(							a__  ):





							'''simple docstring'''






							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							Union[str, Any]:
														lowerCAmelCase__          =	self.config_class(**self.inputs_dict   )
														self.parent.assertTrue(hasattr(lowerCamelCase_					,    '''hidden_sizes'''   )   )
														self.parent.assertTrue(hasattr(lowerCamelCase_					,    '''num_encoder_blocks'''   )   )




class 		a__				:





							'''simple docstring'''






							def __init__(		self					,    lowerCamelCase_					,    lowerCamelCase_=13					,    lowerCamelCase_=64					,    lowerCamelCase_=3					,    lowerCamelCase_=4					,    lowerCamelCase_=[2, 2, 2, 2]					,    lowerCamelCase_=[8, 4, 2, 1]					,    lowerCamelCase_=[16, 32, 64, 1_28]					,    lowerCamelCase_=[1, 4, 8, 16]					,    lowerCamelCase_=False					,    lowerCamelCase_=True					,    lowerCamelCase_="gelu"					,    lowerCamelCase_=0.1					,    lowerCamelCase_=0.02					,    lowerCamelCase_=3					,    lowerCamelCase_=None					,    )       ->							Any:
														lowerCAmelCase__          =	parent
														lowerCAmelCase__          =	batch_size
														lowerCAmelCase__          =	image_size
														lowerCAmelCase__          =	num_channels
														lowerCAmelCase__          =	num_encoder_blocks
														lowerCAmelCase__          =	sr_ratios
														lowerCAmelCase__          =	depths
														lowerCAmelCase__          =	hidden_sizes
														lowerCAmelCase__          =	downsampling_rates
														lowerCAmelCase__          =	is_training
														lowerCAmelCase__          =	use_labels
														lowerCAmelCase__          =	hidden_act
														lowerCAmelCase__          =	hidden_dropout_prob
														lowerCAmelCase__          =	initializer_range
														lowerCAmelCase__          =	num_labels
														lowerCAmelCase__          =	scope







							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							Optional[int]:
														lowerCAmelCase__          =	floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]   )

														lowerCAmelCase__          =	None
														if self.use_labels:
																					lowerCAmelCase__          =	ids_tensor([self.batch_size, self.image_size, self.image_size]					,    self.num_labels   )

														lowerCAmelCase__          =	PoolFormerConfig(
														    image_size=self.image_size					,    num_channels=self.num_channels					,    num_encoder_blocks=self.num_encoder_blocks					,    depths=self.depths					,    hidden_sizes=self.hidden_sizes					,    hidden_act=self.hidden_act					,    hidden_dropout_prob=self.hidden_dropout_prob					,    initializer_range=self.initializer_range					,    )

														return config, pixel_values, labels







							def   __SCREAMING_SNAKE_CASE							(		self					,    lowerCamelCase_					,    lowerCamelCase_					,    lowerCamelCase_   )       ->							Any:
														lowerCAmelCase__          =	PoolFormerModel(config=lowerCamelCase_   )
														model.to(lowerCamelCase_   )
														model.eval()
														lowerCAmelCase__          =	model(lowerCamelCase_   )
														lowerCAmelCase__          =	lowerCAmelCase__          =	self.image_size // 32.0
														self.parent.assertEqual(
														    result.last_hidden_state.shape					,    (self.batch_size, self.hidden_sizes[-1], expected_height, expected_width)   )







							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							str:
														lowerCAmelCase__          =	self.prepare_config_and_inputs()
														lowerCAmelCase__   ,					lowerCAmelCase__   ,					lowerCAmelCase__          =	config_and_inputs
														lowerCAmelCase__          =	{'''pixel_values''': pixel_values}
														return config, inputs_dict




@require_torch
class 		a__				(							a__				,				a__				,				unittest.TestCase  ):





							'''simple docstring'''






							lowercase__			:   Dict							=      (PoolFormerModel, PoolFormerForImageClassification) if is_torch_available() else ()
							lowercase__			:   Any							=      (
							    {"feature-extraction": PoolFormerModel, "image-classification": PoolFormerForImageClassification}
							    if is_torch_available()
							    else {}
							)

							lowercase__			:   Optional[int]							=      False
							lowercase__			:   Optional[int]							=      False
							lowercase__			:   Union[str, Any]							=      False
							lowercase__			:   List[Any]							=      False
							lowercase__			:   Tuple							=      False







							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							Union[str, Any]:
														lowerCAmelCase__          =	PoolFormerModelTester(self   )
														lowerCAmelCase__          =	PoolFormerConfigTester(self					,    config_class=lowerCamelCase_   )







							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							List[Any]:
														self.config_tester.run_common_tests()







							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							List[Any]:
														lowerCAmelCase__          =	self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_model(*lowerCamelCase_   )







							@unittest.skip('''PoolFormer does not use inputs_embeds'''   )
							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							str:
														pass







							@unittest.skip('''PoolFormer does not have get_input_embeddings method and get_output_embeddings methods'''   )
							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							Tuple:
														pass







							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							int:
														def check_hidden_states_output(lowerCamelCase_					,    lowerCamelCase_					,    lowerCamelCase_   ):
																					lowerCAmelCase__          =	model_class(lowerCamelCase_   )
																					model.to(lowerCamelCase_   )
																					model.eval()

																					with torch.no_grad():
																												lowerCAmelCase__          =	model(**self._prepare_for_class(lowerCamelCase_					,    lowerCamelCase_   )   )

																					lowerCAmelCase__          =	outputs.hidden_states

																					lowerCAmelCase__          =	self.model_tester.num_encoder_blocks
																					self.assertEqual(len(lowerCamelCase_   )					,    lowerCamelCase_   )

																					# verify the first hidden states (first block)
																					self.assertListEqual(
																					    list(hidden_states[0].shape[-3:]   )					,    [
																					        self.model_tester.hidden_sizes[0],
																					        self.model_tester.image_size // 4,
																					        self.model_tester.image_size // 4,
																					    ]					,    )

														lowerCAmelCase__   ,					lowerCAmelCase__          =	self.model_tester.prepare_config_and_inputs_for_common()

														for model_class in self.all_model_classes:
																					lowerCAmelCase__          =	True
																					check_hidden_states_output(lowerCamelCase_					,    lowerCamelCase_					,    lowerCamelCase_   )

																					# check that output_hidden_states also work using config
																					del inputs_dict["output_hidden_states"]
																					lowerCAmelCase__          =	True

																					check_hidden_states_output(lowerCamelCase_					,    lowerCamelCase_					,    lowerCamelCase_   )







							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							int:
														if not self.model_tester.is_training:
																					return

														lowerCAmelCase__   ,					lowerCAmelCase__          =	self.model_tester.prepare_config_and_inputs_for_common()
														lowerCAmelCase__          =	True

														for model_class in self.all_model_classes:
																					if model_class in get_values(lowerCamelCase_   ):
																												continue
																					lowerCAmelCase__          =	model_class(lowerCamelCase_   )
																					model.to(lowerCamelCase_   )
																					model.train()
																					lowerCAmelCase__          =	self._prepare_for_class(lowerCamelCase_					,    lowerCamelCase_					,    return_labels=lowerCamelCase_   )
																					lowerCAmelCase__          =	model(**lowerCamelCase_   ).loss
																					loss.backward()







							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							int:
														lowerCAmelCase__   ,					lowerCAmelCase__          =	self.model_tester.prepare_config_and_inputs_for_common()

														for model_class in self.all_model_classes:
																					lowerCAmelCase__          =	model_class(lowerCamelCase_   )
																					lowerCAmelCase__          =	inspect.signature(model.forward   )
																					# signature.parameters is an OrderedDict => so arg_names order is deterministic
																					lowerCAmelCase__          =	[*signature.parameters.keys()]

																					lowerCAmelCase__          =	['''pixel_values''']
																					self.assertListEqual(arg_names[:1]					,    lowerCamelCase_   )







							@slow
							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							int:
														for model_name in POOLFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																					lowerCAmelCase__          =	PoolFormerModel.from_pretrained(lowerCamelCase_   )
																					self.assertIsNotNone(lowerCamelCase_   )


def 					_snake_case   (     )						->    Optional[int]:
							lowerCAmelCase__          =	Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''       )
							return image






@require_torch
class 		a__				(							unittest.TestCase  ):





							'''simple docstring'''






							@slow
							def   __SCREAMING_SNAKE_CASE							(		self   )       ->							Dict:
														lowerCAmelCase__          =	PoolFormerImageProcessor()
														lowerCAmelCase__          =	PoolFormerForImageClassification.from_pretrained('''sail/poolformer_s12'''   ).to(lowerCamelCase_   )

														lowerCAmelCase__          =	image_processor(images=prepare_img()					,    return_tensors='''pt'''   ).to(lowerCamelCase_   )

														# forward pass
														with torch.no_grad():
																					lowerCAmelCase__          =	model(**lowerCamelCase_   )

														# verify the logits
														lowerCAmelCase__          =	torch.Size((1, 10_00)   )
														self.assertEqual(outputs.logits.shape					,    lowerCamelCase_   )

														lowerCAmelCase__          =	torch.tensor([-0.6_113, 0.1_685, -0.0_492]   ).to(lowerCamelCase_   )
														self.assertTrue(torch.allclose(outputs.logits[0, :3]					,    lowerCamelCase_					,    atol=1e-4   )   )