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}, #"Dutch": {"model_id": "jonatasgrosman/wav2vec2-large-xlsr-53-dutch", "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}, #"Dutch": {"model_id": "jonatasgrosman/wav2vec2-xls-r-1b-dutch", "has_lm": False}, "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(LARGE_MODEL_BY_LANGUAGE.keys()) # the container given by HF has 16GB of RAM, so we need to limit the number of models to load 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 = AutoModelForCTC.from_pretrained(model["model_id"]) 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 = "
" for item in history: if item["error_message"] is not None: html_output += f"
{item['error_message']}
" else: url_suffix = " + LM" if item["decoding_type"] == "LM" else "" html_output += "
" html_output += f'{item["model_id"]}{url_suffix}

' html_output += f'{item["transcription"]}
' html_output += "
" html_output += "
" return html_output, history gr.Interface( run, inputs=[ #gr.inputs.Audio(source="microphone", type="filepath", label="Record something..."), gr.Audio(source="microphone", type='filepath', streaming=True), #gr.inputs.Audio(source="microphone", type="filepath", label="Record something...", streaming="True"), gr.inputs.Radio(label="Language", choices=LANGUAGES), gr.inputs.Radio(label="Decoding type", choices=["greedy", "LM"]), # gr.inputs.Radio(label="Model size", choices=["300M", "1B"]), "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 # test1 ).launch(enable_queue=True)