Update asr.py

asr.py

@@ -1,9 +1,18 @@
 import librosa
 from transformers import Wav2Vec2ForCTC, AutoProcessor
 import torch
+import json
+
+from huggingface_hub import hf_hub_download
+from torchaudio.models.decoder import ctc_decoder
 
 ASR_SAMPLING_RATE = 16_000
 
+ASR_LANGUAGES = {}
+with open(f"data/asr/all_langs.tsv") as f:
+    for line in f:
+        iso, name = line.split(" ", 1)
+        ASR_LANGUAGES[iso] = name
 
 MODEL_ID = "facebook/mms-1b-all"
 
@@ -11,21 +20,68 @@ processor = AutoProcessor.from_pretrained(MODEL_ID)
 model = Wav2Vec2ForCTC.from_pretrained(MODEL_ID)
 
 
-def transcribe(microphone, file_upload, lang):
+# lm_decoding_config = {}
+# lm_decoding_configfile = hf_hub_download(
+#     repo_id="facebook/mms-cclms",
+#     filename="decoding_config.json",
+#     subfolder="mms-1b-all",
+# )
 
-    warn_output = ""
-    if (microphone is not None) and (file_upload is not None):
-        warn_output = (
-            "WARNING: You've uploaded an audio file and used the microphone. "
-            "The recorded file from the microphone will be used and the uploaded audio will be discarded.\n"
-        )
-    elif (microphone is None) and (file_upload is None):
-        return "ERROR: You have to either use the microphone or upload an audio file"
+# with open(lm_decoding_configfile) as f:
+#     lm_decoding_config = json.loads(f.read())
+
+# # allow language model decoding for "eng"
+
+# decoding_config = lm_decoding_config["eng"]
 
-    audio_fp = microphone if microphone is not None else file_upload
+# lm_file = hf_hub_download(
+#     repo_id="facebook/mms-cclms",
+#     filename=decoding_config["lmfile"].rsplit("/", 1)[1],
+#     subfolder=decoding_config["lmfile"].rsplit("/", 1)[0],
+# )
+# token_file = hf_hub_download(
+#     repo_id="facebook/mms-cclms",
+#     filename=decoding_config["tokensfile"].rsplit("/", 1)[1],
+#     subfolder=decoding_config["tokensfile"].rsplit("/", 1)[0],
+# )
+# lexicon_file = None
+# if decoding_config["lexiconfile"] is not None:
+#     lexicon_file = hf_hub_download(
+#         repo_id="facebook/mms-cclms",
+#         filename=decoding_config["lexiconfile"].rsplit("/", 1)[1],
+#         subfolder=decoding_config["lexiconfile"].rsplit("/", 1)[0],
+#     )
+    
+# beam_search_decoder = ctc_decoder(
+#     lexicon=lexicon_file,
+#     tokens=token_file,
+#     lm=lm_file,
+#     nbest=1,
+#     beam_size=500,
+#     beam_size_token=50,
+#     lm_weight=float(decoding_config["lmweight"]),
+#     word_score=float(decoding_config["wordscore"]),
+#     sil_score=float(decoding_config["silweight"]),
+#     blank_token="<s>",
+# )
+
+
+def transcribe(
+    audio_source=None, microphone=None, file_upload=None, lang="eng (English)"
+):
+    if type(microphone) is dict:
+        # HACK: microphone variable is a dict when running on examples
+        microphone = microphone["name"]
+    audio_fp = (
+        file_upload if "upload" in str(audio_source or "").lower() else microphone
+    )
+    
+    if audio_fp is None:
+        return "ERROR: You have to either use the microphone or upload an audio file"
+    
     audio_samples = librosa.load(audio_fp, sr=ASR_SAMPLING_RATE, mono=True)[0]
 
-    lang_code = lang.split(":")[0]
+    lang_code = lang.split()[0]
     processor.tokenizer.set_target_lang(lang_code)
     model.load_adapter(lang_code)
 
@@ -33,9 +89,42 @@ def transcribe(microphone, file_upload, lang):
         audio_samples, sampling_rate=ASR_SAMPLING_RATE, return_tensors="pt"
     )
 
+    # set device
+    if torch.cuda.is_available():
+        device = torch.device("cuda")
+    elif (
+        hasattr(torch.backends, "mps")
+        and torch.backends.mps.is_available()
+        and torch.backends.mps.is_built()
+    ):
+        device = torch.device("mps")
+    else:
+        device = torch.device("cpu")
+
+    model.to(device)
+    inputs = inputs.to(device)
+
     with torch.no_grad():
         outputs = model(**inputs).logits
 
-    ids = torch.argmax(outputs, dim=-1)[0]
-    transcription = processor.decode(ids)
-    return warn_output + transcription
+    if lang_code != "eng" or True:
+        ids = torch.argmax(outputs, dim=-1)[0]
+        transcription = processor.decode(ids)
+    else:
+        assert False 
+        # beam_search_result = beam_search_decoder(outputs.to("cpu"))
+        # transcription = " ".join(beam_search_result[0][0].words).strip()
+
+    return transcription
+
+
+ASR_EXAMPLES = [
+    [None, "assets/english.mp3", None, "eng (English)"],
+    # [None, "assets/tamil.mp3", None, "tam (Tamil)"],
+    # [None, "assets/burmese.mp3", None, "mya (Burmese)"],
+]
+
+ASR_NOTE = """
+The above demo doesn't use beam-search decoding using a language model. 
+Checkout the instructions [here](https://huggingface.co/facebook/mms-1b-all) on how to run LM decoding for better accuracy.
+"""