|
import gradio as gr |
|
import logging |
|
import sys |
|
import tempfile |
|
import numpy as np |
|
import datetime |
|
|
|
from transformers import pipeline, AutoModelForCTC, Wav2Vec2Processor, Wav2Vec2ProcessorWithLM |
|
from typing import Optional |
|
from TTS.utils.manage import ModelManager |
|
from TTS.utils.synthesizer import Synthesizer |
|
|
|
logging.basicConfig( |
|
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", |
|
datefmt="%m/%d/%Y %H:%M:%S", |
|
handlers=[logging.StreamHandler(sys.stdout)], |
|
) |
|
logger = logging.getLogger(__name__) |
|
logger.setLevel(logging.DEBUG) |
|
|
|
|
|
LARGE_MODEL_BY_LANGUAGE = { |
|
"Arabic": {"model_id": "jonatasgrosman/wav2vec2-large-xlsr-53-arabic", "has_lm": False}, |
|
"Chinese": {"model_id": "jonatasgrosman/wav2vec2-large-xlsr-53-chinese-zh-cn", "has_lm": False}, |
|
|
|
"English": {"model_id": "jonatasgrosman/wav2vec2-large-xlsr-53-english", "has_lm": True}, |
|
"Finnish": {"model_id": "jonatasgrosman/wav2vec2-large-xlsr-53-finnish", "has_lm": False}, |
|
"French": {"model_id": "jonatasgrosman/wav2vec2-large-xlsr-53-french", "has_lm": True}, |
|
"German": {"model_id": "jonatasgrosman/wav2vec2-large-xlsr-53-german", "has_lm": True}, |
|
"Greek": {"model_id": "jonatasgrosman/wav2vec2-large-xlsr-53-greek", "has_lm": False}, |
|
"Hungarian": {"model_id": "jonatasgrosman/wav2vec2-large-xlsr-53-hungarian", "has_lm": False}, |
|
"Italian": {"model_id": "jonatasgrosman/wav2vec2-large-xlsr-53-italian", "has_lm": True}, |
|
"Japanese": {"model_id": "jonatasgrosman/wav2vec2-large-xlsr-53-japanese", "has_lm": False}, |
|
"Persian": {"model_id": "jonatasgrosman/wav2vec2-large-xlsr-53-persian", "has_lm": False}, |
|
"Polish": {"model_id": "jonatasgrosman/wav2vec2-large-xlsr-53-polish", "has_lm": True}, |
|
"Portuguese": {"model_id": "jonatasgrosman/wav2vec2-large-xlsr-53-portuguese", "has_lm": True}, |
|
"Russian": {"model_id": "jonatasgrosman/wav2vec2-large-xlsr-53-russian", "has_lm": True}, |
|
"Spanish": {"model_id": "jonatasgrosman/wav2vec2-large-xlsr-53-spanish", "has_lm": True}, |
|
} |
|
|
|
XLARGE_MODEL_BY_LANGUAGE = { |
|
"English": {"model_id": "jonatasgrosman/wav2vec2-xls-r-1b-english", "has_lm": True}, |
|
"Spanish": {"model_id": "jonatasgrosman/wav2vec2-xls-r-1b-spanish", "has_lm": True}, |
|
"German": {"model_id": "jonatasgrosman/wav2vec2-xls-r-1b-german", "has_lm": True}, |
|
"Russian": {"model_id": "jonatasgrosman/wav2vec2-xls-r-1b-russian", "has_lm": True}, |
|
"French": {"model_id": "jonatasgrosman/wav2vec2-xls-r-1b-french", "has_lm": True}, |
|
"Italian": {"model_id": "jonatasgrosman/wav2vec2-xls-r-1b-italian", "has_lm": True}, |
|
|
|
"Polish": {"model_id": "jonatasgrosman/wav2vec2-xls-r-1b-polish", "has_lm": True}, |
|
"Portuguese": {"model_id": "jonatasgrosman/wav2vec2-xls-r-1b-portuguese", "has_lm": True}, |
|
} |
|
|
|
|
|
|
|
|
|
|
|
LANGUAGES = sorted(XLARGE_MODEL_BY_LANGUAGE.keys()) |
|
CACHED_MODELS_BY_ID = {} |
|
|
|
|
|
def run(input_file, language, decoding_type, history, model_size="300M"): |
|
|
|
logger.info(f"Running ASR {language}-{model_size}-{decoding_type} for {input_file}") |
|
|
|
history = history or [] |
|
|
|
if model_size == "300M": |
|
model = LARGE_MODEL_BY_LANGUAGE.get(language, None) |
|
else: |
|
model = XLARGE_MODEL_BY_LANGUAGE.get(language, None) |
|
|
|
if model is None: |
|
history.append({ |
|
"error_message": f"Model size {model_size} not found for {language} language :(" |
|
}) |
|
elif decoding_type == "LM" and not model["has_lm"]: |
|
history.append({ |
|
"error_message": f"LM not available for {language} language :(" |
|
}) |
|
else: |
|
|
|
|
|
model_instance = CACHED_MODELS_BY_ID.get(model["model_id"], None) |
|
if model_instance is None: |
|
model_instance = AutoModelForCTC.from_pretrained(model["model_id"]) |
|
CACHED_MODELS_BY_ID[model["model_id"]] = model_instance |
|
|
|
if decoding_type == "LM": |
|
processor = Wav2Vec2ProcessorWithLM.from_pretrained(model["model_id"]) |
|
asr = pipeline("automatic-speech-recognition", model=model_instance, tokenizer=processor.tokenizer, |
|
feature_extractor=processor.feature_extractor, decoder=processor.decoder) |
|
else: |
|
processor = Wav2Vec2Processor.from_pretrained(model["model_id"]) |
|
asr = pipeline("automatic-speech-recognition", model=model_instance, tokenizer=processor.tokenizer, |
|
feature_extractor=processor.feature_extractor, decoder=None) |
|
|
|
transcription = asr(input_file, chunk_length_s=5, stride_length_s=1)["text"] |
|
|
|
logger.info(f"Transcription for {input_file}: {transcription}") |
|
|
|
history.append({ |
|
"model_id": model["model_id"], |
|
"language": language, |
|
"model_size": model_size, |
|
"decoding_type": decoding_type, |
|
"transcription": transcription, |
|
"error_message": None |
|
}) |
|
|
|
html_output = "<div class='result'>" |
|
for item in history: |
|
if item["error_message"] is not None: |
|
html_output += f"<div class='result_item result_item_error'>{item['error_message']}</div>" |
|
else: |
|
url_suffix = " + LM" if item["decoding_type"] == "LM" else "" |
|
html_output += "<div class='result_item result_item_success'>" |
|
html_output += f'<strong><a target="_blank" href="https://huggingface.co/{item["model_id"]}">{item["model_id"]}{url_suffix}</a></strong><br/><br/>' |
|
html_output += f'{item["transcription"]}<br/>' |
|
html_output += "</div>" |
|
html_output += "</div>" |
|
|
|
return html_output, history |
|
|
|
|
|
gr.Interface( |
|
run, |
|
inputs=[ |
|
|
|
gr.Audio(source="microphone", type='filepath', streaming=True), |
|
|
|
gr.inputs.Radio(label="Language", choices=LANGUAGES), |
|
gr.inputs.Radio(label="Decoding type", choices=["greedy", "LM"]), |
|
|
|
"state" |
|
], |
|
outputs=[ |
|
gr.outputs.HTML(label="Outputs"), |
|
"state" |
|
], |
|
title="🗣️NLP ASR Wav2Vec2 GR📄", |
|
description="", |
|
css=""" |
|
.result {display:flex;flex-direction:column} |
|
.result_item {padding:15px;margin-bottom:8px;border-radius:15px;width:100%} |
|
.result_item_success {background-color:mediumaquamarine;color:white;align-self:start} |
|
.result_item_error {background-color:#ff7070;color:white;align-self:start} |
|
""", |
|
allow_screenshot=False, |
|
allow_flagging="never", |
|
theme="grass", |
|
live=True |
|
).launch(enable_queue=True) |
