code: def load_textgrid_wav(
      data_dir: Path,
      token_seq_len: int,
      batch_size,
      pad_wav_len,
      mode: str,
  ):
      """load wav and textgrid files to memory."""
      tg_files = sorted(data_dir.glob("*.TextGrid"))
      random.Random(42).shuffle(tg_files)
      L = len(tg_files) * 95 // 100
      assert mode in ["train", "val", "gta"]
      phonemes = load_phonemes_set()
      if mode == "gta":
          tg_files = tg_files  # all files
      elif mode == "train":
          tg_files = tg_files[:L]
      elif mode == "val":
          tg_files = tg_files[L:]
  
      data = []
      for fn in tg_files:
          ps, ds = zip(*load_textgrid(fn))
          ps = [phonemes.index(p) for p in ps]
          l = len(ps)
          ps = pad_seq(ps, token_seq_len, 0)
          ds = pad_seq(ds, token_seq_len, 0)
  
          wav_file = data_dir / f"{fn.stem}.wav"
          sr, y = wavfile.read(wav_file)
          y = np.copy(y)
          start_time = 0
          for i, (phone_idx, duration) in enumerate(zip(ps, ds)):
              l = int(start_time * sr)
              end_time = start_time + duration
              r = int(end_time * sr)
              if i == len(ps) - 1:
                  r = len(y)
              if phone_idx < len(FLAGS.special_phonemes):
                  y[l:r] = 0
              start_time = end_time
  
          if len(y) > pad_wav_len:
              y = y[:pad_wav_len]
  
          # # normalize to match hifigan preprocessing
          # y = y.astype(np.float32)
          # y = y / np.max(np.abs(y))
          # y = y * 0.95
          # y = y * (2 ** 15)
          # y = y.astype(np.int16)
  
          wav_length = len(y)
          y = np.pad(y, (0, pad_wav_len - len(y)))
          data.append((fn.stem, ps, ds, l, y, wav_length))
  
      batch = []
      while True:
          random.shuffle(data)
          for idx, e in enumerate(data):
              batch.append(e)
              if len(batch) == batch_size or (
                  mode == "gta" and idx == len(data) - 1
              ):
                  names, ps, ds, lengths, wavs, wav_lengths = zip(
                      *batch
                  )
                  ps = np.array(ps, dtype=np.int32)
                  ds = np.array(ds, dtype=np.float32)
                  lengths = np.array(lengths, dtype=np.int32)
                  wavs = np.array(wavs, dtype=np.int16)
                  wav_lengths = np.array(wav_lengths, dtype=np.int32)
                  if mode == "gta":
                      yield names, AcousticInput(
                          ps, lengths, ds, wavs, wav_lengths, None
                      )
                  else:
                      yield AcousticInput(
                          ps, lengths, ds, wavs, wav_lengths, None
                      )
                  batch = []
          if mode == "gta":
              assert len(batch) == 0
              break
  
file_path: underthesea/underthesea/pipeline/say/viettts_/nat/data_loader.py
imports: import random
  from pathlib import Path
  import numpy
  from scipy.io import wavfile
  from .config import FLAGS, AcousticInput, DurationInput
label_desc: Use image node_id 6 for calling a global function with example usage: load_textgrid_wav(data_dir, token_seq_len, batch_size, pad_wav_len, mode) without return types