Include MultitaskASRModel

wav2vecasr/PhonemeASRModel.py

@@ -1,4 +1,5 @@
 import torch
+from wav2vecasr.models import MultiTaskWav2Vec2
 from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor, Wav2Vec2ProcessorWithLM, \
   Wav2Vec2CTCTokenizer, Wav2Vec2FeatureExtractor
 import pyctcdecode
@@ -6,6 +7,7 @@ import json
 import re
 from sys import platform
 
+
 class PhonemeASRModel:
   def get_l2_phoneme_sequence(self, audio):
     """
@@ -38,6 +40,68 @@ class PhonemeASRModel:
     """
     pass
 
+class MultitaskPhonemeASRModel(PhonemeASRModel):
+    def __init__(self, model_path, best_model_vocab_path, device):
+        self.device = device
+        tokenizer = Wav2Vec2CTCTokenizer(best_model_vocab_path, unk_token="[UNK]", pad_token="[PAD]", word_delimiter_token="|")
+        feature_extractor = Wav2Vec2FeatureExtractor(
+          feature_size=1,
+          sampling_rate=16000,
+          padding_value=0.0,
+          do_normalize=True,
+          return_attention_mask=False,
+        )
+        processor = Wav2Vec2Processor(feature_extractor=feature_extractor, tokenizer=tokenizer)
+
+        wav2vec2_backbone = Wav2Vec2ForCTC.from_pretrained(
+          pretrained_model_name_or_path="facebook/wav2vec2-xls-r-300m",
+          ignore_mismatched_sizes=True,
+          ctc_loss_reduction="mean",
+          pad_token_id=processor.tokenizer.pad_token_id,
+          vocab_size=len(processor.tokenizer),
+          output_hidden_states=True,
+        )
+        wav2vec2_backbone = wav2vec2_backbone.to(device)
+
+        model = MultiTaskWav2Vec2(
+          wav2vec2_backbone=wav2vec2_backbone,
+          backbone_hidden_size=1024,
+          projection_hidden_size=256,
+          num_accent_class=3,
+        )
+
+        model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
+        model.to(device)
+        model.eval()
+
+        self.multitask_model = model
+        self.processor = processor
+
+    def get_l2_phoneme_sequence(self, audio):
+        audio = audio.unsqueeze(0)
+        audio = self.processor(audio, sampling_rate=16000).input_values[0]
+        audio = torch.tensor(audio, device=self.device)
+
+        with torch.no_grad():
+            _, lm_logits, _, _ = self.multitask_model(audio)
+        lm_preds = torch.argmax(lm_logits, dim=-1)
+
+        # Decode output results
+        pred_decoded = self.processor.batch_decode(lm_preds)
+
+        pred_phones = pred_decoded[0].split(" ")
+
+        # remove sil and sp
+        pred_phones = [phone for phone in pred_phones if phone != "sil" and phone != "sp"]
+
+        return pred_phones
+
+    def standardise_g2p_phoneme_sequence(self, phones):
+      return phones
+
+    def standardise_l2_artic_groundtruth_phoneme_sequence(self, phones):
+      return phones
+
 class Wav2Vec2PhonemeASRModel(PhonemeASRModel):
   """
   Uses greedy decoding
@@ -98,4 +162,3 @@ class Wav2Vec2OptimisedPhonemeASRModel(PhonemeASRModel):
   def standardise_l2_artic_groundtruth_phoneme_sequence(self, phones):
     return [re.sub(r'\d', "", phone_str) for phone_str in phones]
 
-