Adjust app to support the use of cleaner model interface

app.py

@@ -1,19 +1,19 @@
 import streamlit as st
-import torch
-from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
 from speechbrain.pretrained import GraphemeToPhoneme
 import os
 import torchaudio
 from wav2vecasr.MispronounciationDetector import MispronounciationDetector
+from wav2vecasr.PhonemeASRModel import Wav2Vec2PhonemeASRModel, Wav2Vec2OptimisedPhonemeASRModel
+
 
 @st.cache_resource
 def load_model():
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    path = os.path.join(os.getcwd(), "wav2vecasr", "model", "checkpoint-1200")
-    model = Wav2Vec2ForCTC.from_pretrained(path).to(device)
-    processor = Wav2Vec2Processor.from_pretrained(path)
+    path = os.path.join(os.getcwd(), "wav2vecasr", "model", "checkpoint-600")
+    asr_model = Wav2Vec2OptimisedPhonemeASRModel(path, os.path.join(path, "wav2vec2_vocab_final.json"),
+                                                         os.path.join(os.getcwd(), "wav2vecasr", "pretrained_models",
+                                                                      "ken-lm-ngram"))
     g2p = GraphemeToPhoneme.from_hparams("speechbrain/soundchoice-g2p")
-    mispronounciation_detector = MispronounciationDetector(model, processor, g2p, "cpu")
+    mispronounciation_detector = MispronounciationDetector(asr_model, g2p, "cpu")
     return mispronounciation_detector
 
 

requirements.txt

@@ -4,7 +4,7 @@ altair==5.1.2
 async-timeout==4.0.3
 attrs==23.1.0
 audioread==3.0.1
-backports.zoneinfo;python_version<"3.9"
+backports.zoneinfo==0.2.1
 blinker==1.6.3
 cachetools==5.3.1
 certifi==2023.7.22
@@ -15,6 +15,7 @@ colorama==0.4.6
 datasets==1.18.3
 decorator==5.1.1
 dill==0.3.7
+exceptiongroup==1.1.3
 filelock==3.12.4
 frozenlist==1.4.0
 fsspec==2023.9.2
@@ -22,6 +23,7 @@ gitdb==4.0.10
 GitPython==3.1.37
 huggingface-hub==0.17.3
 HyperPyYAML==1.2.2
+hypothesis==6.88.1
 idna==3.4
 importlib-metadata==6.8.0
 importlib-resources==6.1.0
@@ -30,6 +32,7 @@ jiwer==3.0.3
 joblib==1.3.2
 jsonschema==4.19.1
 jsonschema-specifications==2023.7.1
+kenlm @ git+https://github.com/kpu/kenlm/archive/master.zip
 lazy_loader==0.3
 librosa==0.10.1
 llvmlite==0.41.0
@@ -52,8 +55,10 @@ pooch==1.7.0
 protobuf==4.24.4
 pyarrow==13.0.0
 pycparser==2.21
+pyctcdecode==0.5.0
 pydeck==0.8.1b0
 Pygments==2.16.1
+pygtrie==2.5.0
 python-dateutil==2.8.2
 pytz==2023.3.post1
 PyYAML==6.0.1
@@ -71,6 +76,7 @@ scipy==1.10.1
 sentencepiece==0.1.99
 six==1.16.0
 smmap==5.0.1
+sortedcontainers==2.4.0
 soundfile==0.12.1
 soxr==0.3.7
 speechbrain==0.5.15