code: def main():
      parser = argparse.ArgumentParser(description="Evaluate the model")
      parser.add_argument(
          "--ssl_model_ckpt_path",
          type=str,
          required=True,
      )
      parser.add_argument("--manifest_paths", type=str, required=True)
      parser.add_argument("--sup_data_dir", type=str, default=None)
      parser.add_argument("--batch_size", type=int, default=32)
      parser.add_argument(
          "--ssl_content_emb_type",
          type=str,
          default="embedding_and_probs",
      )
      parser.add_argument("--use_unique_tokens", type=int, default=1)
      parser.add_argument("--num_workers", type=int, default=8)
      parser.add_argument("--pool_workers", type=int, default=30)
      parser.add_argument(
          "--compute_pitch_contours", type=int, default=1
      )
      parser.add_argument(
          "--num_pitch_per_speaker", type=int, default=None
      )  # saves time.
      parser.add_argument("--sample_rate", type=int, default=22050)
      parser.add_argument("--pad_multiple", type=int, default=1024)
      parser.add_argument(
          "--ssl_downsampling_factor", type=int, default=4
      )
      parser.add_argument("--stft_n_fft", type=int, default=1024)
      parser.add_argument("--stft_hop_length", type=int, default=256)
      parser.add_argument("--stft_win_length", type=int, default=1024)
      parser.add_argument("--stft_n_mel", type=int, default=80)
      parser.add_argument("--stft_fmin", type=int, default=0)
      parser.add_argument("--stft_fmax", type=int, default=8000)
      parser.add_argument("--yin_fmax", type=int, default=500)
      parser.add_argument("--segment_length", type=int, default=44100)
      parser.add_argument("--segment_hop_size", type=int, default=22050)
      parser.add_argument(
          "--min_segment_length", type=int, default=22050
      )
  
      args = parser.parse_args()
  
      device = "cuda:0" if torch.cuda.is_available() else "cpu"
  
      manifest_paths = args.manifest_paths.split(",")
      ssl_model_ckpt_path = args.ssl_model_ckpt_path
  
      dataset = AudioDataset(
          manifest_paths,
          pad_multiple=args.pad_multiple,
          sample_rate=args.sample_rate,
          sup_data_dir=args.sup_data_dir,
      )
      dataloader = torch.utils.data.DataLoader(
          dataset,
          batch_size=args.batch_size,
          shuffle=False,
          collate_fn=dataset.pad_collate_fn,
          num_workers=args.num_workers,
      )
  
      ssl_model = ssl_tts.SSLDisentangler.load_from_checkpoint(
          ssl_model_ckpt_path, strict=False
      )
      with open_dict(ssl_model.cfg):
          ssl_model.cfg.preprocessor.exact_pad = True
      ssl_model.preprocessor = hydra.utils.instantiate(
          ssl_model.cfg.preprocessor
      )
      ssl_model.preprocessor_disentangler = ssl_model.preprocessor
      ssl_model.eval()
      ssl_model.to(device)
  
      sample_rate = args.sample_rate
      stft_params = {
          "n_fft": args.stft_n_fft,
          "hop_length": args.stft_hop_length,
          "win_length": args.stft_win_length,
          "n_mel": args.stft_n_mel,
          "sample_rate": sample_rate,
          "pad_multiple": args.pad_multiple,
          "fmin": args.stft_fmin,
          "fmax": args.stft_fmax,
          "yin_fmax": args.yin_fmax,
      }
  
      fb = (
          torch.tensor(
              librosa.filters.mel(
                  sr=sample_rate,
                  n_fft=stft_params["n_fft"],
                  n_mels=stft_params["n_mel"],
                  fmin=stft_params["fmin"],
                  fmax=stft_params["fmax"],
              ),
              dtype=torch.float,
          )
          .unsqueeze(0)
          .to(device)
      )
  
      st = time.time()
      bidx = 0
      wav_and_id_list = []
  
      for batch in tqdm(dataloader):
          bidx += 1
          with torch.no_grad():
              (
                  _,
                  _,
                  batch_content_embedding,
                  batch_content_log_probs,
                  batch_encoded_len,
              ) = ssl_model.forward_for_export(
                  input_signal=batch["audio"].to(device),
                  input_signal_length=batch["audio_len"].to(device),
                  normalize_content=True,
              )
  
              batch_mel_specs = get_mel_spectrogram(
                  fb, batch["audio"][:, :-1].to(device), stft_params
              )
              audio_segmented, segment_indices = segment_batch(
                  batch,
                  args.segment_length,
                  args.segment_hop_size,
                  args.min_segment_length,
              )
              audio_seg_len = (
                  torch.tensor(
                      [len(segment) for segment in audio_segmented]
                  )
                  .to(device)
                  .long()
              )
  
              (
                  _,
                  batch_speaker_embeddings,
                  _,
                  _,
                  _,
              ) = ssl_model.forward_for_export(
                  input_signal=audio_segmented.to(device),
                  input_signal_length=audio_seg_len,
                  normalize_content=True,
              )
  
              for idx in range(batch["audio"].shape[0]):
                  _speaker = batch["speaker"][idx].item()
                  wav_path = batch["wav_path"][idx]
  
                  wav_id = batch["rel_audio_path_as_text_id"][idx]
                  wav_and_id_list.append((wav_path, wav_id, _speaker))
                  content_embedding = batch_content_embedding[
                      idx
                  ].detach()
                  content_log_probs = batch_content_log_probs[
                      :, idx, :
                  ].detach()  # (content lob prob is (t, b, c))
                  encoded_len = batch_encoded_len[idx].detach()
                  content_embedding = content_embedding[
                      : encoded_len.item()
                  ]
                  content_embedding = content_embedding.t()
                  content_log_probs = content_log_probs[
                      : encoded_len.item()
                  ]
                  content_log_probs = content_log_probs.t()
                  content_probs = torch.exp(content_log_probs)
  
                  duration = (
                      torch.ones(content_embedding.shape[1])
                      * args.ssl_downsampling_factor
                  )
  
                  bsi_start = segment_indices[idx][0]
                  bsi_end = segment_indices[idx][1]
                  speaker_embedding = torch.mean(
                      batch_speaker_embeddings[bsi_start : bsi_end + 1],
                      dim=0,
                  )
  
                  l2_norm = torch.norm(speaker_embedding, p=2)
                  speaker_embedding = speaker_embedding / l2_norm
  
                  if args.ssl_content_emb_type == "probs":
                      # content embedding is only character probabilities
                      final_content_embedding = content_probs
                  elif args.ssl_content_emb_type == "embedding":
                      # content embedding is only output of content head of SSL backbone
                      final_content_embedding = content_embedding
                  elif args.ssl_content_emb_type == "log_probs":
                      # content embedding is only log of character probabilities
                      final_content_embedding = content_log_probs
                  elif (
                      args.ssl_content_emb_type == "embedding_and_probs"
                  ):
                      # content embedding is the concatenation of character probabilities and output of content head of SSL backbone
                      final_content_embedding = torch.cat(
                          [content_embedding, content_probs], dim=0
                      )
  
                  if args.use_unique_tokens == 1:
                      # group content embeddings with same predicted token (by averaging) and add the durations of the grouped embeddings
                      # Eg. By default each content embedding corresponds to 4 frames of spectrogram (ssl_downsampling_factor)
                      # If we group 3 content embeddings, the duration of the grouped embedding will be 12 frames.
                      # This is useful for adapting the duration during inference based on the speaker.
                      token_predictions = torch.argmax(
                          content_probs, dim=0
                      )
                      content_buffer = [final_content_embedding[:, 0]]
                      unique_content_embeddings = []
                      unique_tokens = []
                      durations = []
                      for _t in range(
                          1, final_content_embedding.shape[1]
                      ):
                          if (
                              token_predictions[_t]
                              == token_predictions[_t - 1]
                          ):
                              content_buffer.append(
                                  final_content_embedding[:, _t]
                              )
                          else:
                              durations.append(
                                  len(content_buffer)
                                  * args.ssl_downsampling_factor
                              )
                              unique_content_embeddings.append(
                                  torch.mean(
                                      torch.stack(content_buffer), dim=0
                                  )
                              )
                              content_buffer = [
                                  final_content_embedding[:, _t]
                              ]
                              unique_tokens.append(
                                  token_predictions[_t].item()
                              )
  
                      if len(content_buffer) > 0:
                          durations.append(
                              len(content_buffer)
                              * args.ssl_downsampling_factor
                          )
                          unique_content_embeddings.append(
                              torch.mean(
                                  torch.stack(content_buffer), dim=0
                              )
                          )
                          unique_tokens.append(
                              token_predictions[_t].item()
                          )
  
                      unique_content_embedding = torch.stack(
                          unique_content_embeddings
                      )
                      final_content_embedding = (
                          unique_content_embedding.t()
                      )
                      duration = torch.tensor(durations).float()
  
                  mel_len = int(
                      batch["audio_len"][idx].item()
                      / stft_params["hop_length"]
                  )
                  item_mel = batch_mel_specs[idx][:, :mel_len]
  
                  wav_text_id = batch["rel_audio_path_as_text_id"][idx]
                  content_emb_fn = f"{args.ssl_content_emb_type}_content_embedding_{wav_text_id}.pt"
                  speaker_emb_fn = f"speaker_embedding_{wav_text_id}.pt"
                  duration_fn = f"duration_embedding_{wav_text_id}.pt"  # embedding just for namesake
                  content_emb_fp = os.path.join(
                      dataset.sup_data_dir, content_emb_fn
                  )
                  speaker_emb_fp = os.path.join(
                      dataset.sup_data_dir, speaker_emb_fn
                  )
                  duration_fp = os.path.join(
                      dataset.sup_data_dir, duration_fn
                  )
  
                  mel_spec_fn = f"mel_spec_{wav_text_id}.pt"
                  mel_spec_fp = os.path.join(
                      dataset.sup_data_dir, mel_spec_fn
                  )
  
                  torch.save(item_mel.cpu(), mel_spec_fp)
                  torch.save(
                      final_content_embedding.cpu(), content_emb_fp
                  )
                  torch.save(speaker_embedding.cpu(), speaker_emb_fp)
                  torch.save(duration.cpu(), duration_fp)
  
              et = time.time()
              logging.info(
                  "Processed Batch {} of {} | Time per batch: {:.4f} s".format(
                      bidx + 1, len(dataloader), (et - st) / bidx
                  )
              )
  
      if args.compute_pitch_contours == 1:
          speaker_wise_records = {}
          for row in wav_and_id_list:
              wav_path, wav_id, speaker = row
              if speaker not in speaker_wise_records:
                  speaker_wise_records[speaker] = []
              speaker_wise_records[speaker].append(
                  {
                      "wav_path": wav_path,
                      "wav_id": wav_id,
                      "sup_data_dir": dataset.sup_data_dir,
                      "stft_params": stft_params,
                      "speaker": speaker,
                  }
              )
  
          filtered_records = []
          for speaker in speaker_wise_records:
              if args.num_pitch_per_speaker is not None:
                  filtered_records += speaker_wise_records[speaker][
                      : args.num_pitch_per_speaker
                  ]
              else:
                  filtered_records += speaker_wise_records[speaker]
  
          with Pool(args.pool_workers) as p:
              p.map(save_pitch_contour, filtered_records)
  
          compute_pitch_stats(filtered_records)
  
imports: import argparse
  import json
  import os
  import time
  from multiprocessing import Pool
  from pathlib import Path
  import hydra.utils
  import librosa
  import numpy
  import torch
  from omegaconf import open_dict
  from tqdm import tqdm
  from nemo.collections.asr.parts.preprocessing.segment import AudioSegment
  from nemo.collections.tts.models import ssl_tts
  from nemo.collections.tts.parts.utils.tts_dataset_utils import get_base_dir
  from nemo.core.classes import Dataset
  from nemo.utils import logging
label_desc: Use image node_id 7 for calling a global function with example usage: main() without return types