content: from pathlib import Path
  
  import torch
  from torch import nn
  from einops import rearrange
  
  import fairseq
  
  from torchaudio.functional import resample
  
  from audiolm_pytorch.utils import curtail_to_multiple
  
  import logging
  logging.root.setLevel(logging.ERROR)
  
  def exists(val):
      return val is not None
  
  class FairseqVQWav2Vec(nn.Module):
      """
      checkpoint path can be found at https://github.com/facebookresearch/fairseq/blob/main/examples/wav2vec/README.md#vq-wav2vec
      specifically download the kmeans model for now
  
      $ wget https://dl.fbaipublicfiles.com/fairseq/wav2vec/vq-wav2vec_kmeans.pt
      """
  
      def __init__(
          self,
          checkpoint_path,
          target_sample_hz = 24000,
          seq_len_multiple_of = None
      ):
          super().__init__()
          self.target_sample_hz = target_sample_hz
          self.seq_len_multiple_of = seq_len_multiple_of
  
          path = Path(checkpoint_path)
          assert path.exists(), f'path {checkpoint_path} does not exist'
  
          checkpoint = torch.load(checkpoint_path)
          load_model_input = {checkpoint_path: checkpoint}
          model, *_ = fairseq.checkpoint_utils.load_model_ensemble_and_task(load_model_input)
  
          self.model = model[0]
          self.model.eval()
  
          assert hasattr(self.model, 'vector_quantizer') and hasattr(self.model.vector_quantizer, 'embedding'), 'the vq wav2vec model does not seem to be valid'
  
      @property
      def groups(self):
          return self.model.vector_quantizer.groups
  
      @property
      def downsample_factor(self):
          # todo: double check architecture
          return 80
  
      @property
      def codebook_size(self):
          return self.model.vector_quantizer.embedding.shape[0]
  
      @torch.inference_mode()
      def forward(
          self,
          wav_input,
          flatten = True,
          input_sample_hz = None
      ):
          if exists(input_sample_hz):
              wav_input = resample(wav_input, input_sample_hz, self.target_sample_hz)
  
          if exists(self.seq_len_multiple_of):
              wav_input = curtail_to_multiple(wav_input, self.seq_len_multiple_of)
  
          embed = self.model.feature_extractor(wav_input)
          _, codebook_indices = self.model.vector_quantizer.forward_idx(embed)
  
          if not flatten:
              return codebook_indices
  
          return rearrange(codebook_indices, 'b ... -> b (...)')
  