Spaces:
Runtime error
Runtime error
| import logging | |
| import sys | |
| import gradio as gr | |
| from transformers import pipeline, AutoModelForCTC, Wav2Vec2Processor, Wav2Vec2ProcessorWithLM | |
| 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 [] | |
| # the history seems to be not by session anymore, so I'll deactivate this for now | |
| history = [] | |
| 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.name, chunk_length_s=5, stride_length_s=1)["text"] | |
| logger.info(f"Transcription for {language}-{model_size}-{decoding_type} 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.inputs.Audio(source="microphone", type="file", 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) | |