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

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

class 		__magic_name__				( lowercase_					):


  """simple docstring"""





  _UpperCamelCase						 = RetriBertConfig
  _UpperCamelCase						 = None
  _UpperCamelCase						 = "retribert"


  def     _UpperCAmelCase  (					self	,		a__		):
    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		)







_UpperCAmelCase    			=					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."						,lowercase_						,)
class 		__magic_name__				( lowercase_					):


  """simple docstring"""





  def __init__(					self	,		a__		):
    super().__init__(a__		)
    _lowerCamelCase      =						config.projection_dim

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

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

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


  def     _UpperCAmelCase  (					self	,		a__	,		a__	,		a__	,		a__=-1	,		):
    # reproduces BERT forward pass with checkpointing
    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
      _lowerCamelCase      =						input_ids.device
      _lowerCamelCase      =						input_ids.size()
      _lowerCamelCase      =						torch.zeros(a__	,		dtype=torch.long	,		device=a__		)
      _lowerCamelCase      =						[None] * sent_encoder.config.num_hidden_layers
      _lowerCamelCase      =						sent_encoder.get_extended_attention_mask(
          a__	,		a__		)

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

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


  def     _UpperCAmelCase  (					self	,		a__	,		a__=None	,		a__=-1	,		):
    _lowerCamelCase      =						self.embed_sentences_checkpointed(
        a__	,		a__	,		self.bert_query	,		a__	,		)
    return self.project_query(a__		)


  def     _UpperCAmelCase  (					self	,		a__	,		a__=None	,		a__=-1	,		):
    _lowerCamelCase      =						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	,		a__	,		a__	,		a__	,		a__	,		a__ = -1	,		):
    _lowerCamelCase      =						input_ids_query.device
    _lowerCamelCase      =						self.embed_questions(a__	,		a__	,		a__		)
    _lowerCamelCase      =						self.embed_answers(a__	,		a__	,		a__		)
    _lowerCamelCase      =						torch.mm(a__	,		a_reps.t()		)
    _lowerCamelCase      =						self.ce_loss(a__	,		torch.arange(compare_scores.shape[1]		).to(a__		)		)
    _lowerCamelCase      =						self.ce_loss(compare_scores.t()	,		torch.arange(compare_scores.shape[0]		).to(a__		)		)
    _lowerCamelCase      =						(loss_qa + loss_aq) / 2
    return loss






from typing import Dict, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    get_resize_output_image_size,
    normalize,
    pad,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    IMAGENET_STANDARD_MEAN,
    IMAGENET_STANDARD_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    get_image_size,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, logging
from ...utils.import_utils import is_vision_available


_UpperCAmelCase    			=					logging.get_logger(__name__)


if is_vision_available():
  import PIL

