def _compute_scores(
        self,
        batch: MutableMapping,
        outputs: CausalLMOutputWithPast,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Run pairs through the encoder separately in two passes.

        This function splits queries and passages based on the step size, which represents
        the number of hard negatives per passage. It then runs the queries and passages
        through the encoder separately to obtain the encoded representations. The encoded
        representations are used for further computations in the model.

        Args:
            batch (MutableMapping): The input batch containing queries and passages.
            outputs (CausalLMOutputWithPast): The model outputs.

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: The encoded representations of queries and passages.
        """
        hidden_state = self.get_hidden_state(outputs)
        zero = self.handle_language_head(outputs)
        (
            queries_batch,
            queries_last_hidden_state,
        ) = self.format_queries_batch(batch, hidden_state)
        (
            passages_batch,
            passages_last_hidden_state,
        ) = self.format_passages_batch(batch, hidden_state)

        query_attn_mask = queries_batch.get('attention_mask')
        passage_attn_mask = passages_batch.get('attention_mask')
        assert isinstance(query_attn_mask, torch.Tensor)
        assert isinstance(passage_attn_mask, torch.Tensor)
        if self.vector_representation == 'eos':

            def pool_fn(x: torch.Tensor, mask: torch.Tensor):
                row_indices = torch.arange(mask.shape[0])
                flipped_mask = ~mask.bool()
                last_true_indices = flipped_mask.int().argmax(dim=1) - 1
                pooled_outputs = x[row_indices, last_true_indices, :]
                return pooled_outputs
        elif self.vector_representation == 'avg':

            def pool_fn(x: torch.Tensor, mask: torch.Tensor):
                x = x.masked_fill(~mask[..., None].bool(), 0.0)
                pooled_outputs = x.sum(dim=1) / (mask.sum(dim=1)[..., None])
                return pooled_outputs
        else:
            raise ValueError(
                f'Specified vector_representation={self.vector_representation} not recognized. `vector_representation` must be one of [`avg`, `eos`].',
            )

        q_pooled_outputs = pool_fn(queries_last_hidden_state, query_attn_mask)
        p_pooled_outputs = pool_fn(
            passages_last_hidden_state,
            passage_attn_mask,
        )

        if self.normalize_output:
            q_pooled_outputs = F.normalize(q_pooled_outputs, dim=-1)
            p_pooled_outputs = F.normalize(p_pooled_outputs, dim=-1)

        # Use all_gather to include negatives across mini batch
        if self.gather_in_batch_negatives:
            all_q_pooled_outputs = self._cat_gather(
                q_pooled_outputs,
                group=self.infonce_process_group,
            )
            all_p_pooled_outputs = self._cat_gather(
                p_pooled_outputs,
                group=self.infonce_process_group,
            )
        else:
            all_q_pooled_outputs = q_pooled_outputs
            all_p_pooled_outputs = p_pooled_outputs

        assert all_q_pooled_outputs is not None
        assert all_p_pooled_outputs is not None

        all_scores = self._full_contrastive_scores(
            queries=all_q_pooled_outputs,
            passages=all_p_pooled_outputs,
        )
        all_scores = all_scores * (1 / self.temperature) + zero

        all_labels = torch.arange(
            all_scores.size(0),
            device=q_pooled_outputs.device,
            dtype=torch.long,
        )
        all_labels = all_labels * (
            p_pooled_outputs.size(0) // q_pooled_outputs.size(0)
        )

        return all_scores, all_labels
Run pairs through the encoder separately in two passes.

        This function splits queries and passages based on the step size, which represents
        the number of hard negatives per passage. It then runs the queries and passages
        through the encoder separately to obtain the encoded representations. The encoded
        representations are used for further computations in the model.

        Args:
            batch (MutableMapping): The input batch containing queries and passages.
            outputs (CausalLMOutputWithPast): The model outputs.

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: The encoded representations of queries and passages.
        
llmfoundry/models/llm_embed/modeling_llm_embed.py