Update demo to use MultitaskASRModel

wav2vecasr/demo.py

@@ -3,13 +3,17 @@ from speechbrain.pretrained import GraphemeToPhoneme
 import datasets
 import os
 import torchaudio
-from MispronounciationDetector import MispronounciationDetector
+from wav2vecasr.MispronounciationDetector import MispronounciationDetector
+from wav2vecasr.PhonemeASRModel import Wav2Vec2PhonemeASRModel, Wav2Vec2OptimisedPhonemeASRModel, MultitaskPhonemeASRModel
+import jiwer
+import re
 
 # Load sample data
-audio_path, transcript_path = os.path.join(os.getcwd(), "data", "arctic_a0003.wav"), os.path.join(os.getcwd(), "data", "arctic_a0003.txt")
+audio_path, transcript_path = os.path.join(os.getcwd(), "data", "arctic_a0003.wav"), os.path.join(os.getcwd(),"data", "arctic_a0003.txt")
 audio, org_sr = torchaudio.load(audio_path)
 audio = torchaudio.functional.resample(audio, orig_freq=org_sr, new_freq=16000)
 audio = audio.view(audio.shape[1])
+audio = audio.to("cpu")
 with open(transcript_path) as f:
   text = f.read()
 f.close()
@@ -17,16 +21,16 @@ print("Done loading sample data")
 
 # Load processors and models
 device = "cpu"
-path = os.path.join(os.getcwd(), "model", "checkpoint-1200")
-model = Wav2Vec2ForCTC.from_pretrained(path).to(device)
-processor = Wav2Vec2Processor.from_pretrained(path)
+path = os.path.join(os.getcwd(), "model", "multitask_best_ctc.pt")
+vocab_path = os.path.join(os.getcwd(), "model", "vocab")
+asr_model = MultitaskPhonemeASRModel(path, vocab_path, device)
 g2p = GraphemeToPhoneme.from_hparams("speechbrain/soundchoice-g2p")
-mispronounciation_detector = MispronounciationDetector(model, processor, g2p, "cpu")
+mispronounciation_detector = MispronounciationDetector(asr_model, g2p, "cpu")
 print("Done loading models and processors")
 
 # Predict
 raw_info = mispronounciation_detector.detect(audio, text)
-aligned_phoneme_output_delimited_by_words = " ".join(raw_info['ref']) + "\n" + " ".join(raw_info['hyp']) + "\n" +\
-                                            " ".join(raw_info['phoneme_errors'])
+print(raw_info['ref'])
+print(raw_info['hyp'])
+print(raw_info['phoneme_errors'])
 print(f"PER: {raw_info['per']}\n")
-print(f"Phoneme level errors:\n{raw_info['phoneme_output']}\n")