class 		__magic_name__				( lowercase_					):


  """simple docstring"""





  _UpperCamelCase						 = ["pixel_values"]


  def __init__(					self	,		a__ = True	,		a__ = None	,		a__ = PILImageResampling.BILINEAR	,		a__ = True	,		a__ = False	,		a__ = True	,		a__ = True	,		a__ = 1 / 2_55	,		a__ = True	,		a__ = None	,		a__ = None	,		**a__	,		):
    super().__init__(**a__		)

    _lowerCamelCase      =						size if size is not None else {'''height''': 25_60, '''width''': 19_20}
    if isinstance(a__	,		(tuple, list)		):
      # The previous feature extractor size parameter was in (width, height) format
      _lowerCamelCase      =						size[::-1]
    _lowerCamelCase      =						get_size_dict(a__		)

    _lowerCamelCase      =						do_resize
    _lowerCamelCase      =						size
    _lowerCamelCase      =						resample
    _lowerCamelCase      =						do_thumbnail
    _lowerCamelCase      =						do_align_long_axis
    _lowerCamelCase      =						do_pad
    _lowerCamelCase      =						do_rescale
    _lowerCamelCase      =						rescale_factor
    _lowerCamelCase      =						do_normalize
    _lowerCamelCase      =						image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
    _lowerCamelCase      =						image_std if image_std is not None else IMAGENET_STANDARD_STD


  def     _UpperCAmelCase  (					self	,		a__	,		a__	,		a__ = None		):
    _lowerCamelCase    ,    _lowerCamelCase      =						get_image_size(a__		)
    _lowerCamelCase    ,    _lowerCamelCase      =						size['''height'''], size['''width''']

    if (output_width < output_height and input_width > input_height) or (
        output_width > output_height and input_width < input_height
    ):
      _lowerCamelCase      =						np.rotaa(a__	,		3		)

    if data_format is not None:
      _lowerCamelCase      =						to_channel_dimension_format(a__	,		a__		)

    return image


  def     _UpperCAmelCase  (					self	,		a__	,		a__	,		a__ = False	,		a__ = None	,		):
    _lowerCamelCase    ,    _lowerCamelCase      =						size['''height'''], size['''width''']
    _lowerCamelCase    ,    _lowerCamelCase      =						get_image_size(a__		)

    _lowerCamelCase      =						output_width - input_width
    _lowerCamelCase      =						output_height - input_height

    if random_padding:
      _lowerCamelCase      =						np.random.randint(low=0	,		high=delta_height + 1		)
      _lowerCamelCase      =						np.random.randint(low=0	,		high=delta_width + 1		)
    else:
      _lowerCamelCase      =						delta_height // 2
      _lowerCamelCase      =						delta_width // 2

    _lowerCamelCase      =						delta_height - pad_top
    _lowerCamelCase      =						delta_width - pad_left

    _lowerCamelCase      =						((pad_top, pad_bottom), (pad_left, pad_right))
    return pad(a__	,		a__	,		data_format=a__		)


  def     _UpperCAmelCase  (					self	,		*a__	,		**a__		):
    logger.info('''pad is deprecated and will be removed in version 4.27. Please use pad_image instead.'''		)
    return self.pad_image(*a__	,		**a__		)


  def     _UpperCAmelCase  (					self	,		a__	,		a__	,		a__ = PILImageResampling.BICUBIC	,		a__ = None	,		**a__	,		):
    _lowerCamelCase    ,    _lowerCamelCase      =						get_image_size(a__		)
    _lowerCamelCase    ,    _lowerCamelCase      =						size['''height'''], size['''width''']

    # We always resize to the smallest of either the input or output size.
    _lowerCamelCase      =						min(a__	,		a__		)
    _lowerCamelCase      =						min(a__	,		a__		)

    if height == input_height and width == input_width:
      return image

    if input_height > input_width:
      _lowerCamelCase      =						int(input_width * height / input_height		)
    elif input_width > input_height:
      _lowerCamelCase      =						int(input_height * width / input_width		)

    return resize(
        a__	,		size=(height, width)	,		resample=a__	,		reducing_gap=2.0	,		data_format=a__	,		**a__		)


  def     _UpperCAmelCase  (					self	,		a__	,		a__	,		a__ = PILImageResampling.BICUBIC	,		a__ = None	,		**a__	,		):
    _lowerCamelCase      =						get_size_dict(a__		)
    _lowerCamelCase      =						min(size['''height''']	,		size['''width''']		)
    _lowerCamelCase      =						get_resize_output_image_size(a__	,		size=a__	,		default_to_square=a__		)
    _lowerCamelCase      =						resize(a__	,		size=a__	,		resample=a__	,		data_format=a__	,		**a__		)
    return resized_image


  def     _UpperCAmelCase  (					self	,		a__	,		a__	,		a__ = None	,		**a__	,		):
    return rescale(a__	,		scale=a__	,		data_format=a__	,		**a__		)


  def     _UpperCAmelCase  (					self	,		a__	,		a__	,		a__	,		a__ = None	,		**a__	,		):
    return normalize(a__	,		mean=a__	,		std=a__	,		data_format=a__	,		**a__		)


  def     _UpperCAmelCase  (					self	,		a__	,		a__ = None	,		a__ = None	,		a__ = None	,		a__ = None	,		a__ = None	,		a__ = None	,		a__ = False	,		a__ = None	,		a__ = None	,		a__ = None	,		a__ = None	,		a__ = None	,		a__ = None	,		a__ = ChannelDimension.FIRST	,		**a__	,		):
    _lowerCamelCase      =						do_resize if do_resize is not None else self.do_resize
    _lowerCamelCase      =						size if size is not None else self.size
    if isinstance(a__	,		(tuple, list)		):
      # Previous feature extractor had size in (width, height) format
      _lowerCamelCase      =						size[::-1]
    _lowerCamelCase      =						get_size_dict(a__		)
    _lowerCamelCase      =						resample if resample is not None else self.resample
    _lowerCamelCase      =						do_thumbnail if do_thumbnail is not None else self.do_thumbnail
    _lowerCamelCase      =						do_align_long_axis if do_align_long_axis is not None else self.do_align_long_axis
    _lowerCamelCase      =						do_pad if do_pad is not None else self.do_pad
    _lowerCamelCase      =						do_rescale if do_rescale is not None else self.do_rescale
    _lowerCamelCase      =						rescale_factor if rescale_factor is not None else self.rescale_factor
    _lowerCamelCase      =						do_normalize if do_normalize is not None else self.do_normalize
    _lowerCamelCase      =						image_mean if image_mean is not None else self.image_mean
    _lowerCamelCase      =						image_std if image_std is not None else self.image_std

    _lowerCamelCase      =						make_list_of_images(a__		)

    if not valid_images(a__		):
      raise ValueError(
          '''Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, '''
          '''torch.Tensor, tf.Tensor or jax.ndarray.'''		)

    if do_resize and size is None:
      raise ValueError('''Size must be specified if do_resize is True.'''		)

    if do_rescale and rescale_factor is None:
      raise ValueError('''Rescale factor must be specified if do_rescale is True.'''		)

    if do_pad and size is None:
      raise ValueError('''Size must be specified if do_pad is True.'''		)

    if do_normalize and (image_mean is None or image_std is None):
      raise ValueError('''Image mean and std must be specified if do_normalize is True.'''		)

    # All transformations expect numpy arrays.
    _lowerCamelCase      =						[to_numpy_array(a__		) for image in images]

    if do_align_long_axis:
      _lowerCamelCase      =						[self.align_long_axis(a__	,		size=a__		) for image in images]

    if do_resize:
      _lowerCamelCase      =						[self.resize(image=a__	,		size=a__	,		resample=a__		) for image in images]

    if do_thumbnail:
      _lowerCamelCase      =						[self.thumbnail(image=a__	,		size=a__		) for image in images]

    if do_pad:
      _lowerCamelCase      =						[self.pad_image(image=a__	,		size=a__	,		random_padding=a__		) for image in images]

    if do_rescale:
      _lowerCamelCase      =						[self.rescale(image=a__	,		scale=a__		) for image in images]

    if do_normalize:
      _lowerCamelCase      =						[self.normalize(image=a__	,		mean=a__	,		std=a__		) for image in images]

    _lowerCamelCase      =						[to_channel_dimension_format(a__	,		a__		) for image in images]

    _lowerCamelCase      =						{'''pixel_values''': images}
    return BatchFeature(data=a__	,		tensor_type=a__		)




