Spaces:
Running
Running
| from datetime import datetime, timedelta | |
| import os | |
| import gradio as gr | |
| import nemo.collections.asr as nemo_asr | |
| import wandb | |
| MODEL_HISTORY_DAYS = 180 | |
| WANDB_ENTITY = os.environ.get("WANDB_ENTITY", "tarteel") | |
| WANDB_PROJECT_NAME = os.environ.get("WANDB_PROJECT_NAME", "nemo-experiments") | |
| wandb_api = wandb.Api(overrides={"entity": WANDB_ENTITY}) | |
| all_artifacts_versions = [ | |
| version | |
| for version in [ | |
| collection.versions() | |
| for collection in wandb_api.artifact_type( | |
| type_name="model", project=WANDB_PROJECT_NAME | |
| ).collections() | |
| ] | |
| ] | |
| latest_artifacts = [ | |
| artifact | |
| for artifact_versions in all_artifacts_versions | |
| for artifact in artifact_versions | |
| if ( | |
| datetime.fromisoformat(artifact.created_at) | |
| > datetime.now() - timedelta(days=MODEL_HISTORY_DAYS) # last 180 days | |
| and artifact.state != "DELETED" | |
| ) | |
| ] | |
| latest_artifacts.sort(key=lambda a: a.created_at, reverse=True) | |
| models = {artifact.name: None for artifact in latest_artifacts} | |
| def lazy_load_models(models_names): | |
| for model_name in models_names: | |
| model = models[model_name] | |
| if not model: | |
| models[model_name] = nemo_asr.models.EncDecCTCModelBPE.restore_from( | |
| list(filter(lambda x: x.name == model_name, latest_artifacts))[0].file() | |
| ) | |
| models[model_name].eval() | |
| def transcribe(audio_mic, audio_file, models_names): | |
| lazy_load_models(models_names) | |
| # transcribe audio_mic and audio_file separately | |
| # because transcribe() fails is path is empty | |
| transcription_mic = "\n".join( | |
| [ | |
| f"{model_name} => {models[model_name].transcribe([audio_mic])[0]}" | |
| for model_name in models_names | |
| ] | |
| if audio_mic | |
| else "" | |
| ) | |
| transcription_file = "\n".join( | |
| [ | |
| f"{model_name} => {models[model_name].transcribe([audio_file])[0]}" | |
| for model_name in models_names | |
| ] | |
| if audio_file | |
| else "" | |
| ) | |
| return transcription_mic, transcription_file | |
| model_selection = list(models.keys()) | |
| demo = gr.Blocks() | |
| with demo: | |
| gr.Markdown( | |
| """ | |
| # ﷽ | |
| These are the latest* Tarteel models. | |
| Please note that the models are lazy loaded. | |
| This means that the first time you use a model, | |
| it might take some time to be downloaded and loaded for inference. | |
| *: last 180 days since the space was launched. | |
| To update the list, restart the space. | |
| """ | |
| ) | |
| with gr.Row(): | |
| audio_mic = gr.Audio(source="microphone", type="filepath", label="Microphone") | |
| audio_file = gr.Audio(source="upload", type="filepath", label="File") | |
| with gr.Row(): | |
| output_mic = gr.TextArea(label="Microphone Transcription") | |
| output_file = gr.TextArea(label="Audio Transcription") | |
| models_names = gr.CheckboxGroup(model_selection, label="Select Models to Use") | |
| b1 = gr.Button("Transcribe") | |
| b1.click( | |
| transcribe, | |
| inputs=[audio_mic, audio_file, models_names], | |
| outputs=[output_mic, output_file], | |
| ) | |
| demo.launch() | |