def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.FloatTensor] = None,
        token_type_ids: Optional[torch.LongTensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        head_mask: Optional[torch.FloatTensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        relevance_score: Optional[torch.FloatTensor] = None,
        block_mask: Optional[torch.BoolTensor] = None,
        start_positions: Optional[torch.LongTensor] = None,
        end_positions: Optional[torch.LongTensor] = None,
        has_answers: Optional[torch.BoolTensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, RealmReaderOutput]:
        r"""
        relevance_score (`torch.FloatTensor` of shape `(searcher_beam_size,)`, *optional*):
            Relevance score, which must be specified if you want to compute the logits and marginal log loss.
        block_mask (`torch.BoolTensor` of shape `(searcher_beam_size, sequence_length)`, *optional*):
            The mask of the evidence block, which must be specified if you want to compute the logits and marginal log
            loss.
        start_positions (`torch.LongTensor` of shape `(searcher_beam_size,)`, *optional*):
            Labels for position (index) of the start of the labelled span for computing the token classification loss.
            Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
            are not taken into account for computing the loss.
        end_positions (`torch.LongTensor` of shape `(searcher_beam_size,)`, *optional*):
            Labels for position (index) of the end of the labelled span for computing the token classification loss.
            Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
            are not taken into account for computing the loss.
        has_answers (`torch.BoolTensor` of shape `(searcher_beam_size,)`, *optional*):
            Whether or not the evidence block has answer(s).

        Returns:
        """
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        if relevance_score is None:
            raise ValueError("You have to specify `relevance_score` to calculate logits and loss.")
        if block_mask is None:
            raise ValueError("You have to specify `block_mask` to separate question block and evidence block.")
        if token_type_ids.size(1) < self.config.max_span_width:
            raise ValueError("The input sequence length must be greater than or equal to config.max_span_width.")
        outputs = self.realm(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        # [reader_beam_size, joint_seq_len, hidden_size]
        sequence_output = outputs[0]

        # [reader_beam_size, num_candidates], [num_candidates], [num_candidates]
        reader_logits, candidate_starts, candidate_ends = self.qa_outputs(
            sequence_output, block_mask[0 : self.config.reader_beam_size]
        )
        # [searcher_beam_size, 1]
        retriever_logits = torch.unsqueeze(relevance_score[0 : self.config.reader_beam_size], -1)
        # [reader_beam_size, num_candidates]
        reader_logits += retriever_logits
        # []
        predicted_block_index = torch.argmax(torch.max(reader_logits, dim=1).values)
        # []
        predicted_candidate = torch.argmax(torch.max(reader_logits, dim=0).values)
        # [1]
        predicted_start = torch.index_select(candidate_starts, dim=0, index=predicted_candidate)
        # [1]
        predicted_end = torch.index_select(candidate_ends, dim=0, index=predicted_candidate)

        total_loss = None
        retriever_loss = None
        reader_loss = None
        retriever_correct = None
        reader_correct = None
        if start_positions is not None and end_positions is not None and has_answers is not None:

            def compute_correct_candidates(candidate_starts, candidate_ends, gold_starts, gold_ends):
                """Compute correct span."""
                # [reader_beam_size, num_answers, num_candidates]
                is_gold_start = torch.eq(
                    torch.unsqueeze(torch.unsqueeze(candidate_starts, 0), 0), torch.unsqueeze(gold_starts, -1)
                )
                is_gold_end = torch.eq(
                    torch.unsqueeze(torch.unsqueeze(candidate_ends, 0), 0), torch.unsqueeze(gold_ends, -1)
                )

                # [reader_beam_size, num_candidates]
                return torch.any(torch.logical_and(is_gold_start, is_gold_end), 1)

            def marginal_log_loss(logits, is_correct):
                """Loss based on the negative marginal log-likelihood."""

                def mask_to_score(mask, dtype=torch.float32):
                    return (1.0 - mask.type(dtype)) * torch.finfo(dtype).min

                # []
                log_numerator = torch.logsumexp(logits + mask_to_score(is_correct, dtype=logits.dtype), dim=-1)
                log_denominator = torch.logsumexp(logits, dim=-1)
                return log_denominator - log_numerator

            # sometimes the start/end positions are outside our model inputs, we ignore these terms
            # `-1` is reserved for no answer.
            ignored_index = sequence_output.size(1)
            start_positions = start_positions.clamp(-1, ignored_index)
            end_positions = end_positions.clamp(-1, ignored_index)

            retriever_correct = has_answers
            any_retriever_correct = torch.any(retriever_correct)

            reader_correct = compute_correct_candidates(
                candidate_starts=candidate_starts,
                candidate_ends=candidate_ends,
                gold_starts=start_positions[0 : self.config.reader_beam_size],
                gold_ends=end_positions[0 : self.config.reader_beam_size],
            )
            any_reader_correct = torch.any(reader_correct)

            retriever_loss = marginal_log_loss(relevance_score, retriever_correct)
            reader_loss = marginal_log_loss(reader_logits.view(-1), reader_correct.view(-1))
            retriever_loss *= any_retriever_correct.type(torch.float32)
            reader_loss *= any_reader_correct.type(torch.float32)

            total_loss = (retriever_loss + reader_loss).mean()

        if not return_dict:
            output = (predicted_block_index, predicted_candidate, predicted_start, predicted_end) + outputs[2:]
            return (
                ((total_loss, retriever_loss, reader_loss, retriever_correct, reader_correct) + output)
                if total_loss is not None
                else output
            )

        return RealmReaderOutput(
            loss=total_loss,
            retriever_loss=retriever_loss,
            reader_loss=reader_loss,
            retriever_correct=retriever_correct,
            reader_correct=reader_correct,
            block_idx=predicted_block_index,
            candidate=predicted_candidate,
            start_pos=predicted_start,
            end_pos=predicted_end,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )

        relevance_score (`torch.FloatTensor` of shape `(searcher_beam_size,)`, *optional*):
            Relevance score, which must be specified if you want to compute the logits and marginal log loss.
        block_mask (`torch.BoolTensor` of shape `(searcher_beam_size, sequence_length)`, *optional*):
            The mask of the evidence block, which must be specified if you want to compute the logits and marginal log
            loss.
        start_positions (`torch.LongTensor` of shape `(searcher_beam_size,)`, *optional*):
            Labels for position (index) of the start of the labelled span for computing the token classification loss.
            Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
            are not taken into account for computing the loss.
        end_positions (`torch.LongTensor` of shape `(searcher_beam_size,)`, *optional*):
            Labels for position (index) of the end of the labelled span for computing the token classification loss.
            Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
            are not taken into account for computing the loss.
        has_answers (`torch.BoolTensor` of shape `(searcher_beam_size,)`, *optional*):
            Whether or not the evidence block has answer(s).

        Returns:
        
dpkd/transformers/src/transformers/models/realm/modeling_realm.py