test python files

app.py

@@ -0,0 +1,77 @@
+
+from random import sample
+import gradio as gr
+import torchaudio
+import torch
+import torch.nn as nn
+import lightning_module
+import pdb
+import jiwer
+# ASR part
+from transformers import pipeline
+p = pipeline("automatic-speech-recognition")
+
+# WER part
+transformation = jiwer.Compose([
+    jiwer.ToLowerCase(),
+    jiwer.RemoveWhiteSpace(replace_by_space=True),
+    jiwer.RemoveMultipleSpaces(),
+    jiwer.ReduceToListOfListOfWords(word_delimiter=" ")
+])
+
+class ChangeSampleRate(nn.Module):
+    def __init__(self, input_rate: int, output_rate: int):
+        super().__init__()
+        self.output_rate = output_rate
+        self.input_rate = input_rate
+
+    def forward(self, wav: torch.tensor) -> torch.tensor:
+        # Only accepts 1-channel waveform input
+        wav = wav.view(wav.size(0), -1)
+        new_length = wav.size(-1) * self.output_rate // self.input_rate
+        indices = (torch.arange(new_length) * (self.input_rate / self.output_rate))
+        round_down = wav[:, indices.long()]
+        round_up = wav[:, (indices.long() + 1).clamp(max=wav.size(-1) - 1)]
+        output = round_down * (1. - indices.fmod(1.)).unsqueeze(0) + round_up * indices.fmod(1.).unsqueeze(0)
+        return output
+
+model = lightning_module.BaselineLightningModule.load_from_checkpoint("epoch=3-step=7459.ckpt").eval()
+def calc_mos(audio_path, ref):
+    wav, sr = torchaudio.load(audio_path)
+    osr = 16_000
+    batch = wav.unsqueeze(0).repeat(10, 1, 1)
+    csr = ChangeSampleRate(sr, osr)
+    out_wavs = csr(wav)
+    # ASR
+    trans = p(audio_path)["text"]
+    # WER
+    wer = jiwer.wer(ref, trans, truth_transform=transformation, hypothesis_transform=transformation)
+    
+    batch = {
+        'wav': out_wavs,
+        'domains': torch.tensor([0]),
+        'judge_id': torch.tensor([288])
+    }
+    with torch.no_grad():
+        output = model(batch)
+    
+    predic_mos = output.mean(dim=1).squeeze().detach().numpy()*2 + 3
+    
+    return predic_mos, trans, wer
+
+description ="""
+MOS prediction demo using UTMOS-strong w/o phoneme encoder model, which is trained on the main track dataset.
+This demo only accepts .wav format. Best at 16 kHz sampling rate.
+
+Paper is available [here](https://arxiv.org/abs/2204.02152)
+""" 
+
+iface = gr.Interface(
+  fn=calc_mos,
+  inputs=[gr.Audio(type='filepath'), gr.Textbox(placeholder="Insert referance here", label="Referance")],
+  outputs=[gr.Textbox("Predicted MOS"), gr.Textbox("Hypothesis"), gr.Textbox("WER")],
+  title="UTMOS Demo",
+  description=description,
+  allow_flagging="auto",
+)
+iface.launch()

app_record.py

