




import inspect
import unittest

from transformers import ResNetConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_backbone_common import BackboneTesterMixin
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 torch import nn

 from transformers import ResNetBackbone, ResNetForImageClassification, ResNetModel
 from transformers.models.resnet.modeling_resnet import RESNET_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
 from PIL import Image

 from transformers import AutoImageProcessor

class 			_lowerCamelCase    :

      """simple docstring"""

      def __init__(  self       :			List[Any]			,							snake_case       :			int			,							snake_case       :			List[Any]=3			,							snake_case       :			Dict=32			,							snake_case       :			Dict=3			,							snake_case       :			Optional[int]=10			,							snake_case       :			int=[10, 20, 30, 40]			,							snake_case       :			List[Any]=[1, 1, 2, 1]			,							snake_case       :			List[Any]=True			,							snake_case       :			str=True			,							snake_case       :			int="relu"			,							snake_case       :			List[str]=3			,							snake_case       :			List[str]=None			,							snake_case       :			List[Any]=["stage2", "stage3", "stage4"]			,							snake_case       :			Optional[Any]=[2, 3, 4]			,							):
          __UpperCamelCase			=    parent
          __UpperCamelCase			=    batch_size
          __UpperCamelCase			=    image_size
          __UpperCamelCase			=    num_channels
          __UpperCamelCase			=    embeddings_size
          __UpperCamelCase			=    hidden_sizes
          __UpperCamelCase			=    depths
          __UpperCamelCase			=    is_training
          __UpperCamelCase			=    use_labels
          __UpperCamelCase			=    hidden_act
          __UpperCamelCase			=    num_labels
          __UpperCamelCase			=    scope
          __UpperCamelCase			=    len(snake_case					)
          __UpperCamelCase			=    out_features
          __UpperCamelCase			=    out_indices







      def    snake_case						(  self       :			Optional[int]					):
          __UpperCamelCase			=    floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]					)

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

          __UpperCamelCase			=    self.get_config()

          return config, pixel_values, labels







      def    snake_case						(  self       :			Tuple					):
          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			,							out_features=self.out_features			,							out_indices=self.out_indices			,							)







      def    snake_case						(  self       :			List[str]			,							snake_case       :			Optional[int]			,							snake_case       :			Tuple			,							snake_case       :			int					):
          __UpperCamelCase			=    ResNetModel(config=snake_case					)
          model.to(snake_case					)
          model.eval()
          __UpperCamelCase			=    model(snake_case					)
          # expected last hidden states: B, C, H // 32, W // 32
          self.parent.assertEqual(
              result.last_hidden_state.shape			,							(self.batch_size, self.hidden_sizes[-1], self.image_size // 32, self.image_size // 32)			,							)







      def    snake_case						(  self       :			Tuple			,							snake_case       :			Union[str, Any]			,							snake_case       :			List[str]			,							snake_case       :			List[str]					):
          __UpperCamelCase			=    self.num_labels
          __UpperCamelCase			=    ResNetForImageClassification(snake_case					)
          model.to(snake_case					)
          model.eval()
          __UpperCamelCase			=    model(snake_case			,							labels=snake_case					)
          self.parent.assertEqual(result.logits.shape			,							(self.batch_size, self.num_labels)					)







      def    snake_case						(  self       :			int			,							snake_case       :			Dict			,							snake_case       :			List[str]			,							snake_case       :			Optional[int]					):
          __UpperCamelCase			=    ResNetBackbone(config=snake_case					)
          model.to(snake_case					)
          model.eval()
          __UpperCamelCase			=    model(snake_case					)

          # verify feature maps
          self.parent.assertEqual(len(result.feature_maps					)			,							len(config.out_features					)					)
          self.parent.assertListEqual(list(result.feature_maps[0].shape					)			,							[self.batch_size, self.hidden_sizes[1], 4, 4]					)

          # verify channels
          self.parent.assertEqual(len(model.channels					)			,							len(config.out_features					)					)
          self.parent.assertListEqual(model.channels			,							config.hidden_sizes[1:]					)

          # verify backbone works with out_features=None
          __UpperCamelCase			=    None
          __UpperCamelCase			=    ResNetBackbone(config=snake_case					)
          model.to(snake_case					)
          model.eval()
          __UpperCamelCase			=    model(snake_case					)

          # verify feature maps
          self.parent.assertEqual(len(result.feature_maps					)			,							1					)
          self.parent.assertListEqual(list(result.feature_maps[0].shape					)			,							[self.batch_size, self.hidden_sizes[-1], 1, 1]					)

          # verify channels
          self.parent.assertEqual(len(model.channels					)			,							1					)
          self.parent.assertListEqual(model.channels			,							[config.hidden_sizes[-1]]					)







      def    snake_case						(  self       :			List[str]					):
          __UpperCamelCase			=    self.prepare_config_and_inputs()
          __UpperCamelCase		,   __UpperCamelCase		,   __UpperCamelCase			=    config_and_inputs
          __UpperCamelCase			=    {'''pixel_values''': pixel_values}
          return config, inputs_dict







@require_torch
class 			_lowerCamelCase    (      _SCREAMING_SNAKE_CASE				,      _SCREAMING_SNAKE_CASE				,      unittest.TestCase ):

      """simple docstring"""

      lowerCAmelCase__      :			List[str]     =       (
          (
              ResNetModel,
              ResNetForImageClassification,
              ResNetBackbone,
          )
          if is_torch_available()
          else ()
      )
      lowerCAmelCase__      :			Optional[int]     =       (
          {"feature-extraction": ResNetModel, "image-classification": ResNetForImageClassification}
          if is_torch_available()
          else {}
      )

      lowerCAmelCase__      :			Optional[Any]     =       True
      lowerCAmelCase__      :			Union[str, Any]     =       False
      lowerCAmelCase__      :			Optional[int]     =       False
      lowerCAmelCase__      :			Tuple     =       False
      lowerCAmelCase__      :			Any     =       False







      def    snake_case						(  self       :			Tuple					):
          __UpperCamelCase			=    ResNetModelTester(self					)
          __UpperCamelCase			=    ConfigTester(self			,							config_class=snake_case			,							has_text_modality=snake_case					)







      def    snake_case						(  self       :			Optional[Any]					):
          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    snake_case						(  self       :			int					):
          return







      @unittest.skip(reason='''ResNet does not use inputs_embeds'''					)
      def    snake_case						(  self       :			Union[str, Any]					):
          pass







      @unittest.skip(reason='''ResNet does not support input and output embeddings'''					)
      def    snake_case						(  self       :			Optional[int]					):
          pass







      def    snake_case						(  self       :			Optional[int]					):
          __UpperCamelCase		,   __UpperCamelCase			=    self.model_tester.prepare_config_and_inputs_for_common()

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

              __UpperCamelCase			=    ['''pixel_values''']
              self.assertListEqual(arg_names[:1]			,							snake_case					)







      def    snake_case						(  self       :			Union[str, Any]					):
          __UpperCamelCase			=    self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*snake_case					)







      def    snake_case						(  self       :			List[Any]					):
          __UpperCamelCase			=    self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_backbone(*snake_case					)







      def    snake_case						(  self       :			str					):
          __UpperCamelCase		,   __UpperCamelCase			=    self.model_tester.prepare_config_and_inputs_for_common()

          for model_class in self.all_model_classes:
              __UpperCamelCase			=    model_class(config=snake_case					)
              for name, module in model.named_modules():
                  if isinstance(snake_case			,							(nn.BatchNormad, nn.GroupNorm)					):
                      self.assertTrue(
                          torch.all(module.weight == 1					)			,							msg=F"Parameter {name} of model {model_class} seems not properly initialized"			,							)
                      self.assertTrue(
                          torch.all(module.bias == 0					)			,							msg=F"Parameter {name} of model {model_class} seems not properly initialized"			,							)







      def    snake_case						(  self       :			Optional[int]					):
          def check_hidden_states_output(snake_case       :			str			,							snake_case       :			Union[str, Any]			,							snake_case       :			str					):
              __UpperCamelCase			=    model_class(snake_case					)
              model.to(snake_case					)
              model.eval()

              with torch.no_grad():
                  __UpperCamelCase			=    model(**self._prepare_for_class(snake_case			,							snake_case					)					)

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

              __UpperCamelCase			=    self.model_tester.num_stages
              self.assertEqual(len(snake_case					)			,							expected_num_stages + 1					)

              # ResNet's feature maps are of shape (batch_size, num_channels, height, width)
              self.assertListEqual(
                  list(hidden_states[0].shape[-2:]					)			,							[self.model_tester.image_size // 4, self.model_tester.image_size // 4]			,							)

          __UpperCamelCase		,   __UpperCamelCase			=    self.model_tester.prepare_config_and_inputs_for_common()
          __UpperCamelCase			=    ['''basic''', '''bottleneck''']
          for model_class in self.all_model_classes:
              for layer_type in layers_type:
                  __UpperCamelCase			=    layer_type
                  __UpperCamelCase			=    True
                  check_hidden_states_output(snake_case			,							snake_case			,							snake_case					)

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

                  check_hidden_states_output(snake_case			,							snake_case			,							snake_case					)







      @unittest.skip(reason='''ResNet does not use feedforward chunking'''					)
      def    snake_case						(  self       :			Any					):
          pass







      def    snake_case						(  self       :			List[str]					):
          __UpperCamelCase			=    self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_image_classification(*snake_case					)







      @slow
      def    snake_case						(  self       :			Union[str, Any]					):
          for model_name in RESNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
              __UpperCamelCase			=    ResNetModel.from_pretrained(snake_case					)
              self.assertIsNotNone(snake_case					)







def        __SCREAMING_SNAKE_CASE (	)							->						int:

    """simple docstring"""


    __UpperCamelCase			=    Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''   )
    return image

@require_torch
@require_vision
class 			_lowerCamelCase    (      unittest.TestCase ):

      """simple docstring"""

      @cached_property
      def    snake_case						(  self       :			str					):
          return (
              AutoImageProcessor.from_pretrained(RESNET_PRETRAINED_MODEL_ARCHIVE_LIST[0]					)
              if is_vision_available()
              else None
          )







      @slow
      def    snake_case						(  self       :			Dict					):
          __UpperCamelCase			=    ResNetForImageClassification.from_pretrained(RESNET_PRETRAINED_MODEL_ARCHIVE_LIST[0]					).to(snake_case					)

          __UpperCamelCase			=    self.default_image_processor
          __UpperCamelCase			=    prepare_img()
          __UpperCamelCase			=    image_processor(images=snake_case			,							return_tensors='''pt'''					).to(snake_case					)

          # forward pass
          with torch.no_grad():
              __UpperCamelCase			=    model(**snake_case					)

          # verify the logits
          __UpperCamelCase			=    torch.Size((1, 1000)					)
          self.assertEqual(outputs.logits.shape			,							snake_case					)

          __UpperCamelCase			=    torch.tensor([-11.10_69, -9.78_77, -8.37_77]					).to(snake_case					)

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







@require_torch
class 			_lowerCamelCase    (      _SCREAMING_SNAKE_CASE				,      unittest.TestCase ):

      """simple docstring"""

      lowerCAmelCase__      :			Tuple     =       (ResNetBackbone,) if is_torch_available() else ()
      lowerCAmelCase__      :			List[str]     =       False
      lowerCAmelCase__      :			int     =       ResNetConfig







      def    snake_case						(  self       :			Optional[Any]					):
          __UpperCamelCase			=    ResNetModelTester(self					)












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_ 		=  logging.get_logger(__name__)

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

class 			_lowerCamelCase    (      _SCREAMING_SNAKE_CASE ):

      """simple docstring"""

      lowerCAmelCase__      :			Optional[Any]     =       RetriBertConfig
      lowerCAmelCase__      :			Optional[Any]     =       None
      lowerCAmelCase__      :			Dict     =       "retribert"







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







a_ 		=  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."				,      _SCREAMING_SNAKE_CASE				,      )
class 			_lowerCamelCase    (      _SCREAMING_SNAKE_CASE ):

      """simple docstring"""

      def __init__(  self       :			List[Any]			,							snake_case       :			RetriBertConfig					):
          super().__init__(snake_case					)
          __UpperCamelCase			=    config.projection_dim

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

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

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







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

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

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







      def    snake_case						(  self       :			Dict			,							snake_case       :			Any			,							snake_case       :			Union[str, Any]=None			,							snake_case       :			List[str]=-1			,							):
          __UpperCamelCase			=    self.embed_sentences_checkpointed(
              snake_case			,							snake_case			,							self.bert_query			,							snake_case			,							)
          return self.project_query(snake_case					)







      def    snake_case						(  self       :			Optional[Any]			,							snake_case       :			Optional[Any]			,							snake_case       :			Optional[int]=None			,							snake_case       :			Optional[Any]=-1			,							):
          __UpperCamelCase			=    self.embed_sentences_checkpointed(
              snake_case			,							snake_case			,							self.bert_query if self.bert_doc is None else self.bert_doc			,							snake_case			,							)
          return self.project_doc(snake_case					)







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






