


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


snake_case_  :     Any							  =							logging.get_logger(__name__)

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

class   A__							(						UpperCamelCase__					):
						UpperCAmelCase     =  RetriBertConfig
						UpperCAmelCase     =  None
						UpperCAmelCase     =  "retribert"




						def 							__UpperCamelCase			(    self   : Dict			,      _a   : Dict					)							->       int:


										"""simple docstring"""





										if isinstance(_a			,      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(_a			,      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(_a			,      nn.LayerNorm					):
														module.bias.data.zero_()
														module.weight.data.fill_(1.0					)
snake_case_  :     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."				,       UpperCamelCase__				,       )
class   A__							(						UpperCamelCase__					):




						def __init__(    self   : Optional[int]			,      _a   : RetriBertConfig					)							->       None:


										"""simple docstring"""





										super().__init__(_a					)
										_SCREAMING_SNAKE_CASE     					=config.projection_dim

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

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

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




						def 							__UpperCamelCase			(    self   : Union[str, Any]			,      _a   : Optional[Any]			,      _a   : List[str]			,      _a   : Dict			,      _a   : Optional[Any]=-1			,      )							->       int:


										"""simple docstring"""





										if checkpoint_batch_size < 0 or input_ids.shape[0] < checkpoint_batch_size:
														return sent_encoder(_a			,      attention_mask=_a					)[1]
										else:
														# prepare implicit variables
														_SCREAMING_SNAKE_CASE     					=input_ids.device
														_SCREAMING_SNAKE_CASE     					=input_ids.size()
														_SCREAMING_SNAKE_CASE     					=torch.zeros(_a			,      dtype=torch.long			,      device=_a					)
														_SCREAMING_SNAKE_CASE     					=[None] * sent_encoder.config.num_hidden_layers
														_SCREAMING_SNAKE_CASE     					=sent_encoder.get_extended_attention_mask(
														    _a			,      _a					)

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

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




						def 							__UpperCamelCase			(    self   : Optional[Any]			,      _a   : str			,      _a   : Tuple=None			,      _a   : Any=-1			,      )							->       Any:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE     					=self.embed_sentences_checkpointed(
										    _a			,      _a			,      self.bert_query			,      _a			,      )
										return self.project_query(_a					)




						def 							__UpperCamelCase			(    self   : int			,      _a   : Optional[Any]			,      _a   : int=None			,      _a   : int=-1			,      )							->       Any:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE     					=self.embed_sentences_checkpointed(
										    _a			,      _a			,      self.bert_query if self.bert_doc is None else self.bert_doc			,      _a			,      )
										return self.project_doc(_a					)




						def 							__UpperCamelCase			(    self   : Optional[Any]			,      _a   : torch.LongTensor			,      _a   : Optional[torch.FloatTensor]			,      _a   : torch.LongTensor			,      _a   : Optional[torch.FloatTensor]			,      _a   : int = -1			,      )							->       torch.FloatTensor:


										"""simple docstring"""





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



import inspect
import unittest

from transformers import ResNetConfig, is_flax_available
from transformers.testing_utils import require_flax, slow
from transformers.utils import cached_property, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor


if is_flax_available():
							import jax
							import jax.numpy as jnp

							from transformers.models.resnet.modeling_flax_resnet import FlaxResNetForImageClassification, FlaxResNetModel

if is_vision_available():
							from PIL import Image

							from transformers import AutoImageProcessor

class   A__							(						unittest.TestCase					):




						def __init__(    self   : List[Any]			,      _a   : Optional[Any]			,      _a   : int=3			,      _a   : int=32			,      _a   : str=3			,      _a   : List[str]=10			,      _a   : Optional[Any]=[10, 20, 30, 40]			,      _a   : Tuple=[1, 1, 2, 1]			,      _a   : Optional[Any]=True			,      _a   : Optional[int]=True			,      _a   : Any="relu"			,      _a   : Any=3			,      _a   : Tuple=None			,      )							->       Optional[Any]:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE     					=parent
										_SCREAMING_SNAKE_CASE     					=batch_size
										_SCREAMING_SNAKE_CASE     					=image_size
										_SCREAMING_SNAKE_CASE     					=num_channels
										_SCREAMING_SNAKE_CASE     					=embeddings_size
										_SCREAMING_SNAKE_CASE     					=hidden_sizes
										_SCREAMING_SNAKE_CASE     					=depths
										_SCREAMING_SNAKE_CASE     					=is_training
										_SCREAMING_SNAKE_CASE     					=use_labels
										_SCREAMING_SNAKE_CASE     					=hidden_act
										_SCREAMING_SNAKE_CASE     					=num_labels
										_SCREAMING_SNAKE_CASE     					=scope
										_SCREAMING_SNAKE_CASE     					=len(_a					)




						def 							__UpperCamelCase			(    self   : Union[str, Any]					)							->       str:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE     					=floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]					)

										_SCREAMING_SNAKE_CASE     					=self.get_config()

										return config, pixel_values




						def 							__UpperCamelCase			(    self   : List[str]					)							->       Dict:


										"""simple docstring"""





										return ResNetConfig(
										    num_channels=self.num_channels			,      embeddings_size=self.embeddings_size			,      hidden_sizes=self.hidden_sizes			,      depths=self.depths			,      hidden_act=self.hidden_act			,      num_labels=self.num_labels			,      image_size=self.image_size			,      )




						def 							__UpperCamelCase			(    self   : Optional[Any]			,      _a   : Union[str, Any]			,      _a   : List[Any]					)							->       Union[str, Any]:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE     					=FlaxResNetModel(config=_a					)
										_SCREAMING_SNAKE_CASE     					=model(_a					)

										# Output shape (b, c, h, w)
										self.parent.assertEqual(
										    result.last_hidden_state.shape			,      (self.batch_size, self.hidden_sizes[-1], self.image_size // 32, self.image_size // 32)			,      )




						def 							__UpperCamelCase			(    self   : int			,      _a   : str			,      _a   : str					)							->       Optional[Any]:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE     					=self.num_labels
										_SCREAMING_SNAKE_CASE     					=FlaxResNetForImageClassification(config=_a					)
										_SCREAMING_SNAKE_CASE     					=model(_a					)
										self.parent.assertEqual(result.logits.shape			,      (self.batch_size, self.num_labels)					)




						def 							__UpperCamelCase			(    self   : Tuple					)							->       Optional[Any]:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE     					=self.prepare_config_and_inputs()
										_SCREAMING_SNAKE_CASE					,			_SCREAMING_SNAKE_CASE     					=config_and_inputs
										_SCREAMING_SNAKE_CASE     					={'''pixel_values''': pixel_values}
										return config, inputs_dict

@require_flax
class   A__							(						UpperCamelCase__				,       unittest.TestCase					):
						UpperCAmelCase     =  (FlaxResNetModel, FlaxResNetForImageClassification) if is_flax_available() else ()

						UpperCAmelCase     =  False
						UpperCAmelCase     =  False
						UpperCAmelCase     =  False




						def 							__UpperCamelCase			(    self   : str					)							->       None:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE     					=FlaxResNetModelTester(self					)
										_SCREAMING_SNAKE_CASE     					=ConfigTester(self			,      config_class=_a			,      has_text_modality=_a					)




						def 							__UpperCamelCase			(    self   : Optional[Any]					)							->       int:


										"""simple docstring"""





										self.create_and_test_config_common_properties()
										self.config_tester.create_and_test_config_to_json_string()
										self.config_tester.create_and_test_config_to_json_file()
										self.config_tester.create_and_test_config_from_and_save_pretrained()
										self.config_tester.create_and_test_config_with_num_labels()
										self.config_tester.check_config_can_be_init_without_params()
										self.config_tester.check_config_arguments_init()




						def 							__UpperCamelCase			(    self   : Optional[Any]					)							->       Union[str, Any]:


										"""simple docstring"""





										return




						def 							__UpperCamelCase			(    self   : int					)							->       Optional[int]:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE     					=self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*_a					)




						def 							__UpperCamelCase			(    self   : Optional[int]					)							->       List[str]:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE     					=self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_image_classification(*_a					)




						@unittest.skip(reason='''ResNet does not use inputs_embeds'''					)
						def 							__UpperCamelCase			(    self   : List[Any]					)							->       List[Any]:


										"""simple docstring"""





										pass




						@unittest.skip(reason='''ResNet does not support input and output embeddings'''					)
						def 							__UpperCamelCase			(    self   : Dict					)							->       Dict:


										"""simple docstring"""





										pass




						def 							__UpperCamelCase			(    self   : Any					)							->       int:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE					,			_SCREAMING_SNAKE_CASE     					=self.model_tester.prepare_config_and_inputs_for_common()

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

														_SCREAMING_SNAKE_CASE     					=['''pixel_values''']
														self.assertListEqual(arg_names[:1]			,      _a					)




						def 							__UpperCamelCase			(    self   : Dict					)							->       Union[str, Any]:


										"""simple docstring"""





										def check_hidden_states_output(_a   : Optional[Any]			,      _a   : Tuple			,      _a   : Optional[Any]					):
														_SCREAMING_SNAKE_CASE     					=model_class(_a					)

														_SCREAMING_SNAKE_CASE     					=model(**self._prepare_for_class(_a			,      _a					)					)

														_SCREAMING_SNAKE_CASE     					=outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

														_SCREAMING_SNAKE_CASE     					=self.model_tester.num_stages
														self.assertEqual(len(_a					)			,      expected_num_stages + 1					)




						@unittest.skip(reason='''ResNet does not use feedforward chunking'''					)
						def 							__UpperCamelCase			(    self   : Optional[int]					)							->       Dict:


										"""simple docstring"""





										pass




						def 							__UpperCamelCase			(    self   : List[Any]					)							->       Any:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE					,			_SCREAMING_SNAKE_CASE     					=self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
														with self.subTest(model_class.__name__					):
																		_SCREAMING_SNAKE_CASE     					=self._prepare_for_class(_a			,      _a					)
																		_SCREAMING_SNAKE_CASE     					=model_class(_a					)

																		@jax.jit
																		def model_jitted(_a   : Optional[Any]			,      **_a   : Any					):
																						return model(pixel_values=_a			,      **_a					)

																		with self.subTest('''JIT Enabled'''					):
																						_SCREAMING_SNAKE_CASE     					=model_jitted(**_a					).to_tuple()

																		with self.subTest('''JIT Disabled'''					):
																						with jax.disable_jit():
																										_SCREAMING_SNAKE_CASE     					=model_jitted(**_a					).to_tuple()

																		self.assertEqual(len(_a					)			,      len(_a					)					)
																		for jitted_output, output in zip(_a			,      _a					):
																						self.assertEqual(jitted_output.shape			,      output.shape					)







def   lowerCamelCase(				):
				_SCREAMING_SNAKE_CASE     					=Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png''')
				return image


@require_flax
class   A__							(						unittest.TestCase					):




						@cached_property
						def 							__UpperCamelCase			(    self   : List[Any]					)							->       Any:


										"""simple docstring"""





										return AutoImageProcessor.from_pretrained('''microsoft/resnet-50'''					) if is_vision_available() else None




						@slow
						def 							__UpperCamelCase			(    self   : Union[str, Any]					)							->       Dict:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE     					=FlaxResNetForImageClassification.from_pretrained('''microsoft/resnet-50'''					)

										_SCREAMING_SNAKE_CASE     					=self.default_image_processor
										_SCREAMING_SNAKE_CASE     					=prepare_img()
										_SCREAMING_SNAKE_CASE     					=image_processor(images=_a			,      return_tensors='''np'''					)

										_SCREAMING_SNAKE_CASE     					=model(**_a					)

										# verify the logits
										_SCREAMING_SNAKE_CASE     					=(1, 1000)
										self.assertEqual(outputs.logits.shape			,      _a					)

										_SCREAMING_SNAKE_CASE     					=jnp.array([-11.10_69, -9.78_77, -8.37_77]					)

										self.assertTrue(jnp.allclose(outputs.logits[0, :3]			,      _a			,      atol=1E-4					)					)