@@ -0,0 +1,65 @@
+
+from random import sample
+import gradio as gr
+import torchaudio
+import torch
+import torch.nn as nn
+import lightning_module
+import pdb
+
+# ASR part
+from transformers import pipeline
+p = pipeline("automatic-speech-recognition")
+
+class ChangeSampleRate(nn.Module):
+    def __init__(self, input_rate: int, output_rate: int):
+        super().__init__()
+        self.output_rate = output_rate
+        self.input_rate = input_rate
+
+    def forward(self, wav: torch.tensor) -> torch.tensor:
+        # Only accepts 1-channel waveform input
+        wav = wav.view(wav.size(0), -1)
+        new_length = wav.size(-1) * self.output_rate // self.input_rate
+        indices = (torch.arange(new_length) * (self.input_rate / self.output_rate))
+        round_down = wav[:, indices.long()]
+        round_up = wav[:, (indices.long() + 1).clamp(max=wav.size(-1) - 1)]
+        output = round_down * (1. - indices.fmod(1.)).unsqueeze(0) + round_up * indices.fmod(1.).unsqueeze(0)
+        return output
+
+model = lightning_module.BaselineLightningModule.load_from_checkpoint("epoch=3-step=7459.ckpt").eval()
+def calc_mos(audio_path):
+    wav, sr = torchaudio.load(audio_path)
+    osr = 16_000
+    batch = wav.unsqueeze(0).repeat(10, 1, 1)
+    csr = ChangeSampleRate(sr, osr)
+    out_wavs = csr(wav)
+    
+    transcription = p(audio_path)["text"]
+    batch = {
+        'wav': out_wavs,
+        'domains': torch.tensor([0]),
+        'judge_id': torch.tensor([288])
+    }
+    with torch.no_grad():
+        output = model(batch)
+    return output.mean(dim=1).squeeze().detach().numpy()*2 + 3, transcription
+
+
+description ="""
+MOS prediction demo using UTMOS-strong w/o phoneme encoder model, which is trained on the main track dataset.
+This demo only accepts .wav format. Best at 16 kHz sampling rate.
+
+Paper is available [here](https://arxiv.org/abs/2204.02152)
+""" 
+
+# inputs=gr.inputs.Audio(type='filepath'),
+iface = gr.Interface(
+  fn=calc_mos,
+  inputs = gr.Audio(source="microphone", type="filepath"), 
+  outputs=["text","textbox"],
+  title="UTMOS Demo",
+  description=description,
+  allow_flagging=True,
+
+).launch()

app_record_streaming.py

@@ -0,0 +1,63 @@
+
+from random import sample
+import gradio as gr
+import torchaudio
+import torch
+import torch.nn as nn
+import lightning_module
+import pdb
+# ASR part
+from transformers import pipeline
+p = pipeline("automatic-speech-recognition")
+
+class ChangeSampleRate(nn.Module):
+    def __init__(self, input_rate: int, output_rate: int):
+        super().__init__()
+        self.output_rate = output_rate
+        self.input_rate = input_rate
+
+    def forward(self, wav: torch.tensor) -> torch.tensor:
+        # Only accepts 1-channel waveform input
+        wav = wav.view(wav.size(0), -1)
+        new_length = wav.size(-1) * self.output_rate // self.input_rate
+        indices = (torch.arange(new_length) * (self.input_rate / self.output_rate))
+        round_down = wav[:, indices.long()]
+        round_up = wav[:, (indices.long() + 1).clamp(max=wav.size(-1) - 1)]
+        output = round_down * (1. - indices.fmod(1.)).unsqueeze(0) + round_up * indices.fmod(1.).unsqueeze(0)
+        return output
+
+model = lightning_module.BaselineLightningModule.load_from_checkpoint("epoch=3-step=7459.ckpt").eval()
+def calc_mos(audio_path):
+    wav, sr = torchaudio.load(audio_path)
+    osr = 16_000
+    batch = wav.unsqueeze(0).repeat(10, 1, 1)
+    csr = ChangeSampleRate(sr, osr)
+    out_wavs = csr(wav)
+    transcription = p(audio_path)["text"]
+    batch = {
+        'wav': out_wavs,
+        'domains': torch.tensor([0]),
+        'judge_id': torch.tensor([288])
+    }
+    with torch.no_grad():
+        output = model(batch)
+    return output.mean(dim=1).squeeze().detach().numpy()*2 + 3, transcription
+
+
+description ="""
+MOS prediction demo using UTMOS-strong w/o phoneme encoder model, which is trained on the main track dataset.
+This demo only accepts .wav format. Best at 16 kHz sampling rate.
+
+Paper is available [here](https://arxiv.org/abs/2204.02152)
+""" 
+
+# inputs=gr.inputs.Audio(type='filepath'),
+iface = gr.Interface(
+  fn=calc_mos,
+  inputs = gr.Audio(source="microphone", type="filepath", streaming=True), 
+  outputs=["text","textbox"],
+  title="UTMOS Demo",
+  description=description,
+  allow_flagging=False,
+
+).launch()

lightning_module.py

