docstr: Collate function for dataloader
  Can accept mixed batch of text-to-text items and audio-text items (typical for ASR)
docstr_tok: ["Collate","function","for","dataloader","Can","accept","mixed","batch","of","text-to-text","items","and","audio-text","items","(","typical","for","ASR",")"]
code: def collate_fn(
      batch: List[Union[TextToTextItem, tuple]],
      tts_text_pad_id: int,
      asr_pad_id: int,
  ) -> Union[TextToTextBatch, TextOrAudioToTextBatch, tuple]:
      """
      Collate function for dataloader
      Can accept mixed batch of text-to-text items and audio-text items (typical for ASR)
      """
      text_items: List[TextToTextItem] = [
          item for item in batch if isinstance(item, TextToTextItem)
      ]
      if not text_items:
          # pure audio-text batch
          return _speech_collate_fn(batch=batch, pad_id=asr_pad_id)
  
      asr_items = [
          item for item in batch if not isinstance(item, TextToTextItem)
      ]
  
      if not asr_items:
          # pure text-to-text batch
          return TextToTextBatch.collate_fn(
              batch=text_items,
              asr_pad_id=asr_pad_id,
              tts_text_pad_id=tts_text_pad_id,
          )
  
      # mixed batch
  
      # each asr item is a tuple:
      # audio_signal (0), audio_length (1), transcript (2), transcript_length (3), sample_id (4, optional)
      audio_signals = pad_sequence(
          [item[0] for item in asr_items],
          batch_first=True,
          padding_value=0.0,
      )
      audio_signal_lengths = torch.tensor(
          [item[1] for item in asr_items]
      ).long()
  
      tts_texts = pad_sequence(
          [item.tts_text for item in text_items],
          batch_first=True,
          padding_value=tts_text_pad_id,
      )
      tts_text_lengths = torch.tensor(
          [item.tts_text.shape[0] for item in text_items]
      ).long()
      speakers = torch.tensor(
          [item.speaker for item in text_items]
      ).long()
  
      transcripts = pad_sequence(
          [item.transcript for item in text_items]
          + [item[2] for item in asr_items],
          batch_first=True,
          padding_value=asr_pad_id,
      )
      transcript_lengths = torch.tensor(
          [item.transcript.shape[0] for item in text_items]
          + [item[3] for item in asr_items]
      ).long()
  
      return TextOrAudioToTextBatch(
          audio_signals=audio_signals,
          audio_signal_lengths=audio_signal_lengths,
          tts_texts=tts_texts,
          tts_text_lengths=tts_text_lengths,
          speakers=speakers,
          transcripts=transcripts,
          transcript_lengths=transcript_lengths,
      )
  
imports: from __future__ import annotations
  import concurrent.futures
  import copy
  import gc
  import json
  import math
  import random
  from pathlib import Path
  from typing import Any, Callable, Dict, Iterable, List, NamedTuple, Optional, Set, Union
  import numpy
  import torch
  import torch.utils.data
  from torch.nn.utils.rnn import pad_sequence
  from tqdm.auto import tqdm
  from nemo.collections.asr.data.audio_to_text import _speech_collate_fn
  from nemo.collections.common.tokenizers import TokenizerSpec
  from nemo.core.classes import Dataset, IterableDataset
  from nemo.utils import logging
label_desc: Use image node_id 1 for calling the TextOrAudioToTextBatch obj's underlying member method code with example usage: obj.collate_fn(batch, tts_text_pad_id, asr_pad_id) and returns: TextOrAudioToTextBatch, _speech_collate_fn, TextToTextBatch