Add files using upload-large-folder tool

omni_speech/infer/fairseq/examples/textless_nlp/gslm/unit2speech/tacotron2/model.py

@@ -0,0 +1,669 @@
+from math import sqrt
+import torch
+import torch.distributions as distr
+from torch.autograd import Variable
+from torch import nn
+from torch.nn import functional as F
+from .layers import ConvNorm, LinearNorm, GlobalAvgPool
+from .utils import to_gpu, get_mask_from_lengths
+
+
+class LocationLayer(nn.Module):
+    def __init__(self, attention_n_filters, attention_kernel_size,
+                 attention_dim):
+        super(LocationLayer, self).__init__()
+        padding = int((attention_kernel_size - 1) / 2)
+        self.location_conv = ConvNorm(2, attention_n_filters,
+                                      kernel_size=attention_kernel_size,
+                                      padding=padding, bias=False, stride=1,
+                                      dilation=1)
+        self.location_dense = LinearNorm(attention_n_filters, attention_dim,
+                                         bias=False, w_init_gain='tanh')
+
+    def forward(self, attention_weights_cat):
+        processed_attention = self.location_conv(attention_weights_cat)
+        processed_attention = processed_attention.transpose(1, 2)
+        processed_attention = self.location_dense(processed_attention)
+        return processed_attention
+
+
+class Attention(nn.Module):
+    def __init__(self, attention_rnn_dim, embedding_dim, attention_dim,
+                 attention_location_n_filters, attention_location_kernel_size):
+        super(Attention, self).__init__()
+        self.query_layer = LinearNorm(attention_rnn_dim, attention_dim,
+                                      bias=False, w_init_gain='tanh')
+        self.memory_layer = LinearNorm(embedding_dim, attention_dim, bias=False,
+                                       w_init_gain='tanh')
+        self.v = LinearNorm(attention_dim, 1, bias=False)
+        self.location_layer = LocationLayer(attention_location_n_filters,
+                                            attention_location_kernel_size,
+                                            attention_dim)
+        self.score_mask_value = -float("inf")
+
+    def get_alignment_energies(self, query, processed_memory,
+                               attention_weights_cat):
+        """
+        PARAMS
+        ------
+        query: decoder output (batch, n_mel_channels * n_frames_per_step)
+        processed_memory: processed encoder outputs (B, T_in, attention_dim)
+        attention_weights_cat: cumulative and prev. att weights (B, 2, max_time)
+
+        RETURNS
+        -------
+        alignment (batch, max_time)
+        """
+
+        processed_query = self.query_layer(query.unsqueeze(1))
+        processed_attention_weights = self.location_layer(attention_weights_cat)
+        energies = self.v(torch.tanh(
+            processed_query + processed_attention_weights + processed_memory))
+
+        energies = energies.squeeze(-1)
+        return energies
+
+    def forward(self, attention_hidden_state, memory, processed_memory,
+                attention_weights_cat, mask):
+        """
+        PARAMS
+        ------
+        attention_hidden_state: attention rnn last output
+        memory: encoder outputs
+        processed_memory: processed encoder outputs
+        attention_weights_cat: previous and cummulative attention weights
+        mask: binary mask for padded data
+        """
+        alignment = self.get_alignment_energies(
+            attention_hidden_state, processed_memory, attention_weights_cat)
+
+        if mask is not None:
+            alignment.data.masked_fill_(mask, self.score_mask_value)
+
+        attention_weights = F.softmax(alignment, dim=1)
+        attention_context = torch.bmm(attention_weights.unsqueeze(1), memory)
+        attention_context = attention_context.squeeze(1)
+
+        return attention_context, attention_weights
+
+
+class Prenet(nn.Module):
+    def __init__(self, in_dim, sizes):
+        super(Prenet, self).__init__()
+        in_sizes = [in_dim] + sizes[:-1]
+        self.layers = nn.ModuleList(
+            [LinearNorm(in_size, out_size, bias=False)
+             for (in_size, out_size) in zip(in_sizes, sizes)])
+
+    def forward(self, x):
+        for linear in self.layers:
+            x = F.dropout(F.relu(linear(x)), p=0.5, training=True)
+        return x
+
+
+class Postnet(nn.Module):
+    """Postnet
+        - Five 1-d convolution with 512 channels and kernel size 5
+    """
+
+    def __init__(self, hparams):
+        super(Postnet, self).__init__()
+        self.convolutions = nn.ModuleList()
+
+        self.convolutions.append(
+            nn.Sequential(
+                ConvNorm(hparams.n_mel_channels, hparams.postnet_embedding_dim,
+                         kernel_size=hparams.postnet_kernel_size, stride=1,
+                         padding=int((hparams.postnet_kernel_size - 1) / 2),
+                         dilation=1, w_init_gain='tanh'),
+                nn.BatchNorm1d(hparams.postnet_embedding_dim))
+        )
+
+        for i in range(1, hparams.postnet_n_convolutions - 1):
+            self.convolutions.append(
+                nn.Sequential(
+                    ConvNorm(hparams.postnet_embedding_dim,
+                             hparams.postnet_embedding_dim,
+                             kernel_size=hparams.postnet_kernel_size, stride=1,
+                             padding=int((hparams.postnet_kernel_size - 1) / 2),
+                             dilation=1, w_init_gain='tanh'),
+                    nn.BatchNorm1d(hparams.postnet_embedding_dim))
+            )
+
+        self.convolutions.append(
+            nn.Sequential(
+                ConvNorm(hparams.postnet_embedding_dim, hparams.n_mel_channels,
+                         kernel_size=hparams.postnet_kernel_size, stride=1,
+                         padding=int((hparams.postnet_kernel_size - 1) / 2),
+                         dilation=1, w_init_gain='linear'),
+                nn.BatchNorm1d(hparams.n_mel_channels))
+            )
+
+    def forward(self, x):
+        for i in range(len(self.convolutions) - 1):
+            x = F.dropout(torch.tanh(self.convolutions[i](x)), 0.5, self.training)
+        x = F.dropout(self.convolutions[-1](x), 0.5, self.training)
+
+        return x
+
+
+class Encoder(nn.Module):
+    """Encoder module:
+        - Three 1-d convolution banks
+        - Bidirectional LSTM
+    """
+    def __init__(self, hparams):
+        super(Encoder, self).__init__()
+
+        convolutions = []
+        for _ in range(hparams.encoder_n_convolutions):
+            conv_layer = nn.Sequential(
+                ConvNorm(hparams.encoder_embedding_dim,
+                         hparams.encoder_embedding_dim,
+                         kernel_size=hparams.encoder_kernel_size, stride=1,
+                         padding=int((hparams.encoder_kernel_size - 1) / 2),
+                         dilation=1, w_init_gain='relu'),
+                nn.BatchNorm1d(hparams.encoder_embedding_dim))
+            convolutions.append(conv_layer)
+        self.convolutions = nn.ModuleList(convolutions)
+
+        self.lstm = nn.LSTM(hparams.encoder_embedding_dim,
+                            int(hparams.encoder_embedding_dim / 2), 1,
+                            batch_first=True, bidirectional=True)
+
+    def forward(self, x, input_lengths):
+        for conv in self.convolutions:
+            x = F.dropout(F.relu(conv(x)), 0.5, self.training)
+
+        x = x.transpose(1, 2)
+
+        # pytorch tensor are not reversible, hence the conversion
+        input_lengths = input_lengths.cpu().numpy()
+        x = nn.utils.rnn.pack_padded_sequence(
+            x, input_lengths, batch_first=True)
+
+        self.lstm.flatten_parameters()
+        outputs, _ = self.lstm(x)
+
+        outputs, _ = nn.utils.rnn.pad_packed_sequence(
+            outputs, batch_first=True)
+
+        return outputs
+
+    def inference(self, x):
+        for conv in self.convolutions:
+            x = F.dropout(F.relu(conv(x)), 0.5, self.training)
+
+        x = x.transpose(1, 2)
+
+        self.lstm.flatten_parameters()
+        outputs, _ = self.lstm(x)
+
+        return outputs
+
+
+class AudioEncoder(nn.Module):
+    def __init__(self, hparams):
+        super(AudioEncoder, self).__init__()
+
+        assert hparams.lat_dim > 0
+
+        convolutions = []
+        inp_dim = hparams.n_mel_channels
+        for _ in range(hparams.lat_n_convolutions):
+            conv_layer = nn.Sequential(
+                ConvNorm(inp_dim, hparams.lat_n_filters,
+                         kernel_size=hparams.lat_kernel_size, stride=1,
+                         padding=int((hparams.lat_kernel_size - 1) / 2),
+                         dilation=1, w_init_gain='tanh'),
+                nn.BatchNorm1d(hparams.lat_n_filters))
+            inp_dim = hparams.lat_n_filters
+            convolutions.append(conv_layer)
+        self.convolutions = nn.ModuleList(convolutions)
+
+        self.lstm = nn.LSTM(hparams.lat_n_filters,
+                            int(hparams.lat_n_filters / 2),
+                            hparams.lat_n_blstms, batch_first=True,
+                            bidirectional=True)
+        self.pool = GlobalAvgPool()
+
+        self.mu_proj = LinearNorm(hparams.lat_n_filters, hparams.lat_dim)
+        self.logvar_proj = LinearNorm(hparams.lat_n_filters, hparams.lat_dim)
+        self.lat_dim = hparams.lat_dim
+
+    def forward(self, x, lengths):
+        """
+        Args:
+            x (torch.Tensor): (B, F, T)
+        """
+
+        for conv in self.convolutions:
+            x = F.dropout(F.tanh(conv(x)), 0.5, self.training)
+
+        x = x.transpose(1, 2)  # (B, T, D)
+
+        # x may not be sorted by length. Sort->process->unsort
+        max_len = x.size(1)
+        assert max_len == torch.max(lengths).item()
+
+        lengths, perm_idx = lengths.sort(0, descending=True)
+        x = x[perm_idx]
+        x = nn.utils.rnn.pack_padded_sequence(x, lengths, batch_first=True)
+
+        self.lstm.flatten_parameters()
+        outputs, _ = self.lstm(x)
+        outputs, _ = nn.utils.rnn.pad_packed_sequence(outputs, batch_first=True)
+
+        _, unperm_idx = perm_idx.sort(0)
+        outputs = outputs[unperm_idx]  # (B, T, D)
+        lengths = lengths[unperm_idx]  # (B, T, D)
+
+        outputs = self.pool(outputs, lengths)  # (B, D)
+
+        mu = self.mu_proj(outputs)
+        logvar = self.logvar_proj(outputs)
+        z = distr.Normal(mu, logvar).rsample()
+        return z, mu, logvar
+
+
+class Decoder(nn.Module):
+    def __init__(self, hparams):
+        super(Decoder, self).__init__()
+        self.n_mel_channels = hparams.n_mel_channels
+        self.n_frames_per_step = hparams.n_frames_per_step
+        self.encoder_embedding_dim = hparams.encoder_embedding_dim
+        self.obs_dim = hparams.obs_dim
+        self.lat_dim = hparams.lat_dim
+        self.attention_rnn_dim = hparams.attention_rnn_dim
+        self.decoder_rnn_dim = hparams.decoder_rnn_dim
+        self.prenet_dim = hparams.prenet_dim
+        self.max_decoder_steps = hparams.max_decoder_steps
+        self.gate_threshold = hparams.gate_threshold
+        self.p_attention_dropout = hparams.p_attention_dropout
+        self.p_decoder_dropout = hparams.p_decoder_dropout
+
+        self.prenet = Prenet(
+            hparams.n_mel_channels * hparams.n_frames_per_step,
+            [hparams.prenet_dim, hparams.prenet_dim])
+
+        self.attention_rnn = nn.LSTMCell(
+            hparams.prenet_dim + hparams.encoder_embedding_dim,
+            hparams.attention_rnn_dim)
+
+        self.attention_layer = Attention(
+            hparams.attention_rnn_dim, hparams.encoder_embedding_dim,
+            hparams.attention_dim, hparams.attention_location_n_filters,
+            hparams.attention_location_kernel_size)
+
+        encoder_tot_dim = (hparams.encoder_embedding_dim + \
+                           hparams.lat_dim + hparams.obs_dim)
+        self.decoder_rnn = nn.LSTMCell(
+            hparams.attention_rnn_dim + encoder_tot_dim,
+            hparams.decoder_rnn_dim, 1)
+
+        self.linear_projection = LinearNorm(
+            hparams.decoder_rnn_dim + encoder_tot_dim,
+            hparams.n_mel_channels * hparams.n_frames_per_step)
+
+        self.gate_layer = LinearNorm(
+            hparams.decoder_rnn_dim + encoder_tot_dim, 1,
+            bias=True, w_init_gain='sigmoid')
+
+    def get_go_frame(self, memory):
+        """ Gets all zeros frames to use as first decoder input
+        PARAMS
+        ------
+        memory: decoder outputs
+
+        RETURNS
+        -------
+        decoder_input: all zeros frames
+        """
+        B = memory.size(0)
+        decoder_input = Variable(memory.data.new(
+            B, self.n_mel_channels * self.n_frames_per_step).zero_())
+        return decoder_input
+
+    def initialize_decoder_states(self, memory, obs_and_lat, mask):
+        """ Initializes attention rnn states, decoder rnn states, attention
+        weights, attention cumulative weights, attention context, stores memory
+        and stores processed memory
+        PARAMS
+        ------
+        memory: Encoder outputs
+        obs_and_lat: Observed and latent attribute embeddings
+        mask: Mask for padded data if training, expects None for inference
+        """
+        B = memory.size(0)
+        MAX_TIME = memory.size(1)
+
+        self.attention_hidden = Variable(memory.data.new(
+            B, self.attention_rnn_dim).zero_())
+        self.attention_cell = Variable(memory.data.new(
+            B, self.attention_rnn_dim).zero_())
+
+        self.decoder_hidden = Variable(memory.data.new(
+            B, self.decoder_rnn_dim).zero_())
+        self.decoder_cell = Variable(memory.data.new(
+            B, self.decoder_rnn_dim).zero_())
+
+        self.attention_weights = Variable(memory.data.new(
+            B, MAX_TIME).zero_())
+        self.attention_weights_cum = Variable(memory.data.new(
+            B, MAX_TIME).zero_())
+        self.attention_context = Variable(memory.data.new(
+            B, self.encoder_embedding_dim).zero_())
+
+        self.memory = memory
+        self.processed_memory = self.attention_layer.memory_layer(memory)
+        self.obs_and_lat = obs_and_lat
+        self.mask = mask
+
+    def parse_decoder_inputs(self, decoder_inputs):
+        """ Prepares decoder inputs, i.e. mel outputs
+        PARAMS
+        ------
+        decoder_inputs: inputs used for teacher-forced training, i.e. mel-specs
+
+        RETURNS
+        -------
+        inputs: processed decoder inputs
+
+        """
+        # (B, n_mel_channels, T_out) -> (B, T_out, n_mel_channels)
+        decoder_inputs = decoder_inputs.transpose(1, 2)
+        decoder_inputs = decoder_inputs.view(
+            decoder_inputs.size(0),
+            int(decoder_inputs.size(1)/self.n_frames_per_step), -1)
+        # (B, T_out, n_mel_channels) -> (T_out, B, n_mel_channels)
+        decoder_inputs = decoder_inputs.transpose(0, 1)
+        return decoder_inputs
+
+    def parse_decoder_outputs(self, mel_outputs, gate_outputs, alignments):
+        """ Prepares decoder outputs for output
+        PARAMS
+        ------
+        mel_outputs:
+        gate_outputs: gate output energies
+        alignments:
+
+        RETURNS
+        -------
+        mel_outputs:
+        gate_outpust: gate output energies
+        alignments:
+        """
+        # (T_out, B) -> (B, T_out)
+        alignments = torch.stack(alignments).transpose(0, 1)
+        # (T_out, B) -> (B, T_out)
+        gate_outputs = torch.stack(gate_outputs).transpose(0, 1)
+        gate_outputs = gate_outputs.contiguous()
+        # (T_out, B, n_mel_channels) -> (B, T_out, n_mel_channels)
+        mel_outputs = torch.stack(mel_outputs).transpose(0, 1).contiguous()
+        # decouple frames per step
+        mel_outputs = mel_outputs.view(
+            mel_outputs.size(0), -1, self.n_mel_channels)
+        # (B, T_out, n_mel_channels) -> (B, n_mel_channels, T_out)
+        mel_outputs = mel_outputs.transpose(1, 2)
+
+        return mel_outputs, gate_outputs, alignments
+
+    def decode(self, decoder_input):
+        """ Decoder step using stored states, attention and memory
+        PARAMS
+        ------
+        decoder_input: previous mel output
+
+        RETURNS
+        -------
+        mel_output:
+        gate_output: gate output energies
+        attention_weights:
+        """
+        cell_input = torch.cat((decoder_input, self.attention_context), -1)
+        self.attention_hidden, self.attention_cell = self.attention_rnn(
+            cell_input, (self.attention_hidden, self.attention_cell))
+        self.attention_hidden = F.dropout(
+            self.attention_hidden, self.p_attention_dropout, self.training)
+
+        attention_weights_cat = torch.cat(
+            (self.attention_weights.unsqueeze(1),
+             self.attention_weights_cum.unsqueeze(1)), dim=1)
+        self.attention_context, self.attention_weights = self.attention_layer(
+            self.attention_hidden, self.memory, self.processed_memory,
+            attention_weights_cat, self.mask)
+
+        self.attention_weights_cum += self.attention_weights
+        decoder_input = torch.cat(
+            (self.attention_hidden, self.attention_context), -1)
+        if self.obs_and_lat is not None:
+            decoder_input = torch.cat((decoder_input, self.obs_and_lat), -1)
+        self.decoder_hidden, self.decoder_cell = self.decoder_rnn(
+            decoder_input, (self.decoder_hidden, self.decoder_cell))
+        self.decoder_hidden = F.dropout(
+            self.decoder_hidden, self.p_decoder_dropout, self.training)
+
+        decoder_hidden_attention_context = torch.cat(
+            (self.decoder_hidden, self.attention_context), dim=1)
+        if self.obs_and_lat is not None:
+            decoder_hidden_attention_context = torch.cat(
+                    (decoder_hidden_attention_context, self.obs_and_lat), dim=1)
+        decoder_output = self.linear_projection(
+            decoder_hidden_attention_context)
+
+        gate_prediction = self.gate_layer(decoder_hidden_attention_context)
+        return decoder_output, gate_prediction, self.attention_weights
+
+    def forward(self, memory, obs_and_lat, decoder_inputs, memory_lengths):
+        """ Decoder forward pass for training
+        PARAMS
+        ------
+        memory: Encoder outputs
+        obs_and_lat: Observed and latent attribute embeddings
+        decoder_inputs: Decoder inputs for teacher forcing. i.e. mel-specs
+        memory_lengths: Encoder output lengths for attention masking.
+
+        RETURNS
+        -------
+        mel_outputs: mel outputs from the decoder
+        gate_outputs: gate outputs from the decoder
+        alignments: sequence of attention weights from the decoder
+        """
+
+        decoder_input = self.get_go_frame(memory).unsqueeze(0)
+        decoder_inputs = self.parse_decoder_inputs(decoder_inputs)
+        decoder_inputs = torch.cat((decoder_input, decoder_inputs), dim=0)
+        decoder_inputs = self.prenet(decoder_inputs)
+
+        self.initialize_decoder_states(
+            memory, obs_and_lat, mask=~get_mask_from_lengths(memory_lengths))
+
+        mel_outputs, gate_outputs, alignments = [], [], []
+        while len(mel_outputs) < decoder_inputs.size(0) - 1:
+            decoder_input = decoder_inputs[len(mel_outputs)]
+            mel_output, gate_output, attention_weights = self.decode(
+                decoder_input)
+            mel_outputs += [mel_output.squeeze(1)]
+            gate_outputs += [gate_output.squeeze()]
+            alignments += [attention_weights]
+
+        mel_outputs, gate_outputs, alignments = self.parse_decoder_outputs(
+            mel_outputs, gate_outputs, alignments)
+
+        return mel_outputs, gate_outputs, alignments
+
+    def inference(self, memory, obs_and_lat, ret_has_eos=False):
+        """ Decoder inference
+        PARAMS
+        ------
+        memory: Encoder outputs
+        obs_and_lat: Observed and latent attribute embeddings
+
+        RETURNS
+        -------
+        mel_outputs: mel outputs from the decoder
+        gate_outputs: gate outputs from the decoder
+        alignments: sequence of attention weights from the decoder
+        """
+        decoder_input = self.get_go_frame(memory)
+
+        self.initialize_decoder_states(memory, obs_and_lat, mask=None)
+
+        mel_outputs, gate_outputs, alignments = [], [], []
+        has_eos = False
+        while True:
+            decoder_input = self.prenet(decoder_input)
+            mel_output, gate_output, alignment = self.decode(decoder_input)
+
+            mel_outputs += [mel_output.squeeze(1)]
+            gate_outputs += [gate_output]
+            alignments += [alignment]
+
+            if torch.sigmoid(gate_output.data) > self.gate_threshold:
+                has_eos = True
+                break
+            elif len(mel_outputs) == self.max_decoder_steps:
+                # print("Warning! Reached max decoder steps")
+                break
+
+            decoder_input = mel_output
+
+        mel_outputs, gate_outputs, alignments = self.parse_decoder_outputs(
+            mel_outputs, gate_outputs, alignments)
+
+        if ret_has_eos:
+            return mel_outputs, gate_outputs, alignments, has_eos
+        else:
+            return mel_outputs, gate_outputs, alignments
+
+
+class Tacotron2(nn.Module):
+    def __init__(self, hparams):
+        super(Tacotron2, self).__init__()
+        self.mask_padding = hparams.mask_padding
+        self.fp16_run = hparams.fp16_run
+        self.n_mel_channels = hparams.n_mel_channels
+        self.n_frames_per_step = hparams.n_frames_per_step
+
+        # initialize text encoder embedding
+        self.embedding = nn.Embedding(
+            hparams.n_symbols, hparams.symbols_embedding_dim)
+        std = sqrt(2.0 / (hparams.n_symbols + hparams.symbols_embedding_dim))
+        val = sqrt(3.0) * std  # uniform bounds for std
+        self.embedding.weight.data.uniform_(-val, val)
+
+        # initialize observed attribute embedding
+        self.obs_embedding = None
+        if hparams.obs_dim > 0:
+            self.obs_embedding = nn.Embedding(
+                hparams.obs_n_class, hparams.obs_dim)
+            std = sqrt(2.0 / (hparams.obs_n_class + hparams.obs_dim))
+            val = sqrt(3.0) * std  # uniform bounds for std
+            self.obs_embedding.weight.data.uniform_(-val, val)
+
+        self.encoder = Encoder(hparams)
+        self.decoder = Decoder(hparams)
+        self.postnet = Postnet(hparams)
+
+        self.lat_encoder = None
+        if hparams.lat_dim > 0:
+            self.lat_encoder = AudioEncoder(hparams)
+
+    def parse_batch(self, batch):
+        (text_padded, input_lengths, obs_labels,
+         mel_padded, gate_padded, output_lengths) = batch
+        text_padded = to_gpu(text_padded).long()
+        input_lengths = to_gpu(input_lengths).long()
+        obs_labels = to_gpu(obs_labels).long()
+        max_len = torch.max(input_lengths.data).item()
+        mel_padded = to_gpu(mel_padded).float()
+        gate_padded = to_gpu(gate_padded).float()
+        output_lengths = to_gpu(output_lengths).long()
+
+        return (
+            (text_padded, input_lengths, obs_labels,
+             mel_padded, max_len, output_lengths),
+            (mel_padded, gate_padded))
+
+    def parse_output(self, outputs, output_lengths=None):
+        if self.mask_padding and output_lengths is not None:
+            mask = ~get_mask_from_lengths(output_lengths)
+            mask = mask.expand(self.n_mel_channels, mask.size(0), mask.size(1))
+            mask = mask.permute(1, 0, 2)
+
+            outputs[0].data.masked_fill_(mask, 0.0)
+            outputs[1].data.masked_fill_(mask, 0.0)
+            outputs[2].data.masked_fill_(mask[:, 0, :], 1e3)  # gate energies
+
+        return outputs
+
+    def forward(self, inputs):
+        (text_inputs, text_lengths, obs_labels,
+         mels, max_len, output_lengths) = inputs
+        text_lengths, output_lengths = text_lengths.data, output_lengths.data
+
+        embedded_inputs = self.embedding(text_inputs).transpose(1, 2)
+
+        encoder_outputs = self.encoder(embedded_inputs, text_lengths)
+
+        obs = None
+        if self.obs_embedding is not None:
+            obs = self.obs_embedding(obs_labels)
+
+        lat, lat_mu, lat_logvar = None, None, None
+        if self.lat_encoder is not None:
+            (lat, lat_mu, lat_logvar) = self.lat_encoder(mels, output_lengths)
+
+        obs_and_lat = [x for x in [obs, lat] if x is not None]
+        if bool(obs_and_lat):
+            obs_and_lat = torch.cat(obs_and_lat, dim=-1)
+        else:
+            obs_and_lat = None
+
+        mel_outputs, gate_outputs, alignments = self.decoder(
+            encoder_outputs, obs_and_lat, mels, memory_lengths=text_lengths)
+
+        mel_outputs_postnet = self.postnet(mel_outputs)
+        mel_outputs_postnet = mel_outputs + mel_outputs_postnet
+
+        return self.parse_output(
+            [mel_outputs, mel_outputs_postnet, gate_outputs, alignments,
+             lat_mu, lat_logvar],
+            output_lengths)
+
+    def inference(self, inputs, obs_labels=None, lat=None, ret_has_eos=False):
+        embedded_inputs = self.embedding(inputs).transpose(1, 2)
+        encoder_outputs = self.encoder.inference(embedded_inputs)
+
+        if obs_labels is None:
+            obs_labels = torch.LongTensor(len(inputs))
+            obs_labels = obs_labels.to(inputs.device).zero_()
+
+        obs = None
+        if self.obs_embedding is not None:
+            obs = self.obs_embedding(obs_labels)
+
+        if self.lat_encoder is not None:
+            if lat is None:
+                lat = torch.FloatTensor(len(inputs), self.lat_encoder.lat_dim)
+                lat = lat.to(inputs.device).zero_().type(encoder_outputs.type())
+
+        obs_and_lat = [x for x in [obs, lat] if x is not None]
+        if bool(obs_and_lat):
+            obs_and_lat = torch.cat(obs_and_lat, dim=-1)
+        else:
+            obs_and_lat = None
+
+        mel_outputs, gate_outputs, alignments, has_eos = self.decoder.inference(
+            encoder_outputs, obs_and_lat, ret_has_eos=True)
+
+        mel_outputs_postnet = self.postnet(mel_outputs)
+        mel_outputs_postnet = mel_outputs + mel_outputs_postnet
+
+        outputs = self.parse_output(
+            [mel_outputs, mel_outputs_postnet, gate_outputs, alignments])
+
+        if ret_has_eos:
+            return outputs + [has_eos]
+        else:
+            return outputs

omni_speech/infer/fairseq/examples/textless_nlp/gslm/unit2speech/tacotron2/numbers.py

