Update app.py

app.py

@@ -2,8 +2,88 @@ import gradio as gr
 from transformers import pipeline
 from librosa import resample
 import numpy as np
+import os
+import sys
+import glob
+import torch
+from huggingface_hub import snapshot_download
+from fairseq_chunking import infer_fairseq_with_chunking
 
-def transcribe(input_audio):
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+fairseq_cache = {}
+
+
+def load_dictionary(dict_path):
+    """Load the fairseq dictionary file"""
+    dictionary = {}
+    special_tokens = ['<s>', '<pad>', '</s>', '<unk>']
+    for i, token in enumerate(special_tokens):
+        dictionary[i] = token
+    with open(dict_path, 'r', encoding='utf-8') as f:
+        for line in f:
+            parts = line.strip().split()
+            if len(parts) >= 1:
+                token = parts[0]
+                dictionary[len(dictionary)] = token
+    return dictionary
+
+
+def decode_predictions(emissions, dictionary, verbose=False):
+    predictions = emissions.argmax(dim=-1)
+    vocab_size = emissions.shape[-1]
+    blank_idx = 0
+
+    tokens = []
+    prev_token_id = None
+
+    for pred in predictions[0]:
+        token_id = pred.item()
+        if token_id == blank_idx:
+            prev_token_id = None
+            continue
+        if token_id == prev_token_id or token_id < 4:
+            if token_id != prev_token_id and token_id < 4:
+                prev_token_id = None
+            continue
+        token = dictionary.get(token_id, f'<unk_{token_id}>')
+        tokens.append(token)
+        prev_token_id = token_id
+
+    transcription = ''.join(tokens).replace('|', ' ').strip()
+    return transcription
+
+
+def load_fairseq_model(model_id):
+    """Download the fairseq model repo, import the encoders, and load the checkpoint"""
+    if model_id in fairseq_cache:
+        return fairseq_cache[model_id]
+
+    local_path = snapshot_download(model_id)
+    if local_path not in sys.path:
+        sys.path.insert(0, local_path)
+
+    original_load = torch.load
+    torch.load = lambda *args, **kwargs: original_load(*args, **{**kwargs, "weights_only": False})
+    try:
+        import fairseq_extra_encoders
+        import fairseq
+        checkpoint_path = os.path.join(local_path, "fairseq_checkpoint.pt")
+        models, cfg, task = fairseq.checkpoint_utils.load_model_ensemble_and_task([checkpoint_path])
+    finally:
+        torch.load = original_load
+
+    model = models[0]
+    model.eval()
+    model = model.to(device)
+
+    dict_path = os.path.join(local_path, "dict.ltr.txt")
+    dictionary = load_dictionary(dict_path)
+
+    fairseq_cache[model_id] = (model, dictionary, cfg.task.normalize)
+    return fairseq_cache[model_id]
+
+
+def transcribe(input_audio, model_id):
     sr, speech = input_audio
     # Convert to mono if stereo
     if speech.ndim > 1:
@@ -12,23 +92,51 @@ def transcribe(input_audio):
     if speech.dtype != "float32":
         speech = speech.astype(np.float32)
     # Resample if sampling rate is not 16kHz
-    if sr!=16000:
+    if sr != 16000:
         speech = resample(speech, orig_sr=sr, target_sr=16000)
-    output = pipe(speech, chunk_length_s=30, stride_length_s=5)['text']
+        sr = 16000
+
+    if "ebranch" in model_id:
+        model, dictionary, normalize_audio = load_fairseq_model(model_id)
+        output = infer_fairseq_with_chunking(
+            audio=speech,
+            sampling_rate=sr,
+            model=model,
+            dictionary=dictionary,
+            device=device,
+            normalize_audio=normalize_audio,
+            chunk_length_s=30.0,
+            stride_length_s=(5.0, 5.0),
+            model_downsample_ratio=320.0,
+            decode_fn=lambda emissions, dict: decode_predictions(emissions, dict, verbose=False)
+        )
+    else:
+        pipe = pipeline(
+            "automatic-speech-recognition",
+            model=model_id,
+            device="cpu"
+        )
+        output = pipe(speech, chunk_length_s=30, stride_length_s=5)['text']
     return output
 
-pipe = pipeline(
-    "automatic-speech-recognition",
-    model="GetmanY1/wav2vec2-large-sami-cont-pt-22k-finetuned",
-    device="cpu"
-    )
 
+model_ids_list = [
+ "GetmanY1/wav2vec2-large-sami-cont-pt-22k-finetuned",
+ "GetmanY1/wav2vec2-large-ebranch-sami-18k-finetuned-experimental"
+]
 gradio_app = gr.Interface(
     fn=transcribe,
-    inputs=gr.Audio(sources=["upload","microphone"]),
+    inputs=[
+        gr.Audio(sources=["upload","microphone"]),
+        gr.Dropdown(
+            label="Model",
+            value="GetmanY1/wav2vec2-large-sami-cont-pt-22k-finetuned",
+            choices=model_ids_list
+        )
+        ],
     outputs="text",
     title="Sámi Automatic Speech Recognition",
+    description ="Choose a model from the list."
 )
-
 if __name__ == "__main__":
     gradio_app.launch(server_name="0.0.0.0", server_port=7860)