docstr: Args:
      input_ids_query (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
          Indices of input sequence tokens in the vocabulary for the queries in a batch.
  
          Indices can be obtained using [`RetriBertTokenizer`]. See [`PreTrainedTokenizer.encode`] and
          [`PreTrainedTokenizer.__call__`] for details.
  
          [What are input IDs?](../glossary#input-ids)
      attention_mask_query (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):
          Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
  
          - 1 for tokens that are **not masked**,
          - 0 for tokens that are **masked**.
  
          [What are attention masks?](../glossary#attention-mask)
      input_ids_doc (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
          Indices of input sequence tokens in the vocabulary for the documents in a batch.
      attention_mask_doc (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):
          Mask to avoid performing attention on documents padding token indices.
      checkpoint_batch_size (`int`, *optional*, defaults to `-1`):
          If greater than 0, uses gradient checkpointing to only compute sequence representation on
          `checkpoint_batch_size` examples at a time on the GPU. All query representations are still compared to
          all document representations in the batch.
  
  Return:
      `torch.FloatTensor``: The bidirectional cross-entropy loss obtained while trying to match each query to its
      corresponding document and each document to its corresponding query in the batch
docstr_tok: ["Args",":","input_ids_query","(","`","torch.LongTensor","`","of","shape","`","(","batch_size",",","sequence_length",")","`",")",":","Indices","of","input","sequence","tokens","in","the","vocabulary","for","the","queries","in","a","batch",".","Indices","can","be","obtained","using","[","`","RetriBertTokenizer","`","]",".","See","[","`","PreTrainedTokenizer.encode","`","]","and","[","`","PreTrainedTokenizer.__call__","`","]","for","details",".","[","What","are","input","IDs","?","]","(","..","\/glossary","#","input-ids",")","attention_mask_query","(","`","torch.FloatTensor","`","of","shape","`","(","batch_size",",","sequence_length",")","`",",","*","optional","*",")",":","Mask","to","avoid","performing","attention","on","padding","token","indices",".","Mask","values","selected","in","`","[","0",",","1","]","`",":","-","1","for","tokens","that","are","*","*","not","masked","*","*",",","-","0","for","tokens","that","are","*","*","masked","*","*",".","[","What","are","attention","masks","?","]","(","..","\/glossary","#","attention-mask",")","input_ids_doc","(","`","torch.LongTensor","`","of","shape","`","(","batch_size",",","sequence_length",")","`",")",":","Indices","of","input","sequence","tokens","in","the","vocabulary","for","the","documents","in","a","batch",".","attention_mask_doc","(","`","torch.FloatTensor","`","of","shape","`","(","batch_size",",","sequence_length",")","`",",","*","optional","*",")",":","Mask","to","avoid","performing","attention","on","documents","padding","token","indices",".","checkpoint_batch_size","(","`","int","`",",","*","optional","*",",","defaults","to","`","-1","`",")",":","If","greater","than","0",",","uses","gradient","checkpointing","to","only","compute","sequence","representation","on","`","checkpoint_batch_size","`","examples","at","a","time","on","the","GPU",".","All","query","representations","are","still","compared","to","all","document","representations","in","the","batch",".","Return",":","`","torch.FloatTensor","``",":","The","bidirectional","cross-entropy","loss","obtained","while","trying","to","match","each","query","to","its","corresponding","document","and","each","document","to","its","corresponding","query","in","the","batch"]
code: def forward(
      self,
      input_ids_query: torch.LongTensor,
      attention_mask_query: Optional[torch.FloatTensor],
      input_ids_doc: torch.LongTensor,
      attention_mask_doc: Optional[torch.FloatTensor],
      checkpoint_batch_size: int = -1,
  ) -> torch.FloatTensor:
      r"""
      Args:
          input_ids_query (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
              Indices of input sequence tokens in the vocabulary for the queries in a batch.
  
              Indices can be obtained using [`RetriBertTokenizer`]. See [`PreTrainedTokenizer.encode`] and
              [`PreTrainedTokenizer.__call__`] for details.
  
              [What are input IDs?](../glossary#input-ids)
          attention_mask_query (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):
              Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
  
              - 1 for tokens that are **not masked**,
              - 0 for tokens that are **masked**.
  
              [What are attention masks?](../glossary#attention-mask)
          input_ids_doc (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
              Indices of input sequence tokens in the vocabulary for the documents in a batch.
          attention_mask_doc (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):
              Mask to avoid performing attention on documents padding token indices.
          checkpoint_batch_size (`int`, *optional*, defaults to `-1`):
              If greater than 0, uses gradient checkpointing to only compute sequence representation on
              `checkpoint_batch_size` examples at a time on the GPU. All query representations are still compared to
              all document representations in the batch.
  
      Return:
          `torch.FloatTensor``: The bidirectional cross-entropy loss obtained while trying to match each query to its
          corresponding document and each document to its corresponding query in the batch
      """
      device = input_ids_query.device
      q_reps = self.embed_questions(
          input_ids_query, attention_mask_query, checkpoint_batch_size
      )
      a_reps = self.embed_answers(
          input_ids_doc, attention_mask_doc, checkpoint_batch_size
      )
      compare_scores = torch.mm(q_reps, a_reps.t())
      loss_qa = self.ce_loss(
          compare_scores,
          torch.arange(compare_scores.shape[1]).to(device),
      )
      loss_aq = self.ce_loss(
          compare_scores.t(),
          torch.arange(compare_scores.shape[0]).to(device),
      )
      loss = (loss_qa + loss_aq) / 2
      return loss
  
file_path: H2O/h2o_flexgen/benchmark/third_party/transformers/src/transformers/models/retribert/modeling_retribert.py
imports: import math
  from typing import Optional
  import torch
  import torch.utils.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
label_desc: Use image node_id 5 for calling the RetriBertModel obj's underlying member method code with example usage: obj.forward(input_ids_query, attention_mask_query, input_ids_doc, attention_mask_doc, checkpoint_batch_size) and returns: loss