Spaces:
Sleeping
Sleeping
| # Import Libraries to load Models | |
| from transformers import WhisperProcessor, WhisperForConditionalGeneration | |
| from transformers import Speech2TextForConditionalGeneration, Speech2TextProcessor | |
| from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC | |
| # Import Libraries to access Datasets | |
| from datasets import load_dataset | |
| from datasets import Audio | |
| # Helper Libraries | |
| import plotly.graph_objs as go | |
| import evaluate | |
| import librosa | |
| import torch | |
| import numpy as np | |
| import pandas as pd | |
| import time | |
| # This constant determines on how many samples the Models are evaluated on | |
| N_SAMPLES = 50 | |
| # Load the WER Metric | |
| wer_metric = evaluate.load("wer") | |
| def run(data_subset:str, model_1:str, model_2:str, own_audio, own_transcription:str): | |
| """ | |
| Main Function running an entire evaluation cycle | |
| Params: | |
| - data_subset (str) :The name of a valid Dataset to choose from ["Common Voice", "Librispeech ASR clean", "Librispeech ASR other", "OWN Recording/Sample"] | |
| - model_1 (str) :The name of a valid model to choose form ["openai/whisper-tiny.en", "facebook/s2t-medium-librispeech-asr", "facebook/wav2vec2-base-960h","openai/whisper-large-v2"] | |
| - model_2 (str) :The name of a valid model to choose form ["openai/whisper-tiny.en", "facebook/s2t-medium-librispeech-asr", "facebook/wav2vec2-base-960h","openai/whisper-large-v2"] | |
| - own_audio (gr.Audio) :The return value of an gr.Audio component (sr, audio (as numpy array)) | |
| - own_transcription (str) :The paired transcription to the own_audio | |
| """ | |
| # A little bit of Error Handling | |
| if data_subset is None and own_audio is None and own_transcription is None: | |
| raise ValueError("No Dataset selected") | |
| if model_1 is None: | |
| raise ValueError("No Model 1 selected") | |
| if model_2 is None: | |
| raise ValueError("No Model 2 selected") | |
| # Load the selected Dataset but only N_SAMPLES of it | |
| if data_subset == "Common Voice": | |
| dataset, text_column = load_Common_Voice() | |
| elif data_subset == "Librispeech ASR clean": | |
| dataset, text_column = load_Librispeech_ASR_clean() | |
| elif data_subset == "Librispeech ASR other": | |
| dataset, text_column = load_Librispeech_ASR_other() | |
| elif data_subset == "OWN Recording/Sample": | |
| sr, audio = own_audio | |
| audio = audio.astype(np.float32) | |
| print("AUDIO: ", type(audio), audio) | |
| audio = librosa.resample(audio, orig_sr=sr, target_sr=16000) | |
| else: | |
| # if data_subset is None then still load load_Common_Voice | |
| dataset, text_column = load_Common_Voice() | |
| # I have left the print statements because users have access to the logs in Spaces and this might help to understand what's going on | |
| print("Dataset Loaded") | |
| # Load the selected Models | |
| model1, processor1 = load_model(model_1) | |
| model2, processor2 = load_model(model_2) | |
| print("Models Loaded") | |
| # In case a own Recording is selected only a single sample has to be evaluated | |
| if data_subset == "OWN Recording/Sample": | |
| sample = {"audio":{"array":audio,"sampling_rate":16000}} | |
| inference_times1 = [] | |
| inference_times2 = [] | |
| time_start = time.time() | |
| transcription1 = model_compute(model1, processor1, sample, model_1) | |
| time_stop = time.time() | |
| duration = time_stop - time_start | |
| inference_times1.append(duration) | |
| time_start = time.time() | |
| transcription2 = model_compute(model2, processor2, sample, model_2) | |
| time_stop = time.time() | |
| duration = time_stop - time_start | |
| inference_times2.append(duration) | |
| transcriptions1 = [transcription1] | |
| transcriptions2 = [transcription2] | |
| references = [own_transcription.lower()] | |
| wer1 = round(N_SAMPLES * compute_wer(references, transcriptions1), 2) | |
| wer2 = round(N_SAMPLES * compute_wer(references, transcriptions2), 2) | |
| results_md = f""" | |
| #### {model_1} | |
| - WER Score: {wer1} | |
| - Avg. Inference Duration: {round(sum(inference_times1)/len(inference_times1), 4)}s | |
| #### {model_2} | |
| - WER Score: {wer2} | |
| - Avg. Inference Duration: {round(sum(inference_times2)/len(inference_times2), 4)}s""" | |
| # Create the bar plot | |
| fig = go.Figure( | |
| data=[ | |
| go.Bar(x=[f"{model_1}"], y=[wer1], showlegend=False), | |
| go.Bar(x=[f"{model_2}"], y=[wer2], showlegend=False), | |
| ] | |
| ) | |
| # Update the layout for better visualization | |
| fig.update_layout( | |
| title="Comparison of Two Models", | |
| xaxis_title="Models", | |
| yaxis_title="Value", | |
| barmode="group", | |
| ) | |
| df = pd.DataFrame({"references":references, "transcriptions 1":transcriptions1,"WER 1":[wer1],"transcriptions 2":transcriptions2,"WER 2":[wer2]}) | |
| yield results_md, fig, df | |
| # In case a Dataset has been selected | |
| else: | |
| references = [] | |
| transcriptions1 = [] | |
| transcriptions2 = [] | |
| WER1s = [] | |
| WER2s = [] | |
| inference_times1 = [] | |
| inference_times2 = [] | |
| counter = 0 | |
| for i, sample in enumerate(dataset, start=1): | |
| print(counter) | |
| counter += 1 | |
| references.append(sample[text_column]) | |
| if model_1 == model_2: | |
| time_start = time.time() | |
| transcription = model_compute(model1, processor1, sample, model_1) | |
| time_stop = time.time() | |
| duration = time_stop - time_start | |
| inference_times1.append(duration) | |
| inference_times2.append(duration) | |
| transcriptions1.append(transcription) | |
| transcriptions2.append(transcription) | |
| else: | |
| time_start = time.time() | |
| transcription1 = model_compute(model1, processor1, sample, model_1) | |
| time_stop = time.time() | |
| duration = time_stop - time_start | |
| inference_times1.append(duration) | |
| transcriptions1.append(transcription1) | |
| time_start = time.time() | |
| transcription2 = model_compute(model2, processor2, sample, model_2) | |
| time_stop = time.time() | |
| duration = time_stop - time_start | |
| inference_times2.append(duration) | |
| transcriptions2.append(transcription2) | |
| WER1s.append(round(compute_wer([sample[text_column]], [transcription1]),4)) | |
| WER2s.append(round(compute_wer([sample[text_column]], [transcription2]),4)) | |
| wer1 = round(sum(WER1s)/len(WER1s), 4) | |
| wer2 = round(sum(WER2s)/len(WER2s), 4) | |
| results_md = f""" | |
| {i}/{len(dataset)}-{'#'*i}{'_'*(N_SAMPLES-i)} | |
| #### {model_1} | |
| - WER Score: {wer1} | |
| - Avg. Inference Duration: {round(sum(inference_times1)/len(inference_times1), 4)}s | |
| #### {model_2} | |
| - WER Score: {wer2} | |
| - Avg. Inference Duration: {round(sum(inference_times2)/len(inference_times2), 4)}s""" | |
| # Create the bar plot | |
| fig = go.Figure( | |
| data=[ | |
| go.Bar(x=[f"{model_1}"], y=[wer1], showlegend=False), | |
| go.Bar(x=[f"{model_2}"], y=[wer2], showlegend=False), | |
| ] | |
| ) | |
| # Update the layout for better visualization | |
| fig.update_layout( | |
| title="Comparison of Two Models", | |
| xaxis_title="Models", | |
| yaxis_title="Value", | |
| barmode="group", | |
| ) | |
| df = pd.DataFrame({"references":references, f"{model_1}":transcriptions1,"WER 1":WER1s,f"{model_2}":transcriptions2,"WER 2":WER2s}) | |
| yield results_md, fig, df | |
| # DATASET LOADERS | |
| def load_Common_Voice(): | |
| dataset = load_dataset("mozilla-foundation/common_voice_11_0", "en", revision="streaming", split="test", streaming=True, token=True, trust_remote_code=True) | |
| text_column = "sentence" | |
| dataset = dataset.take(N_SAMPLES) | |
| dataset = dataset.cast_column("audio", Audio(sampling_rate=16000)) | |
| dataset = list(dataset) | |
| for sample in dataset: | |
| sample["text"] = sample["text"].lower() | |
| return dataset, text_column | |
| def load_Librispeech_ASR_clean(): | |
| dataset = load_dataset("librispeech_asr", "clean", split="test", streaming=True, token=True, trust_remote_code=True) | |
| print(next(iter(dataset))) | |
| text_column = "text" | |
| dataset = dataset.take(N_SAMPLES) | |
| dataset = dataset.cast_column("audio", Audio(sampling_rate=16000)) | |
| dataset = list(dataset) | |
| for sample in dataset: | |
| sample["text"] = sample["text"].lower() | |
| return dataset, text_column | |
| def load_Librispeech_ASR_other(): | |
| dataset = load_dataset("librispeech_asr", "other", split="test", streaming=True, token=True, trust_remote_code=True) | |
| print(next(iter(dataset))) | |
| text_column = "text" | |
| dataset = dataset.take(N_SAMPLES) | |
| dataset = dataset.cast_column("audio", Audio(sampling_rate=16000)) | |
| dataset = list(dataset) | |
| for sample in dataset: | |
| sample["text"] = sample["text"].lower() | |
| return dataset, text_column | |
| # MODEL LOADERS | |
| def load_model(model_id:str): | |
| if model_id == "openai/whisper-tiny.en": | |
| model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-tiny.en") | |
| processor = WhisperProcessor.from_pretrained("openai/whisper-tiny.en") | |
| elif model_id == "facebook/s2t-medium-librispeech-asr": | |
| model = Speech2TextForConditionalGeneration.from_pretrained("facebook/s2t-medium-librispeech-asr") | |
| processor = Speech2TextProcessor.from_pretrained("facebook/s2t-medium-librispeech-asr") | |
| elif model_id == "facebook/wav2vec2-base-960h": | |
| processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-base-960h") | |
| model = Wav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-base-960h") | |
| elif model_id == "openai/whisper-large-v2": | |
| processor = WhisperProcessor.from_pretrained("openai/whisper-large-v2") | |
| model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-large-v2") | |
| model.config.forced_decoder_ids = None | |
| else: # In case no model has been selected the Whipser-Tiny.En is selected - just for completeness | |
| model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-tiny.en") | |
| processor = WhisperProcessor.from_pretrained("openai/whisper-tiny.en") | |
| return model, processor | |
| # MODEL INFERENCE | |
| def model_compute(model, processor, sample, model_id): | |
| if model_id == "openai/whisper-tiny.en": | |
| sample = sample["audio"] | |
| input_features = processor(sample["array"], sampling_rate=sample["sampling_rate"], return_tensors="pt").input_features | |
| predicted_ids = model.generate(input_features) | |
| transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True) | |
| transcription = processor.tokenizer.normalize(transcription[0]) | |
| return transcription | |
| elif model_id == "facebook/s2t-medium-librispeech-asr": | |
| sample = sample["audio"] | |
| features = processor(sample["array"], sampling_rate=16000, padding=True, return_tensors="pt") | |
| input_features = features.input_features | |
| attention_mask = features.attention_mask | |
| gen_tokens = model.generate(input_features=input_features, attention_mask=attention_mask) | |
| transcription= processor.batch_decode(gen_tokens, skip_special_tokens=True) | |
| return transcription[0] | |
| elif model_id == "facebook/wav2vec2-base-960h": | |
| sample = sample["audio"] | |
| input_values = processor(sample["array"], sampling_rate=16000, return_tensors="pt", padding="longest").input_values # Batch size 1 | |
| logits = model(input_values).logits | |
| predicted_ids = torch.argmax(logits, dim=-1) | |
| transcription = processor.batch_decode(predicted_ids) | |
| return transcription[0].lower() | |
| elif model_id == "openai/whisper-large-v2": | |
| sample = sample["audio"] | |
| input_features = processor(sample["array"], sampling_rate=sample["sampling_rate"], return_tensors="pt").input_features | |
| predicted_ids = model.generate(input_features) | |
| transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True) | |
| transcription = processor.tokenizer.normalize(transcription[0]) | |
| print("TRANSCRIPTION Whisper Large v2: ", transcription) | |
| return transcription | |
| else: # In case no model has been selected the Whipser-Tiny.En is selected - just for completeness | |
| sample = sample["audio"] | |
| input_features = processor(sample["array"], sampling_rate=sample["sampling_rate"], return_tensors="pt").input_features | |
| predicted_ids = model.generate(input_features) | |
| transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True) | |
| return transcription[0] | |
| # UTILS | |
| def compute_wer(references, predictions): | |
| wer = wer_metric.compute(references=references, predictions=predictions) | |
| return wer | |