code: def get_modality_length_grouped_indices(
      lengths, batch_size, world_size, generator=None
  ):
      # We need to use torch for the random part as a distributed sampler will set the random seed for torch.
      assert all(
          l != 0 for l in lengths
      ), "Should not have zero length."
      mm_indices, mm_lengths = zip(
          *[(i, l) for i, l in enumerate(lengths) if l > 0]
      )
      lang_indices, lang_lengths = zip(
          *[(i, -l) for i, l in enumerate(lengths) if l < 0]
      )
  
      assert (
          len(mm_indices) > 0
      ), "Should have at least one multimodal sample."
      assert (
          len(lang_indices) > 0
      ), "Should have at least one language sample."
  
      mm_shuffle = [
          mm_indices[i]
          for i in get_length_grouped_indices(
              mm_lengths, batch_size, world_size, generator=None
          )
      ]
      lang_shuffle = [
          lang_indices[i]
          for i in get_length_grouped_indices(
              lang_lengths, batch_size, world_size, generator=None
          )
      ]
      megabatch_size = world_size * batch_size
      mm_megabatches = [
          mm_shuffle[i : i + megabatch_size]
          for i in range(0, len(mm_shuffle), megabatch_size)
      ]
      lang_megabatches = [
          lang_shuffle[i : i + megabatch_size]
          for i in range(0, len(lang_shuffle), megabatch_size)
      ]
  
      last_mm = mm_megabatches[-1]
      last_lang = lang_megabatches[-1]
      additional_batch = last_mm + last_lang
      megabatches = mm_megabatches[:-1] + lang_megabatches[:-1]
      megabatch_indices = torch.randperm(
          len(megabatches), generator=generator
      )
      megabatches = [megabatches[i] for i in megabatch_indices]
  
      if len(additional_batch) >= megabatch_size:
          megabatches = [
              additional_batch[:megabatch_size]
          ] + megabatches
          additional_batch = additional_batch[megabatch_size:]
  
      if len(additional_batch) > 0:
          megabatches.append(additional_batch)
  
      return [i for megabatch in megabatches for i in megabatch]
  
imports: import os
  import torch
  import torch.nn
  from fairscale.optim import OSS
  from torch.utils.data import Sampler
  from transformers import Trainer
  from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS
  from transformers.trainer_pt_utils import get_parameter_names
  from transformers.trainer_utils import ShardedDDPOption, has_length
  from transformers.utils import is_sagemaker_mp_enabled, logging
  from typing import List, Optional, Union, Dict, Any
label_desc: Use image node_id 4 for calling a global function with example usage: get_modality_length_grouped_indices(lengths, batch_size, world_size, generator) and returns: unknown