def get_tokens_per_batch_func(
    decoder_only: bool = True,
) -> Callable[[Batch], Union[int, dict[str, int]]]:
    """Returns a callable that counts the number of tokens in a batch.

    Args:
        pad_token_id (int): The id of the padding token.
        decoder_only (bool, optional): Whether to expect the batch to just contain ``input_ids`` (decoder only)
            or to also contain ``decoder_input_ids`` (encoder decoder). Defaults to ``True``.

    Returns:
        Callable[[Batch], int]: A callable that counts the number of tokens in a batch.
    """

    def get_num_tokens_in_batch(batch: Batch) -> Union[int, dict[str, int]]:
        if not isinstance(batch, Mapping) or (
            'attention_mask' not in batch and 'input_ids' not in batch
        ):
            raise ValueError(
                'get_tokens_per_batch_func() requires a batch with an attention_mask key or an input_ids key',
            )

        if not decoder_only and 'decoder_attention_mask' not in batch:
            raise ValueError(
                'get_tokens_per_batch_func() for encoder decoder requires a batch with a decoder_attention_mask key',
            )

        # Short cut if the dataloader has already calculated the number of tokens
        if 'total_tokens' in batch and 'loss_generating_tokens' in batch:
            return {
                'total': sum(batch['total_tokens']),
                'loss_generating': sum(batch['loss_generating_tokens']),
            }

        # Count number of non padding tokens in batch
        if 'attention_mask' in batch:
            input_ids_tokens = int(torch.sum(batch['attention_mask']).item())
        else:
            input_ids_tokens = batch['input_ids'].numel()

        loss_generating_tokens = None
        if 'labels' in batch:
            loss_generating_tokens = (
                batch['labels'].shape[0] * (batch['labels'].shape[1] - 1)
            ) - torch.count_nonzero(
                torch.eq(batch['labels'][..., 1:], CROSS_ENTROPY_IGNORE_INDEX),
            )

        # For encoder decoder models only
        decoder_input_ids_tokens = 0
        if not decoder_only:
            decoder_input_ids_tokens = int(
                torch.sum(batch['decoder_attention_mask']).item(),
            )

        if loss_generating_tokens is not None:
            return {
                'total': input_ids_tokens + decoder_input_ids_tokens,
                'loss_generating': loss_generating_tokens,
            }
        return input_ids_tokens + decoder_input_ids_tokens

    return get_num_tokens_in_batch
Returns a callable that counts the number of tokens in a batch.

    Args:
        pad_token_id (int): The id of the padding token.
        decoder_only (bool, optional): Whether to expect the batch to just contain ``input_ids`` (decoder only)
            or to also contain ``decoder_input_ids`` (encoder decoder). Defaults to ``True``.

    Returns:
        Callable[[Batch], int]: A callable that counts the number of tokens in a batch.
    
llmfoundry/data/utils.py