@@ -0,0 +1,71 @@
+""" from https://github.com/keithito/tacotron """
+
+import inflect
+import re
+
+
+_inflect = inflect.engine()
+_comma_number_re = re.compile(r'([0-9][0-9\,]+[0-9])')
+_decimal_number_re = re.compile(r'([0-9]+\.[0-9]+)')
+_pounds_re = re.compile(r'£([0-9\,]*[0-9]+)')
+_dollars_re = re.compile(r'\$([0-9\.\,]*[0-9]+)')
+_ordinal_re = re.compile(r'[0-9]+(st|nd|rd|th)')
+_number_re = re.compile(r'[0-9]+')
+
+
+def _remove_commas(m):
+  return m.group(1).replace(',', '')
+
+
+def _expand_decimal_point(m):
+  return m.group(1).replace('.', ' point ')
+
+
+def _expand_dollars(m):
+  match = m.group(1)
+  parts = match.split('.')
+  if len(parts) > 2:
+    return match + ' dollars'  # Unexpected format
+  dollars = int(parts[0]) if parts[0] else 0
+  cents = int(parts[1]) if len(parts) > 1 and parts[1] else 0
+  if dollars and cents:
+    dollar_unit = 'dollar' if dollars == 1 else 'dollars'
+    cent_unit = 'cent' if cents == 1 else 'cents'
+    return '%s %s, %s %s' % (dollars, dollar_unit, cents, cent_unit)
+  elif dollars:
+    dollar_unit = 'dollar' if dollars == 1 else 'dollars'
+    return '%s %s' % (dollars, dollar_unit)
+  elif cents:
+    cent_unit = 'cent' if cents == 1 else 'cents'
+    return '%s %s' % (cents, cent_unit)
+  else:
+    return 'zero dollars'
+
+
+def _expand_ordinal(m):
+  return _inflect.number_to_words(m.group(0))
+
+
+def _expand_number(m):
+  num = int(m.group(0))
+  if num > 1000 and num < 3000:
+    if num == 2000:
+      return 'two thousand'
+    elif num > 2000 and num < 2010:
+      return 'two thousand ' + _inflect.number_to_words(num % 100)
+    elif num % 100 == 0:
+      return _inflect.number_to_words(num // 100) + ' hundred'
+    else:
+      return _inflect.number_to_words(num, andword='', zero='oh', group=2).replace(', ', ' ')
+  else:
+    return _inflect.number_to_words(num, andword='')
+
+
+def normalize_numbers(text):
+  text = re.sub(_comma_number_re, _remove_commas, text)
+  text = re.sub(_pounds_re, r'\1 pounds', text)
+  text = re.sub(_dollars_re, _expand_dollars, text)
+  text = re.sub(_decimal_number_re, _expand_decimal_point, text)
+  text = re.sub(_ordinal_re, _expand_ordinal, text)
+  text = re.sub(_number_re, _expand_number, text)
+  return text

omni_speech/infer/fairseq/examples/textless_nlp/gslm/unit2speech/tacotron2/stft.py

@@ -0,0 +1,141 @@
+"""
+BSD 3-Clause License
+
+Copyright (c) 2017, Prem Seetharaman
+All rights reserved.
+
+* Redistribution and use in source and binary forms, with or without
+  modification, are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice,
+  this list of conditions and the following disclaimer.
+
+* Redistributions in binary form must reproduce the above copyright notice, this
+  list of conditions and the following disclaimer in the
+  documentation and/or other materials provided with the distribution.
+
+* Neither the name of the copyright holder nor the names of its
+  contributors may be used to endorse or promote products derived from this
+  software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
+ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
+ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+"""
+
+import torch
+import numpy as np
+import torch.nn.functional as F
+from torch.autograd import Variable
+from scipy.signal import get_window
+from librosa.util import pad_center, tiny
+from .audio_processing import window_sumsquare
+
+
+class STFT(torch.nn.Module):
+    """adapted from Prem Seetharaman's https://github.com/pseeth/pytorch-stft"""
+    def __init__(self, filter_length=800, hop_length=200, win_length=800,
+                 window='hann'):
+        super(STFT, self).__init__()
+        self.filter_length = filter_length
+        self.hop_length = hop_length
+        self.win_length = win_length
+        self.window = window
+        self.forward_transform = None
+        scale = self.filter_length / self.hop_length
+        fourier_basis = np.fft.fft(np.eye(self.filter_length))
+
+        cutoff = int((self.filter_length / 2 + 1))
+        fourier_basis = np.vstack([np.real(fourier_basis[:cutoff, :]),
+                                   np.imag(fourier_basis[:cutoff, :])])
+
+        forward_basis = torch.FloatTensor(fourier_basis[:, None, :])
+        inverse_basis = torch.FloatTensor(
+            np.linalg.pinv(scale * fourier_basis).T[:, None, :])
+
+        if window is not None:
+            assert(filter_length >= win_length)
+            # get window and zero center pad it to filter_length
+            fft_window = get_window(window, win_length, fftbins=True)
+            fft_window = pad_center(fft_window, filter_length)
+            fft_window = torch.from_numpy(fft_window).float()
+
+            # window the bases
+            forward_basis *= fft_window
+            inverse_basis *= fft_window
+
+        self.register_buffer('forward_basis', forward_basis.float())
+        self.register_buffer('inverse_basis', inverse_basis.float())
+
+    def transform(self, input_data):
+        num_batches = input_data.size(0)
+        num_samples = input_data.size(1)
+
+        self.num_samples = num_samples
+
+        # similar to librosa, reflect-pad the input
+        input_data = input_data.view(num_batches, 1, num_samples)
+        input_data = F.pad(
+            input_data.unsqueeze(1),
+            (int(self.filter_length / 2), int(self.filter_length / 2), 0, 0),
+            mode='reflect')
+        input_data = input_data.squeeze(1)
+
+        forward_transform = F.conv1d(
+            input_data,
+            Variable(self.forward_basis, requires_grad=False),
+            stride=self.hop_length,
+            padding=0)
+
+        cutoff = int((self.filter_length / 2) + 1)
+        real_part = forward_transform[:, :cutoff, :]
+        imag_part = forward_transform[:, cutoff:, :]
+
+        magnitude = torch.sqrt(real_part**2 + imag_part**2)
+        phase = torch.autograd.Variable(
+            torch.atan2(imag_part.data, real_part.data))
+
+        return magnitude, phase
+
+    def inverse(self, magnitude, phase):
+        recombine_magnitude_phase = torch.cat(
+            [magnitude*torch.cos(phase), magnitude*torch.sin(phase)], dim=1)
+
+        inverse_transform = F.conv_transpose1d(
+            recombine_magnitude_phase,
+            Variable(self.inverse_basis, requires_grad=False),
+            stride=self.hop_length,
+            padding=0)
+
+        if self.window is not None:
+            window_sum = window_sumsquare(
+                self.window, magnitude.size(-1), hop_length=self.hop_length,
+                win_length=self.win_length, n_fft=self.filter_length,
+                dtype=np.float32)
+            # remove modulation effects
+            approx_nonzero_indices = torch.from_numpy(
+                np.where(window_sum > tiny(window_sum))[0])
+            window_sum = torch.autograd.Variable(
+                torch.from_numpy(window_sum), requires_grad=False)
+            window_sum = window_sum.cuda() if magnitude.is_cuda else window_sum
+            inverse_transform[:, :, approx_nonzero_indices] /= window_sum[approx_nonzero_indices]
+
+            # scale by hop ratio
+            inverse_transform *= float(self.filter_length) / self.hop_length
+
+        inverse_transform = inverse_transform[:, :, int(self.filter_length/2):]
+        inverse_transform = inverse_transform[:, :, :-int(self.filter_length/2):]
+
+        return inverse_transform
+
+    def forward(self, input_data):
+        self.magnitude, self.phase = self.transform(input_data)
+        reconstruction = self.inverse(self.magnitude, self.phase)
+        return reconstruction

omni_speech/infer/fairseq/examples/textless_nlp/gslm/unit2speech/tacotron2/symbols.py

@@ -0,0 +1,18 @@
+""" from https://github.com/keithito/tacotron """
+
+'''
+Defines the set of symbols used in text input to the model.
+
+The default is a set of ASCII characters that works well for English or text that has been run through Unidecode. For other data, you can modify _characters. See TRAINING_DATA.md for details. '''
+from . import cmudict
+
+_pad        = '_'
+_punctuation = '!\'(),.:;? '
+_special = '-'
+_letters = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz'
+
+# Prepend "@" to ARPAbet symbols to ensure uniqueness (some are the same as uppercase letters):
+_arpabet = ['@' + s for s in cmudict.valid_symbols]
+
+# Export all symbols:
+symbols = [_pad] + list(_special) + list(_punctuation) + list(_letters) + _arpabet

omni_speech/infer/fairseq/examples/textless_nlp/gslm/unit2speech/tacotron2/text.py

@@ -0,0 +1,107 @@
+""" from https://github.com/keithito/tacotron """
+import numpy as np
+import re
+from . import cleaners
+from .symbols import symbols
+
+
+# Mappings from symbol to numeric ID and vice versa:
+_symbol_to_id = {s: i for i, s in enumerate(symbols)}
+_id_to_symbol = {i: s for i, s in enumerate(symbols)}
+
+# Regular expression matching text enclosed in curly braces:
+_curly_re = re.compile(r'(.*?)\{(.+?)\}(.*)')
+
+# Special symbols
+SOS_TOK = '<s>'
+EOS_TOK = '</s>'
+
+def text_to_sequence(text, cleaner_names):
+  '''Converts a string of text to a sequence of IDs corresponding to the symbols in the text.
+
+    The text can optionally have ARPAbet sequences enclosed in curly braces embedded
+    in it. For example, "Turn left on {HH AW1 S S T AH0 N} Street."
+
+    Args:
+      text: string to convert to a sequence
+      cleaner_names: names of the cleaner functions to run the text through
+
+    Returns:
+      List of integers corresponding to the symbols in the text
+  '''
+  sequence = []
+
+  # Check for curly braces and treat their contents as ARPAbet:
+  while len(text):
+    m = _curly_re.match(text)
+    if not m:
+      sequence += _symbols_to_sequence(_clean_text(text, cleaner_names))
+      break
+    sequence += _symbols_to_sequence(_clean_text(m.group(1), cleaner_names))
+    sequence += _arpabet_to_sequence(m.group(2))
+    text = m.group(3)
+
+  return sequence
+
+
+def sample_code_chunk(code, size):
+    assert(size > 0 and size <= len(code))
+    start = np.random.randint(len(code) - size + 1)
+    end = start + size
+    return code[start:end], start, end
+
+
+def code_to_sequence(code, code_dict, collapse_code):
+    if collapse_code:
+        prev_c = None
+        sequence = []
+        for c in code:
+            if c in code_dict and c != prev_c:
+                sequence.append(code_dict[c])
+                prev_c = c
+    else:
+        sequence = [code_dict[c] for c in code if c in code_dict]
+        if len(sequence) < 0.95 * len(code):
+            print('WARNING : over 5%% codes are OOV')
+
+    return sequence
+
+
+def sequence_to_text(sequence):
+  '''Converts a sequence of IDs back to a string'''
+  result = ''
+  for symbol_id in sequence:
+    if symbol_id in _id_to_symbol:
+      s = _id_to_symbol[symbol_id]
+      # Enclose ARPAbet back in curly braces:
+      if len(s) > 1 and s[0] == '@':
+        s = '{%s}' % s[1:]
+      result += s
+  return result.replace('}{', ' ')
+
+
+def sequence_to_code(sequence, code_dict):
+    '''Analogous to sequence_to_text'''
+    id_to_code = {i: c for c, i in code_dict.items()}
+    return ' '.join([id_to_code[i] for i in sequence])
+
+
+def _clean_text(text, cleaner_names):
+  for name in cleaner_names:
+    cleaner = getattr(cleaners, name)
+    if not cleaner:
+      raise Exception('Unknown cleaner: %s' % name)
+    text = cleaner(text)
+  return text
+
+
+def _symbols_to_sequence(symbols):
+  return [_symbol_to_id[s] for s in symbols if _should_keep_symbol(s)]
+
+
+def _arpabet_to_sequence(text):
+  return _symbols_to_sequence(['@' + s for s in text.split()])
+
+
+def _should_keep_symbol(s):
+  return s in _symbol_to_id and s != '_' and s != '~'

omni_speech/infer/fairseq/examples/textless_nlp/gslm/unit2speech/tacotron2/utils.py

@@ -0,0 +1,171 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import collections
+import io
+import json
+import librosa
+import numpy as np
+import soundfile as sf
+import time
+import torch
+from scipy.io.wavfile import read
+from .text import SOS_TOK, EOS_TOK
+
+
+def get_mask_from_lengths(lengths):
+    max_len = torch.max(lengths).item()
+    ids = torch.arange(0, max_len, out=torch.cuda.LongTensor(max_len))
+    mask = (ids < lengths.unsqueeze(1))
+    return mask
+
+
+def load_wav_to_torch(full_path, sr=None):
+    data, sr = librosa.load(full_path, sr=sr)
+    data = np.clip(data, -1, 1)  # potentially out of [-1, 1] due to resampling
+    data = data * 32768.0  # match values loaded by scipy
+    return torch.FloatTensor(data.astype(np.float32)), sr
+
+
+def read_binary_audio(bin_data, tar_sr=None):
+    """
+    read binary audio (`bytes` or `uint8` `numpy.ndarray`) to `float32`
+    `numpy.ndarray`
+
+    RETURNS:
+        data (np.ndarray) : audio of shape (n,) or (2, n)
+        tar_sr (int) : sample rate
+    """
+    data, ori_sr = sf.read(io.BytesIO(bin_data), dtype='float32')
+    data = data.T
+    if (tar_sr is not None) and (ori_sr != tar_sr):
+        data = librosa.resample(data, ori_sr, tar_sr)
+    else:
+        tar_sr = ori_sr
+    data = np.clip(data, -1, 1)
+    data = data * 32768.0
+    return torch.FloatTensor(data.astype(np.float32)), tar_sr
+
+
+def load_filepaths_and_text(filename):
+    with open(filename, encoding='utf-8') as f:
+        data = [json.loads(line.rstrip()) for line in f]
+    return data
+
+
+def to_gpu(x):
+    x = x.contiguous()
+
+    if torch.cuda.is_available():
+        x = x.cuda(non_blocking=True)
+    return torch.autograd.Variable(x)
+
+
+def load_code_dict(path, add_sos=False, add_eos=False):
+    if not path:
+        return {}
+
+    with open(path, 'r') as f:
+        codes = ['_'] + [line.rstrip() for line in f]  # '_' for pad
+    code_dict = {c: i for i, c in enumerate(codes)}
+
+    if add_sos:
+        code_dict[SOS_TOK] = len(code_dict)
+    if add_eos:
+        code_dict[EOS_TOK] = len(code_dict)
+    assert(set(code_dict.values()) == set(range(len(code_dict))))
+
+    return code_dict
+
+
+def load_obs_label_dict(path):
+    if not path:
+        return {}
+    with open(path, 'r') as f:
+        obs_labels = [line.rstrip() for line in f]
+    return {c: i for i, c in enumerate(obs_labels)}
+
+
+# A simple timer class inspired from `tnt.TimeMeter`
+class CudaTimer:
+    def __init__(self, keys):
+        self.keys = keys
+        self.reset()
+
+    def start(self, key):
+        s = torch.cuda.Event(enable_timing=True)
+        s.record()
+        self.start_events[key].append(s)
+        return self
+
+    def stop(self, key):
+        e = torch.cuda.Event(enable_timing=True)
+        e.record()
+        self.end_events[key].append(e)
+        return self
+
+    def reset(self):
+        self.start_events = collections.defaultdict(list)
+        self.end_events = collections.defaultdict(list)
+        self.running_times = collections.defaultdict(float)
+        self.n = collections.defaultdict(int)
+        return self
+
+    def value(self):
+        self._synchronize()
+        return {k: self.running_times[k] / self.n[k] for k in self.keys}
+
+    def _synchronize(self):
+        torch.cuda.synchronize()
+        for k in self.keys:
+            starts = self.start_events[k]
+            ends = self.end_events[k]
+            if len(starts) == 0:
+                raise ValueError("Trying to divide by zero in TimeMeter")
+            if len(ends) != len(starts):
+                raise ValueError("Call stop before checking value!")
+            time = 0
+            for start, end in zip(starts, ends):
+                time += start.elapsed_time(end)
+            self.running_times[k] += time * 1e-3
+            self.n[k] += len(starts)
+        self.start_events = collections.defaultdict(list)
+        self.end_events = collections.defaultdict(list)
+
+
+# Used to measure the time taken for multiple events
+class Timer:
+    def __init__(self, keys):
+        self.keys = keys
+        self.n = {}
+        self.running_time = {}
+        self.total_time = {}
+        self.reset()
+
+    def start(self, key):
+        self.running_time[key] = time.time()
+        return self
+
+    def stop(self, key):
+        self.total_time[key] = time.time() - self.running_time[key]
+        self.n[key] += 1
+        self.running_time[key] = None
+        return self
+
+    def reset(self):
+        for k in self.keys:
+            self.total_time[k] = 0
+            self.running_time[k] = None
+            self.n[k] = 0
+        return self
+
+    def value(self):
+        vals = {}
+        for k in self.keys:
+            if self.n[k] == 0:
+                raise ValueError("Trying to divide by zero in TimeMeter")
+            else:
+                vals[k] = self.total_time[k] / self.n[k]
+        return vals

omni_speech/infer/fairseq/examples/textless_nlp/gslm/unit2speech/tacotron2/waveglow_denoiser.py

@@ -0,0 +1,40 @@
+# import sys
+# sys.path.append('tacotron2')
+import torch
+from .layers import STFT
+
+
+class Denoiser(torch.nn.Module):
+    """ Removes model bias from audio produced with waveglow """
+
+    def __init__(self, waveglow, filter_length=1024, n_overlap=4,
+                 win_length=1024, mode='zeros'):
+        super(Denoiser, self).__init__()
+        self.stft = STFT(filter_length=filter_length,
+                         hop_length=int(filter_length/n_overlap),
+                         win_length=win_length).cuda()
+        if mode == 'zeros':
+            mel_input = torch.zeros(
+                (1, 80, 88),
+                dtype=waveglow.upsample.weight.dtype,
+                device=waveglow.upsample.weight.device)
+        elif mode == 'normal':
+            mel_input = torch.randn(
+                (1, 80, 88),
+                dtype=waveglow.upsample.weight.dtype,
+                device=waveglow.upsample.weight.device)
+        else:
+            raise Exception("Mode {} if not supported".format(mode))
+
+        with torch.no_grad():
+            bias_audio = waveglow.infer(mel_input, sigma=0.0).float()
+            bias_spec, _ = self.stft.transform(bias_audio)
+
+        self.register_buffer('bias_spec', bias_spec[:, :, 0][:, :, None])
+
+    def forward(self, audio, strength=0.1):
+        audio_spec, audio_angles = self.stft.transform(audio.cuda().float())
+        audio_spec_denoised = audio_spec - self.bias_spec * strength
+        audio_spec_denoised = torch.clamp(audio_spec_denoised, 0.0)
+        audio_denoised = self.stft.inverse(audio_spec_denoised, audio_angles)
+        return audio_denoised

omni_speech/infer/fairseq/examples/textless_nlp/pgslm/README.md

@@ -0,0 +1,318 @@
+# Text-Free Prosody-Aware Generative Spoken Language Modeling
+
+This folder contains code and recipes to reproduce results reported in a paper _Text-Free Prosody-Aware Generative Spoken Language Modeling_,
+Eugene Kharitonov*, Ann Lee*, Adam Polyak, Yossi Adi, Jade Copet, Kushal Lakhotia, Tu-Anh Nguyen, Morgane Rivière, Abdelrahman Mohamed, Emmanuel Dupoux, Wei-Ning Hsu, 2021. arxiv/2109.03264 [[arxiv]](https://arxiv.org/abs/2109.03264).
+
+`*` denotes equal contribution.
+
+You can find demo samples [[here]](https://speechbot.github.io/pgslm/index.html).
+
+<details>
+  <summary>If you find this code useful, please consider citing our work using this bibtex </summary>
+  
+```
+  @misc{Kharitonov2021,
+      title={Text-Free Prosody-Aware Generative Spoken Language Modeling}, 
+      author={Eugene Kharitonov and Ann Lee and Adam Polyak and Yossi Adi and Jade Copet and Kushal Lakhotia and Tu-Anh Nguyen and Morgane Rivière and Abdelrahman Mohamed and Emmanuel Dupoux and Wei-Ning Hsu},
+      year={2021},
+      eprint={2109.03264},
+      archivePrefix={arXiv},
+      primaryClass={cs.CL}
+}
+```
+</details>
+
+
+## Additional requirements
+Three packages are required in addition to fairseq, they are installable with pip:
+```bash 
+pip install AMFM-decompy SoundFile scipy sklearn torchaudio npy-append-array
+```
+
+## Data preprocessing
+
+### Prepare unit pseudo-text transcriptions of the audio
+To get unit trascripts of the speech data we rely on the preprocessing steps of [GSLM](https://github.com/pytorch/fairseq/tree/main/examples/textless_nlp/gslm/speech2unit/) work.
+
+Firstly, we will need to prepare manifest files for the dataset we want to preprocess
+```
+mkdir manifests/
+python examples/wav2vec/wav2vec_manifest.py --valid-percent=0.0 $DATA_PATH --dest=manifests/train/
+```
+Next, we need a pre-trained HuBERT-base-ls960 model [[download]](https://dl.fbaipublicfiles.com/hubert/hubert_base_ls960.pt) and a corresponding kmeans-100 quantizer [[download]](https://dl.fbaipublicfiles.com/textless_nlp/gslm/hubert/km100/km.bin). Having those we can quantize the dataset:
+```
+python examples/textless_nlp/gslm/speech2unit/clustering/quantize_with_kmeans.py \
+    --feature_type hubert \
+    --kmeans_model_path km.bin \
+    --acoustic_model_path hubert_base_ls960.pt \
+    --layer 6 \
+    --manifest_path manifests/train/train.tsv \
+    --out_quantized_file_path manifests/train/units
+```
+
+Finally, by running
+```
+python examples/textless_nlp/pgslm/scripts/join_units_manifest.py --manifest=manifests/train/train.tsv --units=manifests/train/units --output=train.txt
+```
+We will get the training data description `train.txt` in the format that pGSLM expects. The above steps have to be repeated for 
+dev/test sets. Importantly, we rely on an assumption that the directories are structured as in LibriSpeech, i.e. the file paths follow the
+`<spk_id>/<session_id>/<sample_id>.wav` format.
+
+### Preprocess data for pGSLM
+The very first step is to obtain the F0 quantization bins.
+Assume the vocoder training manifest is `vocoder_train.txt` (in pGSLM data format prepared with the same process above).
+We prepare the quantized F0 from the vocoder training data by running
+```sh
+bash examples/textless_nlp/pgslm/scripts/prepare_f0_quantization.sh \
+  vocoder_train.txt <sample_rate> 32 <preprocessed_dir> <output_prefix> # we use 32 bins in the paper
+```
+- `<sample_rate>`: sampling rate of the audio files in the manifest
+- `<preprocessed_dir>`: where to output the output files
+- `<output_prefix>`: prefix of the output files
+
+The script will generate 
+- `<output_prefix>.f0_stat.pt`: the speaker-level F0 statistics, which can be used in vocoder training
+- `<output_prefix>_mean_norm_log_f0_bin.th`: the quantized F0, which should be used in `prepare_data.sh` below
+
+**Note:** See "Pre-trained models" for the pre-computed speaker-level F0 statistics and quantized F0 bins. We suggest using the pre-computed statistics for the data preparation below in order to take advantage of the pre-trained vocoder for waveform generation.
+
+Next prepare the pGSLM data.
+Assume train/valid/test manifests are `{train,valid,test}.txt`.
+Here is an example of how to preprocess data:
+
+```sh
+bash examples/textless_nlp/pgslm/scripts/prepare_data.sh \
+  train.txt valid.txt test.txt <n_unit> <hop_size> <sample_rate> \
+  <preprocessed_dir>/<output_prefix>_mean_norm_log_f0_bin.th <preprocessed_dir>
+```
+- `<n_unit>`: discrete unit vocabulary size (we used a kmeans quantizer with the number of units equal to 100 in the example above)
+- `<hop_size>`: downsampling rate relative to the waveform (e.g., 320 for HuBERT units)
+- `<sample_rate>`: sampling rate of the audio files in the manifest
+- `<preprocessed_dir>`: where to output the preprocessed files
+
+This will create the dataset json config used for the next section at
+`<preprocessed_dir>/data_config.json`.
+
+Note that the example script uses only one thread to compute F0, which can take
+_very long_ for preprocessing large datasets. It is suggested to distribute
+jobs over multiple nodes/processes with `--nshards=x` and `--rank=z` (where z is
+in [1, x]) in `preprocess_f0.py`, and set `--nshards_list=x` in
+`prepare_data.py` correspondingly to collect sharded F0 data.
+
+Now, everything is ready for training a model.
+
+## Training Multi-Stream Transformer Unit Language Model (MS-TLM)
+
+Below is an example command that trains Multi-Stream Transformer Language Model (MS-TLM) on a prepared dataset:
+```bash
+DATASET=data_config.json
+
+fairseq-train $DATASET \
+  --task=speech_unit_modeling \
+  --arch="transformer_ulm_tiny" \
+  --criterion=speech_unit_lm_criterion \
+  --share-decoder-input-output-embed \
+  --dropout=0.1 \
+  --attention-dropout=0.1 \
+  --optimizer="adam" \
+  --adam-betas="(0.9, 0.98)" \
+  --clip-norm=1.0 \
+  --lr=0.0005 \
+  --lr-scheduler="inverse_sqrt" \
+  --warmup-updates=4000 \
+  --warmup-init-lr=1e-07 \
+  --tokens-per-sample=3072 \
+  --max-tokens=3072 \
+  --update-freq=4 \
+  --max-epoch=70 \
+  --num-workers=0 \
+  --skip-invalid-size-inputs-valid-test \
+  --loss-weights="1.0;0.5;0.0" \
+  --ignore-f0-input \
+  --checkpoint-activations \
+  --fp16 \
+  --max-target-positions=4096 \
+  --stream-shifts="1,1" \
+  --log-f0 --normalize-f0-mean --interpolate-f0 \
+  --ignore-unused-valid-subsets \
+  --discrete-duration --discrete-f0
+```
+
+Some of the important parameters that are specific to MS-TLM:
+ *  `arch`: specifies the Transformer architecture used. Supported options are:
+    * `transformer_ulm_tiny` - a tiny model that can be used for debugging; it has 2 layers, 1 attention head, FFN and embedding dimensions of 64,
+    * `transformer_ulm` - a base model with 6 layers, 8 heads, embedding dimension 512, and FFN dimensionality of 2048,
+    * `transformer_ulm_big` - the largest model we experiment with in the paper: 12-layer/16 heads, 1024/4096 embedding and FFN dimensions;
+ * `loss-weights`: this parameter sets importance weights (must be non-negative) for the components of the loss that correspond to unit, duration, and F0 streams. To turn off a component of the loss, its weight has to be set to 0. For instance, to predict only unit stream the parameter should be set to "1;0;0";
+ * `stream-shifts`: specifies relative shifts of the two prosodic streams w.r.t. the unit stream (duration and F0, respectively). No shift corresponds to "0,0";
+ * `ignore-duration-input`/`ignore-f0-input`: setting these flags would zero-out correpsonding input streams;
+ * `max-token-duration`: duration values would be max-capped by the specified value;
+ * `discrete-duration`/`discrete-f0`: whether duration and F0 streams should be quantized;
+ * `log_f0`, `normalize-f0-mean`, `normalize-f0-std`, `interpolate-f0`: configure how F0 stream is treated. `log_f0` sets up modelling in the log-space, `normalize-f0-mean`/`normalize-f0-std` control per-speaker normalization, and `interpolate-f0` enables F0 interpolation for unvoiced regions where F0 was set to 0,
+ * `mask-dur-prob`, `mask-f0-prob`, `mask-dur-seg-prob`, `mask-f0-seg-prob`, `mask-unit-seg-prob`, `mask-unit-seg-leng`: this family of parameters sets the probababilities of masking individual steps and spans on each stream as well as lengths of the maked spans.
+
+
+## Pre-trained models
+### MS-TLM
+Below you can find checkpoints for four best-performing models from the paper (IDs 9..12 in Table 1). These models are trained on Hubert-100 transcripts of the LibriLight-6K dataset. They have the prosody streams shifted by 1 w.r.t. the unit stream. All models predict all three streams (units, duration, and F0), but two
+of them only have unit steam in their input.
+
+|                   | Continuous prosody | Quantized prosody |
+|-------------------|--------------------|-------------------|
+| No prosody input  | [[download]](https://dl.fbaipublicfiles.com/textless_nlp/pgslm/ulm_checkpoints/continuous_no_prosody_shift_1_1.pt) | [[download]](https://dl.fbaipublicfiles.com/textless_nlp/pgslm/ulm_checkpoints/discrete_no_prosody_shift_1_1.pt)  |
+| Has prosody input | [[download]](https://dl.fbaipublicfiles.com/textless_nlp/pgslm/ulm_checkpoints/continuous_prosody_shift_1_1.pt) | [[download]](https://dl.fbaipublicfiles.com/textless_nlp/pgslm/ulm_checkpoints/discrete_prosody_shift_1_1.pt)|
+
+The optimal per-stream sampling temperatures/scaling parameters that we have identified for these models, in the (`T-token, T-duration, T-f0`) format:
+
+|                   | Continuous prosody | Quantized prosody |
+|-------------------|--------------------|-------------------|
+| No prosody input  |  0.7, 0.125, 0.0003125|    0.7, 0.25, 0.5 |
+| Has prosody input |  0.7, 0.125, 0.00125  |   0.7, 0.25, 0.7  |
+
+## Vocoder
+|       Units       | Prosody | F0 stats     | Checkpoint | Config |
+|-------------------|---------|--------------|------------|--------|
+| HuBERT-base-ls960, kmeans-100 | [[Quantized 32 bins]](https://dl.fbaipublicfiles.com/textless_nlp/pgslm/vocoder/blizzard2013/mean_norm_log_f0_seg_bin.th) | [[download]](https://dl.fbaipublicfiles.com/textless_nlp/pgslm/vocoder/blizzard2013/f0_stats.pt) | [[download]](https://dl.fbaipublicfiles.com/textless_nlp/pgslm/vocoder/blizzard2013/naive_quant_32_norm_log_seg_hubert/checkpoint.pt) | [[download]](https://dl.fbaipublicfiles.com/textless_nlp/pgslm/vocoder/blizzard2013/naive_quant_32_norm_log_seg_hubert/config.json) |
+| HuBERT-base-ls960, kmeans-100 | Continuous | [[download]](https://dl.fbaipublicfiles.com/textless_nlp/pgslm/vocoder/blizzard2013/f0_stats.pt) | [[download]](https://dl.fbaipublicfiles.com/textless_nlp/pgslm/vocoder/blizzard2013/mean_norm_log_f0_hubert/checkpoint.pt) | [[download]](https://dl.fbaipublicfiles.com/textless_nlp/pgslm/vocoder/blizzard2013/mean_norm_log_f0_hubert/config.json) |
+
+
+## Evaluating a trained model
+Evaluation is done with the `eval/cont_metrics.py` scripts. As described in the paper, there are several metrics used.
+
+**Teacher-forced metrics**
+```bash
+SET=valid
+CHECKPOINT_PATH=discrete_prosody_shift_1_1.pt
+DATA=data_config.json
+
+python examples/textless_nlp/pgslm/eval/cont_metrics.py $DATA \
+  --metric=teacher_force_everything \
+  --path=$CHECKPOINT_PATH \
+  --batch-size=16 \
+  --fp16 \
+  --seed=111 \
+  --eval-subset=$SET \
+  --f0-discretization-bounds=mean_norm_log_f0_seg_bin.th --dequantize-prosody 
+```
+(Using this command, our provided `discrete_prosody_shift_1_1.pt` checkpoint should produce `{'token_loss': 1.408..., 'duration_loss': 0.5424..., 'f0_loss': 0.0474...}` on LibriSpeech dev-clean).
+
+The parameters `--f0-discretization-bounds=mean_norm_log_f0_seg_bin.th --dequantize-prosody` are specific for quantized-prosody models. They signal that the prosody streams must be decoded into the continuous domain before calculating correlation. It is the same `*_mean_norm_log_f0_bin.th` file as we prepared before.
+The `mean_norm_log_f0_seg_bin.th` file we used with the pre-trained models can be downloaded [[here]](https://dl.fbaipublicfiles.com/textless_nlp/pgslm/vocoder/blizzard2013/mean_norm_log_f0_seg_bin.th).
+
+
+**Consistency (aka Correlation) metrics**
+
+The following command estimates correlation between mean values of the F0 stream in the prompt and in the generated continuation (unit and duration steams are fixed).
+
+```bash
+T_F0=0.7
+EXPLOSION=20
+SET=test
+CHECKPOINT_PATH=discrete_prosody_shift_1_1.pt
+DATA=data_config.json
+
+python examples/textless_nlp/pgslm/eval/cont_metrics.py $DATA \
+    --prefix-length=150 \
+    --metric=correlation \
+    --path=$CHECKPOINT_PATH \
+    --batch-size=16 \
+    --fp16 \
+    --seed=111 \
+    --teacher-force-tokens \
+    --teacher-force-duration  \
+    --min-length=300  \
+    --batch-explosion-rate=$EXPLOSION \
+    --T-f0=$T_F0 \
+    --eval-subset=$SET \
+    --f0-discretization-bounds=mean_norm_log_f0_seg_bin.th \
+    --dequantize-prosody --n-workers=8
+```
+(Using this command, our provided `discrete_prosody_shift_1_1.pt` checkpoint should produce `{...'F0 corr': 0.315 ..}` on LibriSpeech test-clean).
+
+ * By using flags `--teacher-force-tokens, --teacher-force-duration, --teacher-force-f0` one can calculate correlations along each stream while having other two streams fixed to ground-truth values (or freeze all three streams to get ground-truth correlation values);
+ * The parameters `T-f0`, `T-duration`, and `T-token` specify per-stream temperatures and, in the case of continuous-valued prosody, scaling parameter of the corresponding Laplace distribution (setting a temperature to 0 will enforce greedy sampling);
+ * `min-length` filters out sequences that are shorter then 300 duration units (i.e. 6s in the case of Hubert units);
+ * `prefix-length` specifies that we want to use first 150 duration units are prompt (i.e. 3s in the case of Hubert units)
+
+
+**Correctness (aka Continuation) and Expressiveness (aka Std) metrics**
+
+By running the following command, we can get minMAE and Std for the log-F0 stream for the model with quantized prosody.
+```bash
+DATA=data_config.json
+EXPLOSION=20
+SET=test
+CHECKPOINT_PATH=discrete_prosody_shift_1_1.pt
+T_F0=0.7
+
+python examples/textless_nlp/pgslm/eval/cont_metrics.py $DATA \
+  --prefix-length=150 \
+  --metric=continuation \
+  --path=$CHECKPOINT_PATH \
+  --batch-size=16 \
+  --fp16 \
+  --seed=111 \
+  --batch-explosion-rate=$EXPLOSION \
+  --teacher-force-tokens \
+  --teacher-force-duration \
+  --T-f0=$T_F0 \
+  --eval-subset=$SET \
+  --f0-discretization-bounds=mean_norm_log_f0_seg_bin.th --dequantize-prosody
+```
+(Using this command, our provided `discrete_prosody_shift_1_1.pt` checkpoint should produce `{...'F0 MAE': 0.0772, 'F0 Std': 0.1489...}` on LibriSpeech test-clean).
+
+Again, by setting `--teacher-force-tokens, --teacher-force-duration, --teacher-force-f0` we can calculate Token BLEU for the token stream (when `--teacher-force-duration` &  `--teacher-force-f0` are on) and per-stream min MAE for each prosody stream individually.
+
+Finally, `cont_metrics.py` allows to specify the number of workers (e.g., `n-workers=8`) which allows to speed up the computation by spreading multiple worker processes 
+over the available GPUs.
+
+**Cont Word BLEU**
+
+We used the code and the evaluation protocol of [(Lakhotia et al., 2021)](https://arxiv.org/abs/2102.01192).
+
+## Sampling from a trained model
+
+To get (prompted or not) samples from a trained model it is enough to run `sample.py`:
+```bash
+CHECKPOINT_PATH=checkpoints/checkpoint_best.pt
+DATASET=examples/textless_nlp/pgslm/repro/dataset/data_config.json 
+python examples/textless_nlp/pgslm/sample/sample.py $DATASET \
+  --output=$SAMPLES \
+  --path=$CHECKPOINT_PATH \
+  --sampling \
+  --T-token=0.7 \
+  --T-duration=0.25 \
+  --T-f0=0.7 \
+  --max-length=500 \
+  --prefix-length=150 \
+  --subset=valid \
+  --seed=1 \
+  --match-duration \
+  --code-type=hubert \
+  --batch-explosion-rate=2
+```
+
+Some useful parameters:
+ * `T-token`, `T-duration`, `T-f0` specify sampling temperature for the three streams. Setting a temperature to `0` switches sample to the greedy (argmax) one;
+ * `prefix-length`: length of the prompt, measured in timesteps (e.g. for Hubert (CPC) each timestep is 20 (10) ms);
+ * `subset`: which subset of the dataset to use as prompts (can be `train`, `valid`, `test`);
+ * `teacher-force-tokens`, `teacher-force-duration`, `teacher-force-f0`: if set, at each autoregressive step, ground-truth values replace the produced one;
+ * `short-curcuit`: replace sampling by ground-truth inputs;
+ * `match-duration`: forces the produced sample to have the same duration (in time), as the entire sequence (beyond the prompt if there is any);
+ * `batch-explosion-rate`: number of samples per prompt;
+ * `f0-discretization-bounds`: path to a file with quantization boundaries. If it is set, F0 values are de-quantized back to the continuous domain
+      (the model must be a quanized one);
+  * `max-length` sets the maximal number of segment steps to be produced.
+
+Note that `sample.py` automatically uses all available GPUs, to avoid that please use environment variable `CUDA_VISIBLE_DEVICES`.
+
+## Vocoding samples
+To generate audios for output from `sample.py` (`$IN_FILE`):
+```bash
+python examples/textless_nlp/pgslm/generate_waveform.py \
+  --in-file=$IN_FILE \
+  --vocoder=$VODOER \
+  --vocoder-cfg=$VOCODER_CFG \
+  --results-path=$RESULTS_PATH
+```
+See "Pre-trained model" for `$VOCODER` and `VOCODER_CFG`.

omni_speech/infer/fairseq/examples/textless_nlp/pgslm/data_utils.py

@@ -0,0 +1,107 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+import torch
+
+from tqdm import tqdm
+
+
+class Stat:
+    def __init__(self, keep_raw=False):
+        self.x = 0.0
+        self.x2 = 0.0
+        self.z = 0.0  # z = logx
+        self.z2 = 0.0
+        self.n = 0.0
+        self.u = 0.0
+        self.keep_raw = keep_raw
+        self.raw = []
+
+    def update(self, new_x):
+        new_z = new_x.log()
+
+        self.x += new_x.sum()
+        self.x2 += (new_x**2).sum()
+        self.z += new_z.sum()
+        self.z2 += (new_z**2).sum()
+        self.n += len(new_x)
+        self.u += 1
+
+        if self.keep_raw:
+            self.raw.append(new_x)
+
+    @property
+    def mean(self):
+        return self.x / self.n
+
+    @property
+    def std(self):
+        return (self.x2 / self.n - self.mean**2) ** 0.5
+
+    @property
+    def mean_log(self):
+        return self.z / self.n
+
+    @property
+    def std_log(self):
+        return (self.z2 / self.n - self.mean_log**2) ** 0.5
+
+    @property
+    def n_frms(self):
+        return self.n
+
+    @property
+    def n_utts(self):
+        return self.u
+
+    @property
+    def raw_data(self):
+        assert self.keep_raw, "does not support storing raw data!"
+        return torch.cat(self.raw)
+
+
+class F0Stat(Stat):
+    def update(self, new_x):
+        # assume unvoiced frames are 0 and consider only voiced frames
+        if new_x is not None:
+            super().update(new_x[new_x != 0])
+
+
+def dump_speaker_f0_stat(speaker_to_f0_stat, out_prefix):
+    path = f"{out_prefix}.f0_stat.pt"
+    assert not os.path.exists(path)
+
+    d = {
+        speaker: {
+            "f0_mean": speaker_to_f0_stat[speaker].mean,
+            "f0_std": speaker_to_f0_stat[speaker].std,
+            "logf0_mean": speaker_to_f0_stat[speaker].mean_log,
+            "logf0_std": speaker_to_f0_stat[speaker].std_log,
+        }
+        for speaker in speaker_to_f0_stat
+    }
+    torch.save(d, path)
+
+    return d
+
+
+def load_audio_path(path):
+    audio_paths = []
+    with open(path) as f:
+        for line in f.readlines():
+            sample = eval(line.strip())
+            audio_paths.append(sample["audio"])
+
+    return audio_paths
+
+
+def load_f0(f0_dir, nshards):
+    path_to_f0 = {}
+    for rank in tqdm(range(1, nshards + 1), desc=f"load f0"):
+        f0_shard_path = f"{f0_dir}/f0_{rank}_{nshards}.pt"
+        shard_path_to_f0 = torch.load(f0_shard_path)
+        path_to_f0.update(shard_path_to_f0)
+    return path_to_f0

omni_speech/infer/fairseq/examples/textless_nlp/pgslm/eval/__init__.py

@@ -0,0 +1,4 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.

omni_speech/infer/fairseq/examples/textless_nlp/pgslm/eval/cont_metrics.py

@@ -0,0 +1,730 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+import numpy as np
+import scipy
+
+import torch
+import torch.multiprocessing as mp
+from fairseq import checkpoint_utils, options
+from fairseq.data.codedataset import CodeDataset, ExpressiveCodeDataConfig
+from fairseq.dataclass.utils import convert_namespace_to_omegaconf
+from torch.utils.data import DataLoader, DistributedSampler
+from fairseq.utils import move_to_cuda
+from fairseq import utils
+from fairseq.criterions.speech_ulm_criterion import nll_loss, mae_loss
+
+import time
+from types import SimpleNamespace
+
+import sys, pathlib
+
+sys.path.append(str(pathlib.Path(__file__).parent.parent.resolve()))
+
+from naive_decoder import Naive_F0_Decoder
+from inference_dataset import InferenceDataset, explode_batch
+from sample.sample import do_sampling, TemperatureDecoder, FilterNamesDataset
+
+try:
+    from nltk.translate.bleu_score import sentence_bleu
+except ImportError:
+    print("Please install nltk: `pip install --user -U nltk`")
+    raise
+
+
+@torch.no_grad()
+def teacher_force_everything(
+    args, dataset, model, criterion, tgt_dict, rank, world_size
+):
+    prefix = args.prefix_length
+
+    f0_decoder = None
+    if args.dequantize_prosody:
+        assert dataset.discrete_f0
+        print("Reporting MAE for a discrete model")
+        f0_decoder = Naive_F0_Decoder(
+            args.f0_discretization_bounds, dataset.config.f0_vq_n_units
+        ).cuda()
+
+    dataset = InferenceDataset(
+        dataset,
+        prefix=args.prefix_length,
+        only_prefix=False,
+        filter_short=True,
+        presort_by_length=True,
+    )
+    sampler = (
+        None
+        if world_size == 1
+        else DistributedSampler(
+            dataset, num_replicas=world_size, rank=rank, shuffle=False
+        )
+    )
+    dataloader = DataLoader(
+        dataset,
+        args.batch_size,
+        shuffle=False,
+        collate_fn=dataset.collater,
+        sampler=sampler,
+    )
+
+    total_token_loss, total_duration_loss, total_f0_loss, total_tokens = (
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+    )
+
+    i = 0
+    for batch in dataloader:
+        i += 1
+        batch = move_to_cuda(batch)
+        output = model(**batch["net_input"])
+
+        tokens, durations, f0 = output["token"], output["duration"], output["f0"]
+        durations, f0 = durations.squeeze(), f0.squeeze()
+
+        token_loss = nll_loss(
+            tokens[:, prefix - 1 :],
+            batch["target"][:, prefix - 1 :].contiguous(),
+            batch["mask"][:, prefix - 1 :].contiguous(),
+            reduce=True,
+        )
+
+        if args.dequantize_prosody:
+            durations = durations.argmax(dim=-1)
+            duration_loss = mae_loss(
+                durations[:, prefix - 1 :].contiguous().float(),
+                batch["dur_target"][:, prefix - 1 :].contiguous().float(),
+                batch["dur_mask"][:, prefix - 1 :].contiguous(),
+                reduce=True,
+            )
+        else:
+            duration_loss = criterion.dur_loss_fn(
+                durations[:, prefix - 1 :].contiguous(),
+                batch["dur_target"][:, prefix - 1 :].contiguous(),
+                batch["dur_mask"][:, prefix - 1 :].contiguous(),
+                reduce=True,
+            )
+
+        if f0_decoder:
+            f0 = f0.argmax(dim=-1)
+            f0 = f0_decoder(f0).squeeze(-1)
+
+            f0_target = batch["raw_f0"]
+            f0_loss = mae_loss(
+                f0[:, prefix - 1 :].contiguous(),
+                f0_target[:, prefix - 1 :].contiguous(),
+                batch["f0_mask"][:, prefix - 1 :].contiguous(),
+                reduce=True,
+            )
+        else:
+            f0_loss = criterion.f0_loss_fn(
+                f0[:, prefix - 1 :].contiguous(),
+                batch["f0_target"][:, prefix - 1 :].contiguous(),
+                batch["f0_mask"][:, prefix - 1 :].contiguous(),
+                reduce=True,
+            )
+
+        n_tokens = (~batch["dur_mask"])[:, prefix - 1 :].sum()
+
+        total_token_loss += token_loss.item()
+        total_duration_loss += duration_loss.item()
+        total_f0_loss += f0_loss.item()
+
+        total_tokens += n_tokens.item()
+        if args.debug and i > 5:
+            break
+
+    values = torch.tensor([total_token_loss, total_duration_loss, total_f0_loss])
+    normalizers = torch.tensor([total_tokens for _ in range(3)])
+
+    return values, normalizers
+
+
+def get_bleu(produced_tokens, target_tokens, tgt_dict):
+    assert target_tokens.ndim == 1
+    assert produced_tokens.size(1) == target_tokens.size(0)
+
+    # we can have padding due to shifted channels
+    shift = 0
+    for token in reversed(target_tokens.cpu().tolist()):
+        if token in [tgt_dict.pad(), tgt_dict.eos()]:
+            shift += 1
+        else:
+            break
+    target_tokens = target_tokens[:-shift]
+    produced_tokens = produced_tokens[:, :-shift]
+
+    string_target = tgt_dict.string(target_tokens).split()
+    string_candidates = [
+        tgt_dict.string(produced_tokens[i, :]).split()
+        for i in range(produced_tokens.size(0))
+    ]
+
+    bleu3 = sentence_bleu(
+        references=string_candidates,
+        hypothesis=string_target,
+        weights=(1.0 / 3, 1.0 / 3, 1.0 / 3),
+    )
+    return bleu3
+
+
+@torch.no_grad()
+def continuation(args, dataset, model, criterion, tgt_dict, rank, world_size):
+    is_discrete_duration = dataset.discrete_dur
+    is_discrete_f0 = dataset.discrete_f0
+
+    f0_decoder = None
+    if args.dequantize_prosody:
+        assert dataset.discrete_f0
+        print("Reporting MAE F0 for a discrete model")
+        f0_decoder = Naive_F0_Decoder(
+            args.f0_discretization_bounds, dataset.config.f0_vq_n_units
+        ).cuda()
+
+    dataset = InferenceDataset(
+        dataset, args.prefix_length, filter_short=True, presort_by_length=True
+    )
+    sampler = (
+        None
+        if world_size == 1
+        else DistributedSampler(
+            dataset, num_replicas=world_size, rank=rank, shuffle=False
+        )
+    )
+    dataloader = DataLoader(
+        dataset,
+        batch_size=1,
+        shuffle=False,
+        collate_fn=dataset.collater,
+        sampler=sampler,
+    )
+
+    Ts = args.T_token, args.T_duration, args.T_f0
+    decoder = TemperatureDecoder(
+        Ts, discrete_dur=is_discrete_duration, discrete_f0=is_discrete_f0
+    )
+
+    running_stats = SimpleNamespace(
+        token_bleu=0.0,
+        duration_nll=0.0,
+        duration_mae=0.0,
+        f0_nll=0.0,
+        f0_mae=0.0,
+        n_tokens=0.0,
+        n_sentences=0.0,
+        f0_sum=0.0,
+        f0_sum_sq=0.0,
+        dur_sum=0.0,
+        dur_sum_sq=0.0,
+    )
+
+    for i, batch in enumerate(dataloader):
+        batch = explode_batch(batch, args.batch_explosion_rate)
+        bsz = batch["target"].size(0)
+
+        batch = move_to_cuda(batch)
+        prefix = batch["prefix"][0]
+
+        max_length_to_unroll = batch["target"].size(1)
+        prefix_length = batch["net_input"]["src_tokens"].size(1)
+        steps = max_length_to_unroll - prefix_length + 1
+
+        assert steps > 0
+        produced_tokens, produced_durations, produced_f0, outputs = do_sampling(
+            model,
+            batch,
+            tgt_dict.eos(),
+            decoder,
+            autoregressive_steps=steps,
+            teacher_force_tokens=args.teacher_force_tokens,
+            teacher_force_duration=args.teacher_force_duration,
+            teacher_force_f0=args.teacher_force_f0,
+        )
+
+        if args.teacher_force_tokens:
+            assert (produced_tokens[:, 1:] == batch["target"]).all()
+        if args.teacher_force_duration:
+            assert (produced_durations[:, 1:] == batch["dur_target"]).all()
+        if args.teacher_force_f0:
+            assert (produced_f0[:, 1:] == batch["f0_target"]).all()
+
+        dur_target = batch["dur_target"][:, prefix - 1 :].contiguous()
+        f0_target = batch["f0_target"][:, prefix - 1 :].contiguous()
+
+        f0_mask = batch["f0_mask"][:, prefix - 1 :].contiguous()
+        dur_mask = batch["dur_mask"][:, prefix - 1 :].contiguous()
+
+        duration_mae = mae_loss(
+            produced_durations[:, prefix:].float(),
+            dur_target.float(),
+            dur_mask,
+            reduce=False,
+        )
+        min_duration_mae = duration_mae.view(bsz, -1).sum(dim=-1).min(dim=0)[0]
+        running_stats.duration_mae += min_duration_mae
+
+        running_stats.dur_sum += (
+            produced_durations[:, prefix:].float() * (~dur_mask)
+        ).sum() / args.batch_explosion_rate
+        running_stats.dur_sum_sq += (
+            produced_durations[:, prefix:].float() * (~dur_mask)
+        ).pow(2.0).sum() / args.batch_explosion_rate
+
+        if is_discrete_duration:
+            duration_loss = criterion.dur_loss_fn(
+                torch.stack([x[1] for x in outputs], dim=1),
+                dur_target,
+                dur_mask,
+                reduce=False,
+            )
+            min_duration_loss = duration_loss.view(bsz, -1).sum(dim=-1).min(dim=0)[0]
+            running_stats.duration_nll += min_duration_loss
+
+        if f0_decoder:  # can only exist for discrete F0 models
+            decoded_produced_f0 = f0_decoder(produced_f0[:, prefix:])
+            decoded_f0_target = batch["raw_f0"][:, prefix - 1 :].contiguous()
+
+            if produced_f0.ndim == 3:
+                decoded_produced_f0 = decoded_produced_f0.squeeze(2)
+                decoded_f0_target = decoded_f0_target.squeeze(2)
+
+            f0_mae = mae_loss(
+                decoded_produced_f0, decoded_f0_target, f0_mask, reduce=False
+            )
+            f0_mae = f0_mae.view(bsz, -1).sum(dim=-1).min(dim=0)[0]
+            running_stats.f0_mae += f0_mae
+
+            f0_loss = criterion.f0_loss_fn(
+                torch.stack([x[2] for x in outputs], dim=1),
+                f0_target.long(),
+                f0_mask,
+                reduce=False,
+            )
+            f0_loss = f0_loss.view(bsz, -1).sum(dim=-1).min(dim=0)[0]
+            running_stats.f0_nll += f0_loss
+
+            running_stats.f0_sum += (
+                decoded_produced_f0 * (~f0_mask)
+            ).sum() / args.batch_explosion_rate
+            running_stats.f0_sum_sq += (decoded_produced_f0 * (~f0_mask)).pow(
+                2.0
+            ).sum() / args.batch_explosion_rate
+
+        else:
+            assert not is_discrete_duration
+
+            f0_loss = mae_loss(
+                produced_f0[:, prefix:], f0_target, f0_mask, reduce=False
+            )
+            f0_loss = f0_loss.view(bsz, -1).sum(dim=-1).min(dim=0)[0]
+            running_stats.f0_mae += f0_loss
+
+            running_stats.f0_sum += (
+                produced_f0[:, prefix:].sum() / args.batch_explosion_rate
+            )
+            running_stats.f0_sum_sq += (
+                produced_f0[:, prefix:].pow(2.0).sum() / args.batch_explosion_rate
+            )
+
+        running_stats.n_tokens += (~dur_mask)[0, ...].sum()
+
+        token_loss = get_bleu(
+            produced_tokens[:, prefix:], batch["target"][0, prefix - 1 :], tgt_dict
+        )
+        running_stats.token_bleu += token_loss
+        running_stats.n_sentences += 1
+
+        if args.debug:
+            break
+
+    values = torch.tensor(
+        [
+            running_stats.token_bleu,
+            running_stats.duration_nll,
+            running_stats.duration_mae,
+            running_stats.f0_nll,
+            running_stats.f0_mae,
+            running_stats.f0_sum,
+            running_stats.f0_sum_sq,
+            running_stats.dur_sum,
+            running_stats.dur_sum_sq,
+        ]
+    )
+    normalizers = torch.tensor(
+        [running_stats.n_sentences] + [running_stats.n_tokens] * 8
+    )
+
+    return values, normalizers
+
+
+@torch.no_grad()
+def correlation(args, dataset, model, criterion, tgt_dict, rank, world_size):
+    is_discrete_duration = dataset.discrete_dur
+    is_discrete_f0 = dataset.discrete_f0
+
+    f0_decoder = None
+    if is_discrete_f0:
+        assert dataset.discrete_f0
+        f0_decoder = Naive_F0_Decoder(
+            args.f0_discretization_bounds, dataset.config.f0_vq_n_units
+        ).cuda()
+
+    if is_discrete_f0:
+        assert f0_decoder  # correlation on tokens is meaningless
+
+    dataset = InferenceDataset(
+        dataset,
+        args.prefix_length,
+        filter_short=True,
+        presort_by_length=True,
+        min_length=args.min_length,
+    )
+    sampler = (
+        None
+        if world_size == 1
+        else DistributedSampler(
+            dataset, num_replicas=world_size, rank=rank, shuffle=False
+        )
+    )
+    dataloader = DataLoader(
+        dataset,
+        batch_size=1,
+        shuffle=False,
+        collate_fn=dataset.collater,
+        sampler=sampler,
+    )
+
+    Ts = args.T_token, args.T_duration, args.T_f0
+    decoder = TemperatureDecoder(
+        Ts, discrete_dur=is_discrete_duration, discrete_f0=is_discrete_f0
+    )
+
+    mean_dur_prefix, mean_dur_cont = [], []
+    mean_f0_prefix, mean_f0_cont = [], []
+
+    for batch in dataloader:
+        batch = explode_batch(batch, args.batch_explosion_rate)
+        batch = move_to_cuda(batch)
+
+        assert len(batch["prefix"]) == 1
+
+        if args.teacher_force_tokens:
+            autoregressive_steps = batch["target"].size(1) - args.prefix_length - 1
+        else:
+            autoregressive_steps = args.max_length - args.prefix_length  # + max_shift?
+
+        if args.copy_target:
+            produced_durations, produced_f0 = batch["dur_target"], batch["f0_target"]
+        else:
+            _, produced_durations, produced_f0, outputs = do_sampling(
+                model,
+                batch,
+                tgt_dict.eos(),
+                decoder,
+                autoregressive_steps=autoregressive_steps,
+                teacher_force_tokens=args.teacher_force_tokens,
+                teacher_force_duration=args.teacher_force_duration,
+                teacher_force_f0=args.teacher_force_f0,
+            )
+
+            # first tokens actually correspond to BOS
+            produced_durations = produced_durations[:, 1:]
+            produced_f0 = produced_f0[:, 1:]
+
+        dur_target = batch["dur_target"]
+        if is_discrete_duration:
+            produced_durations = produced_durations.float()
+            dur_target = dur_target.float()
+
+        if is_discrete_f0:
+            produced_f0 = f0_decoder(produced_f0).squeeze(-1)
+            f0_target = batch["raw_f0"]
+        else:
+            f0_target = batch["f0_target"]
+
+        # prefix values
+        prefix = batch["prefix"][0]
+        dur_prefix_mean = dur_target[:, :prefix].sum(dim=-1) / (
+            (~batch["dur_mask"][:, :prefix]).sum(dim=-1)
+        )
+
+        non_voiced = f0_target[:, :prefix] == 0.0
+        f0_mask = batch["f0_mask"][:, :prefix].logical_or(non_voiced)
+        f0_prefix_mean = f0_target[:, :prefix].sum(dim=-1) / ((~f0_mask).sum(dim=-1))
+
+        # continuation values
+        dur_cont_mean = produced_durations[:, prefix:].sum(dim=-1) / (
+            (~batch["dur_mask"][:, prefix:]).sum(dim=-1)
+        )
+
+        non_voiced = produced_f0[:, prefix:] == 0.0
+        f0_mask = non_voiced
+        f0_cont_mean = produced_f0[:, prefix:].sum(dim=-1) / ((~f0_mask).sum(dim=-1))
+
+        assert not f0_cont_mean.isnan().any()
+
+        mean_dur_prefix.append(dur_prefix_mean.cpu())
+        mean_dur_cont.append(dur_cont_mean.cpu())
+
+        mean_f0_prefix.append(f0_prefix_mean.cpu())
+        mean_f0_cont.append(f0_cont_mean.cpu())
+
+        if args.debug and len(mean_dur_prefix) > 10:
+            break
+
+    mean_dur_prefix, mean_dur_cont = torch.cat(mean_dur_prefix), torch.cat(
+        mean_dur_cont
+    )
+    mean_f0_prefix, mean_f0_cont = torch.cat(mean_f0_prefix), torch.cat(mean_f0_cont)
+
+    return mean_dur_prefix, mean_dur_cont, mean_f0_prefix, mean_f0_cont
+
+
+def main(rank, world_size, args):
+    start = time.time()
+
+    if world_size > 1:
+        torch.distributed.init_process_group(
+            backend="gloo", init_method="env://", world_size=world_size, rank=rank
+        )
+        torch.cuda.set_device(rank % torch.cuda.device_count())
+
+    raw_args = args
+
+    args = convert_namespace_to_omegaconf(args)
+    if args.common.seed is not None:
+        np.random.seed(args.common.seed)
+        utils.set_torch_seed(args.common.seed)
+
+    models, model_args, task = checkpoint_utils.load_model_ensemble_and_task(
+        [raw_args.path], arg_overrides={"data": args.task.data}
+    )
+
+    tgt_dict = task.target_dictionary
+
+    for model in models:
+        model.prepare_for_inference_(args)
+        model.cuda().eval()
+        if raw_args.fp16:
+            model = model.half()
+    model = models[0]
+
+    config = ExpressiveCodeDataConfig(args.task.data)
+
+    dataset = CodeDataset(
+        manifest=config.manifests[raw_args.eval_subset],
+        dictionary=task.source_dictionary,
+        dur_dictionary=task.source_duration_dictionary,
+        f0_dictionary=task.source_f0_dictionary,
+        config=config,
+        discrete_dur=task.cfg.discrete_duration,
+        discrete_f0=task.cfg.discrete_f0,
+        log_f0=task.cfg.log_f0,
+        normalize_f0_mean=task.cfg.normalize_f0_mean,
+        normalize_f0_std=task.cfg.normalize_f0_std,
+        interpolate_f0=task.cfg.interpolate_f0,
+        shifts=task.cfg.stream_shifts,
+        return_filename=True,
+        strip_filename=False,
+        return_continuous_f0=raw_args.dequantize_prosody,
+    )
+
+    if raw_args.filter_names:
+        dataset = FilterNamesDataset(dataset, raw_args.filter_names)
+
+    criterion = task.build_criterion(model_args.criterion)
+
+    name2metric = {
+        "continuation": continuation,
+        "teacher_force_everything": teacher_force_everything,
+        "correlation": correlation,
+    }
+
+    name2keys = {
+        "continuation": (
+            "Token BLEU3",
+            "Duration NLL",
+            "Duration MAE",
+            "F0 NLL",
+            "F0 MAE",
+            "F0 sum",
+            "F0 sum_sq",
+            "Dur sum",
+            "Dur sum_sq",
+        ),
+        "teacher_force_everything": ("token_loss", "duration_loss", "f0_loss"),
+        "correlation": ("Duration corr", "F0 corr"),
+    }
+    metric_name = raw_args.metric
+
+    metric = name2metric[metric_name]
+    results = metric(raw_args, dataset, model, criterion, tgt_dict, rank, world_size)
+
+    values = None
+
+    if metric_name not in [
+        "correlation",
+    ]:
+        values, normalizers = results
+        values = maybe_aggregate_normalize(values, normalizers, world_size)
+    elif metric_name == "correlation":
+        values = maybe_aggregate_correlations(results, world_size)
+    else:
+        assert False
+
+    assert values is not None
+    summary = dict(zip(name2keys[raw_args.metric], values.tolist()))
+    if metric_name == "continuation":
+        summary["F0 Std"] = np.sqrt(-summary["F0 sum"] ** 2 + summary["F0 sum_sq"])
+        summary["Dur Std"] = np.sqrt(-summary["Dur sum"] ** 2 + summary["Dur sum_sq"])
+        del summary["F0 sum"]
+        del summary["F0 sum_sq"]
+        del summary["Dur sum"]
+        del summary["Dur sum_sq"]
+
+    summary["metric"] = metric_name
+
+    if rank == 0:
+        print(summary)
+        if raw_args.wandb:
+            wandb_results(summary, raw_args)
+        print("# finished in ", time.time() - start, "seconds")
+
+
+def wandb_results(summary, raw_args):
+    import wandb
+
+    run = wandb.init(
+        project=raw_args.wandb_project_name, tags=raw_args.wandb_tags.split(",")
+    )
+    run.config.metric = raw_args.metric
+    run.config.model = raw_args.path
+    run.config.data = raw_args.data
+
+    if raw_args.wandb_run_name:
+        run.name = raw_args.wandb_run_name
+        run.save()
+
+    wandb.log(summary)
+    wandb.finish()
+
+
+def maybe_aggregate_normalize(values, normalizers, world_size):
+    if world_size > 1:
+        torch.distributed.barrier()
+
+        torch.distributed.all_reduce_multigpu([values])
+        torch.distributed.all_reduce_multigpu([normalizers])
+
+    return values / normalizers
+
+
+def maybe_aggregate_correlations(results, world_size):
+    if world_size > 1:
+        output = [None for _ in range(world_size)]
+        torch.distributed.all_gather_object(output, results)
+        mean_dur_prefix, mean_dur_cont, mean_f0_prefix, mean_f0_cont = [
+            torch.cat([x[i] for x in output]) for i in range(4)
+        ]
+    else:
+        mean_dur_prefix, mean_dur_cont, mean_f0_prefix, mean_f0_cont = results
+
+    corr_dur = scipy.stats.pearsonr(mean_dur_prefix.numpy(), mean_dur_cont.numpy())[0]
+    corr_f0 = scipy.stats.pearsonr(mean_f0_prefix.numpy(), mean_f0_cont.numpy())[0]
+    values = torch.tensor([corr_dur, corr_f0])
+
+    return values
+
+
+def cli_main():
+    parser = options.get_interactive_generation_parser()
+    parser.add_argument(
+        "--prefix-length",
+        type=int,
+        default=1,
+        help="Prompt prefix length (including <s>)",
+    )
+    parser.add_argument(
+        "--duration-scale",
+        type=float,
+        default=1,
+        help="Multiply durations by the given scaler",
+    )
+    parser.add_argument(
+        "--debug", action="store_true", help="Process only the first batch"
+    )
+    parser.add_argument("--n_hypotheses", type=int, default=1)
+    parser.add_argument("--filter-names", type=str, default=None)
+    parser.add_argument(
+        "--max-length", type=int, default=200, help="Maximal produced length"
+    )
+
+    parser.add_argument("--teacher-force-tokens", action="store_true", default=False)
+    parser.add_argument("--teacher-force-duration", action="store_true", default=False)
+    parser.add_argument("--teacher-force-f0", action="store_true", default=False)
+
+    parser.add_argument("--copy-target", action="store_true", default=False)
+    parser.add_argument("--min-length", type=int, default=None)
+    parser.add_argument("--f0-discretization-bounds", type=str, default=None)
+    parser.add_argument("--dequantize-prosody", action="store_true")
+    parser.add_argument("--batch-explosion-rate", type=int, default=1)
+
+    parser.add_argument(
+        "--metric",
+        choices=["continuation", "teacher_force_everything", "correlation"],
+        required=True,
+    )
+
+    parser.add_argument("--wandb", action="store_true")
+    parser.add_argument("--wandb-project-name", type=str, default="eslm")
+    parser.add_argument("--wandb-tags", type=str, default="")
+    parser.add_argument("--wandb-run-name", type=str, default="")
+
+    parser.add_argument("--T-token", type=float, default=1.0)
+    parser.add_argument("--T-duration", type=float, default=1.0)
+    parser.add_argument("--T-f0", type=float, default=1.0)
+
+    parser.add_argument("--n-workers", type=int, default=1)
+
+    parser.add_argument(
+        "--eval-subset", type=str, default="valid", choices=["valid", "test"]
+    )
+
+    args = options.parse_args_and_arch(parser)
+
+    assert (
+        args.prefix_length >= 1
+    ), "Prefix length includes bos token <s>, hence the minimum is 1."
+    assert args.temperature >= 0.0, "T must be non-negative!"
+
+    if args.dequantize_prosody:
+        assert args.f0_discretization_bounds
+
+    world_size = args.n_workers or torch.cuda.device_count()
+    if world_size > 1:
+        import random
+
+        mp.set_start_method("spawn", force=True)
+        os.environ["MASTER_ADDR"] = "localhost"
+        os.environ["MASTER_PORT"] = str(random.randint(10_000, 50_000))
+
+        mp.spawn(
+            main,
+            nprocs=world_size,
+            args=(
+                world_size,
+                args,
+            ),
+            join=True,
+        )
+    else:
+        main(rank=0, world_size=world_size, args=args)
+
+
+if __name__ == "__main__":
+    cli_main()

omni_speech/infer/fairseq/examples/textless_nlp/pgslm/generate_waveform.py

@@ -0,0 +1,120 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import ast
+import argparse
+import json
+import logging
+from pathlib import Path
+import soundfile as sf
+import torch
+
+from tqdm import tqdm
+
+from fairseq import utils
+from fairseq.models.text_to_speech.vocoder import CodeHiFiGANVocoder
+
+
+logging.basicConfig()
+logging.root.setLevel(logging.INFO)
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+def dump_result(args, data, sample_id, pred_wav):
+    assert "audio" in data or args.results_path is not None
+    if args.results_path:
+        fname = Path(data["audio"]).name if "audio" in data else f"{sample_id}_pred.wav"
+        out_file = Path(args.results_path) / fname
+
+    sf.write(
+        out_file.as_posix(),
+        pred_wav.detach().cpu().numpy(),
+        args.sample_rate,
+    )
+
+
+def load_data(in_file):
+    with open(in_file) as f:
+        data = [ast.literal_eval(line.strip()) for line in f]
+
+    return data
+
+
+def get_f0_upsample_ratio(code_hop_size, f_hop_size):
+    ratio = (code_hop_size // 160) // (f_hop_size // 256) * 2
+    return ratio
+
+
+def main(args):
+    logger.info(args)
+
+    use_cuda = torch.cuda.is_available() and not args.cpu
+
+    with open(args.vocoder_cfg) as f:
+        vocoder_cfg = json.load(f)
+    vocoder = CodeHiFiGANVocoder(args.vocoder, vocoder_cfg)
+    if use_cuda:
+        vocoder = vocoder.cuda()
+
+    data = load_data(args.in_file)
+
+    if args.results_path:
+        Path(args.results_path).mkdir(exist_ok=True, parents=True)
+
+    for i, d in tqdm(enumerate(data), total=len(data)):
+        code_key = "cpc_km100" if "cpc_km100" in d else "hubert"
+        code = list(map(int, d[code_key].split()))
+
+        x = {
+            "code": torch.LongTensor(code).view(1, -1),
+            "f0": torch.Tensor(d["f0"]).view(1, -1),
+        }
+
+        f0_up_ratio = get_f0_upsample_ratio(
+            vocoder_cfg["code_hop_size"], vocoder_cfg["hop_size"]
+        )
+        if f0_up_ratio > 1:
+            bsz, cond_length = x["f0"].size()
+            x["f0"] = x["f0"].unsqueeze(2).repeat(1, 1, f0_up_ratio).view(bsz, -1)
+
+        x = utils.move_to_cuda(x) if use_cuda else x
+        wav = vocoder(x)
+        dump_result(args, d, i, wav)
+
+
+def cli_main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--in-file",
+        type=str,
+        required=True,
+        help="Input file following the same format of the output from sample.py ('f0' and 'cpc_km100/hubert' are required fields)",
+    )
+    parser.add_argument(
+        "--vocoder", type=str, required=True, help="path to the vocoder"
+    )
+    parser.add_argument(
+        "--vocoder-cfg",
+        type=str,
+        required=True,
+        help="path to the vocoder config",
+    )
+    parser.add_argument("--sample-rate", type=int, default=16_000)
+    parser.add_argument(
+        "--results-path",
+        type=str,
+        default=None,
+        help="Output directory. If not set, the audios will be stored following the 'audio' field specified in the input file.",
+    )
+    parser.add_argument("--cpu", action="store_true", help="run on CPU")
+
+    args = parser.parse_args()
+
+    main(args)
+
+
+if __name__ == "__main__":
+    cli_main()

omni_speech/infer/fairseq/examples/textless_nlp/pgslm/inference_dataset.py

@@ -0,0 +1,103 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+import torch
+
+
+class InferenceDataset:
+    def __init__(
+        self,
+        dataset,
+        prefix,
+        only_prefix=True,
+        presort_by_length=True,
+        filter_short=False,
+        min_length=None,
+    ):
+        self.dataset = dataset
+        self.collater = self.dataset.collater
+        self.prefix = prefix
+        self.only_prefix = only_prefix
+        self.filter_short = filter_short
+
+        self.remapping = list(range(len(self.dataset)))
+        if min_length:
+            assert min_length >= prefix + 1
+
+        length_thr = prefix + 1 if not min_length else min_length
+
+        if filter_short:
+            self.remapping = list(
+                filter(
+                    lambda i: self.dataset[i]["dur_source"].sum() > length_thr,
+                    self.remapping,
+                )
+            )
+            print(
+                f"# the initial dataset of {len(self.dataset)} examples became {len(self.remapping)} after filtering"
+                f" examples shorter than {length_thr} (in duration units)"
+            )
+
+        if presort_by_length:
+            lengths = {index: dataset.size(index) for index in self.remapping}
+            self.remapping.sort(key=lambda i: lengths[i])
+
+    @property
+    def pads(self):
+        return self.dataset.pads
+
+    def __len__(self):
+        return len(self.remapping)
+
+    def original_size(self, k):
+        k = self.remapping[k]
+        return self.dataset.size(k)
+
+    def __getitem__(self, k):
+        k = self.remapping[k]
+        channels = self.dataset[k]
+
+        if self.prefix and self.only_prefix:
+            dur_channel = channels["dur_source"]
+            assert dur_channel.sum() >= self.prefix
+
+            token_times = dur_channel.cumsum(dim=-1)
+            cut_after = torch.searchsorted(token_times, torch.tensor(self.prefix))
+
+            r = {}
+            for channel_name, value in channels.items():
+                if isinstance(value, torch.Tensor) and "source" in channel_name:
+                    # if self.filter_short: assert value.size(0) >= self.prefix
+                    r[channel_name] = value[: cut_after + 1]
+                else:
+                    r[channel_name] = value
+
+            r["prefix"] = cut_after + 1
+        else:
+            r = channels
+
+        return r
+
+
+def explode_batch(batch, times):
+    if times == 1:
+        return batch
+
+    new_batch = {}
+
+    for key, value in batch.items():
+        if isinstance(value, torch.Tensor):
+            assert value.size(0) == 1
+            new_batch[key] = torch.cat([value] * times)
+        elif key in ["ntokens", "nsentences"]:
+            new_batch[key] = value * times
+        elif key in ["prefix", "filename"]:
+            new_batch[key] = value
+        elif key == "net_input":
+            new_batch[key] = explode_batch(value, times)
+        else:
+            assert False, key
+    return new_batch

omni_speech/infer/fairseq/examples/textless_nlp/pgslm/naive_decoder.py

@@ -0,0 +1,40 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+import warnings
+
+
+class Naive_F0_Decoder(torch.nn.Module):
+    def __init__(self, bounds_path, n_units=32):
+        super().__init__()
+
+        bounds = torch.load(bounds_path)
+        bounds = torch.from_numpy(bounds[n_units])
+        assert bounds.ndim == 1
+
+        pad = torch.tensor([-5.0, -5.0])  # bos, eos, pad are in the dictionary
+        centers = torch.cat(
+            [bounds[0:1], 0.5 * (bounds[1:] + bounds[:-1]), bounds[-1:], pad[:]]
+        )
+
+        self.embedding = torch.nn.Embedding.from_pretrained(
+            centers.unsqueeze(-1), freeze=True
+        )
+        self.max_n = self.embedding.weight.numel()
+
+    def forward(self, discrete_f0: torch.Tensor):
+        in_bounds = (0 <= discrete_f0).all() and (discrete_f0 < self.max_n).all()
+        if not in_bounds:
+            warnings.warn(
+                f"F0 contains some weird outputs: discrete_f0.max().item()={discrete_f0.max().item()} discrete_f0.min().item()={discrete_f0.min().item()}; "
+                f"while we have embeddings for {self.max_n} values. "
+                "Assuming this is a no-prosody model -- but be careful!"
+            )
+
+            mask = discrete_f0 >= self.max_n
+            discrete_f0 = discrete_f0.masked_fill(mask, self.max_n - 1)
+
+        return self.embedding(discrete_f0).squeeze(-1)

omni_speech/infer/fairseq/examples/textless_nlp/pgslm/prepare_dataset.py

@@ -0,0 +1,143 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from multiprocessing import Pool
+
+import os
+from collections import defaultdict
+from itertools import starmap
+
+import torch
+from npy_append_array import NpyAppendArray
+from tqdm import tqdm
+
+from data_utils import dump_speaker_f0_stat, F0Stat, load_f0
+from fairseq.data.codedataset import (
+    ExpressiveCodeDataConfig,
+    parse_manifest,
+    F0_FRAME_SPACE,
+    align_f0_to_durations,
+)
+from fairseq.tasks.speech_ulm_task import UnitDictionary
+
+
+def load_meta(meta_path, split):
+    config = ExpressiveCodeDataConfig(meta_path)
+    manifest_path = config.manifests[split]
+    dictionary = UnitDictionary(n_units=config.n_units)
+    audio_paths, codes, durs, speakers = parse_manifest(manifest_path, dictionary)
+    return config, audio_paths, codes, durs, speakers
+
+
+def _align_f0(f0, dur, ratio, frm_tol=5):
+    if f0 is None:
+        seg_f0 = torch.zeros_like(dur, dtype=torch.float)
+    else:
+        seg_f0 = align_f0_to_durations(f0, dur, ratio, tol=frm_tol * ratio)
+    return seg_f0.numpy()  # try a hacky stuff
+
+
+def align_f0(path_to_f0, audio_paths, durs, ratio, mp=False):
+    chunk_size = 2000
+    num_procs = 40
+    iterable = ((path_to_f0[p], d, ratio) for p, d in zip(audio_paths, durs))
+
+    seg_f0s = []
+    if mp:
+        with Pool(num_procs) as pool:
+            iterator = tqdm(
+                pool.istarmap(_align_f0, iterable, chunk_size),
+                desc="align f0",
+                total=len(durs),
+            )
+            for seg_f0 in iterator:
+                seg_f0s.append(torch.from_numpy(seg_f0).float())
+    else:
+        iterator = tqdm(starmap(_align_f0, iterable), desc="align f0", total=len(durs))
+        for seg_f0 in iterator:
+            seg_f0s.append(torch.from_numpy(seg_f0).float())
+
+    return seg_f0s
+
+
+def prepare_seg_data(config, audio_paths, codes, durs, speakers, path_to_f0):
+    ratio = config.code_hop_size / (config.sampling_rate * F0_FRAME_SPACE)
+    seg_f0s = align_f0(path_to_f0, audio_paths, durs, ratio)
+    data = {
+        "codes": codes,
+        "duration": durs,
+        "f0": seg_f0s,
+        "speaker": speakers,
+        "path": audio_paths,
+    }
+    return data
+
+
+def dump_seg_data(data, out_prefix):
+    key_targs = {
+        "codes": f"{out_prefix}.code.npy",
+        "duration": f"{out_prefix}.dur.npy",
+        "f0": f"{out_prefix}.f0.npy",
+    }
+    for key, targ in key_targs.items():
+        assert not os.path.exists(targ)
+        npaa = NpyAppendArray(targ)
+        for utt_data in tqdm(data[key], desc=f"dumping {key}"):
+            npaa.append(utt_data.numpy())
+
+    assert not os.path.exists(f"{out_prefix}.path.txt")
+    with open(f"{out_prefix}.path.txt", "w") as f:
+        for x in data["path"]:
+            f.write(f"{str(x)}\n")
+
+    assert not os.path.exists(f"{out_prefix}.leng.txt")
+    with open(f"{out_prefix}.leng.txt", "w") as f:
+        for x in data["codes"]:
+            f.write(f"{len(x)}\n")
+
+    assert not os.path.exists(f"{out_prefix}.speaker.txt")
+    with open(f"{out_prefix}.speaker.txt", "w") as f:
+        for x in data["speaker"]:
+            f.write(f"{str(x)}\n")
+
+    print(f"wrote to files with prefix {out_prefix}")
+
+
+def main(meta_path, f0_dir, splits, nshards_list):
+    speaker_to_stat = defaultdict(F0Stat)
+    if len(nshards_list) == 1:
+        nshards_list = nshards_list * len(splits)
+    else:
+        assert len(nshards_list) == len(splits)
+
+    for split, nshards in zip(splits, nshards_list):
+        config, audio_paths, codes, durs, speakers = load_meta(meta_path, split)
+        path_to_f0 = load_f0(f"{f0_dir}/{split}", nshards)
+
+        # segment-level data
+        data = prepare_seg_data(config, audio_paths, codes, durs, speakers, path_to_f0)
+        dump_seg_data(data, config.manifests[split])
+
+        # speaker f0
+        for audio_path, speaker in tqdm(zip(audio_paths, speakers)):
+            f0 = path_to_f0[audio_path]
+            speaker_to_stat[speaker].update(f0)
+        dump_speaker_f0_stat(speaker_to_stat, config.manifests[split])
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("meta_path")
+    parser.add_argument("f0_dir", help="out_dir from preprocess_f0")
+    parser.add_argument("--splits", nargs="+", default=["train", "valid"])
+    parser.add_argument(
+        "--nshards_list", type=int, nargs="+", default=[20], help="number of f0 shards"
+    )
+    args = parser.parse_args()
+    print(args)
+
+    main(**vars(args))

omni_speech/infer/fairseq/examples/textless_nlp/pgslm/preprocess_f0.py

@@ -0,0 +1,65 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+import torch
+from tqdm import tqdm
+from data_utils import load_audio_path
+from fairseq.data.codedataset import get_f0_by_filename
+
+
+def process_one(path, sr):
+    """
+    Args:
+        path: audio file path
+        sr: sampling rate
+    """
+    try:
+        # YAAPT throws errors in some rare cases
+        f0 = get_f0_by_filename(path, sr)
+    except Exception as e:
+        print(
+            f"WARNING: error when processing {path}. set f0 to zero. original error message:\n{e}"
+        )
+        f0 = None
+    return f0
+
+
+def main(file_path, out_dir, nshards, rank, sampling_rate):
+    # load data
+    audio_paths = load_audio_path(file_path)
+
+    # shard
+    assert nshards <= len(audio_paths) and nshards > 0
+    shard_size = len(audio_paths) / nshards
+    s = int(round((rank - 1) * shard_size))
+    e = int(round(rank * shard_size))
+    audio_paths = audio_paths[s:e]
+
+    # process
+    path_to_f0 = {}
+    for i, audio_path in enumerate(tqdm(audio_paths)):
+        f0 = process_one(audio_path, sampling_rate)
+        path_to_f0[audio_path] = f0
+    print(f"finished processing {len(path_to_f0)} utterances ({s}-{e})")
+
+    f0_path = f"{out_dir}/f0_{rank}_{nshards}.pt"
+    os.makedirs(out_dir, exist_ok=True)
+    torch.save(path_to_f0, f0_path)
+    print(f"saved to {f0_path}")
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("file_path")
+    parser.add_argument("out_dir")
+    parser.add_argument("--nshards", type=int, default=20)
+    parser.add_argument("--rank", type=int, default=1)
+    parser.add_argument("--sampling_rate", type=int, default=16000)
+    args = parser.parse_args()
+
+    main(**vars(args))

omni_speech/infer/fairseq/examples/textless_nlp/pgslm/quantize_f0.py

@@ -0,0 +1,94 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from collections import defaultdict
+from functools import partial
+
+import numpy as np
+import torch
+from tqdm import tqdm
+
+from data_utils import dump_speaker_f0_stat, F0Stat, load_audio_path, load_f0
+
+
+def load_speaker(path):
+    speakers = []
+    with open(path) as f:
+        for line in f.readlines():
+            sample = eval(line.strip())
+            assert "speaker" in sample
+            speakers.append(sample["speaker"])
+    return speakers
+
+
+def quantize_f0(speaker_to_f0, f0_stats, nbins, normalize, log):
+    f0_all = []
+    for speaker, f0 in speaker_to_f0.items():
+        f0 = f0.raw_data
+        if log:
+            f0 = f0.log()
+        mean = f0_stats[speaker]["logf0_mean"] if log else f0_stats[speaker]["f0_mean"]
+        std = f0_stats[speaker]["logf0_std"] if log else f0_stats[speaker]["f0_std"]
+        if normalize == "mean":
+            f0 = f0 - mean
+        elif normalize == "meanstd":
+            f0 = (f0 - mean) / std
+        f0_all.extend(f0.tolist())
+
+    hist, bin_x = np.histogram(f0_all, 100000)
+    cum_hist = np.cumsum(hist) / len(f0_all) * 100
+
+    f0_bin = {}
+    for num_bin in nbins:
+        bin_offset = []
+        bin_size = 100 / num_bin
+        threshold = bin_size
+        for i in range(num_bin - 1):
+            index = (np.abs(cum_hist - threshold)).argmin()
+            bin_offset.append(bin_x[index])
+            threshold += bin_size
+        f0_bin[num_bin] = np.array(bin_offset)
+
+    return f0_bin
+
+
+def main(file_path, f0_dir, out_dir, out_prefix, nbins, nshards, normalize, log):
+    audio_paths = load_audio_path(file_path)
+    path_to_f0 = load_f0(f0_dir, nshards)
+
+    speakers = load_speaker(file_path)
+    speaker_to_f0 = defaultdict(partial(F0Stat, True))
+
+    # speaker f0 stats
+    for audio_path, speaker in tqdm(zip(audio_paths, speakers)):
+        f0 = path_to_f0[audio_path]
+        speaker_to_f0[speaker].update(f0)
+    f0_stats = dump_speaker_f0_stat(speaker_to_f0, f"{out_dir}/{out_prefix}")
+
+    # quantize
+    f0_bin = quantize_f0(speaker_to_f0, f0_stats, nbins, normalize, log)
+    log_suffix = "_log" if log else ""
+    f0_bin_out_file = f"{out_dir}/{out_prefix}_{normalize}_norm{log_suffix}_f0_bin.th"
+    torch.save(f0_bin, f0_bin_out_file)
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("file_path")
+    parser.add_argument("f0_dir", help="out_dir from preprocess_f0")
+    parser.add_argument("out_dir")
+    parser.add_argument("out_prefix")
+    parser.add_argument("--nbins", nargs="+", type=int, default=[32])
+    parser.add_argument("--nshards", type=int, default=20, help="number of f0 shards")
+    parser.add_argument(
+        "--normalize", type=str, choices=["meanstd", "mean", "none"], default="mean"
+    )
+    parser.add_argument("--log", action="store_true")
+    args = parser.parse_args()
+    print(args)
+
+    main(**vars(args))

omni_speech/infer/fairseq/examples/textless_nlp/pgslm/sample/__init__.py

@@ -0,0 +1,4 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.

omni_speech/infer/fairseq/examples/textless_nlp/pgslm/sample/sample.py

@@ -0,0 +1,612 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+import torch.multiprocessing as mp
+import numpy as np
+import json
+
+import torch
+from torch.distributions.categorical import Categorical
+
+from fairseq import checkpoint_utils, options, utils
+from fairseq.data.codedataset import CodeDataset, ExpressiveCodeDataConfig
+from fairseq.dataclass.utils import convert_namespace_to_omegaconf
+from torch.utils.data import DataLoader, DistributedSampler
+from fairseq.utils import move_to_cuda
+
+import tqdm
+import random
+import pathlib
+
+import sys, pathlib
+
+sys.path.append(str(pathlib.Path(__file__).parent.parent))
+from inference_dataset import InferenceDataset, explode_batch
+from naive_decoder import Naive_F0_Decoder
+from truncated_laplace import truncated_laplace
+
+CODETYPE_TO_FRAMETIME = {"cpc_km100": 0.01, "hubert": 0.02}  # 10ms  # 20ms
+
+
+class TemperatureDecoder:
+    def __init__(self, Ts, discrete_dur=False, discrete_f0=False):
+        self.T_token, self.T_dur, self.T_f0 = Ts
+        self.discrete_dur = discrete_dur
+        self.discrete_f0 = discrete_f0
+
+    def __call__(self, output):
+        def sample_multinomial(key, T):
+            logits = output[key][:, -1, :].float()
+            return Categorical(logits=logits / T).sample().unsqueeze(-1)
+
+        def sample_laplace(key, T, truncate_at_zero):
+            mean = output[key][:, -1, :].float()
+            return truncated_laplace(mean=mean, T=T, truncate_by_zero=truncate_at_zero)
+
+        if self.T_token > 0:
+            new_tokens = sample_multinomial("token", self.T_token)
+        else:
+            new_tokens = output["token"][:, -1, :].argmax(dim=-1, keepdim=True)
+
+        if not self.discrete_dur and self.T_dur == 0:
+            new_durations = output["duration"][:, -1].round().int()
+        elif not self.discrete_dur and self.T_dur > 0:
+            new_durations = (
+                sample_laplace("duration", self.T_dur, truncate_at_zero=True)
+                .round()
+                .int()
+            )
+        elif self.discrete_dur and self.T_dur > 0:
+            new_durations = sample_multinomial("duration", self.T_dur)
+        elif self.discrete_dur and self.T_dur == 0:
+            new_durations = output["duration"][:, -1, :].argmax(dim=-1, keepdim=True)
+        else:
+            assert False
+
+        if not self.discrete_f0 and self.T_f0 == 0:
+            new_f0 = output["f0"][:, -1]
+        elif not self.discrete_f0 and self.T_f0 > 0:
+            new_f0 = sample_laplace("f0", self.T_f0, truncate_at_zero=False)
+        elif self.discrete_f0 and self.T_f0 > 0:
+            new_f0 = sample_multinomial("f0", self.T_f0)
+        elif self.discrete_f0 and self.T_f0 == 0:
+            new_f0 = output["f0"][:, -1, :].argmax(dim=-1, keepdim=True)
+        else:
+            assert False
+
+        return new_tokens, new_durations, new_f0
+
+
+class FilterNamesDataset:
+    def __init__(self, dataset, fnames_path):
+        self.dataset = dataset
+
+        with open(fnames_path, "r") as fin:
+            fnames = set((eval(line)["audio"] for line in fin))
+            print(f"# will retrict the dataset for {len(fnames)} files")
+
+        self.indexes = []
+
+        for i, datapoint in enumerate(dataset):
+            if datapoint["filename"] in fnames:
+                self.indexes.append(i)
+        assert len(self.indexes) == len(fnames), f"{len(self.indexes)} {len(fnames)}"
+
+        self.collater = self.dataset.collater
+        self.discrete_dur = self.dataset.discrete_dur
+        self.discrete_f0 = self.dataset.discrete_f0
+
+    def __len__(self):
+        return len(self.indexes)
+
+    def __getitem__(self, k):
+        k = self.indexes[k]
+        return self.dataset[k]
+
+    def size(self, k):
+        k = self.indexes[k]
+        return self.dataset.size(k)
+
+
+@torch.no_grad()
+def do_sampling(
+    model,
+    batch,
+    eos_token,
+    decoder,
+    autoregressive_steps=100,
+    teacher_force_tokens=False,
+    teacher_force_duration=False,
+    teacher_force_f0=False,
+    match_duration=False,
+):
+    def autoregressive_step_(output, autoregressive_steps):
+        new_tokens, new_durations, new_f0 = decoder(output)
+
+        n = output["token"].size(1) if output["token"].ndim == 3 else 1
+
+        if teacher_force_tokens:
+            new_tokens = batch["target"][:, n - 1].unsqueeze(-1)
+        if teacher_force_duration:
+            new_durations = batch["dur_target"][:, n - 1].unsqueeze(-1)
+        if teacher_force_f0:
+            new_f0 = batch["f0_target"][:, n - 1].unsqueeze(-1)
+
+        batch["net_input"]["src_tokens"] = torch.cat(
+            [batch["net_input"]["src_tokens"], new_tokens], dim=1
+        )
+        batch["net_input"]["dur_src"] = torch.cat(
+            [batch["net_input"]["dur_src"], new_durations], dim=1
+        )
+        batch["net_input"]["f0_src"] = torch.cat(
+            [batch["net_input"]["f0_src"], new_f0], dim=1
+        )
+
+    outputs = []
+
+    if teacher_force_tokens or teacher_force_duration or teacher_force_f0:
+        max_time = batch["target"].size(1)
+        prefix_time = batch["net_input"]["src_tokens"].size(1)
+
+        autoregressive_steps = max_time - prefix_time + 1  # should be 0
+
+    for _ in range(autoregressive_steps):
+        output = model(**batch["net_input"])
+
+        last_steps = (
+            output["token"][:, -1, ...],
+            output["duration"][:, -1, ...],
+            output["f0"][:, -1, ...],
+        )
+        outputs.append(last_steps)
+
+        autoregressive_step_(output, autoregressive_steps)
+        tokens, duration, f0 = (
+            batch["net_input"]["src_tokens"],
+            batch["net_input"]["dur_src"],
+            batch["net_input"]["f0_src"],
+        )
+
+        if (
+            match_duration
+            and (batch["dur_target"].sum(dim=-1) < duration.sum(dim=-1)).all()
+        ):
+            break
+
+    return tokens, duration, f0, outputs
+
+
+def unroll_duration(token_stream, duration_stream):
+    assert len(token_stream) == len(
+        duration_stream
+    ), f"{len(token_stream)} != {len(duration_stream)}"
+    non_positive_durations = sum(d <= 0 for d in duration_stream)
+    if non_positive_durations > 0:
+        print(
+            f"# {non_positive_durations} durations are non-positive, they will be capped to 1"
+        )
+
+    result = []
+
+    duration_stream_rounded_capped = [max(1, int(round(x))) for x in duration_stream]
+    for t, d in zip(token_stream, duration_stream_rounded_capped):
+        result.extend([t] * d)
+
+    return result
+
+
+def realign_shifted_streams(tokens, durations, F0s, shifts):
+    """
+    Durations are shifted by 1, F0 by 2
+    >>> tokens = ["<s>", "t1",  "t2", "t3", "</s>", "x", "x"]
+    >>> durations = ["<0>", "<0>", "d1", "d2", "d3", "<0>", "x"]
+    >>> F0s    = ["<0>", "<0>", "<0>", "f1", "f2", "f3", "<0>"]
+    >>> shifts = [1,2]
+    >>> realign_shifted_streams(tokens, durations, F0s, shifts)
+    (['<s>', 't1', 't2', 't3', '</s>'], ['<0>', 'd1', 'd2', 'd3', '<0>'], ['<0>', 'f1', 'f2', 'f3', '<0>'])
+    """
+    max_shift = max(shifts)
+    if max_shift > 0:
+        shift_durations, shift_F0s = shifts
+
+        tokens = tokens[:-max_shift]
+        durations = durations[shift_durations:]
+        if shift_durations < max_shift:
+            durations = durations[: -(max_shift - shift_durations)]
+
+        if F0s is not None:
+            F0s = F0s[shift_F0s:]
+            if shift_F0s < max_shift:
+                F0s = F0s[: -(max_shift - shift_F0s)]
+
+    assert len(tokens) == len(durations), f"{len(tokens)} =! {len(durations)}"
+    if F0s is not None:
+        assert len(tokens) == len(F0s), f"{len(tokens)} =! {len(F0s)}"
+
+    return tokens, durations, F0s
+
+
+def maybe_cut_eos(produced_tokens, produced_duration, produced_f0, eos_idx):
+    if eos_idx in produced_tokens:
+        eos_index = produced_tokens.index(eos_idx)
+        produced_tokens = produced_tokens[:eos_index]
+        produced_duration = produced_duration[:eos_index]
+        produced_f0 = produced_f0[:eos_index]
+    return produced_tokens, produced_duration, produced_f0
+
+
+def maybe_filter_pad(produced_tokens, produced_duration, produced_f0, pad_idx):
+    if pad_idx not in produced_tokens:
+        return produced_tokens, produced_duration, produced_f0
+
+    assert len(produced_tokens) == len(produced_duration) == len(produced_f0)
+
+    print("<pad> is detected in the output!")
+    filtered_tokens, filtered_duration, filtered_f0 = [], [], []
+
+    for t, d, f in zip(produced_tokens, produced_duration, produced_f0):
+        if t != pad_idx:
+            filtered_tokens.append(t)
+            filtered_duration.append(d)
+            filtered_f0.append(f)
+    return filtered_tokens, filtered_duration, filtered_f0
+
+
+def match_duration(produced_tokens, produced_duration, produced_f0, target_duration):
+    """
+    >>> tokens = ['t'] * 4
+    >>> F0s    = ['f0'] * 4
+    >>> produced_duration = [1, 10, 10, 10]
+    >>> match_duration(tokens, produced_duration, F0s, target_duration=100)
+    (['t', 't', 't', 't'], [1, 10, 10, 10], ['f0', 'f0', 'f0', 'f0'])
+    >>> match_duration(tokens, produced_duration, F0s, target_duration=5)
+    (['t', 't'], [1, 4], ['f0', 'f0'])
+    """
+    if sum(produced_duration) <= target_duration:
+        return produced_tokens, produced_duration, produced_f0
+
+    running_duration = 0
+    filtered_duration = []
+
+    for next_tok_duration in produced_duration:
+        if running_duration + next_tok_duration < target_duration:
+            filtered_duration.append(next_tok_duration)
+            running_duration += next_tok_duration
+        else:
+            to_add = target_duration - running_duration
+            assert to_add <= next_tok_duration
+            filtered_duration.append(to_add)
+            break
+
+    produced_duration = filtered_duration
+    assert sum(produced_duration) == target_duration
+
+    n_tok = len(filtered_duration)
+
+    return produced_tokens[:n_tok], produced_duration, produced_f0[:n_tok]
+
+
+def main(rank, world_size, args):
+    if world_size > 1:
+        torch.distributed.init_process_group(
+            backend="gloo", init_method="env://", world_size=world_size, rank=rank
+        )
+        torch.cuda.set_device(rank)
+
+    raw_args = args
+    args = convert_namespace_to_omegaconf(args)
+    if args.common.seed is not None:
+        random.seed(args.common.seed)
+        np.random.seed(args.common.seed)
+        utils.set_torch_seed(args.common.seed)
+
+    models, model_args, task = checkpoint_utils.load_model_ensemble_and_task(
+        [raw_args.path], arg_overrides={"data": args.task.data}
+    )
+    tgt_dict = task.target_dictionary
+
+    for model in models:
+        model.prepare_for_inference_(args)
+        model.cuda().eval()
+        if raw_args.fp16:
+            model = model.half()
+    model = models[0]
+
+    config = ExpressiveCodeDataConfig(args.task.data)
+
+    dataset = CodeDataset(
+        manifest=config.manifests[raw_args.subset],
+        dictionary=task.source_dictionary,
+        dur_dictionary=task.source_duration_dictionary,
+        f0_dictionary=task.source_f0_dictionary,
+        config=config,
+        discrete_dur=task.cfg.discrete_duration,
+        discrete_f0=task.cfg.discrete_f0,
+        log_f0=task.cfg.log_f0,
+        normalize_f0_mean=task.cfg.normalize_f0_mean,
+        normalize_f0_std=task.cfg.normalize_f0_std,
+        interpolate_f0=task.cfg.interpolate_f0,
+        shifts=task.cfg.stream_shifts,
+        return_filename=True,
+        strip_filename=False,
+    )
+    tgt_dict = task.target_dictionary
+    shifts = dataset.shifts.dur, dataset.shifts.f0
+    max_shift = max(shifts)
+
+    fname = raw_args.output
+    if world_size > 1:
+        fname += f"_{rank}"
+    output_file = open(fname, "w")
+
+    if raw_args.filter_names:
+        dataset = FilterNamesDataset(dataset, raw_args.filter_names)
+
+    dataset = InferenceDataset(dataset, raw_args.prefix_length, filter_short=True)
+    print(f"Dataset size {len(dataset)}")
+    sampler = (
+        None
+        if world_size == 1
+        else DistributedSampler(
+            dataset, num_replicas=world_size, rank=rank, shuffle=False
+        )
+    )
+    dataloader = DataLoader(
+        dataset,
+        batch_size=1,
+        shuffle=False,
+        collate_fn=dataset.collater,
+        sampler=sampler,
+    )
+
+    Ts = raw_args.T_token, raw_args.T_duration, raw_args.T_f0
+    decoder = TemperatureDecoder(
+        Ts, discrete_dur=task.cfg.discrete_duration, discrete_f0=task.cfg.discrete_f0
+    )
+
+    dataset_size = len(dataset)
+
+    f0_decoder = None
+    if raw_args.f0_discretization_bounds:
+        assert task.cfg.discrete_f0
+        f0_decoder = Naive_F0_Decoder(raw_args.f0_discretization_bounds).cuda()
+
+    pbar = (
+        tqdm.tqdm(
+            total=dataset_size
+            if raw_args.max_samples is None
+            else min(raw_args.max_samples, dataset_size)
+        )
+        if world_size == 1
+        else None
+    )
+
+    samples_produced = 0
+
+    for batch in dataloader:
+        if (
+            raw_args.max_samples is not None
+            and samples_produced >= raw_args.max_samples
+        ):
+            break
+
+        prefix = batch["prefix"][0]
+
+        batch = explode_batch(batch, raw_args.batch_explosion_rate)
+        batch = move_to_cuda(batch)
+
+        if not raw_args.short_curcuit:
+            produced_tokens, produced_durations, produced_f0, _ = do_sampling(
+                models[0],
+                batch,
+                tgt_dict.eos(),
+                decoder,
+                autoregressive_steps=raw_args.max_length - prefix + max_shift,
+                teacher_force_tokens=raw_args.teacher_force_tokens,
+                match_duration=raw_args.match_duration,
+                teacher_force_duration=raw_args.teacher_force_duration,
+                teacher_force_f0=raw_args.teacher_force_f0,
+            )
+
+            # stip entries corresponding to <s>
+            produced_tokens = produced_tokens[:, 1:]
+            produced_durations = produced_durations[:, 1:]
+            produced_f0 = produced_f0[:, 1:]
+
+        else:
+            max_length = raw_args.max_length + max_shift
+            produced_tokens, produced_durations, produced_f0 = (
+                batch["target"][:, :max_length],
+                batch["dur_target"][:, :max_length],
+                batch["f0_target"][:, :max_length],
+            )
+
+        if f0_decoder is not None:
+            produced_f0 = f0_decoder(produced_f0)
+
+        produced_tokens, produced_durations, produced_f0 = (
+            produced_tokens.cpu().tolist(),
+            produced_durations.cpu().tolist(),
+            produced_f0.cpu().tolist(),
+        )
+
+        bsz = batch["target"].size(0)
+        assert bsz == raw_args.batch_explosion_rate
+
+        for i in range(bsz):
+            if (
+                raw_args.max_samples is not None
+                and samples_produced >= raw_args.max_samples
+            ):
+                break
+
+            produced_tokens_i = produced_tokens[i]
+            produced_durations_i = produced_durations[i]
+            produced_f0_i = produced_f0[i]
+
+            (
+                produced_tokens_i,
+                produced_durations_i,
+                produced_f0_i,
+            ) = realign_shifted_streams(
+                produced_tokens_i, produced_durations_i, produced_f0_i, shifts
+            )
+
+            produced_tokens_i, produced_durations_i, produced_f0_i = maybe_cut_eos(
+                produced_tokens_i, produced_durations_i, produced_f0_i, tgt_dict.eos()
+            )
+
+            produced_tokens_i, produced_durations_i, produced_f0_i = maybe_filter_pad(
+                produced_tokens_i, produced_durations_i, produced_f0_i, tgt_dict.pad()
+            )
+
+            if raw_args.match_duration:
+                # NB: here we cheat a bit and use that padding has duration 0
+                # so no need to re-align and remove padding
+                dur_target_i = batch["dur_target"][i, :].sum().item()
+                produced_tokens_i, produced_durations_i, produced_f0_i = match_duration(
+                    produced_tokens_i, produced_durations_i, produced_f0_i, dur_target_i
+                )
+
+            if raw_args.cut_prompt:
+                produced_tokens_i, produced_durations_i, produced_f0_i = (
+                    produced_tokens_i[prefix:],
+                    produced_durations_i[prefix:],
+                    produced_f0_i[prefix:],
+                )
+
+            prompt_fname = batch["filename"][0]
+            fname = str(pathlib.Path(prompt_fname).with_suffix("")) + f"__{i}.wav"
+
+            token_stream = unroll_duration(produced_tokens_i, produced_durations_i)
+            f0_stream = unroll_duration(produced_f0_i, produced_durations_i)
+            output_line = json.dumps(
+                {
+                    "audio": fname,
+                    "prompt": prompt_fname,
+                    raw_args.code_type: " ".join(map(str, token_stream)),
+                    "duration": round(
+                        sum(produced_durations_i)
+                        * CODETYPE_TO_FRAMETIME[raw_args.code_type],
+                        3,
+                    ),
+                    "raw_duration": produced_durations_i,
+                    "raw_f0": produced_f0_i,
+                    "f0": [round(f0, 3) for f0 in f0_stream],
+                }
+            )
+            print(output_line, file=output_file)
+
+            if pbar:
+                pbar.update(1)
+            samples_produced += 1
+
+        if raw_args.debug:
+            break
+
+    output_file.close()
+
+    if world_size > 1:
+        # important that everything is flushed before aggregating
+        torch.distributed.barrier()
+
+    if world_size > 1 and rank == 0:
+        with open(raw_args.output, "w") as fout:
+            for i in range(world_size):
+                f = raw_args.output + f"_{i}"
+                with open(f, "r") as fin:
+                    fout.write(fin.read())
+                os.remove(f)
+
+
+def cli_main():
+    parser = options.get_interactive_generation_parser()
+    parser.add_argument(
+        "--prefix-length",
+        type=int,
+        default=1,
+        help="Prompt prefix length (including <s>)",
+    )
+    parser.add_argument("--output", type=str, default=None, required=True)
+    parser.add_argument(
+        "--debug", action="store_true", help="Process only the first batch"
+    )
+    parser.add_argument(
+        "--ignore-durations",
+        action="store_true",
+        help="If set, the duration stream is ignored",
+    )
+    parser.add_argument(
+        "--max-length", type=int, default=200, help="Maximal produced length"
+    )
+    parser.add_argument(
+        "--code-type", choices=["cpc_km100", "hubert"], default="cpc_km100"
+    )
+    parser.add_argument("--max-samples", type=int, default=None)
+    parser.add_argument("--prompt-duration-scaler", type=float, default=1.0)
+    parser.add_argument("--teacher-force-tokens", action="store_true", default=False)
+    parser.add_argument("--teacher-force-duration", action="store_true", default=False)
+    parser.add_argument("--teacher-force-f0", action="store_true", default=False)
+    parser.add_argument("--filter-names", type=str, default=None)
+    parser.add_argument(
+        "--match-duration",
+        action="store_true",
+        help="Do not produce sequences longer that ground-truth",
+    )
+    parser.add_argument(
+        "--cut-prompt",
+        action="store_true",
+        help="Remove prompt from the produced audio",
+    )
+    parser.add_argument(
+        "--short-curcuit", action="store_true", help="Use 'target' as a sample"
+    )
+    parser.add_argument("--f0-discretization-bounds", type=str, default=None)
+
+    parser.add_argument("--batch-explosion-rate", type=int, default=1)
+
+    parser.add_argument("--T-token", type=float, default=1.0)
+    parser.add_argument("--T-duration", type=float, default=1.0)
+    parser.add_argument("--T-f0", type=float, default=1.0)
+
+    parser.add_argument(
+        "--subset", type=str, default="valid", choices=["test", "valid"]
+    )
+
+    args = options.parse_args_and_arch(parser)
+
+    assert (
+        args.prefix_length >= 1
+    ), "Prefix length includes bos token <s>, hence the minimum is 1."
+    assert all(
+        t >= 0 for t in [args.T_token, args.T_f0, args.T_duration]
+    ), "T must be non-negative!"
+
+    world_size = torch.cuda.device_count()
+    if world_size > 1:
+        import random
+
+        mp.set_start_method("spawn", force=True)
+        os.environ["MASTER_ADDR"] = "localhost"
+        os.environ["MASTER_PORT"] = str(random.randint(10_000, 50_000))
+
+        print(f"Using {world_size} devices, master port {os.environ['MASTER_PORT']}")
+
+        mp.spawn(
+            main,
+            nprocs=world_size,
+            args=(
+                world_size,
+                args,
+            ),
+            join=True,
+        )
+    else:
+        main(rank=0, world_size=world_size, args=args)
+
+
+if __name__ == "__main__":
+    cli_main()

omni_speech/infer/fairseq/examples/textless_nlp/pgslm/scripts/join_units_manifest.py

@@ -0,0 +1,48 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import json
+import argparse
+import pathlib
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--manifest", required=True)
+    parser.add_argument("--units", required=True)
+    parser.add_argument("--output", required=True)
+    parser.add_argument("--sample_rate", type=int, default=16_000)
+
+    args = parser.parse_args()
+
+    with open(args.manifest, "r") as manifest, open(args.units, "r") as units, open(
+        args.output, "w"
+    ) as outp:
+        root = manifest.readline().strip()
+        root = pathlib.Path(root)
+
+        for manifest_line, unit_line in zip(manifest.readlines(), units.readlines()):
+            path, frames = manifest_line.split()
+            duration = int(frames) / float(args.sample_rate)
+            fname = root / path
+            speaker = fname.parent.parent.name
+
+            units = unit_line.split("|")[1]
+
+            print(
+                json.dumps(
+                    dict(
+                        audio=str(root / path),
+                        duration=duration,
+                        hubert_km100=units.strip(),
+                        speaker=speaker,
+                    )
+                ),
+                file=outp,
+            )
+
+
+if __name__ == "__main__":
+    main()

omni_speech/infer/fairseq/examples/textless_nlp/pgslm/scripts/prepare_data.sh

@@ -0,0 +1,57 @@
+#!/bin/bash
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+set -eu
+
+train_json=$1
+valid_json=$2
+test_json=$3
+n_units=$4
+hop_size=$5
+sr=$6
+f0_quantizer=$7
+out_dir=$8
+
+meta_path="$out_dir/data_config.json"
+f0_dir="$out_dir/f0"
+
+mkdir -p $out_dir
+ln -sf $train_json $out_dir/train.txt
+ln -sf $valid_json $out_dir/valid.txt
+ln -sf $test_json $out_dir/test.txt
+
+cat <<EOF >$meta_path
+{
+    "manifests": {
+      "train": "$out_dir/train.txt",
+      "valid": "$out_dir/valid.txt",
+      "test": "$out_dir/test.txt"
+    },
+    "n_units": $n_units,
+    "code_hop_size": $hop_size,
+    "sampling_rate": $sr,
+    "multispkr": "parent_parent_name",
+
+    "f0_vq_type": "naive",
+    "f0_vq_naive_quantizer": {
+      "log_mean_norm": "$f0_quantizer"
+    },
+    "f0_vq_n_units": 32
+}
+EOF
+
+for split in train valid test; do
+  python examples/textless_nlp/pgslm/preprocess_f0.py \
+    $out_dir/$split.txt $f0_dir/$split --nshards=1 --rank=1 --sampling_rate=$sr
+
+  #NSHARDS=16
+  #seq 1 $NSHARDS | parallel -j $NSHARDS python examples/textless_nlp/pgslm/preprocess_f0.py \
+  #  $out_dir/$split.txt $f0_dir/$split --nshards=$NSHARDS --sampling_rate=$sr --rank
+done
+
+# Please make sure that the number of shards (--nshards_list) is consistent across commands
+python examples/textless_nlp/pgslm/prepare_dataset.py \
+  $meta_path $f0_dir --splits test valid train --nshards_list 1

omni_speech/infer/fairseq/examples/textless_nlp/pgslm/scripts/prepare_f0_quantization.sh

@@ -0,0 +1,27 @@
+#!/bin/bash
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+set -eu
+
+train_json=$1
+sr=$2
+nbins=$3
+out_dir=$4
+out_prefix=$5
+
+f0_dir="$out_dir/f0"
+
+python examples/textless_nlp/pgslm/preprocess_f0.py \
+    $train_json $f0_dir/${out_prefix}_f0_quant --nshards 1 --rank 1 --sampling_rate $sr
+
+# NB: one can use parallel here:
+# NSHARDS=16
+#
+#seq 1 $NSHARDS | parallel -j $NSHARDS python examples/textless_nlp/pgslm/preprocess_f0.py \
+#    $train_json $f0_dir/${out_prefix}_f0_quant --nshards $NSHARDS --sampling_rate $sr --rank
+
+python examples/textless_nlp/pgslm/quantize_f0.py \
+    $train_json $f0_dir/${out_prefix}_f0_quant $out_dir $out_prefix --nbins $nbins --nshards 1 --normalize mean --log

omni_speech/infer/fairseq/examples/textless_nlp/pgslm/truncated_laplace.py

@@ -0,0 +1,29 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+import warnings
+
+
+def truncated_laplace(mean, T, truncate_by_zero=False):
+    """Generating a sample from a Laplace distribution, possible left-truncated at zero.
+    A bit of explanation here https://stats.stackexchange.com/a/357598 .
+    """
+    assert isinstance(mean, torch.Tensor)
+
+    if not truncate_by_zero:
+        percentile = 0.0
+    else:
+        if not (mean >= 0.0).all():
+            warnings.warn(f"means are supposed to be non-negative, but got {mean}")
+            mean = torch.clamp_min(mean, 0.0)
+
+        lower_bound = mean.new_tensor([0.0])
+        percentile = 0.5 + 0.5 * torch.sign(lower_bound - mean) * (
+            1.0 - torch.exp(-1.0 / T * torch.abs(mean - lower_bound))
+        )
+
+    p = torch.empty_like(mean).uniform_() * (1.0 - percentile) + percentile
+    return mean - T * torch.sign(p - 0.5) * torch.log(1 - 2 * torch.abs(p - 0.5))

omni_speech/infer/fairseq/examples/textless_nlp/speech-resynth/README.md

@@ -0,0 +1,28 @@
+
+# Speech Resynthesis from Discrete Disentangled Self-Supervised Representations
+Landing page with usfull resources for the [Speech Resynthesis from Discrete Disentangled Self-Supervised Representations](https://arxiv.org/abs/2104.00355) paper.
+
+<p align="center"><img width="70%" src="img/fig.png" /></p>
+
+__Abstract__: We propose using self-supervised discrete representations for the task of speech resynthesis. To generate disentangled representation, we separately extract low-bitrate representations for speech content, prosodic information, and speaker identity. This allows to synthesize speech in a controllable manner. We analyze various state-of-the-art, self-supervised representation learning methods and shed light on the advantages of each method while considering reconstruction quality and disentanglement properties. Specifically, we evaluate the F0 reconstruction, speaker identification performance (for both resynthesis and voice conversion), recordings' intelligibility, and overall quality using subjective human evaluation. Lastly, we demonstrate how these representations can be used for an ultra-lightweight speech codec. Using the obtained representations, we can get to a rate of 365 bits per second while providing better speech quality than the baseline methods.
+
+
+## Quick Links
+- [Paper](https://arxiv.org/pdf/2104.00355.pdf)
+- [Samples](https://speechbot.github.io/resynthesis/index.html)
+- [Code](https://github.com/facebookresearch/speech-resynthesis)
+
+The codebase for the [Speech Resynthesis from Discrete Disentangled Self-Supervised Representations](https://arxiv.org/abs/2104.00355) paper can be found under the following [repository](https://github.com/facebookresearch/speech-resynthesis). 
+
+
+## Citation
+```
+@inproceedings{polyak21_interspeech,
+  author={Adam Polyak and Yossi Adi and Jade Copet and 
+          Eugene Kharitonov and Kushal Lakhotia and 
+          Wei-Ning Hsu and Abdelrahman Mohamed and Emmanuel Dupoux},
+  title={{Speech Resynthesis from Discrete Disentangled Self-Supervised Representations}},
+  year=2021,
+  booktitle={Proc. Interspeech 2021},
+}
+``` 

omni_speech/infer/fairseq/examples/translation/README.md

@@ -0,0 +1,301 @@
+# Neural Machine Translation
+
+This README contains instructions for [using pretrained translation models](#example-usage-torchhub)
+as well as [training new models](#training-a-new-model).
+
+## Pre-trained models
+
+Model | Description | Dataset | Download
+---|---|---|---
+`conv.wmt14.en-fr` | Convolutional <br> ([Gehring et al., 2017](https://arxiv.org/abs/1705.03122)) | [WMT14 English-French](http://statmt.org/wmt14/translation-task.html#Download) | model: <br> [download (.tar.bz2)](https://dl.fbaipublicfiles.com/fairseq/models/wmt14.v2.en-fr.fconv-py.tar.bz2) <br> newstest2014: <br> [download (.tar.bz2)](https://dl.fbaipublicfiles.com/fairseq/data/wmt14.v2.en-fr.newstest2014.tar.bz2) <br> newstest2012/2013: <br> [download (.tar.bz2)](https://dl.fbaipublicfiles.com/fairseq/data/wmt14.v2.en-fr.ntst1213.tar.bz2)
+`conv.wmt14.en-de` | Convolutional <br> ([Gehring et al., 2017](https://arxiv.org/abs/1705.03122)) | [WMT14 English-German](http://statmt.org/wmt14/translation-task.html#Download) | model: <br> [download (.tar.bz2)](https://dl.fbaipublicfiles.com/fairseq/models/wmt14.en-de.fconv-py.tar.bz2) <br> newstest2014: <br> [download (.tar.bz2)](https://dl.fbaipublicfiles.com/fairseq/data/wmt14.en-de.newstest2014.tar.bz2)
+`conv.wmt17.en-de` | Convolutional <br> ([Gehring et al., 2017](https://arxiv.org/abs/1705.03122)) | [WMT17 English-German](http://statmt.org/wmt17/translation-task.html#Download) | model: <br> [download (.tar.bz2)](https://dl.fbaipublicfiles.com/fairseq/models/wmt17.v2.en-de.fconv-py.tar.bz2) <br> newstest2014: <br> [download (.tar.bz2)](https://dl.fbaipublicfiles.com/fairseq/data/wmt17.v2.en-de.newstest2014.tar.bz2)
+`transformer.wmt14.en-fr` | Transformer <br> ([Ott et al., 2018](https://arxiv.org/abs/1806.00187)) | [WMT14 English-French](http://statmt.org/wmt14/translation-task.html#Download) | model: <br> [download (.tar.bz2)](https://dl.fbaipublicfiles.com/fairseq/models/wmt14.en-fr.joined-dict.transformer.tar.bz2) <br> newstest2014: <br> [download (.tar.bz2)](https://dl.fbaipublicfiles.com/fairseq/data/wmt14.en-fr.joined-dict.newstest2014.tar.bz2)
+`transformer.wmt16.en-de` | Transformer <br> ([Ott et al., 2018](https://arxiv.org/abs/1806.00187)) | [WMT16 English-German](https://drive.google.com/uc?export=download&id=0B_bZck-ksdkpM25jRUN2X2UxMm8) | model: <br> [download (.tar.bz2)](https://dl.fbaipublicfiles.com/fairseq/models/wmt16.en-de.joined-dict.transformer.tar.bz2) <br> newstest2014: <br> [download (.tar.bz2)](https://dl.fbaipublicfiles.com/fairseq/data/wmt16.en-de.joined-dict.newstest2014.tar.bz2)
+`transformer.wmt18.en-de` | Transformer <br> ([Edunov et al., 2018](https://arxiv.org/abs/1808.09381)) <br> WMT'18 winner | [WMT'18 English-German](http://www.statmt.org/wmt18/translation-task.html) | model: <br> [download (.tar.gz)](https://dl.fbaipublicfiles.com/fairseq/models/wmt18.en-de.ensemble.tar.gz) <br> See NOTE in the archive
+`transformer.wmt19.en-de` | Transformer <br> ([Ng et al., 2019](https://arxiv.org/abs/1907.06616)) <br> WMT'19 winner | [WMT'19 English-German](http://www.statmt.org/wmt19/translation-task.html) | model: <br> [download (.tar.gz)](https://dl.fbaipublicfiles.com/fairseq/models/wmt19.en-de.joined-dict.ensemble.tar.gz)
+`transformer.wmt19.de-en` | Transformer <br> ([Ng et al., 2019](https://arxiv.org/abs/1907.06616)) <br> WMT'19 winner | [WMT'19 German-English](http://www.statmt.org/wmt19/translation-task.html) | model: <br> [download (.tar.gz)](https://dl.fbaipublicfiles.com/fairseq/models/wmt19.de-en.joined-dict.ensemble.tar.gz)
+`transformer.wmt19.en-ru` | Transformer <br> ([Ng et al., 2019](https://arxiv.org/abs/1907.06616)) <br> WMT'19 winner | [WMT'19 English-Russian](http://www.statmt.org/wmt19/translation-task.html) | model: <br> [download (.tar.gz)](https://dl.fbaipublicfiles.com/fairseq/models/wmt19.en-ru.ensemble.tar.gz)
+`transformer.wmt19.ru-en` | Transformer <br> ([Ng et al., 2019](https://arxiv.org/abs/1907.06616)) <br> WMT'19 winner | [WMT'19 Russian-English](http://www.statmt.org/wmt19/translation-task.html) | model: <br> [download (.tar.gz)](https://dl.fbaipublicfiles.com/fairseq/models/wmt19.ru-en.ensemble.tar.gz)
+
+## Example usage (torch.hub)
+
+We require a few additional Python dependencies for preprocessing:
+```bash
+pip install fastBPE sacremoses subword_nmt
+```
+
+Interactive translation via PyTorch Hub:
+```python
+import torch
+
+# List available models
+torch.hub.list('pytorch/fairseq')  # [..., 'transformer.wmt16.en-de', ... ]
+
+# Load a transformer trained on WMT'16 En-De
+# Note: WMT'19 models use fastBPE instead of subword_nmt, see instructions below
+en2de = torch.hub.load('pytorch/fairseq', 'transformer.wmt16.en-de',
+                       tokenizer='moses', bpe='subword_nmt')
+en2de.eval()  # disable dropout
+
+# The underlying model is available under the *models* attribute
+assert isinstance(en2de.models[0], fairseq.models.transformer.TransformerModel)
+
+# Move model to GPU for faster translation
+en2de.cuda()
+
+# Translate a sentence
+en2de.translate('Hello world!')
+# 'Hallo Welt!'
+
+# Batched translation
+en2de.translate(['Hello world!', 'The cat sat on the mat.'])
+# ['Hallo Welt!', 'Die Katze saß auf der Matte.']
+```
+
+Loading custom models:
+```python
+from fairseq.models.transformer import TransformerModel
+zh2en = TransformerModel.from_pretrained(
+  '/path/to/checkpoints',
+  checkpoint_file='checkpoint_best.pt',
+  data_name_or_path='data-bin/wmt17_zh_en_full',
+  bpe='subword_nmt',
+  bpe_codes='data-bin/wmt17_zh_en_full/zh.code'
+)
+zh2en.translate('你好 世界')
+# 'Hello World'
+```
+
+If you are using a `transformer.wmt19` models, you will need to set the `bpe`
+argument to `'fastbpe'` and (optionally) load the 4-model ensemble:
+```python
+en2de = torch.hub.load('pytorch/fairseq', 'transformer.wmt19.en-de',
+                       checkpoint_file='model1.pt:model2.pt:model3.pt:model4.pt',
+                       tokenizer='moses', bpe='fastbpe')
+en2de.eval()  # disable dropout
+```
+
+## Example usage (CLI tools)
+
+Generation with the binarized test sets can be run in batch mode as follows, e.g. for WMT 2014 English-French on a GTX-1080ti:
+```bash
+mkdir -p data-bin
+curl https://dl.fbaipublicfiles.com/fairseq/models/wmt14.v2.en-fr.fconv-py.tar.bz2 | tar xvjf - -C data-bin
+curl https://dl.fbaipublicfiles.com/fairseq/data/wmt14.v2.en-fr.newstest2014.tar.bz2 | tar xvjf - -C data-bin
+fairseq-generate data-bin/wmt14.en-fr.newstest2014  \
+    --path data-bin/wmt14.en-fr.fconv-py/model.pt \
+    --beam 5 --batch-size 128 --remove-bpe | tee /tmp/gen.out
+# ...
+# | Translated 3003 sentences (96311 tokens) in 166.0s (580.04 tokens/s)
+# | Generate test with beam=5: BLEU4 = 40.83, 67.5/46.9/34.4/25.5 (BP=1.000, ratio=1.006, syslen=83262, reflen=82787)
+
+# Compute BLEU score
+grep ^H /tmp/gen.out | cut -f3- > /tmp/gen.out.sys
+grep ^T /tmp/gen.out | cut -f2- > /tmp/gen.out.ref
+fairseq-score --sys /tmp/gen.out.sys --ref /tmp/gen.out.ref
+# BLEU4 = 40.83, 67.5/46.9/34.4/25.5 (BP=1.000, ratio=1.006, syslen=83262, reflen=82787)
+```
+
+## Training a new model
+
+### IWSLT'14 German to English (Transformer)
+
+The following instructions can be used to train a Transformer model on the [IWSLT'14 German to English dataset](http://workshop2014.iwslt.org/downloads/proceeding.pdf).
+
+First download and preprocess the data:
+```bash
+# Download and prepare the data
+cd examples/translation/
+bash prepare-iwslt14.sh
+cd ../..
+
+# Preprocess/binarize the data
+TEXT=examples/translation/iwslt14.tokenized.de-en
+fairseq-preprocess --source-lang de --target-lang en \
+    --trainpref $TEXT/train --validpref $TEXT/valid --testpref $TEXT/test \
+    --destdir data-bin/iwslt14.tokenized.de-en \
+    --workers 20
+```
+
+Next we'll train a Transformer translation model over this data:
+```bash
+CUDA_VISIBLE_DEVICES=0 fairseq-train \
+    data-bin/iwslt14.tokenized.de-en \
+    --arch transformer_iwslt_de_en --share-decoder-input-output-embed \
+    --optimizer adam --adam-betas '(0.9, 0.98)' --clip-norm 0.0 \
+    --lr 5e-4 --lr-scheduler inverse_sqrt --warmup-updates 4000 \
+    --dropout 0.3 --weight-decay 0.0001 \
+    --criterion label_smoothed_cross_entropy --label-smoothing 0.1 \
+    --max-tokens 4096 \
+    --eval-bleu \
+    --eval-bleu-args '{"beam": 5, "max_len_a": 1.2, "max_len_b": 10}' \
+    --eval-bleu-detok moses \
+    --eval-bleu-remove-bpe \
+    --eval-bleu-print-samples \
+    --best-checkpoint-metric bleu --maximize-best-checkpoint-metric
+```
+
+Finally we can evaluate our trained model:
+```bash
+fairseq-generate data-bin/iwslt14.tokenized.de-en \
+    --path checkpoints/checkpoint_best.pt \
+    --batch-size 128 --beam 5 --remove-bpe
+```
+
+### WMT'14 English to German (Convolutional)
+
+The following instructions can be used to train a Convolutional translation model on the WMT English to German dataset.
+See the [Scaling NMT README](../scaling_nmt/README.md) for instructions to train a Transformer translation model on this data.
+
+The WMT English to German dataset can be preprocessed using the `prepare-wmt14en2de.sh` script.
+By default it will produce a dataset that was modeled after [Attention Is All You Need (Vaswani et al., 2017)](https://arxiv.org/abs/1706.03762), but with additional news-commentary-v12 data from WMT'17.
+
+To use only data available in WMT'14 or to replicate results obtained in the original [Convolutional Sequence to Sequence Learning (Gehring et al., 2017)](https://arxiv.org/abs/1705.03122) paper, please use the `--icml17` option.
+
+```bash
+# Download and prepare the data
+cd examples/translation/
+# WMT'17 data:
+bash prepare-wmt14en2de.sh
+# or to use WMT'14 data:
+# bash prepare-wmt14en2de.sh --icml17
+cd ../..
+
+# Binarize the dataset
+TEXT=examples/translation/wmt17_en_de
+fairseq-preprocess \
+    --source-lang en --target-lang de \
+    --trainpref $TEXT/train --validpref $TEXT/valid --testpref $TEXT/test \
+    --destdir data-bin/wmt17_en_de --thresholdtgt 0 --thresholdsrc 0 \
+    --workers 20
+
+# Train the model
+mkdir -p checkpoints/fconv_wmt_en_de
+fairseq-train \
+    data-bin/wmt17_en_de \
+    --arch fconv_wmt_en_de \
+    --dropout 0.2 \
+    --criterion label_smoothed_cross_entropy --label-smoothing 0.1 \
+    --optimizer nag --clip-norm 0.1 \
+    --lr 0.5 --lr-scheduler fixed --force-anneal 50 \
+    --max-tokens 4000 \
+    --save-dir checkpoints/fconv_wmt_en_de
+
+# Evaluate
+fairseq-generate data-bin/wmt17_en_de \
+    --path checkpoints/fconv_wmt_en_de/checkpoint_best.pt \
+    --beam 5 --remove-bpe
+```
+
+### WMT'14 English to French
+```bash
+# Download and prepare the data
+cd examples/translation/
+bash prepare-wmt14en2fr.sh
+cd ../..
+
+# Binarize the dataset
+TEXT=examples/translation/wmt14_en_fr
+fairseq-preprocess \
+    --source-lang en --target-lang fr \
+    --trainpref $TEXT/train --validpref $TEXT/valid --testpref $TEXT/test \
+    --destdir data-bin/wmt14_en_fr --thresholdtgt 0 --thresholdsrc 0 \
+    --workers 60
+
+# Train the model
+mkdir -p checkpoints/fconv_wmt_en_fr
+fairseq-train \
+    data-bin/wmt14_en_fr \
+    --arch fconv_wmt_en_fr \
+    --dropout 0.1 \
+    --criterion label_smoothed_cross_entropy --label-smoothing 0.1 \
+    --optimizer nag --clip-norm 0.1 \
+    --lr 0.5 --lr-scheduler fixed --force-anneal 50 \
+    --max-tokens 3000 \
+    --save-dir checkpoints/fconv_wmt_en_fr
+
+# Evaluate
+fairseq-generate \
+    data-bin/fconv_wmt_en_fr \
+    --path checkpoints/fconv_wmt_en_fr/checkpoint_best.pt \
+    --beam 5 --remove-bpe
+```
+
+## Multilingual Translation
+
+We also support training multilingual translation models. In this example we'll
+train a multilingual `{de,fr}-en` translation model using the IWSLT'17 datasets.
+
+Note that we use slightly different preprocessing here than for the IWSLT'14
+En-De data above. In particular we learn a joint BPE code for all three
+languages and use fairseq-interactive and sacrebleu for scoring the test set.
+
+```bash
+# First install sacrebleu and sentencepiece
+pip install sacrebleu sentencepiece
+
+# Then download and preprocess the data
+cd examples/translation/
+bash prepare-iwslt17-multilingual.sh
+cd ../..
+
+# Binarize the de-en dataset
+TEXT=examples/translation/iwslt17.de_fr.en.bpe16k
+fairseq-preprocess --source-lang de --target-lang en \
+    --trainpref $TEXT/train.bpe.de-en \
+    --validpref $TEXT/valid0.bpe.de-en,$TEXT/valid1.bpe.de-en,$TEXT/valid2.bpe.de-en,$TEXT/valid3.bpe.de-en,$TEXT/valid4.bpe.de-en,$TEXT/valid5.bpe.de-en \
+    --destdir data-bin/iwslt17.de_fr.en.bpe16k \
+    --workers 10
+
+# Binarize the fr-en dataset
+# NOTE: it's important to reuse the en dictionary from the previous step
+fairseq-preprocess --source-lang fr --target-lang en \
+    --trainpref $TEXT/train.bpe.fr-en \
+    --validpref $TEXT/valid0.bpe.fr-en,$TEXT/valid1.bpe.fr-en,$TEXT/valid2.bpe.fr-en,$TEXT/valid3.bpe.fr-en,$TEXT/valid4.bpe.fr-en,$TEXT/valid5.bpe.fr-en \
+    --tgtdict data-bin/iwslt17.de_fr.en.bpe16k/dict.en.txt \
+    --destdir data-bin/iwslt17.de_fr.en.bpe16k \
+    --workers 10
+
+# Train a multilingual transformer model
+# NOTE: the command below assumes 1 GPU, but accumulates gradients from
+#       8 fwd/bwd passes to simulate training on 8 GPUs
+mkdir -p checkpoints/multilingual_transformer
+CUDA_VISIBLE_DEVICES=0 fairseq-train data-bin/iwslt17.de_fr.en.bpe16k/ \
+    --max-epoch 50 \
+    --ddp-backend=legacy_ddp \
+    --task multilingual_translation --lang-pairs de-en,fr-en \
+    --arch multilingual_transformer_iwslt_de_en \
+    --share-decoders --share-decoder-input-output-embed \
+    --optimizer adam --adam-betas '(0.9, 0.98)' \
+    --lr 0.0005 --lr-scheduler inverse_sqrt \
+    --warmup-updates 4000 --warmup-init-lr '1e-07' \
+    --label-smoothing 0.1 --criterion label_smoothed_cross_entropy \
+    --dropout 0.3 --weight-decay 0.0001 \
+    --save-dir checkpoints/multilingual_transformer \
+    --max-tokens 4000 \
+    --update-freq 8
+
+# Generate and score the test set with sacrebleu
+SRC=de
+sacrebleu --test-set iwslt17 --language-pair ${SRC}-en --echo src \
+    | python scripts/spm_encode.py --model examples/translation/iwslt17.de_fr.en.bpe16k/sentencepiece.bpe.model \
+    > iwslt17.test.${SRC}-en.${SRC}.bpe
+cat iwslt17.test.${SRC}-en.${SRC}.bpe \
+    | fairseq-interactive data-bin/iwslt17.de_fr.en.bpe16k/ \
+      --task multilingual_translation --lang-pairs de-en,fr-en \
+      --source-lang ${SRC} --target-lang en \
+      --path checkpoints/multilingual_transformer/checkpoint_best.pt \
+      --buffer-size 2000 --batch-size 128 \
+      --beam 5 --remove-bpe=sentencepiece \
+    > iwslt17.test.${SRC}-en.en.sys
+grep ^H iwslt17.test.${SRC}-en.en.sys | cut -f3 \
+    | sacrebleu --test-set iwslt17 --language-pair ${SRC}-en
+```
+
+##### Argument format during inference
+
+During inference it is required to specify a single `--source-lang` and
+`--target-lang`, which indicates the inference langauge direction.
+`--lang-pairs`, `--encoder-langtok`, `--decoder-langtok` have to be set to
+the same value as training.

omni_speech/infer/fairseq/examples/translation/prepare-iwslt14.sh

@@ -0,0 +1,115 @@
+#!/usr/bin/env bash
+#
+# Adapted from https://github.com/facebookresearch/MIXER/blob/master/prepareData.sh
+
+echo 'Cloning Moses github repository (for tokenization scripts)...'
+git clone https://github.com/moses-smt/mosesdecoder.git
+
+echo 'Cloning Subword NMT repository (for BPE pre-processing)...'
+git clone https://github.com/rsennrich/subword-nmt.git
+
+SCRIPTS=mosesdecoder/scripts
+TOKENIZER=$SCRIPTS/tokenizer/tokenizer.perl
+LC=$SCRIPTS/tokenizer/lowercase.perl
+CLEAN=$SCRIPTS/training/clean-corpus-n.perl
+BPEROOT=subword-nmt/subword_nmt
+BPE_TOKENS=10000
+
+URL="http://dl.fbaipublicfiles.com/fairseq/data/iwslt14/de-en.tgz"
+GZ=de-en.tgz
+
+if [ ! -d "$SCRIPTS" ]; then
+    echo "Please set SCRIPTS variable correctly to point to Moses scripts."
+    exit
+fi
+
+src=de
+tgt=en
+lang=de-en
+prep=iwslt14.tokenized.de-en
+tmp=$prep/tmp
+orig=orig
+
+mkdir -p $orig $tmp $prep
+
+echo "Downloading data from ${URL}..."
+cd $orig
+wget "$URL"
+
+if [ -f $GZ ]; then
+    echo "Data successfully downloaded."
+else
+    echo "Data not successfully downloaded."
+    exit
+fi
+
+tar zxvf $GZ
+cd ..
+
+echo "pre-processing train data..."
+for l in $src $tgt; do
+    f=train.tags.$lang.$l
+    tok=train.tags.$lang.tok.$l
+
+    cat $orig/$lang/$f | \
+    grep -v '<url>' | \
+    grep -v '<talkid>' | \
+    grep -v '<keywords>' | \
+    sed -e 's/<title>//g' | \
+    sed -e 's/<\/title>//g' | \
+    sed -e 's/<description>//g' | \
+    sed -e 's/<\/description>//g' | \
+    perl $TOKENIZER -threads 8 -l $l > $tmp/$tok
+    echo ""
+done
+perl $CLEAN -ratio 1.5 $tmp/train.tags.$lang.tok $src $tgt $tmp/train.tags.$lang.clean 1 175
+for l in $src $tgt; do
+    perl $LC < $tmp/train.tags.$lang.clean.$l > $tmp/train.tags.$lang.$l
+done
+
+echo "pre-processing valid/test data..."
+for l in $src $tgt; do
+    for o in `ls $orig/$lang/IWSLT14.TED*.$l.xml`; do
+    fname=${o##*/}
+    f=$tmp/${fname%.*}
+    echo $o $f
+    grep '<seg id' $o | \
+        sed -e 's/<seg id="[0-9]*">\s*//g' | \
+        sed -e 's/\s*<\/seg>\s*//g' | \
+        sed -e "s/\’/\'/g" | \
+    perl $TOKENIZER -threads 8 -l $l | \
+    perl $LC > $f
+    echo ""
+    done
+done
+
+
+echo "creating train, valid, test..."
+for l in $src $tgt; do
+    awk '{if (NR%23 == 0)  print $0; }' $tmp/train.tags.de-en.$l > $tmp/valid.$l
+    awk '{if (NR%23 != 0)  print $0; }' $tmp/train.tags.de-en.$l > $tmp/train.$l
+
+    cat $tmp/IWSLT14.TED.dev2010.de-en.$l \
+        $tmp/IWSLT14.TEDX.dev2012.de-en.$l \
+        $tmp/IWSLT14.TED.tst2010.de-en.$l \
+        $tmp/IWSLT14.TED.tst2011.de-en.$l \
+        $tmp/IWSLT14.TED.tst2012.de-en.$l \
+        > $tmp/test.$l
+done
+
+TRAIN=$tmp/train.en-de
+BPE_CODE=$prep/code
+rm -f $TRAIN
+for l in $src $tgt; do
+    cat $tmp/train.$l >> $TRAIN
+done
+
+echo "learn_bpe.py on ${TRAIN}..."
+python $BPEROOT/learn_bpe.py -s $BPE_TOKENS < $TRAIN > $BPE_CODE
+
+for L in $src $tgt; do
+    for f in train.$L valid.$L test.$L; do
+        echo "apply_bpe.py to ${f}..."
+        python $BPEROOT/apply_bpe.py -c $BPE_CODE < $tmp/$f > $prep/$f
+    done
+done

omni_speech/infer/fairseq/examples/translation/prepare-iwslt17-multilingual.sh

@@ -0,0 +1,133 @@
+#!/bin/bash
+# Copyright (c) Facebook, Inc. and its affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+SRCS=(
+    "de"
+    "fr"
+)
+TGT=en
+
+ROOT=$(dirname "$0")
+SCRIPTS=$ROOT/../../scripts
+SPM_TRAIN=$SCRIPTS/spm_train.py
+SPM_ENCODE=$SCRIPTS/spm_encode.py
+
+BPESIZE=16384
+ORIG=$ROOT/iwslt17_orig
+DATA=$ROOT/iwslt17.de_fr.en.bpe16k
+mkdir -p "$ORIG" "$DATA"
+
+TRAIN_MINLEN=1  # remove sentences with <1 BPE token
+TRAIN_MAXLEN=250  # remove sentences with >250 BPE tokens
+
+URLS=(
+    "https://wit3.fbk.eu/archive/2017-01-trnted/texts/de/en/de-en.tgz"
+    "https://wit3.fbk.eu/archive/2017-01-trnted/texts/fr/en/fr-en.tgz"
+)
+ARCHIVES=(
+    "de-en.tgz"
+    "fr-en.tgz"
+)
+VALID_SETS=(
+    "IWSLT17.TED.dev2010.de-en IWSLT17.TED.tst2010.de-en IWSLT17.TED.tst2011.de-en IWSLT17.TED.tst2012.de-en IWSLT17.TED.tst2013.de-en IWSLT17.TED.tst2014.de-en IWSLT17.TED.tst2015.de-en"
+    "IWSLT17.TED.dev2010.fr-en IWSLT17.TED.tst2010.fr-en IWSLT17.TED.tst2011.fr-en IWSLT17.TED.tst2012.fr-en IWSLT17.TED.tst2013.fr-en IWSLT17.TED.tst2014.fr-en IWSLT17.TED.tst2015.fr-en"
+)
+
+# download and extract data
+for ((i=0;i<${#URLS[@]};++i)); do
+    ARCHIVE=$ORIG/${ARCHIVES[i]}
+    if [ -f "$ARCHIVE" ]; then
+        echo "$ARCHIVE already exists, skipping download"
+    else
+        URL=${URLS[i]}
+        wget -P "$ORIG" "$URL"
+        if [ -f "$ARCHIVE" ]; then
+            echo "$URL successfully downloaded."
+        else
+            echo "$URL not successfully downloaded."
+            exit 1
+        fi
+    fi
+    FILE=${ARCHIVE: -4}
+    if [ -e "$FILE" ]; then
+        echo "$FILE already exists, skipping extraction"
+    else
+        tar -C "$ORIG" -xzvf "$ARCHIVE"
+    fi
+done
+
+echo "pre-processing train data..."
+for SRC in "${SRCS[@]}"; do
+    for LANG in "${SRC}" "${TGT}"; do
+        cat "$ORIG/${SRC}-${TGT}/train.tags.${SRC}-${TGT}.${LANG}" \
+            | grep -v '<url>' \
+            | grep -v '<talkid>' \
+            | grep -v '<keywords>' \
+            | grep -v '<speaker>' \
+            | grep -v '<reviewer' \
+            | grep -v '<translator' \
+            | grep -v '<doc' \
+            | grep -v '</doc>' \
+            | sed -e 's/<title>//g' \
+            | sed -e 's/<\/title>//g' \
+            | sed -e 's/<description>//g' \
+            | sed -e 's/<\/description>//g' \
+            | sed 's/^\s*//g' \
+            | sed 's/\s*$//g' \
+            > "$DATA/train.${SRC}-${TGT}.${LANG}"
+    done
+done
+
+echo "pre-processing valid data..."
+for ((i=0;i<${#SRCS[@]};++i)); do
+    SRC=${SRCS[i]}
+    VALID_SET=(${VALID_SETS[i]})
+    for ((j=0;j<${#VALID_SET[@]};++j)); do
+        FILE=${VALID_SET[j]}
+        for LANG in "$SRC" "$TGT"; do
+            grep '<seg id' "$ORIG/${SRC}-${TGT}/${FILE}.${LANG}.xml" \
+                | sed -e 's/<seg id="[0-9]*">\s*//g' \
+                | sed -e 's/\s*<\/seg>\s*//g' \
+                | sed -e "s/\’/\'/g" \
+                > "$DATA/valid${j}.${SRC}-${TGT}.${LANG}"
+        done
+    done
+done
+
+# learn BPE with sentencepiece
+TRAIN_FILES=$(for SRC in "${SRCS[@]}"; do echo $DATA/train.${SRC}-${TGT}.${SRC}; echo $DATA/train.${SRC}-${TGT}.${TGT}; done | tr "\n" ",")
+echo "learning joint BPE over ${TRAIN_FILES}..."
+python "$SPM_TRAIN" \
+    --input=$TRAIN_FILES \
+    --model_prefix=$DATA/sentencepiece.bpe \
+    --vocab_size=$BPESIZE \
+    --character_coverage=1.0 \
+    --model_type=bpe
+
+# encode train/valid
+echo "encoding train with learned BPE..."
+for SRC in "${SRCS[@]}"; do
+    python "$SPM_ENCODE" \
+        --model "$DATA/sentencepiece.bpe.model" \
+        --output_format=piece \
+        --inputs $DATA/train.${SRC}-${TGT}.${SRC} $DATA/train.${SRC}-${TGT}.${TGT} \
+        --outputs $DATA/train.bpe.${SRC}-${TGT}.${SRC} $DATA/train.bpe.${SRC}-${TGT}.${TGT} \
+        --min-len $TRAIN_MINLEN --max-len $TRAIN_MAXLEN
+done
+
+echo "encoding valid with learned BPE..."
+for ((i=0;i<${#SRCS[@]};++i)); do
+    SRC=${SRCS[i]}
+    VALID_SET=(${VALID_SETS[i]})
+    for ((j=0;j<${#VALID_SET[@]};++j)); do
+        python "$SPM_ENCODE" \
+            --model "$DATA/sentencepiece.bpe.model" \
+            --output_format=piece \
+            --inputs $DATA/valid${j}.${SRC}-${TGT}.${SRC} $DATA/valid${j}.${SRC}-${TGT}.${TGT} \
+            --outputs $DATA/valid${j}.bpe.${SRC}-${TGT}.${SRC} $DATA/valid${j}.bpe.${SRC}-${TGT}.${TGT}
+    done
+done

omni_speech/infer/fairseq/examples/translation/prepare-wmt14en2de.sh

@@ -0,0 +1,142 @@
+#!/bin/bash
+# Adapted from https://github.com/facebookresearch/MIXER/blob/master/prepareData.sh
+
+echo 'Cloning Moses github repository (for tokenization scripts)...'
+git clone https://github.com/moses-smt/mosesdecoder.git
+
+echo 'Cloning Subword NMT repository (for BPE pre-processing)...'
+git clone https://github.com/rsennrich/subword-nmt.git
+
+SCRIPTS=mosesdecoder/scripts
+TOKENIZER=$SCRIPTS/tokenizer/tokenizer.perl
+CLEAN=$SCRIPTS/training/clean-corpus-n.perl
+NORM_PUNC=$SCRIPTS/tokenizer/normalize-punctuation.perl
+REM_NON_PRINT_CHAR=$SCRIPTS/tokenizer/remove-non-printing-char.perl
+BPEROOT=subword-nmt/subword_nmt
+BPE_TOKENS=40000
+
+URLS=(
+    "http://statmt.org/wmt13/training-parallel-europarl-v7.tgz"
+    "http://statmt.org/wmt13/training-parallel-commoncrawl.tgz"
+    "http://data.statmt.org/wmt17/translation-task/training-parallel-nc-v12.tgz"
+    "http://data.statmt.org/wmt17/translation-task/dev.tgz"
+    "http://statmt.org/wmt14/test-full.tgz"
+)
+FILES=(
+    "training-parallel-europarl-v7.tgz"
+    "training-parallel-commoncrawl.tgz"
+    "training-parallel-nc-v12.tgz"
+    "dev.tgz"
+    "test-full.tgz"
+)
+CORPORA=(
+    "training/europarl-v7.de-en"
+    "commoncrawl.de-en"
+    "training/news-commentary-v12.de-en"
+)
+
+# This will make the dataset compatible to the one used in "Convolutional Sequence to Sequence Learning"
+# https://arxiv.org/abs/1705.03122
+if [ "$1" == "--icml17" ]; then
+    URLS[2]="http://statmt.org/wmt14/training-parallel-nc-v9.tgz"
+    FILES[2]="training-parallel-nc-v9.tgz"
+    CORPORA[2]="training/news-commentary-v9.de-en"
+    OUTDIR=wmt14_en_de
+else
+    OUTDIR=wmt17_en_de
+fi
+
+if [ ! -d "$SCRIPTS" ]; then
+    echo "Please set SCRIPTS variable correctly to point to Moses scripts."
+    exit
+fi
+
+src=en
+tgt=de
+lang=en-de
+prep=$OUTDIR
+tmp=$prep/tmp
+orig=orig
+dev=dev/newstest2013
+
+mkdir -p $orig $tmp $prep
+
+cd $orig
+
+for ((i=0;i<${#URLS[@]};++i)); do
+    file=${FILES[i]}
+    if [ -f $file ]; then
+        echo "$file already exists, skipping download"
+    else
+        url=${URLS[i]}
+        wget "$url"
+        if [ -f $file ]; then
+            echo "$url successfully downloaded."
+        else
+            echo "$url not successfully downloaded."
+            exit -1
+        fi
+        if [ ${file: -4} == ".tgz" ]; then
+            tar zxvf $file
+        elif [ ${file: -4} == ".tar" ]; then
+            tar xvf $file
+        fi
+    fi
+done
+cd ..
+
+echo "pre-processing train data..."
+for l in $src $tgt; do
+    rm $tmp/train.tags.$lang.tok.$l
+    for f in "${CORPORA[@]}"; do
+        cat $orig/$f.$l | \
+            perl $NORM_PUNC $l | \
+            perl $REM_NON_PRINT_CHAR | \
+            perl $TOKENIZER -threads 8 -a -l $l >> $tmp/train.tags.$lang.tok.$l
+    done
+done
+
+echo "pre-processing test data..."
+for l in $src $tgt; do
+    if [ "$l" == "$src" ]; then
+        t="src"
+    else
+        t="ref"
+    fi
+    grep '<seg id' $orig/test-full/newstest2014-deen-$t.$l.sgm | \
+        sed -e 's/<seg id="[0-9]*">\s*//g' | \
+        sed -e 's/\s*<\/seg>\s*//g' | \
+        sed -e "s/\’/\'/g" | \
+    perl $TOKENIZER -threads 8 -a -l $l > $tmp/test.$l
+    echo ""
+done
+
+echo "splitting train and valid..."
+for l in $src $tgt; do
+    awk '{if (NR%100 == 0)  print $0; }' $tmp/train.tags.$lang.tok.$l > $tmp/valid.$l
+    awk '{if (NR%100 != 0)  print $0; }' $tmp/train.tags.$lang.tok.$l > $tmp/train.$l
+done
+
+TRAIN=$tmp/train.de-en
+BPE_CODE=$prep/code
+rm -f $TRAIN
+for l in $src $tgt; do
+    cat $tmp/train.$l >> $TRAIN
+done
+
+echo "learn_bpe.py on ${TRAIN}..."
+python $BPEROOT/learn_bpe.py -s $BPE_TOKENS < $TRAIN > $BPE_CODE
+
+for L in $src $tgt; do
+    for f in train.$L valid.$L test.$L; do
+        echo "apply_bpe.py to ${f}..."
+        python $BPEROOT/apply_bpe.py -c $BPE_CODE < $tmp/$f > $tmp/bpe.$f
+    done
+done
+
+perl $CLEAN -ratio 1.5 $tmp/bpe.train $src $tgt $prep/train 1 250
+perl $CLEAN -ratio 1.5 $tmp/bpe.valid $src $tgt $prep/valid 1 250
+
+for L in $src $tgt; do
+    cp $tmp/bpe.test.$L $prep/test.$L
+done

omni_speech/infer/fairseq/examples/translation/prepare-wmt14en2fr.sh

@@ -0,0 +1,136 @@
+#!/bin/bash
+# Adapted from https://github.com/facebookresearch/MIXER/blob/master/prepareData.sh
+
+echo 'Cloning Moses github repository (for tokenization scripts)...'
+git clone https://github.com/moses-smt/mosesdecoder.git
+
+echo 'Cloning Subword NMT repository (for BPE pre-processing)...'
+git clone https://github.com/rsennrich/subword-nmt.git
+
+SCRIPTS=mosesdecoder/scripts
+TOKENIZER=$SCRIPTS/tokenizer/tokenizer.perl
+CLEAN=$SCRIPTS/training/clean-corpus-n.perl
+NORM_PUNC=$SCRIPTS/tokenizer/normalize-punctuation.perl
+REM_NON_PRINT_CHAR=$SCRIPTS/tokenizer/remove-non-printing-char.perl
+BPEROOT=subword-nmt/subword_nmt
+BPE_TOKENS=40000
+
+URLS=(
+    "http://statmt.org/wmt13/training-parallel-europarl-v7.tgz"
+    "http://statmt.org/wmt13/training-parallel-commoncrawl.tgz"
+    "http://statmt.org/wmt13/training-parallel-un.tgz"
+    "http://statmt.org/wmt14/training-parallel-nc-v9.tgz"
+    "http://statmt.org/wmt10/training-giga-fren.tar"
+    "http://statmt.org/wmt14/test-full.tgz"
+)
+FILES=(
+    "training-parallel-europarl-v7.tgz"
+    "training-parallel-commoncrawl.tgz"
+    "training-parallel-un.tgz"
+    "training-parallel-nc-v9.tgz"
+    "training-giga-fren.tar"
+    "test-full.tgz"
+)
+CORPORA=(
+    "training/europarl-v7.fr-en"
+    "commoncrawl.fr-en"
+    "un/undoc.2000.fr-en"
+    "training/news-commentary-v9.fr-en"
+    "giga-fren.release2.fixed"
+)
+
+if [ ! -d "$SCRIPTS" ]; then
+    echo "Please set SCRIPTS variable correctly to point to Moses scripts."
+    exit
+fi
+
+src=en
+tgt=fr
+lang=en-fr
+prep=wmt14_en_fr
+tmp=$prep/tmp
+orig=orig
+
+mkdir -p $orig $tmp $prep
+
+cd $orig
+
+for ((i=0;i<${#URLS[@]};++i)); do
+    file=${FILES[i]}
+    if [ -f $file ]; then
+        echo "$file already exists, skipping download"
+    else
+        url=${URLS[i]}
+        wget "$url"
+        if [ -f $file ]; then
+            echo "$url successfully downloaded."
+        else
+            echo "$url not successfully downloaded."
+            exit -1
+        fi
+        if [ ${file: -4} == ".tgz" ]; then
+            tar zxvf $file
+        elif [ ${file: -4} == ".tar" ]; then
+            tar xvf $file
+        fi
+    fi
+done
+
+gunzip giga-fren.release2.fixed.*.gz
+cd ..
+
+echo "pre-processing train data..."
+for l in $src $tgt; do
+    rm $tmp/train.tags.$lang.tok.$l
+    for f in "${CORPORA[@]}"; do
+        cat $orig/$f.$l | \
+            perl $NORM_PUNC $l | \
+            perl $REM_NON_PRINT_CHAR | \
+            perl $TOKENIZER -threads 8 -a -l $l >> $tmp/train.tags.$lang.tok.$l
+    done
+done
+
+echo "pre-processing test data..."
+for l in $src $tgt; do
+    if [ "$l" == "$src" ]; then
+        t="src"
+    else
+        t="ref"
+    fi
+    grep '<seg id' $orig/test-full/newstest2014-fren-$t.$l.sgm | \
+        sed -e 's/<seg id="[0-9]*">\s*//g' | \
+        sed -e 's/\s*<\/seg>\s*//g' | \
+        sed -e "s/\’/\'/g" | \
+    perl $TOKENIZER -threads 8 -a -l $l > $tmp/test.$l
+    echo ""
+done
+
+echo "splitting train and valid..."
+for l in $src $tgt; do
+    awk '{if (NR%1333 == 0)  print $0; }' $tmp/train.tags.$lang.tok.$l > $tmp/valid.$l
+    awk '{if (NR%1333 != 0)  print $0; }' $tmp/train.tags.$lang.tok.$l > $tmp/train.$l
+done
+
+TRAIN=$tmp/train.fr-en
+BPE_CODE=$prep/code
+rm -f $TRAIN
+for l in $src $tgt; do
+    cat $tmp/train.$l >> $TRAIN
+done
+
+echo "learn_bpe.py on ${TRAIN}..."
+python $BPEROOT/learn_bpe.py -s $BPE_TOKENS < $TRAIN > $BPE_CODE
+
+for L in $src $tgt; do
+    for f in train.$L valid.$L test.$L; do
+        echo "apply_bpe.py to ${f}..."
+        python $BPEROOT/apply_bpe.py -c $BPE_CODE < $tmp/$f > $tmp/bpe.$f
+    done
+done
+
+perl $CLEAN -ratio 1.5 $tmp/bpe.train $src $tgt $prep/train 1 250
+perl $CLEAN -ratio 1.5 $tmp/bpe.valid $src $tgt $prep/valid 1 250
+
+for L in $src $tgt; do
+    cp $tmp/bpe.test.$L $prep/test.$L
+done

omni_speech/infer/fairseq/examples/translation_moe/README.md

@@ -0,0 +1,89 @@
+# Mixture Models for Diverse Machine Translation: Tricks of the Trade (Shen et al., 2019)
+
+This page includes instructions for reproducing results from the paper [Mixture Models for Diverse Machine Translation: Tricks of the Trade (Shen et al., 2019)](https://arxiv.org/abs/1902.07816).
+
+## Download data
+
+First, follow the [instructions to download and preprocess the WMT'17 En-De dataset](../translation#prepare-wmt14en2desh).
+Make sure to learn a joint vocabulary by passing the `--joined-dictionary` option to `fairseq-preprocess`.
+
+## Train a model
+
+Then we can train a mixture of experts model using the `translation_moe` task.
+Use the `--method` flag to choose the MoE variant; we support hard mixtures with a learned or uniform prior (`--method hMoElp` and `hMoEup`, respectively) and soft mixures (`--method sMoElp` and `sMoEup`).
+The model is trained with online responsibility assignment and shared parameterization.
+
+The following command will train a `hMoElp` model with `3` experts:
+```bash
+fairseq-train --ddp-backend='legacy_ddp' \
+    data-bin/wmt17_en_de \
+    --max-update 100000 \
+    --task translation_moe --user-dir examples/translation_moe/translation_moe_src \
+    --method hMoElp --mean-pool-gating-network \
+    --num-experts 3 \
+    --arch transformer_wmt_en_de --share-all-embeddings \
+    --optimizer adam --adam-betas '(0.9, 0.98)' --clip-norm 0.0 \
+    --lr-scheduler inverse_sqrt --warmup-init-lr 1e-07 --warmup-updates 4000 \
+    --lr 0.0007 \
+    --dropout 0.1 --weight-decay 0.0 --criterion cross_entropy \
+    --max-tokens 3584
+```
+
+## Translate
+
+Once a model is trained, we can generate translations from different experts using the `--gen-expert` option.
+For example, to generate from expert 0:
+```bash
+fairseq-generate data-bin/wmt17_en_de \
+    --path checkpoints/checkpoint_best.pt \
+    --beam 1 --remove-bpe \
+    --task translation_moe --user-dir examples/translation_moe/translation_moe_src \
+    --method hMoElp --mean-pool-gating-network \
+    --num-experts 3 \
+    --gen-expert 0
+```
+
+## Evaluate
+
+First download a tokenized version of the WMT'14 En-De test set with multiple references:
+```bash
+wget dl.fbaipublicfiles.com/fairseq/data/wmt14-en-de.extra_refs.tok
+```
+
+Next apply BPE on the fly and run generation for each expert:
+```bash
+BPE_CODE=examples/translation/wmt17_en_de/code
+for EXPERT in $(seq 0 2); do \
+    cat wmt14-en-de.extra_refs.tok \
+    | grep ^S | cut -f 2 \
+    | fairseq-interactive data-bin/wmt17_en_de \
+        --path checkpoints/checkpoint_best.pt \
+        --beam 1 \
+        --bpe subword_nmt --bpe-codes $BPE_CODE \
+        --buffer-size 500 --max-tokens 6000 \
+        --task translation_moe --user-dir examples/translation_moe/translation_moe_src \
+        --method hMoElp --mean-pool-gating-network \
+        --num-experts 3 \
+        --gen-expert $EXPERT ; \
+done > wmt14-en-de.extra_refs.tok.gen.3experts
+```
+
+Finally use `score_moe.py` to compute pairwise BLUE and average oracle BLEU:
+```bash
+python examples/translation_moe/score.py --sys wmt14-en-de.extra_refs.tok.gen.3experts --ref wmt14-en-de.extra_refs.tok
+# pairwise BLEU: 48.26
+# #refs covered: 2.11
+# multi-reference BLEU (leave-one-out): 59.46
+```
+This matches row 3 from Table 7 in the paper.
+
+## Citation
+
+```bibtex
+@article{shen2019mixture,
+  title = {Mixture Models for Diverse Machine Translation: Tricks of the Trade},
+  author = {Tianxiao Shen and Myle Ott and Michael Auli and Marc'Aurelio Ranzato},
+  journal = {International Conference on Machine Learning},
+  year = 2019,
+}
+```

omni_speech/infer/fairseq/examples/translation_moe/score.py

@@ -0,0 +1,197 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+"""
+Scoring script for computing pairwise BLEU and multi-ref BLEU over a set of
+candidate hypotheses.
+
+See `"Mixture Models for Diverse Machine Translation: Tricks of the Trade"
+(Shen et al., 2019) <https://arxiv.org/abs/1902.07816>`_.
+"""
+
+import argparse
+import random
+import sys
+from itertools import chain
+
+import numpy as np
+import sacrebleu
+from sacrebleu import corpus_bleu as _corpus_bleu
+
+def main():
+    parser = argparse.ArgumentParser(sys.argv[0])
+    parser.add_argument(
+        "--sys", nargs="*", default="", metavar="FILE", help="path to system output"
+    )
+    parser.add_argument("--ref", default="", metavar="FILE", help="path to references")
+    parser.add_argument(
+        "--output",
+        default="",
+        metavar="FILE",
+        help="print outputs into a pretty format",
+    )
+    args = parser.parse_args()
+
+    if args.sys:
+        src, tgt, hypos, log_probs = load_sys(args.sys)
+        print("pairwise BLEU: %.2f" % pairwise(hypos))
+        if args.output:
+            merge(src, tgt, hypos, log_probs, args.output)
+
+    if args.ref:
+        _, _, refs = load_ref(args.ref)
+        if args.sys:
+            multi_ref(refs, hypos)
+        else:
+            intra_ref(refs)
+
+
+def dictolist(d):
+    a = sorted(d.items(), key=lambda i: i[0])
+    return [i[1] for i in a]
+
+
+def load_sys(paths):
+    src, tgt, hypos, log_probs = {}, {}, {}, {}
+    for path in paths:
+        with open(path) as f:
+            for line in f:
+                line = line.rstrip()
+                # S: source
+                # T: target
+                # D: detokenized system output
+                if line.startswith(("S-", "T-", "D-")):
+                    i = int(line[line.find("-") + 1 : line.find("\t")])
+                    if line.startswith("S-"):
+                        src[i] = line.split("\t")[1]
+                    if line.startswith("T-"):
+                        tgt[i] = line.split("\t")[1]
+                    if line.startswith("D-"):
+                        if i not in hypos:
+                            hypos[i] = []
+                            log_probs[i] = []
+                        hypos[i].append(line.split("\t")[2])
+                        log_probs[i].append(float(line.split("\t")[1]))
+    return dictolist(src), dictolist(tgt), dictolist(hypos), dictolist(log_probs)
+
+
+def load_ref(path):
+    with open(path) as f:
+        lines = f.readlines()
+    src, tgt, refs = [], [], []
+    i = 0
+    while i < len(lines):
+        if lines[i].startswith("S-"):
+            src.append(lines[i].split("\t")[1].rstrip())
+            i += 1
+        elif lines[i].startswith("T-"):
+            tgt.append(lines[i].split("\t")[1].rstrip())
+            i += 1
+        else:
+            a = []
+            while i < len(lines) and lines[i].startswith("R"):
+                a.append(lines[i].split("\t")[1].rstrip())
+                i += 1
+            refs.append(a)
+    return src, tgt, refs
+
+
+def merge(src, tgt, hypos, log_probs, path):
+    with open(path, "w") as f:
+        for s, t, hs, lps in zip(src, tgt, hypos, log_probs):
+            f.write(s + "\n")
+            f.write(t + "\n")
+            f.write("\n")
+            for h, lp in zip(hs, lps):
+                f.write("\t%f\t%s\n" % (lp, h.strip()))
+            f.write("------------------------------------------------------\n")
+
+
+def corpus_bleu(sys_stream, ref_streams):
+    bleu = _corpus_bleu(sys_stream, ref_streams, tokenize="none")
+    return bleu.score
+
+
+def sentence_bleu(hypothesis, reference):
+    bleu = _corpus_bleu(hypothesis, reference)
+    for i in range(1, 4):
+        bleu.counts[i] += 1
+        bleu.totals[i] += 1
+    bleu = sacrebleu.BLEU.compute_bleu(
+        bleu.counts,
+        bleu.totals,
+        bleu.sys_len,
+        bleu.ref_len,
+        smooth_method="exp",
+    )
+    return bleu.score
+
+
+def pairwise(sents):
+    _ref, _hypo = [], []
+    for s in sents:
+        for i in range(len(s)):
+            for j in range(len(s)):
+                if i != j:
+                    _ref.append(s[i])
+                    _hypo.append(s[j])
+    return corpus_bleu(_hypo, [_ref])
+
+
+def multi_ref(refs, hypos):
+    _ref, _hypo = [], []
+    ref_cnt = 0
+    assert len(refs) == len(hypos)
+
+    # count number of refs covered
+    for rs, hs in zip(refs, hypos):
+        a = set()
+        for h in hs:
+            s = [sentence_bleu(h, r) for r in rs]
+            j = np.argmax(s)
+            _ref.append(rs[j])
+            _hypo.append(h)
+            best = [k for k in range(len(rs)) if s[k] == s[j]]
+            a.add(random.choice(best))
+        ref_cnt += len(a)
+    print("#refs covered: %.2f" % (ref_cnt / len(refs)))
+
+    # transpose refs and hypos
+    refs = list(zip(*refs))
+    hypos = list(zip(*hypos))
+
+    # compute multi-ref corpus BLEU (leave-one-out to be comparable to intra_ref)
+    k = len(hypos)
+    m = len(refs)
+    flat_hypos = [hypos[j][i] for i in range(len(hypos[0])) for j in range(k)]
+    duplicated_refs = [[ref for ref in refs_i for _ in range(k)] for refs_i in refs]
+    loo_bleus = []
+    for held_out_ref in range(m):
+        remaining_refs = (
+            duplicated_refs[:held_out_ref] + duplicated_refs[held_out_ref + 1 :]
+        )
+        assert len(remaining_refs) == m - 1
+        loo_bleus.append(corpus_bleu(flat_hypos, remaining_refs))
+    print("average multi-reference BLEU (leave-one-out): %.2f" % np.mean(loo_bleus))
+
+
+def intra_ref(refs):
+    print("ref pairwise BLEU: %.2f" % pairwise(refs))
+    refs = list(zip(*refs))
+    m = len(refs)
+    concat_h = []
+    concat_rest = [[] for j in range(m - 1)]
+    for i, h in enumerate(refs):
+        rest = refs[:i] + refs[i + 1 :]
+        concat_h.append(h)
+        for j in range(m - 1):
+            concat_rest[j].extend(rest[j])
+    concat_h = list(chain.from_iterable(concat_h))
+    bleu = corpus_bleu(concat_h, concat_rest)
+    print("multi-reference BLEU (leave-one-out): %.2f" % bleu)
+
+
+if __name__ == "__main__":
+    main()

omni_speech/infer/fairseq/examples/translation_moe/translation_moe_src/__init__.py

@@ -0,0 +1,6 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from . import translation_moe  # noqa

omni_speech/infer/fairseq/examples/translation_moe/translation_moe_src/logsumexp_moe.py

@@ -0,0 +1,26 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+
+
+class LogSumExpMoE(torch.autograd.Function):
+    """Standard LogSumExp forward pass, but use *posterior* for the backward.
+
+    See `"Mixture Models for Diverse Machine Translation: Tricks of the Trade"
+    (Shen et al., 2019) <https://arxiv.org/abs/1902.07816>`_.
+    """
+
+    @staticmethod
+    def forward(ctx, logp, posterior, dim=-1):
+        ctx.save_for_backward(posterior)
+        ctx.dim = dim
+        return torch.logsumexp(logp, dim=dim)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        (posterior,) = ctx.saved_tensors
+        grad_logp = grad_output.unsqueeze(ctx.dim) * posterior
+        return grad_logp, None, None

omni_speech/infer/fairseq/examples/translation_moe/translation_moe_src/mean_pool_gating_network.py

@@ -0,0 +1,50 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+import torch.nn.functional as F
+
+
+class MeanPoolGatingNetwork(torch.nn.Module):
+    """A simple mean-pooling gating network for selecting experts.
+
+    This module applies mean pooling over an encoder's output and returns
+    reponsibilities for each expert. The encoder format is expected to match
+    :class:`fairseq.models.transformer.TransformerEncoder`.
+    """
+
+    def __init__(self, embed_dim, num_experts, dropout=None):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.num_experts = num_experts
+
+        self.fc1 = torch.nn.Linear(embed_dim, embed_dim)
+        self.dropout = torch.nn.Dropout(dropout) if dropout is not None else None
+        self.fc2 = torch.nn.Linear(embed_dim, num_experts)
+
+    def forward(self, encoder_out):
+        if not (
+            "encoder_out" in encoder_out
+            and "encoder_padding_mask" in encoder_out
+            and encoder_out["encoder_out"][0].size(2) == self.embed_dim
+        ):
+            raise ValueError("Unexpected format for encoder_out")
+
+        # mean pooling over time
+        encoder_padding_mask = encoder_out["encoder_padding_mask"][0]  # B x T
+        encoder_out = encoder_out["encoder_out"][0].transpose(0, 1)    # B x T x C
+        if encoder_padding_mask is not None:
+            encoder_out = encoder_out.clone()  # required because of transpose above
+            encoder_out[encoder_padding_mask] = 0
+            ntokens = torch.sum(~encoder_padding_mask, dim=1, keepdim=True)
+            x = torch.sum(encoder_out, dim=1) / ntokens.type_as(encoder_out)
+        else:
+            x = torch.mean(encoder_out, dim=1)
+
+        x = torch.tanh(self.fc1(x))
+        if self.dropout is not None:
+            x = self.dropout(x)
+        x = self.fc2(x)
+        return F.log_softmax(x, dim=-1, dtype=torch.float32).type_as(x)

omni_speech/infer/fairseq/examples/translation_moe/translation_moe_src/translation_moe.py

@@ -0,0 +1,259 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from dataclasses import dataclass, field
+import torch
+from omegaconf import II
+
+from fairseq import utils
+from fairseq.logging import metrics
+from fairseq.dataclass import ChoiceEnum
+from fairseq.tasks import register_task
+from fairseq.tasks.translation import TranslationConfig, TranslationTask
+
+from .logsumexp_moe import LogSumExpMoE
+from .mean_pool_gating_network import MeanPoolGatingNetwork
+
+
+METHOD_CHOICES = ChoiceEnum(["sMoElp", "sMoEup", "hMoElp", "hMoEup"])
+
+
+@dataclass
+class TranslationMoEConfig(TranslationConfig):
+    method: METHOD_CHOICES = field(
+        default="hMoEup",
+        metadata={"help": "MoE method"},
+    )
+    num_experts: int = field(
+        default=3,
+        metadata={"help": "number of experts"},
+    )
+    mean_pool_gating_network: bool = field(
+        default=False,
+        metadata={"help": "use a simple mean-pooling gating network"},
+    )
+    mean_pool_gating_network_dropout: float = field(
+        default=0,
+        metadata={"help": "dropout for mean-pooling gating network"},
+    )
+    mean_pool_gating_network_encoder_dim: int = field(
+        default=0,
+        metadata={"help": "encoder output dim for mean-pooling gating network"},
+    )
+    gen_expert: int = field(
+        default=0,
+        metadata={"help": "which expert to use for generation"},
+    )
+    sentence_avg: bool = II("optimization.sentence_avg")
+
+
+@register_task("translation_moe", dataclass=TranslationMoEConfig)
+class TranslationMoETask(TranslationTask):
+    """
+    Translation task for Mixture of Experts (MoE) models.
+
+    See `"Mixture Models for Diverse Machine Translation: Tricks of the Trade"
+    (Shen et al., 2019) <https://arxiv.org/abs/1902.07816>`_.
+
+    Args:
+        src_dict (~fairseq.data.Dictionary): dictionary for the source language
+        tgt_dict (~fairseq.data.Dictionary): dictionary for the target language
+
+    .. note::
+
+        The translation task is compatible with :mod:`fairseq-train`,
+        :mod:`fairseq-generate` and :mod:`fairseq-interactive`.
+
+    The translation task provides the following additional command-line
+    arguments:
+
+    .. argparse::
+        :ref: fairseq.tasks.translation_parser
+        :prog:
+    """
+
+    cfg: TranslationMoEConfig
+
+    def __init__(self, cfg: TranslationMoEConfig, src_dict, tgt_dict):
+        if cfg.method == "sMoElp":
+            # soft MoE with learned prior
+            self.uniform_prior = False
+            self.hard_selection = False
+        elif cfg.method == "sMoEup":
+            # soft MoE with uniform prior
+            self.uniform_prior = True
+            self.hard_selection = False
+        elif cfg.method == "hMoElp":
+            # hard MoE with learned prior
+            self.uniform_prior = False
+            self.hard_selection = True
+        elif cfg.method == "hMoEup":
+            # hard MoE with uniform prior
+            self.uniform_prior = True
+            self.hard_selection = True
+
+        # add indicator tokens for each expert
+        for i in range(cfg.num_experts):
+            # add to both dictionaries in case we're sharing embeddings
+            src_dict.add_symbol("<expert_{}>".format(i))
+            tgt_dict.add_symbol("<expert_{}>".format(i))
+
+        super().__init__(cfg, src_dict, tgt_dict)
+
+    def build_model(self, cfg, from_checkpoint=False):
+        from fairseq import models
+
+        model = models.build_model(cfg, self)
+        if not self.uniform_prior and not hasattr(model, "gating_network"):
+            if self.cfg.mean_pool_gating_network:
+                if self.cfg.mean_pool_gating_network_encoder_dim > 0:
+                    encoder_dim = self.cfg.mean_pool_gating_network_encoder_dim
+                elif getattr(cfg, "encoder_embed_dim", None):
+                    # assume that encoder_embed_dim is the encoder's output dimension
+                    encoder_dim = cfg.encoder_embed_dim
+                else:
+                    raise ValueError(
+                        "Must specify --mean-pool-gating-network-encoder-dim"
+                    )
+
+                if self.cfg.mean_pool_gating_network_dropout > 0:
+                    dropout = self.cfg.mean_pool_gating_network_dropout
+                elif getattr(cfg, "dropout", None):
+                    dropout = cfg.dropout
+                else:
+                    raise ValueError("Must specify task.mean_pool_gating_network_dropout")
+
+                model.gating_network = MeanPoolGatingNetwork(
+                    encoder_dim,
+                    self.cfg.num_experts,
+                    dropout,
+                )
+            else:
+                raise ValueError(
+                    "translation_moe task with learned prior requires the model to "
+                    "have a gating network; try using --mean-pool-gating-network"
+                )
+        return model
+
+    def expert_index(self, i):
+        return i + self.tgt_dict.index("<expert_0>")
+
+    def _get_loss(self, sample, model, criterion):
+        assert hasattr(
+            criterion, "compute_loss"
+        ), "translation_moe task requires the criterion to implement the compute_loss() method"
+
+        k = self.cfg.num_experts
+        bsz = sample["target"].size(0)
+
+        def get_lprob_y(encoder_out, prev_output_tokens_k):
+            net_output = model.decoder(
+                prev_output_tokens=prev_output_tokens_k,
+                encoder_out=encoder_out,
+            )
+            loss, _ = criterion.compute_loss(model, net_output, sample, reduce=False)
+            loss = loss.view(bsz, -1)
+            return -loss.sum(dim=1, keepdim=True)  # -> B x 1
+
+        def get_lprob_yz(winners=None):
+            encoder_out = model.encoder(
+                src_tokens=sample["net_input"]["src_tokens"],
+                src_lengths=sample["net_input"]["src_lengths"],
+            )
+
+            if winners is None:
+                lprob_y = []
+                for i in range(k):
+                    prev_output_tokens_k = sample["net_input"][
+                        "prev_output_tokens"
+                    ].clone()
+                    assert not prev_output_tokens_k.requires_grad
+                    prev_output_tokens_k[:, 0] = self.expert_index(i)
+                    lprob_y.append(get_lprob_y(encoder_out, prev_output_tokens_k))
+                lprob_y = torch.cat(lprob_y, dim=1)  # -> B x K
+            else:
+                prev_output_tokens_k = sample["net_input"]["prev_output_tokens"].clone()
+                prev_output_tokens_k[:, 0] = self.expert_index(winners)
+                lprob_y = get_lprob_y(encoder_out, prev_output_tokens_k)  # -> B
+
+            if self.uniform_prior:
+                lprob_yz = lprob_y
+            else:
+                lprob_z = model.gating_network(encoder_out)  # B x K
+                if winners is not None:
+                    lprob_z = lprob_z.gather(dim=1, index=winners.unsqueeze(-1))
+                lprob_yz = lprob_y + lprob_z.type_as(lprob_y)  # B x K
+
+            return lprob_yz
+
+        # compute responsibilities without dropout
+        with utils.model_eval(model):  # disable dropout
+            with torch.no_grad():  # disable autograd
+                lprob_yz = get_lprob_yz()  # B x K
+                prob_z_xy = torch.nn.functional.softmax(lprob_yz, dim=1)
+        assert not prob_z_xy.requires_grad
+
+        # compute loss with dropout
+        if self.hard_selection:
+            winners = prob_z_xy.max(dim=1)[1]
+            loss = -get_lprob_yz(winners)
+        else:
+            lprob_yz = get_lprob_yz()  # B x K
+            loss = -LogSumExpMoE.apply(lprob_yz, prob_z_xy, 1)
+
+        loss = loss.sum()
+        sample_size = (
+            sample["target"].size(0) if self.cfg.sentence_avg else sample["ntokens"]
+        )
+        logging_output = {
+            "loss": utils.item(loss.data),
+            "ntokens": sample["ntokens"],
+            "nsentences": bsz,
+            "sample_size": sample_size,
+            "posterior": prob_z_xy.float().sum(dim=0).cpu(),
+        }
+        return loss, sample_size, logging_output
+
+    def train_step(
+        self, sample, model, criterion, optimizer, update_num, ignore_grad=False
+    ):
+        model.train()
+        loss, sample_size, logging_output = self._get_loss(sample, model, criterion)
+        if ignore_grad:
+            loss *= 0
+        optimizer.backward(loss)
+        return loss, sample_size, logging_output
+
+    def valid_step(self, sample, model, criterion):
+        model.eval()
+        with torch.no_grad():
+            loss, sample_size, logging_output = self._get_loss(sample, model, criterion)
+        return loss, sample_size, logging_output
+
+    def inference_step(
+        self,
+        generator,
+        models,
+        sample,
+        prefix_tokens=None,
+        expert=None,
+        constraints=None,
+    ):
+        expert = expert or self.cfg.gen_expert
+        with torch.no_grad():
+            return generator.generate(
+                models,
+                sample,
+                prefix_tokens=prefix_tokens,
+                constraints=constraints,
+                bos_token=self.expert_index(expert),
+            )
+
+    def reduce_metrics(self, logging_outputs, criterion):
+        super().reduce_metrics(logging_outputs, criterion)
+        metrics.log_scalar(
+            "posterior",
+            sum(log["posterior"] for log in logging_outputs if "posterior" in log),
+        )

omni_speech/infer/fairseq/examples/truncated_bptt/README.md

@@ -0,0 +1,70 @@
+# Truncated Backpropagation Through Time (BPTT)
+
+Truncated BPTT is a useful technique for training language models on very long
+sequences. Typically a long sequences is split into chunks and a language model
+is trained over the chunks sequentially. The LM may condition on previous
+chunks, but gradients only flow through the current chunk. This technique was
+the basis for the paper: [Transformer-XL: Attentive Language Models Beyond a
+Fixed-Length Context](https://arxiv.org/abs/1901.02860), which achieved
+state-of-the-art language modeling results at the time of publication.
+
+It is slightly tricky to implement Truncated BPTT efficiently in fairseq, since
+we need to iterate over the data sequentially and disable any batch shuffling
+logic. The code provided in this example illustrates how to implement Truncated
+BPTT in fairseq by overriding ``FairseqTask::get_batch_iterator`` to iterate
+over the data sequentially. Crucially, this example supports batching and
+multi-GPU (data parallel) training.
+
+##### 0. Setup
+
+First, see the general [language modeling README](README.md) for instructions on
+preprocessing the WikiText-103 data.
+
+##### 1. Train a Transformer-XL model on WikiText-103
+
+We will train a 16-layer Transformer-XL model following the [hyperparameters
+used in the original
+paper](https://github.com/kimiyoung/transformer-xl/blob/master/pytorch/run_wt103_base.sh).
+
+The following command assumes 4 GPUs, so that the total batch size is 60
+sequences (15 x 4). Training should take ~24 hours on 4 V100 GPUs:
+```bash
+CUDA_VISIBLE_DEVICES=0,1,2,3 fairseq-train \
+    --user-dir examples/truncated_bptt \
+    data-bin/wikitext-103/ \
+    --task truncated_bptt_lm --tokens-per-sample 150 \
+    --batch-size 15 --max-update 200000 \
+    --arch transformer_xl --n-layer 16 --d-model 410 --n-head 10 \
+    --d-head 41 --d-inner 2100 --dropout 0.1 --dropatt 0.0 --mem-len 150 \
+    --optimizer adam --clip-norm 0.25 \
+    --lr-scheduler cosine --warmup-updates 0 --min-lr 0.0 --lr 0.00025  \
+    --log-format json --log-interval 25 \
+    --fp16
+```
+
+If training on a single GPU, set `--update-freq=4` to accumulate 4x gradients
+and simulate training on 4 GPUs.
+
+##### 2. Evaluate
+
+```bash
+fairseq-eval-lm data-bin/wikitext-103/ \
+    --path checkpoints/checkpoint_best.pt \
+    --user-dir examples/truncated_bptt/ \
+    --task truncated_bptt_lm \
+    --batch-size 1 --required-batch-size-multiple 1 \
+    --model-overrides '{"mem_len":640,"clamp_len":400,"same_length":True}' \
+    --tokens-per-sample 64
+# ... | INFO | fairseq_cli.eval_lm | num. model params: 151123537
+# ... | INFO | fairseq_cli.eval_lm | Evaluated 245569 tokens in 83.1s (2956.82 tokens/s)
+# ... | INFO | fairseq_cli.eval_lm | Loss (base 2): 4.5668, Perplexity: 23.70
+# Compare to 24.0 test perplexity from the paper
+```
+
+*Note:* During training the model saw 150 tokens of context
+(``--tokens-per-sample=150``) and 150 extra memory tokens (``--mem-len=150``).
+During evaluation we measure perplexity on sequences of 64 tokens
+(``--tokens-per-sample=64``) and increase the memory length
+(``--model-overrides='{"mem_len":640}'``). These settings match the evaluation
+settings from [the original
+paper](https://github.com/kimiyoung/transformer-xl/blob/master/pytorch/run_wt103_base.sh).

omni_speech/infer/fairseq/examples/truncated_bptt/__init__.py

@@ -0,0 +1,6 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from . import transformer_xl_model, truncated_bptt_lm_task  # noqa

omni_speech/infer/fairseq/examples/truncated_bptt/transformer_xl_model.py

@@ -0,0 +1,143 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import logging
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional
+
+import torch
+from fairseq.dataclass import FairseqDataclass
+from fairseq.models import (
+    FairseqIncrementalDecoder,
+    FairseqLanguageModel,
+    register_model,
+)
+from fairseq.modules.checkpoint_activations import checkpoint_wrapper
+from omegaconf import II
+
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class TransformerXLConfig(FairseqDataclass):
+    # defaults come from the original Transformer-XL code
+    cutoffs: List[int] = field(default_factory=lambda: [20000, 40000, 200000])
+    d_model: int = 500
+    n_head: int = 10
+    d_head: int = 50
+    d_inner: int = 1000
+    div_val: int = 1
+    n_layer: int = 12
+    mem_len: int = 0
+    clamp_len: int = -1
+    same_length: bool = False
+    dropout: float = 0.0
+    dropatt: float = 0.0
+    checkpoint_activations: bool = False
+    offload_activations: bool = False
+    max_target_positions: int = II("task.max_target_positions")
+
+
+@register_model("transformer_xl", dataclass=TransformerXLConfig)
+class TransformerXLLanguageModel(FairseqLanguageModel):
+    @classmethod
+    def build_model(cls, cfg: TransformerXLConfig, task):
+        return cls(TransformerXLDecoder(cfg, task))
+
+
+class TransformerXLDecoder(FairseqIncrementalDecoder):
+    def __init__(self, cfg, task):
+        try:
+            from transformers.models.transfo_xl import (
+                TransfoXLConfig,
+                TransfoXLLMHeadModel,
+            )
+        except ImportError:
+            from transformers.configuration_transfo_xl import TransfoXLConfig
+            from transformers.modeling_transfo_xl import TransfoXLLMHeadModel
+
+        super().__init__(task.target_dictionary)
+        self.cfg = cfg
+
+        # remove any cutoffs larger than the vocab size
+        cutoffs = [
+            cutoff for cutoff in cfg.cutoffs if cutoff < len(task.target_dictionary)
+        ]
+
+        config = TransfoXLConfig(
+            vocab_size=len(task.target_dictionary),
+            cutoffs=cutoffs,
+            d_model=cfg.d_model,
+            d_embed=cfg.d_model,
+            n_head=cfg.n_head,
+            d_head=cfg.d_head,
+            d_inner=cfg.d_inner,
+            div_val=cfg.div_val,
+            n_layer=cfg.n_layer,
+            mem_len=cfg.mem_len,
+            clamp_len=cfg.clamp_len,
+            same_length=cfg.same_length,
+            dropout=cfg.dropout,
+            dropatt=cfg.dropatt,
+        )
+        logger.info(config)
+        self.model = TransfoXLLMHeadModel(config)
+
+        if cfg.checkpoint_activations or cfg.offload_activations:
+            for i in range(len(self.model.transformer.layers)):
+                self.model.transformer.layers[i] = checkpoint_wrapper(
+                    self.model.transformer.layers[i],
+                    offload_to_cpu=cfg.offload_activations,
+                )
+                # TODO: may save mem to wrap(layer.pos_ff.CoreNet[3])
+
+        self._mems = None
+
+    def forward(
+        self,
+        src_tokens,
+        src_lengths=None,  # unused
+        incremental_state: Optional[Dict[str, List[torch.Tensor]]] = None,
+        encoder_out=None,
+    ):
+        if incremental_state is not None:  # used during inference
+            mems = self.get_incremental_state(incremental_state, "mems")
+            src_tokens = src_tokens[:, -1:]  # only keep the most recent token
+        else:
+            mems = self._mems
+
+        output = self.model(
+            input_ids=src_tokens,
+            mems=mems,
+            return_dict=False,
+        )
+
+        if len(output) >= 2:
+            if incremental_state is not None:
+                self.set_incremental_state(incremental_state, "mems", output[1])
+            else:
+                self._mems = output[1]
+
+        return (output[0],)
+
+    def max_positions(self):
+        return self.cfg.max_target_positions
+
+    def reorder_incremental_state(
+        self,
+        incremental_state: Dict[str, Dict[str, Optional[torch.Tensor]]],
+        new_order: torch.Tensor,
+    ):
+        """Reorder incremental state.
+
+        This will be called when the order of the input has changed from the
+        previous time step. A typical use case is beam search, where the input
+        order changes between time steps based on the selection of beams.
+        """
+        mems = self.get_incremental_state(incremental_state, "mems")
+        if mems is not None:
+            new_mems = [mems_i.index_select(1, new_order) for mems_i in mems]
+            self.set_incremental_state(incremental_state, "mems", new_mems)

omni_speech/infer/fairseq/examples/truncated_bptt/truncated_bptt_lm_task.py

@@ -0,0 +1,285 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import logging
+import os
+from dataclasses import dataclass, field
+from typing import List, Optional, Tuple
+
+import torch
+from fairseq import utils
+from fairseq.data import (
+    Dictionary,
+    TokenBlockDataset,
+    data_utils,
+    iterators,
+)
+from fairseq.dataclass import FairseqDataclass
+from fairseq.distributed import utils as dist_utils
+from fairseq.tasks import FairseqTask, register_task
+from omegaconf import II
+
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class TruncatedBPTTLMConfig(FairseqDataclass):
+    data: str = field(default="???", metadata={"help": "path to data directory"})
+    tokens_per_sample: int = field(
+        default=1024, metadata={"help": "max number of tokens per sequence"},
+    )
+    batch_size: int = II("dataset.batch_size")
+    # Some models use *max_target_positions* to know how many positional
+    # embeddings to learn. We use II(...) to make it default to
+    # *tokens_per_sample*, but in principle there could be more positional
+    # embeddings than tokens in a single batch. This may also be irrelevant for
+    # custom model implementations.
+    max_target_positions: int = II("task.tokens_per_sample")
+    # these will be populated automatically if not provided
+    data_parallel_rank: Optional[int] = None
+    data_parallel_size: Optional[int] = None
+
+
+@register_task("truncated_bptt_lm", dataclass=TruncatedBPTTLMConfig)
+class TruncatedBPTTLMTask(FairseqTask):
+    def __init__(self, cfg: TruncatedBPTTLMConfig):
+        super().__init__(cfg)
+
+        if cfg.data_parallel_rank is None or cfg.data_parallel_size is None:
+            if torch.distributed.is_initialized():
+                cfg.data_parallel_rank = dist_utils.get_data_parallel_rank()
+                cfg.data_parallel_size = dist_utils.get_data_parallel_world_size()
+            else:
+                cfg.data_parallel_rank = 0
+                cfg.data_parallel_size = 1
+
+        # load the dictionary
+        paths = utils.split_paths(cfg.data)
+        assert len(paths) > 0
+        self.dictionary = Dictionary.load(os.path.join(paths[0], "dict.txt"))
+        logger.info("dictionary: {} types".format(len(self.dictionary)))
+
+    def load_dataset(self, split, epoch=1, combine=False, **kwargs):
+        """Load a given dataset split (e.g., train, valid, test)"""
+
+        # support sharded datasets
+        paths = utils.split_paths(self.cfg.data)
+        assert len(paths) > 0
+        data_path = paths[(epoch - 1) % len(paths)]
+        split_path = os.path.join(data_path, split)
+
+        # each element of *data* will be a tensorized line from the original
+        # text dataset, similar to ``open(split_path).readlines()``
+        data = data_utils.load_indexed_dataset(
+            split_path, self.dictionary, combine=combine
+        )
+        if data is None:
+            raise FileNotFoundError(
+                "Dataset not found: {} ({})".format(split, split_path)
+            )
+
+        # this is similar to ``data.view(-1).split(tokens_per_sample)``
+        data = TokenBlockDataset(
+            data,
+            data.sizes,
+            block_size=self.cfg.tokens_per_sample,
+            pad=None,  # unused
+            eos=None,  # unused
+            break_mode="none",
+        )
+
+        self.datasets[split] = TruncatedBPTTDataset(
+            data=data,
+            bsz_per_shard=self.cfg.batch_size,
+            shard_id=self.cfg.data_parallel_rank,
+            num_shards=self.cfg.data_parallel_size,
+        )
+
+    def dataset(self, split):
+        return self.datasets[split]
+
+    def get_batch_iterator(
+        self,
+        dataset,
+        num_workers=0,
+        epoch=1,
+        data_buffer_size=0,
+        skip_remainder_batch=False,
+        **kwargs
+    ):
+        return iterators.EpochBatchIterator(
+            dataset=dataset,
+            collate_fn=self._collate_fn,
+            num_workers=num_workers,
+            epoch=epoch,
+            buffer_size=data_buffer_size,
+            # we don't use the batching functionality from EpochBatchIterator;
+            # instead every item in *dataset* is a whole batch
+            batch_sampler=[[i] for i in range(len(dataset))],
+            disable_shuffling=True,
+            skip_remainder_batch=skip_remainder_batch,
+        )
+
+    def _collate_fn(self, items: List[List[torch.Tensor]]):
+        # we don't use fairseq's batching functionality, so we expect a single
+        # Tensor of type List[torch.Tensor]
+        assert len(items) == 1
+
+        # item will have shape B x T (the last batch may have length < T)
+        id, item = items[0]
+        item = data_utils.collate_tokens(item, pad_idx=self.source_dictionary.pad())
+        B, T = item.size()
+
+        # shift item one position over and append a padding token for the target
+        target = torch.nn.functional.pad(
+            item[:, 1:], (0, 1, 0, 0), value=self.target_dictionary.pad()
+        )
+
+        # fairseq expects batches to have the following structure
+        return {
+            "id": torch.tensor([id] * item.size(0)),
+            "net_input": {"src_tokens": item,},
+            "target": target,
+            "nsentences": item.size(0),
+            "ntokens": item.numel(),
+        }
+
+    def build_dataset_for_inference(
+        self, src_tokens: List[torch.Tensor], src_lengths: List[int], **kwargs
+    ) -> torch.utils.data.Dataset:
+        eos = self.source_dictionary.eos()
+        dataset = TokenBlockDataset(
+            src_tokens,
+            src_lengths,
+            block_size=None,  # ignored for "eos" break mode
+            pad=self.source_dictionary.pad(),
+            eos=eos,
+            break_mode="eos",
+        )
+
+        class Dataset(torch.utils.data.Dataset):
+            def __getitem__(self, i):
+                item = dataset[i]
+                if item[-1] == eos:
+                    # remove eos to support generating with a prefix
+                    item = item[:-1]
+                return (i, [item])
+
+            def __len__(self):
+                return len(dataset)
+
+        return Dataset()
+
+    def inference_step(
+        self, generator, models, sample, prefix_tokens=None, constraints=None
+    ):
+        with torch.no_grad():
+            if constraints is not None:
+                raise NotImplementedError
+
+            # SequenceGenerator doesn't use *src_tokens* directly, we need to
+            # pass the *prefix_tokens* argument instead.
+            if prefix_tokens is None and sample["net_input"]["src_tokens"].nelement():
+                prefix_tokens = sample["net_input"]["src_tokens"]
+
+            # begin generation with the end-of-sentence token
+            bos_token = self.source_dictionary.eos()
+
+            return generator.generate(
+                models, sample, prefix_tokens=prefix_tokens, bos_token=bos_token
+            )
+
+    def eval_lm_dataloader(
+        self,
+        dataset,
+        max_tokens: Optional[int] = 36000,
+        batch_size: Optional[int] = None,
+        max_positions: Optional[int] = None,
+        num_shards: int = 1,
+        shard_id: int = 0,
+        num_workers: int = 1,
+        data_buffer_size: int = 10,
+        context_window: int = 0,
+    ):
+        if context_window > 0:
+            raise NotImplementedError(
+                "Transformer-XL doesn't need --context-window, try "
+                "--model-overrides '{\"mem_len\":42}' instead "
+            )
+        return self.get_batch_iterator(
+            dataset=dataset,
+            max_tokens=max_tokens,
+            max_sentences=batch_size,
+            max_positions=max_positions,
+            ignore_invalid_inputs=True,
+            num_shards=num_shards,
+            shard_id=shard_id,
+            num_workers=num_workers,
+            data_buffer_size=data_buffer_size,
+        ).next_epoch_itr(shuffle=False)
+
+    @property
+    def source_dictionary(self):
+        return self.dictionary
+
+    @property
+    def target_dictionary(self):
+        return self.dictionary
+
+
+class TruncatedBPTTDataset(torch.utils.data.Dataset):
+    def __init__(
+        self,
+        data: List[torch.Tensor],  # ordered list of items
+        bsz_per_shard,  # number of items processed per GPUs per forward
+        shard_id,  # current GPU ID
+        num_shards,  # number of GPUs
+    ):
+        super().__init__()
+        self.data = data
+
+        def batchify(data, bsz):
+            # Work out how cleanly we can divide the dataset into bsz parts.
+            nbatch = data.size(0) // bsz
+            # Trim off any extra elements that wouldn't cleanly fit (remainders).
+            data = data.narrow(0, 0, nbatch * bsz)
+            # Evenly divide the data across the bsz batches.
+            data = data.view(bsz, -1).contiguous()
+            return data
+
+        # total number of sequences processed by all GPUs in each forward pass
+        global_batch_size = bsz_per_shard * num_shards
+
+        """
+        With a 16 item dataset, bsz_per_shard=2 and num_shards=3,
+        *indices* might look like:
+
+            indices = [[0, 1],
+                       [2, 3],
+                       [4, 5],
+                       [6, 7],
+                       [8, 9],
+                       [10, 11]]
+
+        The size of the TruncatedBPTTDataset instance will be 2,
+        and shard 1 will see items:
+
+            [(0, [data[4], data[6]]),
+             (1, [data[5], data[7]])]
+        """
+        indices = batchify(torch.arange(len(data)), global_batch_size)
+        assert indices.size(0) == global_batch_size
+
+        self.my_indices = indices[
+            shard_id * bsz_per_shard : (shard_id + 1) * bsz_per_shard
+        ]
+        assert self.my_indices.size(0) == bsz_per_shard
+
+    def __len__(self):
+        return self.my_indices.size(1)
+
+    def __getitem__(self, i) -> Tuple[int, List[torch.Tensor]]:
+        return (i, [self.data[idx] for idx in self.my_indices[:, i]])

omni_speech/infer/fairseq/examples/unsupervised_quality_estimation/README.md

@@ -0,0 +1,126 @@
+# Unsupervised Quality Estimation for Neural Machine Translation (Fomicheva et al., 2020)
+
+This page includes instructions for reproducing results from the paper [Unsupervised Quality Estimation for Neural
+Machine Translation (Fomicheva et al., 2020)](https://arxiv.org/abs/2005.10608)
+
+## Requirements:
+
+* mosesdecoder: https://github.com/moses-smt/mosesdecoder
+* subword-nmt: https://github.com/rsennrich/subword-nmt
+* flores: https://github.com/facebookresearch/flores
+
+## Download Models and Test Data
+
+Download translation models and test data from [MLQE dataset repository](https://github.com/facebookresearch/mlqe).
+
+## Set up:
+
+Given a testset consisting of source sentences and reference translations:
+
+* `SRC_LANG`: source language
+* `TGT_LANG`: target language
+* `INPUT`: input prefix, such that the file `$INPUT.$SRC_LANG` contains source sentences and `$INPUT.$TGT_LANG`
+contains the reference sentences
+* `OUTPUT_DIR`: output path to store results
+* `MOSES_DECODER`: path to mosesdecoder installation
+* `BPE_ROOT`: path to subword-nmt installation
+* `BPE`: path to BPE model
+* `MODEL_DIR`: directory containing the NMT model `.pt` file as well as the source and target vocabularies.
+* `TMP`: directory for intermediate temporary files
+* `GPU`: if translating with GPU, id of the GPU to use for inference
+* `DROPOUT_N`: number of stochastic forward passes
+
+`$DROPOUT_N` is set to 30 in the experiments reported in the paper. However, we observed that increasing it beyond 10
+does not bring substantial improvements.
+
+## Translate the data using standard decoding
+
+Preprocess the input data:
+```
+for LANG in $SRC_LANG $TGT_LANG; do
+  perl $MOSES_DECODER/scripts/tokenizer/tokenizer.perl -threads 80 -a -l $LANG < $INPUT.$LANG > $TMP/preprocessed.tok.$LANG
+  python $BPE_ROOT/apply_bpe.py -c ${BPE} < $TMP/preprocessed.tok.$LANG > $TMP/preprocessed.tok.bpe.$LANG
+done
+```
+
+Binarize the data for faster translation:
+
+```
+fairseq-preprocess --srcdict $MODEL_DIR/dict.$SRC_LANG.txt --tgtdict $MODEL_DIR/dict.$TGT_LANG.txt
+--source-lang ${SRC_LANG} --target-lang ${TGT_LANG} --testpref $TMP/preprocessed.tok.bpe --destdir $TMP/bin --workers 4
+```
+
+Translate
+
+```
+CUDA_VISIBLE_DEVICES=$GPU fairseq-generate $TMP/bin --path ${MODEL_DIR}/${SRC_LANG}-${TGT_LANG}.pt --beam 5
+--source-lang $SRC_LANG --target-lang $TGT_LANG --no-progress-bar --unkpen 5 > $TMP/fairseq.out
+grep ^H $TMP/fairseq.out | cut -d- -f2- | sort -n | cut -f3- > $TMP/mt.out
+```
+
+Post-process
+
+```
+sed -r 's/(@@ )| (@@ ?$)//g' < $TMP/mt.out | perl $MOSES_DECODER/scripts/tokenizer/detokenizer.perl
+-l $TGT_LANG > $OUTPUT_DIR/mt.out
+```
+
+## Produce uncertainty estimates
+
+### Scoring
+
+Make temporary files to store the translations repeated N times.
+
+```
+python ${SCRIPTS}/scripts/uncertainty/repeat_lines.py -i $TMP/preprocessed.tok.bpe.$SRC_LANG -n $DROPOUT_N
+-o $TMP/repeated.$SRC_LANG
+python ${SCRIPTS}/scripts/uncertainty/repeat_lines.py -i $TMP/mt.out -n $DROPOUT_N -o $TMP/repeated.$TGT_LANG
+
+fairseq-preprocess --srcdict ${MODEL_DIR}/dict.${SRC_LANG}.txt $TGT_DIC --source-lang ${SRC_LANG}
+--target-lang ${TGT_LANG} --testpref ${TMP}/repeated --destdir ${TMP}/bin-repeated
+```
+
+Produce model scores for the generated translations using `--retain-dropout` option to apply dropout at inference time:
+
+```
+CUDA_VISIBLE_DEVICES=${GPU} fairseq-generate ${TMP}/bin-repeated --path ${MODEL_DIR}/${LP}.pt --beam 5
+ --source-lang $SRC_LANG --target-lang $TGT_LANG --no-progress-bar --unkpen 5 --score-reference --retain-dropout
+ --retain-dropout-modules '["TransformerModel","TransformerEncoder","TransformerDecoder","TransformerEncoderLayer"]'
+ TransformerDecoderLayer --seed 46 > $TMP/dropout.scoring.out
+
+grep ^H $TMP/dropout.scoring.out | cut -d- -f2- | sort -n | cut -f2 > $TMP/dropout.scores
+
+```
+
+Use `--retain-dropout-modules` to specify the modules. By default, dropout is applied in the same places
+as for training.
+
+Compute the mean of the resulting output distribution:
+
+```
+python $SCRIPTS/scripts/uncertainty/aggregate_scores.py -i $TMP/dropout.scores -o $OUTPUT_DIR/dropout.scores.mean
+-n $DROPOUT_N
+```
+
+### Generation
+
+Produce multiple translation hypotheses for the same source using `--retain-dropout` option:
+
+```
+CUDA_VISIBLE_DEVICES=${GPU} fairseq-generate ${TMP}/bin-repeated --path ${MODEL_DIR}/${LP}.pt
+ --beam 5 --source-lang $SRC_LANG --target-lang $TGT_LANG --no-progress-bar --retain-dropout
+ --unkpen 5 --retain-dropout-modules TransformerModel TransformerEncoder TransformerDecoder
+TransformerEncoderLayer TransformerDecoderLayer --seed 46 > $TMP/dropout.generation.out
+
+grep ^H $TMP/dropout.generation.out | cut -d- -f2- | sort -n | cut -f3- > $TMP/dropout.hypotheses_
+
+sed -r 's/(@@ )| (@@ ?$)//g' < $TMP/dropout.hypotheses_ | perl $MOSES_DECODER/scripts/tokenizer/detokenizer.perl
+-l $TGT_LANG > $TMP/dropout.hypotheses
+```
+
+Compute similarity between multiple hypotheses corresponding to the same source sentence using Meteor
+evaluation metric:
+```
+python meteor.py -i $TMP/dropout.hypotheses -m <path_to_meteor_installation> -n $DROPOUT_N -o
+$OUTPUT_DIR/dropout.gen.sim.meteor
+```

omni_speech/infer/fairseq/examples/unsupervised_quality_estimation/aggregate_scores.py

@@ -0,0 +1,41 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import argparse
+import sys
+
+import numpy as np
+
+
+aggregate_funcs = {
+    "std": np.std,
+    "var": np.var,
+    "median": np.median,
+    "mean": np.mean,
+    "min": np.min,
+    "max": np.max,
+}
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-i", "--input_file", required=True, type=str)
+    parser.add_argument("-n", "--repeat_times", required=True, type=int)
+    parser.add_argument("-o", "--output_file", required=False)
+    parser.add_argument("-f", "--func", required=False, default="mean")
+    args = parser.parse_args()
+
+    stream = open(args.output_file, "w") if args.output_file else sys.stdout
+
+    segment_scores = []
+    for line in open(args.input_file):
+        segment_scores.append(float(line.strip()))
+        if len(segment_scores) == args.repeat_times:
+            stream.write("{}\n".format(aggregate_funcs[args.func](segment_scores)))
+            segment_scores = []
+
+
+if __name__ == "__main__":
+    main()

omni_speech/infer/fairseq/examples/unsupervised_quality_estimation/meteor.py

@@ -0,0 +1,109 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import argparse
+import math
+import os
+import subprocess
+import sys
+import tempfile
+from collections import defaultdict
+from itertools import combinations
+
+
+def read_translations(path, n_repeats):
+    segment_counter = 0
+    segment_translations = []
+    translations = defaultdict(list)
+    for line in open(path):
+        segment_translations.append(" ".join(line.split()))
+        if len(segment_translations) == n_repeats:
+            translations[segment_counter] = segment_translations
+            segment_translations = []
+            segment_counter += 1
+    return translations
+
+
+def generate_input(translations, n_repeats):
+    _, ref_path = tempfile.mkstemp()
+    _, mt_path = tempfile.mkstemp()
+    ref_fh = open(ref_path, "w")
+    mt_fh = open(mt_path, "w")
+    for segid in sorted(translations.keys()):
+        assert len(translations[segid]) == n_repeats
+        indexes = combinations(range(n_repeats), 2)
+        for idx1, idx2 in indexes:
+            mt_fh.write(translations[segid][idx1].strip() + "\n")
+            ref_fh.write(translations[segid][idx2].strip() + "\n")
+    sys.stderr.write("\nSaved translations to %s and %s" % (ref_path, mt_path))
+    return ref_path, mt_path
+
+
+def run_meteor(ref_path, mt_path, metric_path, lang="en"):
+    _, out_path = tempfile.mkstemp()
+    subprocess.call(
+        [
+            "java",
+            "-Xmx2G",
+            "-jar",
+            metric_path,
+            mt_path,
+            ref_path,
+            "-p",
+            "0.5 0.2 0.6 0.75",  # default parameters, only changed alpha to give equal weight to P and R
+            "-norm",
+            "-l",
+            lang,
+        ],
+        stdout=open(out_path, "w"),
+    )
+    os.remove(ref_path)
+    os.remove(mt_path)
+    sys.stderr.write("\nSaved Meteor output to %s" % out_path)
+    return out_path
+
+
+def read_output(meteor_output_path, n_repeats):
+    n_combinations = math.factorial(n_repeats) / (
+        math.factorial(2) * math.factorial(n_repeats - 2)
+    )
+    raw_scores = []
+    average_scores = []
+    for line in open(meteor_output_path):
+        if not line.startswith("Segment "):
+            continue
+        score = float(line.strip().split("\t")[1])
+        raw_scores.append(score)
+        if len(raw_scores) == n_combinations:
+            average_scores.append(sum(raw_scores) / n_combinations)
+            raw_scores = []
+    os.remove(meteor_output_path)
+    return average_scores
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-i", "--infile")
+    parser.add_argument("-n", "--repeat_times", type=int)
+    parser.add_argument("-m", "--meteor")
+    parser.add_argument("-o", "--output")
+    args = parser.parse_args()
+
+    translations = read_translations(args.infile, args.repeat_times)
+    sys.stderr.write("\nGenerating input for Meteor...")
+    ref_path, mt_path = generate_input(translations, args.repeat_times)
+    sys.stderr.write("\nRunning Meteor...")
+    out_path = run_meteor(ref_path, mt_path, args.meteor)
+    sys.stderr.write("\nReading output...")
+    scores = read_output(out_path, args.repeat_times)
+    sys.stderr.write("\nWriting results...")
+    with open(args.output, "w") as o:
+        for scr in scores:
+            o.write("{}\n".format(scr))
+    o.close()
+
+
+if __name__ == "__main__":
+    main()

omni_speech/infer/fairseq/examples/unsupervised_quality_estimation/repeat_lines.py

@@ -0,0 +1,28 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import argparse
+import sys
+
+
+def _normalize_spaces(line):
+    return " ".join(line.split())
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-i", "--input_file", required=True, type=str)
+    parser.add_argument("-n", "--repeat_times", required=True, type=int)
+    parser.add_argument("-o", "--output_file", required=False, type=str)
+    args = parser.parse_args()
+    stream = open(args.output_file, "w") if args.output_file else sys.stdout
+
+    for line in open(args.input_file):
+        for _ in range(args.repeat_times):
+            stream.write(_normalize_spaces(line) + "\n")
+
+
+if __name__ == "__main__":
+    main()

omni_speech/infer/fairseq/examples/wav2vec/README.md

@@ -0,0 +1,426 @@
+# wav2vec 2.0
+
+wav2vec 2.0 learns speech representations on unlabeled data as described in [wav2vec 2.0: A Framework for Self-Supervised Learning of Speech Representations (Baevski et al., 2020)](https://arxiv.org/abs/2006.11477).
+
+We learned speech representations in multiple languages as well in [Unsupervised Cross-lingual Representation Learning for Speech Recognition (Conneau et al., 2020)](https://arxiv.org/abs/2006.13979).
+
+We also combined wav2vec 2.0 with self-training in [Self-training and Pre-training are Complementary for Speech Recognition (Xu et al., 2020)](https://arxiv.org/abs/2010.11430).
+
+We combined speech data from multiple domains in [Robust wav2vec 2.0: Analyzing Domain Shift in Self-Supervised Pre-Training (Hsu, et al., 2021)](https://arxiv.org/abs/2104.01027).
+
+We finetuned XLSR-53 on multiple languages to transcribe unseen languages in [Simple and Effective Zero-shot Cross-lingual Phoneme Recognition (Xu et al., 2021)](https://arxiv.org/abs/2109.11680).
+
+## Pre-trained models
+
+Model | Finetuning split | Dataset | Model
+|---|---|---|---
+Wav2Vec 2.0 Base | No finetuning | [Librispeech](http://www.openslr.org/12) | [download](https://dl.fbaipublicfiles.com/fairseq/wav2vec/wav2vec_small.pt)
+Wav2Vec 2.0 Base | 10 minutes | [Librispeech](http://www.openslr.org/12) | [download](https://dl.fbaipublicfiles.com/fairseq/wav2vec/wav2vec_small_10m.pt)
+Wav2Vec 2.0 Base | 100 hours | [Librispeech](http://www.openslr.org/12) | [download](https://dl.fbaipublicfiles.com/fairseq/wav2vec/wav2vec_small_100h.pt)
+Wav2Vec 2.0 Base | 960 hours | [Librispeech](http://www.openslr.org/12) | [download](https://dl.fbaipublicfiles.com/fairseq/wav2vec/wav2vec_small_960h.pt)
+Wav2Vec 2.0 Large | No finetuning | [Librispeech](http://www.openslr.org/12)  | [download](https://dl.fbaipublicfiles.com/fairseq/wav2vec/libri960_big.pt)
+Wav2Vec 2.0 Large | 10 minutes | [Librispeech](http://www.openslr.org/12)  | [download](https://dl.fbaipublicfiles.com/fairseq/wav2vec/wav2vec_big_10m.pt)
+Wav2Vec 2.0 Large | 100 hours | [Librispeech](http://www.openslr.org/12)  | [download](https://dl.fbaipublicfiles.com/fairseq/wav2vec/wav2vec_big_100h.pt)
+Wav2Vec 2.0 Large | 960 hours | [Librispeech](http://www.openslr.org/12)  | [download](https://dl.fbaipublicfiles.com/fairseq/wav2vec/wav2vec_big_960h.pt)
+Wav2Vec 2.0 Large (LV-60)* | No finetuning | [Libri-Light](https://github.com/facebookresearch/libri-light) | [download](https://dl.fbaipublicfiles.com/fairseq/wav2vec/wav2vec_vox_new.pt)
+Wav2Vec 2.0 Large conformer - rel_pos (LV-60)* | No finetuning | [Libri-Light](https://github.com/facebookresearch/libri-light) | [download](s3://dl.fbaipublicfiles.com/fairseq/conformer/wav2vec2/librilight/LL_relpos_PT_no_FT)
+Wav2Vec 2.0 Large conformer - rope (LV-60)* | No finetuning | [Libri-Light](https://github.com/facebookresearch/libri-light) | [download](s3://dl.fbaipublicfiles.com/fairseq/conformer/wav2vec2/librilight/LL_rope_PT_no_FT)
+Wav2Vec 2.0 Large (LV-60)* | 10 minutes | [Libri-Light](https://github.com/facebookresearch/libri-light) + [Librispeech](http://www.openslr.org/12) | [download](https://dl.fbaipublicfiles.com/fairseq/wav2vec/wav2vec_vox_10m_new.pt)
+Wav2Vec 2.0 Large (LV-60)* | 100 hours | [Libri-Light](https://github.com/facebookresearch/libri-light) + [Librispeech](http://www.openslr.org/12) | [download](https://dl.fbaipublicfiles.com/fairseq/wav2vec/wav2vec_vox_100h_new.pt)
+Wav2Vec 2.0 Large conformer - rel_pos (LV-60)* | 100 hours | [Libri-Light](https://github.com/facebookresearch/libri-light) | [download](s3://dl.fbaipublicfiles.com/fairseq/conformer/wav2vec2/librilight/LL_relpos_PT_100h_FT.pt)
+Wav2Vec 2.0 Large conformer - rope (LV-60)* | 100 hours | [Libri-Light](https://github.com/facebookresearch/libri-light) | [download](s3://dl.fbaipublicfiles.com/fairseq/conformer/wav2vec2/librilight/LL_rope_PT_100h_FT.pt)
+Wav2Vec 2.0 Large (LV-60)* | 960 hours | [Libri-Light](https://github.com/facebookresearch/libri-light) + [Librispeech](http://www.openslr.org/12) | [download](https://dl.fbaipublicfiles.com/fairseq/wav2vec/wav2vec2_vox_960h_new.pt)
+Wav2Vec 2.0 Large conformer - rel_pos (LV-60)* | 960 hours | [Libri-Light](https://github.com/facebookresearch/libri-light) | [download](s3://dl.fbaipublicfiles.com/fairseq/conformer/wav2vec2/librilight/LL_relpos_PT_960h_FT.pt)
+Wav2Vec 2.0 Large conformer - rope (LV-60)* | 960 hours | [Libri-Light](https://github.com/facebookresearch/libri-light) | [download](s3://dl.fbaipublicfiles.com/fairseq/conformer/wav2vec2/librilight/LL_rope_PT_960h_FT.pt)
+Wav2Vec 2.0 Large (LV-60) + Self Training * | 10 minutes | [Libri-Light](https://github.com/facebookresearch/libri-light) + [Librispeech](http://www.openslr.org/12) | [download](https://dl.fbaipublicfiles.com/fairseq/wav2vec/wav2vec_vox_10m_pl.pt)
+Wav2Vec 2.0 Large (LV-60) + Self Training * | 100 hours | [Libri-Light](https://github.com/facebookresearch/libri-light) + [Librispeech](http://www.openslr.org/12) | [download](https://dl.fbaipublicfiles.com/fairseq/wav2vec/wav2vec_vox_100h_pl.pt)
+Wav2Vec 2.0 Large (LV-60) + Self Training * | 960 hours | [Libri-Light](https://github.com/facebookresearch/libri-light) + [Librispeech](http://www.openslr.org/12) | [download](https://dl.fbaipublicfiles.com/fairseq/wav2vec/wav2vec_vox_960h_pl.pt)
+Wav2Vec 2.0 Large (LV-60 + CV + SWBD + FSH) ** | No finetuning | [Libri-Light](https://github.com/facebookresearch/libri-light) + [CommonVoice](https://commonvoice.mozilla.org/en/languages) + [Switchboard](https://catalog.ldc.upenn.edu/LDC97S62) + [Fisher](https://catalog.ldc.upenn.edu/LDC2004T19) | [download](https://dl.fbaipublicfiles.com/fairseq/wav2vec/w2v_large_lv_fsh_swbd_cv.pt)
+Wav2Vec 2.0 Large (LV-60 + CV + SWBD + FSH) ** | 960 hours Librispeech | [Libri-Light](https://github.com/facebookresearch/libri-light) + [CommonVoice](https://commonvoice.mozilla.org/en/languages) + [Switchboard](https://catalog.ldc.upenn.edu/LDC97S62) + [Fisher](https://catalog.ldc.upenn.edu/LDC2004T19) | [download](https://dl.fbaipublicfiles.com/fairseq/wav2vec/w2v_large_lv_fsh_swbd_cv_ftls960_updated.pt)
+Wav2Vec 2.0 Large (LV-60 + CV + SWBD + FSH) ** | 300 hours Switchboard | [Libri-Light](https://github.com/facebookresearch/libri-light) + [CommonVoice](https://commonvoice.mozilla.org/en/languages) + [Switchboard](https://catalog.ldc.upenn.edu/LDC97S62) + [Fisher](https://catalog.ldc.upenn.edu/LDC2004T19) | [download](https://dl.fbaipublicfiles.com/fairseq/wav2vec/w2v_large_lv_fsh_swbd_cv_ftsb300_updated.pt)
+
+\* updated (Oct. 24, 2020)\
+** updated (Nov. 13, 2021)
+
+We also release multilingual pre-trained wav2vec 2.0 (XLSR) models:
+
+Model | Architecture | Hours | Languages | Datasets | Model
+|---|---|---|---|---|---
+XLSR-53 | Large | 56k | 53 | MLS, CommonVoice, BABEL | [download](https://dl.fbaipublicfiles.com/fairseq/wav2vec/xlsr_53_56k.pt)
+
+The XLSR model uses the following datasets for multilingual pretraining:
+
+* **[MLS: Multilingual LibriSpeech](https://indico2.conference4me.psnc.pl/event/35/contributions/3585/attachments/1060/1101/Wed-2-6-10.pdf)** (8 languages, 50.7k hours): *Dutch, English, French, German, Italian, Polish, Portuguese, Spanish*
+
+* **[CommonVoice](https://commonvoice.mozilla.org/en/languages)** (36 languages, 3.6k hours): *Arabic, Basque, Breton, Chinese (CN), Chinese (HK), Chinese (TW), Chuvash, Dhivehi, Dutch, English, Esperanto, Estonian, French, German, Hakh-Chin, Indonesian, Interlingua, Irish, Italian, Japanese, Kabyle, Kinyarwanda, Kyrgyz, Latvian, Mongolian, Persian, Portuguese, Russian, Sakha, Slovenian, Spanish, Swedish, Tamil, Tatar, Turkish, Welsh* (see also [finetuning splits]([https://dl.fbaipublicfiles.com/cpc_audio/common_voices_splits.tar.gz]) from [this paper](https://arxiv.org/abs/2002.02848)).
+
+* **[Babel](https://catalog.ldc.upenn.edu/byyear)** (17 languages, 1.7k hours): *Assamese, Bengali, Cantonese, Cebuano, Georgian, Haitian, Kazakh, Kurmanji, Lao, Pashto, Swahili, Tagalog, Tamil, Tok, Turkish, Vietnamese, Zulu*
+
+We also finetuned several models on languages from [CommonVoice](https://commonvoice.mozilla.org/en/languages) (version 6.1) and [Babel](https://catalog.ldc.upenn.edu/byyear). Please refer to [our paper](https://arxiv.org/abs/2109.11680) for details about which languages are used.
+
+Pretrained Model | Fintune Dataset | # Languages | Phonemizer | Model | Dictionary
+|---|---|---|---|---|---
+LV-60 | CommonVoice | 26 | [Espeak](https://github.com/espeak-ng/espeak-ng/blob/master/docs/languages.md) | [download](https://dl.fbaipublicfiles.com/fairseq/wav2vec/zero_shot/espeak_en_26lang_m10.pt) | [download](https://dl.fbaipublicfiles.com/fairseq/wav2vec/zero_shot/espeak_dict.txt)
+XLSR-53 | CommonVoice | 26 | [Espeak](https://github.com/espeak-ng/espeak-ng/blob/master/docs/languages.md) | [download](https://dl.fbaipublicfiles.com/fairseq/wav2vec/zero_shot/espeak_26lang_m10.pt) | [download](https://dl.fbaipublicfiles.com/fairseq/wav2vec/zero_shot/espeak_dict.txt)
+XLSR-53 | CommonVoice | 21 | [Phonetisaurus](https://github.com/AdolfVonKleist/Phonetisaurus) | [download](https://dl.fbaipublicfiles.com/fairseq/wav2vec/zero_shot/phonetisaurus_21lang_m10.pt) | [download](https://dl.fbaipublicfiles.com/fairseq/wav2vec/zero_shot/phonetisaurus_dict.txt)
+XLSR-53 | CommonVoice, BABEL | 21, 19 | [Phonetisaurus](https://github.com/AdolfVonKleist/Phonetisaurus) | [download](https://dl.fbaipublicfiles.com/fairseq/wav2vec/zero_shot/phonetisaurus_40lang_m10.pt) | [download](https://dl.fbaipublicfiles.com/fairseq/wav2vec/zero_shot/phonetisaurus_40lang.dict.txt)
+
+We release 2 models that are finetuned on data from 2 different phonemizers. Although the phonemes are all [IPA](https://en.wikipedia.org/wiki/International_Phonetic_Alphabet) symbols, there are still subtle differences between the phonemized transcriptions from the 2 phonemizers. Thus, it's better to use the corresponding model, if your data is phonemized by either phonemizer above.
+
+## Training a new model with the CLI tools
+
+Given a directory containing wav files to be used for pretraining (we recommend splitting each file into separate file 10 to 30 seconds in length)
+
+### Prepare training data manifest
+
+First, install the `soundfile` library:
+
+```shell script
+pip install soundfile
+```
+
+Next, run:
+
+```shell script
+python examples/wav2vec/wav2vec_manifest.py /path/to/waves --dest /manifest/path --ext $ext --valid-percent $valid
+```
+
+$ext should be set to flac, wav, or whatever format your dataset happens to use that soundfile can read.
+
+$valid should be set to some reasonable percentage (like 0.01) of training data to use for validation.
+To use a pre-defined validation set (like dev-other from librispeech), set to it 0 and then overwrite valid.tsv with a
+separately pre-processed manifest file.
+
+### Train a wav2vec 2.0 base model
+
+This configuration was used for the base model trained on the Librispeech dataset in the wav2vec 2.0 paper
+
+Note that the input is expected to be single channel, sampled at 16 kHz
+
+```shell script
+$ fairseq-hydra-train \
+    task.data=/path/to/data \
+    --config-dir /path/to/fairseq-py/examples/wav2vec/config/pretraining \
+    --config-name wav2vec2_base_librispeech
+```
+
+Note: you can simulate 64 GPUs by using k GPUs and adding command line parameters (before `--config-dir`)
+`distributed_training.distributed_world_size=k` `+optimization.update_freq='[x]'` where x = 64/k
+
+### Train a wav2vec 2.0 large model
+
+This configuration was used for the large model trained on the Libri-light dataset in the wav2vec 2.0 paper
+
+```shell script
+$ fairseq-hydra-train \
+    task.data=/path/to/data \
+    --config-dir /path/to/fairseq-py/examples/wav2vec/config/pretraining \
+    --config-name wav2vec2_large_librivox
+```
+
+Note: you can simulate 128 GPUs by using k GPUs and adding command line parameters (before `--config-dir`)
+`distributed_training.distributed_world_size=k` `+optimization.update_freq='[x]'` where x = 128/k
+
+### Train a wav2vec 2.0 model with conformer backbone
+
+To replace the transformer layers in the encoder with the conformer layers, set `--layer-type conformer --attn-type espnet --pos-enc-type ${POS_ENC_TYPE}`. `POS_ENC_TYPE` refers to positional encoding to be used in the conformer encoder.
+Set it to `abs`, `rope` or `rel_pos` to use the absolute positional encoding, rotary positional encoding or relative positional encoding in the conformer layer respectively.
+
+To train a base model with conformer:
+
+```shell script
+$ fairseq-hydra-train \
+    task.data=/path/to/data \
+    --config-dir /path/to/fairseq-py/examples/wav2vec/config/pretraining \
+    --config-name wav2vec2_conformer_base_librispeech \
+    --attn-type espnet --pos-enc-type ${POS_ENC_TYPE}
+```
+
+To train a large model with conformer:
+
+```shell script
+$ fairseq-hydra-train \
+    task.data=/path/to/data \
+    --config-dir /path/to/fairseq-py/examples/wav2vec/config/pretraining \
+    --config-name wav2vec2_conformer_large_librivox
+    --attn-type espnet --pos-enc-type ${POS_ENC_TYPE}
+
+```
+
+### Fine-tune a pre-trained model with CTC
+
+Fine-tuning a model requires parallel audio and labels file, as well as a vocabulary file in fairseq format.
+A letter vocabulary can be downloaded [here](https://dl.fbaipublicfiles.com/fairseq/wav2vec/dict.ltr.txt).
+An example [script](libri_labels.py) that generates labels for the Librispeech dataset from the tsv file produced by wav2vec_manifest.py can be used as follows:
+
+```shell script
+split=train
+$ python libri_labels.py /path/to/tsv --output-dir /output/dir --output-name $split
+```
+
+Fine-tuning on 100h of Librispeech with letter targets:
+
+```shell script
+$ fairseq-hydra-train \
+    distributed_training.distributed_port=$PORT \
+    task.data=/path/to/data \
+    model.w2v_path=/path/to/model.pt \
+    --config-dir /path/to/fairseq-py/examples/wav2vec/config/finetuning \
+    --config-name base_100h
+```
+
+There are other config files in the config/finetuning directory that can be used to fine-tune on other splits.
+You can specify the right config via the `--config-name` parameter.
+
+Note: you can simulate 24 GPUs by using k GPUs and adding command line parameters (before `--config-dir`)
+`distributed_training.distributed_world_size=k` `+optimization.update_freq='[x]'` where x = 24/k
+
+Decoding with a language model during training requires flashlight [python bindings](https://github.com/facebookresearch/flashlight/tree/master/bindings/python) (previously called [wav2letter](https://github.com/facebookresearch/wav2letter).
+If you want to use a language model, add `+criterion.wer_args='[/path/to/kenlm, /path/to/lexicon, 2, -1]'` to the command line.
+
+### Evaluating a CTC model
+
+Evaluating a CTC model with a language model requires [flashlight python bindings](https://github.com/facebookresearch/flashlight/tree/master/bindings/python) (previously called [wav2letter](https://github.com/facebookresearch/wav2letter) to be installed.
+
+Fairseq transformer language model used in the wav2vec 2.0 paper can be obtained from the [wav2letter model repository](https://github.com/facebookresearch/wav2letter/tree/master/recipes/sota/2019).
+Be sure to upper-case the language model vocab after downloading it.
+
+Letter dictionary for pre-trained models can be found [here](https://dl.fbaipublicfiles.com/fairseq/wav2vec/dict.ltr.txt).
+
+Next, run the evaluation command:
+
+```shell script
+$subset=dev_other
+python examples/speech_recognition/infer.py /checkpoint/abaevski/data/speech/libri/10h/wav2vec/raw --task audio_finetuning \
+--nbest 1 --path /path/to/model --gen-subset $subset --results-path /path/to/save/results/for/sclite --w2l-decoder kenlm \
+--lm-model /path/to/kenlm.bin --lm-weight 2 --word-score -1 --sil-weight 0 --criterion ctc --labels ltr --max-tokens 4000000 \
+--post-process letter
+```
+
+To get raw numbers, use --w2l-decoder viterbi and omit the lexicon. To use the transformer language model, use --w2l-decoder fairseqlm.
+
+## Use wav2vec 2.0 with 🤗Transformers
+
+Wav2Vec2 is also available in the [🤗Transformers library](https://github.com/huggingface/transformers) since version 4.4.
+
+Pretrained Models can be found on the [hub](https://huggingface.co/models?filter=wav2vec2)
+and documentation can be found [here](https://huggingface.co/transformers/master/model_doc/wav2vec2.html).
+
+Usage example:
+
+```python
+# !pip install transformers
+# !pip install datasets
+import soundfile as sf
+import torch
+from datasets import load_dataset
+from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
+
+# load pretrained model
+processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-base-960h")
+model = Wav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-base-960h")
+
+
+librispeech_samples_ds = load_dataset("patrickvonplaten/librispeech_asr_dummy", "clean", split="validation")
+
+# load audio
+audio_input, sample_rate = sf.read(librispeech_samples_ds[0]["file"])
+
+# pad input values and return pt tensor
+input_values = processor(audio_input, sampling_rate=sample_rate, return_tensors="pt").input_values
+
+# INFERENCE
+
+# retrieve logits & take argmax
+logits = model(input_values).logits
+predicted_ids = torch.argmax(logits, dim=-1)
+
+# transcribe
+transcription = processor.decode(predicted_ids[0])
+
+# FINE-TUNE
+
+target_transcription = "A MAN SAID TO THE UNIVERSE I EXIST"
+
+# encode labels
+with processor.as_target_processor():
+  labels = processor(target_transcription, return_tensors="pt").input_ids
+
+# compute loss by passing labels
+loss = model(input_values, labels=labels).loss
+loss.backward()
+```
+
+# wav2vec
+
+Example to train a wav2vec model as described in [wav2vec: Unsupervised Pre-training for Speech Recognition (Schneider et al., 2019)](https://arxiv.org/abs/1904.05862).
+
+## Pre-trained models
+
+Description | Dataset | Model
+---|---|---
+Wav2Vec large | [Librispeech](http://www.openslr.org/12) | [download](https://dl.fbaipublicfiles.com/fairseq/wav2vec/wav2vec_large.pt)
+
+#### Example usage
+
+```python
+import torch
+import fairseq
+
+cp_path = '/path/to/wav2vec.pt'
+model, cfg, task = fairseq.checkpoint_utils.load_model_ensemble_and_task([cp_path])
+model = model[0]
+model.eval()
+
+wav_input_16khz = torch.randn(1,10000)
+z = model.feature_extractor(wav_input_16khz)
+c = model.feature_aggregator(z)
+```
+
+## Training a new model with the CLI tools
+
+Given a directory containing wav files to be used for pretraining (we recommend splitting each file into separate files 10 to 30 seconds in length)
+
+### Prepare training data manifest
+
+```
+python examples/wav2vec/wav2vec_manifest.py /path/to/waves --dest /manifest/path --ext wav
+```
+
+### Train a wav2vec model
+
+```
+$ python train.py /manifest/path --save-dir /model/path --num-workers 6 --fp16 --max-update 400000 --save-interval 1 --no-epoch-checkpoints \
+--arch wav2vec --task audio_pretraining --min-lr 1e-06 --stop-min-lr 1e-09 --optimizer adam --lr 0.005 --lr-scheduler cosine \
+--conv-feature-layers [(512, 10, 5), (512, 8, 4), (512, 4, 2), (512, 4, 2), (512, 4, 2), (512, 1, 1), (512, 1, 1)] \
+--conv-aggregator-layers [(512, 2, 1), (512, 3, 1), (512, 4, 1), (512, 5, 1), (512, 6, 1), (512, 7, 1), (512, 8, 1), (512, 9, 1), (512, 10, 1), (512, 11, 1), (512, 12, 1), (512, 13, 1)] \
+--skip-connections-agg --residual-scale 0.5 --log-compression --warmup-updates 500 --warmup-init-lr 1e-07 --criterion wav2vec --num-negatives 10 \
+--max-sample-size 150000 --max-tokens 1500000 --skip-invalid-size-inputs-valid-test
+```
+
+### Run wav2vec2 pre-training on Google Cloud TPUs
+
+Wav2Vec2 is now supported on TPUs! It's currently pre-training only.
+
+#### Using hydra on a v3-8
+
+```
+$ OMP_NUM_THREADS=1 fairseq-hydra-train \
+  task.data=/manifest/path \
+  --config-dir /PATH/TO/FAIRSEQ/examples/wav2vec/config/pretraining \
+  --config-name wav2vec2_large_librivox_tpu.yaml
+```
+
+#### Using command line arguments on a v3-8
+
+Note: Commandline arguments way of execution has a [known-problem](https://github.com/pytorch/fairseq/issues/3741) currently.
+
+```
+$ OMP_NUM_THREADS=1 python train.py /manifest/path --save-dir /model/path --num-workers 6 --fp16 --max-update 400000 --save-interval 1 --no-epoch-checkpoints \
+--arch wav2vec2 --task audio_pretraining --min-lr 1e-06 --stop-min-lr 1e-09 --optimizer adam --lr 0.005 --lr-scheduler cosine \
+--conv-feature-layers [(512, 10, 5), (512, 8, 4), (512, 4, 2), (512, 4, 2), (512, 4, 2), (512, 1, 1), (512, 1, 1)] \
+--conv-aggregator-layers [(512, 2, 1), (512, 3, 1), (512, 4, 1), (512, 5, 1), (512, 6, 1), (512, 7, 1), (512, 8, 1), (512, 9, 1), (512, 10, 1), (512, 11, 1), (512, 12, 1), (512, 13, 1)] \
+--skip-connections-agg --residual-scale 0.5 --log-compression --warmup-updates 500 --warmup-init-lr 1e-07 --criterion wav2vec --num-negatives 10 \
+--max-sample-size 150000 --max-tokens 1500000 --skip-invalid-size-inputs-valid-test \
+--tpu --distributed-world-size 8 --num-batch-buckets 3 --enable-padding \
+--encoder-layerdrop 0 --mask-channel-prob 0.1
+```
+
+#### Using hydra on a pod slice (v3-N with N > 8)
+
+```
+$ OMP_NUM_THREADS=1 fairseq-hydra-train \
+  task.data=/manifest/path \
+  --config-dir /PATH/TO/FAIRSEQ/examples/wav2vec/config/pretraining \
+  --config-name wav2vec2_large_librivox_tpu-pod.yaml  # edit distributed-world-size accordingly
+```
+
+#### Using command line arguments on a pod slice (v3-N with N > 8)
+
+Note: Commandline arguments way of execution has a [known-problem](https://github.com/pytorch/fairseq/issues/3741) currently.
+
+```
+$ python -m torch_xla.distributed.xla_dist \
+  --tpu ${TPUNAME} --conda-env=torch-xla-${TORCH_XLA_VERSION} --env OMP_NUM_THREADS=1 \
+  -- \
+python train.py /manifest/path --save-dir /model/path --num-workers 6 --fp16 --max-update 400000 --save-interval 1 --no-epoch-checkpoints \
+--arch wav2vec2 --task audio_pretraining --min-lr 1e-06 --stop-min-lr 1e-09 --optimizer adam --lr 0.005 --lr-scheduler cosine \
+--conv-feature-layers [(512, 10, 5), (512, 8, 4), (512, 4, 2), (512, 4, 2), (512, 4, 2), (512, 1, 1), (512, 1, 1)] \
+--conv-aggregator-layers [(512, 2, 1), (512, 3, 1), (512, 4, 1), (512, 5, 1), (512, 6, 1), (512, 7, 1), (512, 8, 1), (512, 9, 1), (512, 10, 1), (512, 11, 1), (512, 12, 1), (512, 13, 1)] \
+--skip-connections-agg --residual-scale 0.5 --log-compression --warmup-updates 500 --warmup-init-lr 1e-07 --criterion wav2vec --num-negatives 10 \
+--max-sample-size 150000 --max-tokens 1500000 --skip-invalid-size-inputs-valid-test \
+--tpu --distributed-world-size ${WORLD_SIZE} --num-batch-buckets 3 --enable-padding \
+--encoder-layerdrop 0 --mask-channel-prob 0.1
+```
+
+### Extract embeddings from the downstream task data
+
+```
+$ PYTHONPATH=/path/to/fairseq python examples/wav2vec/wav2vec_featurize.py --input /path/to/task/waves --output /path/to/output \
+--model /model/path/checkpoint_best.pt --split train valid test
+```
+
+# vq-wav2vec
+
+Example to train a vq-wav2vec model as described in [vq-wav2vec: Self-Supervised Learning of Discrete Speech Representations (Baevski et al., 2019)](https://arxiv.org/abs/1910.05453).
+
+These models are also used in [Effectiveness of self-supervised pre-training for speech recognition (Baevski et al., 2019)](https://arxiv.org/abs/1911.03912).
+
+## Pre-trained models
+
+Description | Dataset | Model
+---|---|---
+vq-wav2vec Gumbel | [Librispeech](http://www.openslr.org/12) | [download](https://dl.fbaipublicfiles.com/fairseq/wav2vec/vq-wav2vec.pt)
+vq-wav2vec K-means | [Librispeech](http://www.openslr.org/12) | [download](https://dl.fbaipublicfiles.com/fairseq/wav2vec/vq-wav2vec_kmeans.pt)
+Roberta on K-means codes | [Librispeech](http://www.openslr.org/12) | [download](https://dl.fbaipublicfiles.com/fairseq/wav2vec/bert_kmeans.tar)
+
+#### Example usage
+
+```python
+import torch
+import fairseq
+
+cp = torch.load('/path/to/vq-wav2vec.pt')
+model, cfg, task = fairseq.checkpoint_utils.load_model_ensemble_and_task([cp])
+model = model[0]
+model.eval()
+
+wav_input_16khz = torch.randn(1,10000)
+z = model.feature_extractor(wav_input_16khz)
+_, idxs = model.vector_quantizer.forward_idx(z)
+print(idxs.shape) # output: torch.Size([1, 60, 2]), 60 timesteps with 2 indexes corresponding to 2 groups in the model
+```
+
+## Training a new model with the CLI tools
+
+Given a directory containing wav files to be used for pretraining (we recommend splitting each file into separate file 10 to 30 seconds in length)
+
+### Prepare training data manifest
+
+```
+python examples/wav2vec/wav2vec_manifest.py /path/to/waves --dest /manifest/path --ext wav
+```
+
+### Train a gumbel vq-wav2vec model
+
+```
+$ python train.py /manifest/path --save-dir /model/path --num-workers 6 --fp16 --max-update 400000 \
+--save-interval 1 --no-epoch-checkpoints --arch wav2vec --task audio_pretraining --min-lr 1e-06 --stop-min-lr 1e-09 \
+--optimizer adam --lr 1e-05 --lr-scheduler cosine \
+--conv-feature-layers [(512, 10, 5), (512, 8, 4), (512, 4, 2), (512, 4, 2), (512, 4, 2), (512, 1, 1), (512, 1, 1), (512, 1, 1)] \
+--conv-aggregator-layers [(512, 2, 1), (512, 3, 1), (512, 4, 1), (512, 5, 1), (512, 6, 1), (512, 7, 1), (512, 8, 1), (512, 9, 1), (512, 10, 1), (512, 11, 1), (512, 12, 1), (512, 13, 1)] \
+--activation gelu --offset auto --skip-connections-agg --residual-scale 0.5 \
+--log-keys ["prob_perplexity","code_perplexity","temp"] --vq-type gumbel --vq-groups 2 --vq-depth 2 \
+--combine-groups --vq-vars 320 --vq-temp (2,0.5,0.999995) --prediction-steps 12 --warmup-updates 1000 \
+--warmup-init-lr 1e-07 --criterion wav2vec --num-negatives 10 --max-sample-size 150000 \
+--max-tokens 300000 --cross-sample-negatives 0 --update-freq 1 --seed 2 --skip-invalid-size-inputs-valid-test
+```
+
+for k-means training, set vq-type with "kmeans" and add --loss-weights [1] argument. Pre-trained models were trained on 16 GPUs.
+
+### Tokenize audio data (e.g. for BERT training)
+
+```
+$ PYTHONPATH=/path/to/fairseq python examples/wav2vec/vq-wav2vec_featurize.py --data-dir /manifest/path --output-dir /path/to/output \
+--checkpoint /model/path/checkpoint_best.pt --split train valid test --extension tsv
+```

omni_speech/infer/fairseq/examples/wav2vec/config/finetuning/base_100h.yaml

@@ -0,0 +1,58 @@
+# @package _group_
+
+common:
+  fp16: true
+  log_format: json
+  log_interval: 200
+
+checkpoint:
+  no_epoch_checkpoints: true
+  best_checkpoint_metric: wer
+
+task:
+  _name: audio_finetuning
+  data: ???
+  normalize: false
+  labels: ltr
+
+dataset:
+  num_workers: 6
+  max_tokens: 3200000
+  skip_invalid_size_inputs_valid_test: true
+  valid_subset: dev_other
+
+distributed_training:
+  ddp_backend: legacy_ddp
+  distributed_world_size: 2
+
+criterion:
+  _name: ctc
+  zero_infinity: true
+
+optimization:
+  max_update: 80000
+  lr: [0.00003]
+  sentence_avg: true
+  update_freq: [4]
+
+optimizer:
+  _name: adam
+  adam_betas: (0.9,0.98)
+  adam_eps: 1e-08
+
+lr_scheduler:
+  _name: tri_stage
+  phase_ratio: [0.1, 0.4, 0.5]
+  final_lr_scale: 0.05
+
+model:
+  _name: wav2vec_ctc
+  w2v_path: ???
+  apply_mask: true
+  mask_prob: 0.65
+  mask_channel_prob: 0.5
+  mask_channel_length: 64
+  layerdrop: 0.1
+  activation_dropout: 0.1
+  feature_grad_mult: 0.0
+  freeze_finetune_updates: 0

omni_speech/infer/fairseq/examples/wav2vec/config/finetuning/base_10h.yaml

@@ -0,0 +1,63 @@
+# @package _group_
+
+common:
+  fp16: true
+  log_format: json
+  log_interval: 200
+
+checkpoint:
+  save_interval: 50
+  save_interval_updates: 10000
+  keep_interval_updates: 1
+  no_epoch_checkpoints: true
+  best_checkpoint_metric: wer
+
+task:
+  _name: audio_finetuning
+  data: ???
+  normalize: false
+  labels: ltr
+
+dataset:
+  num_workers: 6
+  max_tokens: 3200000
+  skip_invalid_size_inputs_valid_test: true
+  validate_after_updates: 10000
+  validate_interval: 50
+  valid_subset: dev_other
+
+distributed_training:
+  ddp_backend: legacy_ddp
+  distributed_world_size: 2
+
+criterion:
+  _name: ctc
+  zero_infinity: true
+
+optimization:
+  max_update: 20000
+  lr: [0.00005]
+  sentence_avg: true
+  update_freq: [4]
+
+optimizer:
+  _name: adam
+  adam_betas: (0.9,0.98)
+  adam_eps: 1e-08
+
+lr_scheduler:
+  _name: tri_stage
+  phase_ratio: [0.1, 0.4, 0.5]
+  final_lr_scale: 0.05
+
+model:
+  _name: wav2vec_ctc
+  w2v_path: ???
+  apply_mask: true
+  mask_prob: 0.65
+  mask_channel_prob: 0.5
+  mask_channel_length: 64
+  layerdrop: 0.05
+  activation_dropout: 0.1
+  feature_grad_mult: 0.0
+  freeze_finetune_updates: 10000

omni_speech/infer/fairseq/examples/wav2vec/config/finetuning/base_10m.yaml

@@ -0,0 +1,63 @@
+# @package _group_
+
+common:
+  fp16: true
+  log_format: json
+  log_interval: 200
+
+checkpoint:
+  save_interval: 1000
+  save_interval_updates: 50
+  keep_interval_updates: 1
+  no_epoch_checkpoints: true
+  best_checkpoint_metric: wer
+
+task:
+  _name: audio_finetuning
+  data: ???
+  normalize: false
+  labels: ltr
+
+dataset:
+  num_workers: 6
+  max_tokens: 3200000
+  skip_invalid_size_inputs_valid_test: true
+  validate_after_updates: 10000
+  validate_interval: 1000
+  valid_subset: dev_other
+
+distributed_training:
+  ddp_backend: legacy_ddp
+  distributed_world_size: 2
+
+criterion:
+  _name: ctc
+  zero_infinity: true
+
+optimization:
+  max_update: 13000
+  lr: [0.00005]
+  sentence_avg: true
+  update_freq: [4]
+
+optimizer:
+  _name: adam
+  adam_betas: (0.9,0.98)
+  adam_eps: 1e-08
+
+lr_scheduler:
+  _name: tri_stage
+  phase_ratio: [0.1, 0.4, 0.5]
+  final_lr_scale: 0.05
+
+model:
+  _name: wav2vec_ctc
+  w2v_path: ???
+  apply_mask: true
+  mask_prob: 0.65
+  mask_channel_prob: 0.25
+  mask_channel_length: 64
+  layerdrop: 0.1
+  activation_dropout: 0.1
+  feature_grad_mult: 0.0
+  freeze_finetune_updates: 10000

omni_speech/infer/fairseq/examples/wav2vec/config/finetuning/base_1h.yaml

@@ -0,0 +1,63 @@
+# @package _group_
+
+common:
+  fp16: true
+  log_format: json
+  log_interval: 200
+
+checkpoint:
+  save_interval: 50
+  save_interval_updates: 1000
+  keep_interval_updates: 1
+  no_epoch_checkpoints: true
+  best_checkpoint_metric: wer
+
+task:
+  _name: audio_finetuning
+  data: ???
+  normalize: false
+  labels: ltr
+
+dataset:
+  num_workers: 6
+  max_tokens: 3200000
+  skip_invalid_size_inputs_valid_test: true
+  validate_after_updates: 10000
+  validate_interval: 1000
+  valid_subset: dev_other
+
+distributed_training:
+  ddp_backend: legacy_ddp
+  distributed_world_size: 2
+
+criterion:
+  _name: ctc
+  zero_infinity: true
+
+optimization:
+  max_update: 13000
+  lr: [0.00005]
+  sentence_avg: true
+  update_freq: [4]
+
+optimizer:
+  _name: adam
+  adam_betas: (0.9,0.98)
+  adam_eps: 1e-08
+
+lr_scheduler:
+  _name: tri_stage
+  phase_ratio: [0.1, 0.4, 0.5]
+  final_lr_scale: 0.05
+
+model:
+  _name: wav2vec_ctc
+  w2v_path: ???
+  apply_mask: true
+  mask_prob: 0.65
+  mask_channel_prob: 0.25
+  mask_channel_length: 64
+  layerdrop: 0.1
+  activation_dropout: 0.1
+  feature_grad_mult: 0.0
+  freeze_finetune_updates: 10000