@@ -0,0 +1,41 @@
+import pytorch_lightning as pl
+import torch
+import torch.nn as nn
+import os
+import numpy as np
+import hydra
+from model import load_ssl_model, PhonemeEncoder, DomainEmbedding, LDConditioner, Projection
+
+
+class BaselineLightningModule(pl.LightningModule):
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.construct_model()
+        self.save_hyperparameters()
+    
+    def construct_model(self):
+        self.feature_extractors = nn.ModuleList([
+            load_ssl_model(cp_path='wav2vec_small.pt'),
+            DomainEmbedding(3,128),
+        ])
+        output_dim = sum([ feature_extractor.get_output_dim() for feature_extractor in self.feature_extractors])
+        output_layers = [
+            LDConditioner(judge_dim=128,num_judges=3000,input_dim=output_dim)
+        ]
+        output_dim = output_layers[-1].get_output_dim()
+        output_layers.append(
+            Projection(hidden_dim=2048,activation=torch.nn.ReLU(),range_clipping=False,input_dim=output_dim)
+
+        )
+
+        self.output_layers = nn.ModuleList(output_layers)
+
+    def forward(self, inputs):
+        outputs = {}
+        for feature_extractor in self.feature_extractors:
+            outputs.update(feature_extractor(inputs))
+        x = outputs
+        for output_layer in self.output_layers:
+            x = output_layer(x,inputs)
+        return x

model.py

