def format_queries_batch(
        self,
        batch: MutableMapping,
        last_hidden_state: torch.Tensor,
    ) -> tuple[dict[str, torch.Tensor], torch.Tensor]:
        """Format `queries` by selecting every ``n``th entry from the batch.

        Here ``n`` is the step size, which represents the number of hard
        negatives per passage.
        """
        assert self.step_size
        queries = {}
        indices = list(range(0, batch['input_ids'].size(0), self.step_size))
        for key in batch:
            queries[key] = batch[key][indices, :]
        return queries, last_hidden_state[indices, :, :]
Format `queries` by selecting every ``n``th entry from the batch.

        Here ``n`` is the step size, which represents the number of hard
        negatives per passage.
        
llmfoundry/models/llm_embed/modeling_llm_embed.py