Spaces:
Running
Running
File size: 4,444 Bytes
7e6e084 cf20a06 0b960f8 eeadd68 7e6e084 eeadd68 cf20a06 d4dffb9 cf20a06 eeadd68 d08907e eeadd68 7e6e084 eeadd68 7e6e084 ca8157c eeadd68 ca8157c eeadd68 7e6e084 200b48f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 | 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
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:
speech = speech.mean(axis=1)
# Convert to float32 if needed
if speech.dtype != "float32":
speech = speech.astype(np.float32)
# Resample if sampling rate is not 16kHz
if sr != 16000:
speech = resample(speech, orig_sr=sr, target_sr=16000)
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
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"]),
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) |