code: def _batch_by_token_count(
      idx_target_lengths, max_tokens, batch_size=None
  ):
      batches = []
      current_batch = []
      current_token_count = 0
      for idx, target_length in idx_target_lengths:
          if current_token_count + target_length > max_tokens or (
              batch_size and len(current_batch) == batch_size
          ):
              batches.append(current_batch)
              current_batch = [idx]
              current_token_count = target_length
          else:
              current_batch.append(idx)
              current_token_count += target_length
  
      if current_batch:
          batches.append(current_batch)
  
      return batches
  
imports: import os
  import random
  import torch
  import torchaudio
  from pytorch_lightning import LightningDataModule
label_desc: Use image node_id 1 for calling a global function with example usage: _batch_by_token_count(idx_target_lengths, max_tokens, batch_size) and returns: batches