import gradio as gr import streamlit as st import logging import sys import firebase_admin import tempfile import numpy as np import datetime from firebase_admin import credentials from firebase_admin import firestore #from transformers import pipeline 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": True}, "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 = { "Dutch": {"model_id": "jonatasgrosman/wav2vec2-xls-r-1b-dutch", "has_lm": True}, "English": {"model_id": "jonatasgrosman/wav2vec2-xls-r-1b-english", "has_lm": True}, "French": {"model_id": "jonatasgrosman/wav2vec2-xls-r-1b-french", "has_lm": True}, "German": {"model_id": "jonatasgrosman/wav2vec2-xls-r-1b-german", "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}, "Russian": {"model_id": "jonatasgrosman/wav2vec2-xls-r-1b-russian", "has_lm": True}, "Spanish": {"model_id": "jonatasgrosman/wav2vec2-xls-r-1b-spanish", "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.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="Automatic Speech Recognition", 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" ).launch(enable_queue=True)