docstr: Wraps an arbitrary index-style PyTorch dataset and provides a "textless" quantized view of it.
  It makes the following assumptions about `dataset[i]`:
  * it returns a tuple where the first element is a raw waveform,
  * all waveforms have the same sample rate as speech_encoder expects.
  
  Args:
      dataset (Dataset): Dataset to be wrapped
      speech_encoder (Callable): SpeechEncoder to encode the audio
      device (Optional[Union[str, torch.torch.device]], optional):
          Sets a device to be used for encoding speech. If is set to None, torch.cuda.current_device() will be used.
          When used in a dataloader with multiple workers, it might be useful to set `device` to "auto".
          This way, on the first call of __getitem__(), the QuantizeDataset checks if it runs in a dataloader
          worker. If this is the case, it will grab one of the available GPUs and place its copy of
          SpeechEncoder there. This could be useful as SpeechEncoder is typically GPU-intensive and it is a good idea to parallelize
          across multiple GPUs. Defaults to None.
      speaker_extractor (Optional[Callable], optional): An optional callable that extracts speaker id from the
      dataset[i] output. Only needed if speech_encoder uses per-speaker F0 normalization. Defaults to None.
docstr_tok: ["Wraps","an","arbitrary","index-style","PyTorch","dataset","and","provides","a","``","textless","''","quantized","view","of","it",".","It","makes","the","following","assumptions","about","`","dataset","[","i","]","`",":","*","it","returns","a","tuple","where","the","first","element","is","a","raw","waveform",",","*","all","waveforms","have","the","same","sample","rate","as","speech_encoder","expects",".","Args",":","dataset","(","Dataset",")",":","Dataset","to","be","wrapped","speech_encoder","(","Callable",")",":","SpeechEncoder","to","encode","the","audio","device","(","Optional","[","Union","[","str",",","torch.torch.device","]","]",",","optional",")",":","Sets","a","device","to","be","used","for","encoding","speech",".","If","is","set","to","None",",","torch.cuda.current_device","(",")","will","be","used",".","When","used","in","a","dataloader","with","multiple","workers",",","it","might","be","useful","to","set","`","device","`","to","``","auto","''",".","This","way",",","on","the","first","call","of","__getitem__","(",")",",","the","QuantizeDataset","checks","if","it","runs","in","a","dataloader","worker",".","If","this","is","the","case",",","it","will","grab","one","of","the","available","GPUs","and","place","its","copy","of","SpeechEncoder","there",".","This","could","be","useful","as","SpeechEncoder","is","typically","GPU-intensive","and","it","is","a","good","idea","to","parallelize","across","multiple","GPUs",".","Defaults","to","None",".","speaker_extractor","(","Optional","[","Callable","]",",","optional",")",":","An","optional","callable","that","extracts","speaker","id","from","the","dataset","[","i","]","output",".","Only","needed","if","speech_encoder","uses","per-speaker","F0","normalization",".","Defaults","to","None","."]
code: def __init__(
      self,
      dataset: Dataset,
      speech_encoder: Callable,
      device: Optional[Union[str, torch.torch.device]] = None,
      speaker_extractor: Optional[Callable] = None,
  ):
      """Wraps an arbitrary index-style PyTorch dataset and provides a "textless" quantized view of it.
      It makes the following assumptions about `dataset[i]`:
      * it returns a tuple where the first element is a raw waveform,
      * all waveforms have the same sample rate as speech_encoder expects.
  
      Args:
          dataset (Dataset): Dataset to be wrapped
          speech_encoder (Callable): SpeechEncoder to encode the audio
          device (Optional[Union[str, torch.torch.device]], optional):
              Sets a device to be used for encoding speech. If is set to None, torch.cuda.current_device() will be used.
              When used in a dataloader with multiple workers, it might be useful to set `device` to "auto".
              This way, on the first call of __getitem__(), the QuantizeDataset checks if it runs in a dataloader
              worker. If this is the case, it will grab one of the available GPUs and place its copy of
              SpeechEncoder there. This could be useful as SpeechEncoder is typically GPU-intensive and it is a good idea to parallelize
              across multiple GPUs. Defaults to None.
          speaker_extractor (Optional[Callable], optional): An optional callable that extracts speaker id from the
          dataset[i] output. Only needed if speech_encoder uses per-speaker F0 normalization. Defaults to None.
      """
  
      self.dataset = dataset
      self.speech_encoder = speech_encoder
  
      self.randomize_device_on_next_call = False
      # TODO: allow specifying a list of allowed devices
      # this way we can preprocess and train a model on different GPUs.
      if device == "auto":
          self.device = None
          self.randomize_device_on_next_call = True
      elif device is None:
          self.device = torch.cuda.current_device()
          self.speech_encoder.to(self.device)
      else:
          self.device = torch.device(device)
          self.speech_encoder.to(self.device)
  
      self.speaker_extractor = (
          speaker_extractor
          if speaker_extractor is not None
          else default_speaker_ls
      )
  
      self.unit_vocab_size = torch.tensor(
          [self.speech_encoder.vocab_size]
      )
  
      self.unit_pad = 1 + max(
          self.unit_vocab_size - 1,
          self.speech_encoder.bos.item(),
          self.speech_encoder.eos.item(),
      )
  
imports: from typing import Callable, Union, Optional, Any, Dict, List
  import pathlib
  import torch
  import logging
  from torchaudio.datasets import LIBRISPEECH, LJSPEECH, COMMONVOICE, VCTK_092, YESNO
  from torch.utils.data import Dataset
  from textless.data.speech_encoder import SpeechEncoder
  from textless.data.collater_utils import collate_tensors
label_desc: Use image node_id 1 to create a new QuantizeDataset object  with example: obj = QuantizeDataset(dataset, speech_encoder, device, speaker_extractor)