@@ -0,0 +1,191 @@
+import torch
+import torch.nn as nn
+import fairseq
+import os
+import hydra
+
+def load_ssl_model(cp_path):
+    ssl_model_type = cp_path.split("/")[-1]
+    wavlm =  "WavLM" in ssl_model_type
+    if wavlm:
+        checkpoint = torch.load(cp_path)
+        cfg = WavLMConfig(checkpoint['cfg'])
+        ssl_model = WavLM(cfg)
+        ssl_model.load_state_dict(checkpoint['model'])
+        if 'Large' in ssl_model_type:
+            SSL_OUT_DIM = 1024
+        else:
+            SSL_OUT_DIM = 768
+    else:
+        if ssl_model_type == "wav2vec_small.pt":
+            SSL_OUT_DIM = 768
+        elif ssl_model_type in ["w2v_large_lv_fsh_swbd_cv.pt", "xlsr_53_56k.pt"]:
+            SSL_OUT_DIM = 1024
+        else:
+            print("*** ERROR *** SSL model type " + ssl_model_type + " not supported.")
+            exit()
+        model, cfg, task = fairseq.checkpoint_utils.load_model_ensemble_and_task(
+            [cp_path]
+        )
+        ssl_model = model[0]
+        ssl_model.remove_pretraining_modules()
+    return SSL_model(ssl_model, SSL_OUT_DIM, wavlm)
+
+class SSL_model(nn.Module):
+    def __init__(self,ssl_model,ssl_out_dim,wavlm) -> None:
+        super(SSL_model,self).__init__()
+        self.ssl_model, self.ssl_out_dim = ssl_model, ssl_out_dim
+        self.WavLM = wavlm
+
+    def forward(self,batch):
+        wav = batch['wav'] 
+        wav = wav.squeeze(1) # [batches, audio_len]
+        if self.WavLM:
+            x = self.ssl_model.extract_features(wav)[0]
+        else:
+            res = self.ssl_model(wav, mask=False, features_only=True)
+            x = res["x"]
+        return {"ssl-feature":x}
+    def get_output_dim(self):
+        return self.ssl_out_dim
+
+
+class PhonemeEncoder(nn.Module):
+    '''
+    PhonemeEncoder consists of an embedding layer, an LSTM layer, and a linear layer.
+    Args:
+        vocab_size: the size of the vocabulary
+        hidden_dim: the size of the hidden state of the LSTM
+        emb_dim: the size of the embedding layer
+        out_dim: the size of the output of the linear layer
+        n_lstm_layers: the number of LSTM layers
+    '''
+    def __init__(self, vocab_size, hidden_dim, emb_dim, out_dim,n_lstm_layers,with_reference=True) -> None:
+        super().__init__()
+        self.with_reference = with_reference
+        self.embedding = nn.Embedding(vocab_size, emb_dim)
+        self.encoder = nn.LSTM(emb_dim, hidden_dim,
+                               num_layers=n_lstm_layers, dropout=0.1, bidirectional=True)
+        self.linear = nn.Sequential(
+                nn.Linear(hidden_dim + hidden_dim*self.with_reference, out_dim),
+                nn.ReLU()
+                )
+        self.out_dim = out_dim
+
+    def forward(self,batch):
+        seq = batch['phonemes']
+        lens = batch['phoneme_lens']
+        reference_seq = batch['reference']
+        reference_lens = batch['reference_lens']
+        emb = self.embedding(seq)
+        emb = torch.nn.utils.rnn.pack_padded_sequence(
+            emb, lens, batch_first=True, enforce_sorted=False)
+        _, (ht, _) = self.encoder(emb)
+        feature = ht[-1] + ht[0]
+        if self.with_reference:
+            if reference_seq==None or reference_lens ==None:
+                raise ValueError("reference_batch and reference_lens should not be None when with_reference is True")
+            reference_emb = self.embedding(reference_seq)
+            reference_emb = torch.nn.utils.rnn.pack_padded_sequence(
+                reference_emb, reference_lens, batch_first=True, enforce_sorted=False)
+            _, (ht_ref, _) = self.encoder(emb)
+            reference_feature = ht_ref[-1] + ht_ref[0]
+            feature = self.linear(torch.cat([feature,reference_feature],1))
+        else:
+            feature = self.linear(feature)
+        return {"phoneme-feature": feature}
+    def get_output_dim(self):
+        return self.out_dim
+
+class DomainEmbedding(nn.Module):
+    def __init__(self,n_domains,domain_dim) -> None:
+        super().__init__()
+        self.embedding = nn.Embedding(n_domains,domain_dim)
+        self.output_dim = domain_dim
+    def forward(self, batch):
+        return {"domain-feature": self.embedding(batch['domains'])}
+    def get_output_dim(self):
+        return self.output_dim
+
+
+class LDConditioner(nn.Module):
+    '''
+    Conditions ssl output by listener embedding
+    '''
+    def __init__(self,input_dim, judge_dim, num_judges=None):
+        super().__init__()
+        self.input_dim = input_dim
+        self.judge_dim = judge_dim
+        self.num_judges = num_judges
+        assert num_judges !=None
+        self.judge_embedding = nn.Embedding(num_judges, self.judge_dim)
+        # concat [self.output_layer, phoneme features]
+        
+        self.decoder_rnn = nn.LSTM(
+            input_size = self.input_dim + self.judge_dim,
+            hidden_size = 512,
+            num_layers = 1,
+            batch_first = True,
+            bidirectional = True
+        ) # linear?
+        self.out_dim = self.decoder_rnn.hidden_size*2
+
+    def get_output_dim(self):
+        return self.out_dim
+
+
+    def forward(self, x, batch):
+        judge_ids = batch['judge_id']
+        if 'phoneme-feature' in x.keys():
+            concatenated_feature = torch.cat((x['ssl-feature'], x['phoneme-feature'].unsqueeze(1).expand(-1,x['ssl-feature'].size(1) ,-1)),dim=2)
+        else:
+            concatenated_feature = x['ssl-feature']
+        if 'domain-feature' in x.keys():
+            concatenated_feature = torch.cat(
+                (
+                    concatenated_feature,
+                    x['domain-feature']
+                    .unsqueeze(1)
+                    .expand(-1, concatenated_feature.size(1), -1),
+                ),
+                dim=2,
+            )
+        if judge_ids != None:
+            concatenated_feature = torch.cat(
+                (
+                    concatenated_feature,
+                    self.judge_embedding(judge_ids)
+                    .unsqueeze(1)
+                    .expand(-1, concatenated_feature.size(1), -1),
+                ),
+                dim=2,
+            )
+            decoder_output, (h, c) = self.decoder_rnn(concatenated_feature)
+        return decoder_output
+
+class Projection(nn.Module):
+    def __init__(self, input_dim, hidden_dim, activation, range_clipping=False):
+        super(Projection, self).__init__()
+        self.range_clipping = range_clipping
+        output_dim = 1
+        if range_clipping:
+            self.proj = nn.Tanh()
+        
+        self.net = nn.Sequential(
+            nn.Linear(input_dim, hidden_dim),
+            activation,
+            nn.Dropout(0.3),
+            nn.Linear(hidden_dim, output_dim),
+        )
+        self.output_dim = output_dim
+    
+    def forward(self, x, batch):
+        output = self.net(x)
+
+        # range clipping
+        if self.range_clipping:
+            return self.proj(output) * 2.0 + 3
+        else:
+            return output
+    def get_output_dim(self):
+        return self.output_dim