Spaces:
Runtime error
Runtime error
import os | |
import sys | |
import torch | |
import logging | |
import speechbrain as sb | |
from speechbrain.utils.distributed import run_on_main | |
from hyperpyyaml import load_hyperpyyaml | |
from pathlib import Path | |
import torchaudio.transforms as T | |
from cv_train import ASRCV | |
import torchaudio | |
import numpy as np | |
import kenlm | |
from pyctcdecode import build_ctcdecoder | |
import re | |
from torch.nn.utils.rnn import pad_sequence | |
import torch.optim as optim | |
import torch.nn as nn | |
# Commented out IPython magic to ensure Python compatibility. | |
hparams_file, run_opts, overrides = sb.parse_arguments(["TunisianASR/semi_trained.yaml"]) | |
# If distributed_launch=True then | |
# create ddp_group with the right communication protocol | |
sb.utils.distributed.ddp_init_group(run_opts) | |
with open(hparams_file) as fin: | |
hparams = load_hyperpyyaml(fin, overrides) | |
# Create experiment directory | |
sb.create_experiment_directory( | |
experiment_directory=hparams["output_folder"], | |
hyperparams_to_save=hparams_file, | |
overrides=overrides, | |
) | |
# Dataset prep (parsing Librispeech) | |
def dataio_prepare(hparams): | |
"""This function prepares the datasets to be used in the brain class. | |
It also defines the data processing pipeline through user-defined functions.""" | |
# 1. Define datasets | |
data_folder = hparams["data_folder"] | |
train_data = sb.dataio.dataset.DynamicItemDataset.from_csv( | |
csv_path=hparams["train_csv"], replacements={"data_root": data_folder}, | |
) | |
if hparams["sorting"] == "ascending": | |
# we sort training data to speed up training and get better results. | |
train_data = train_data.filtered_sorted( | |
sort_key="duration", | |
key_max_value={"duration": hparams["avoid_if_longer_than"]}, | |
) | |
# when sorting do not shuffle in dataloader ! otherwise is pointless | |
hparams["dataloader_options"]["shuffle"] = False | |
elif hparams["sorting"] == "descending": | |
train_data = train_data.filtered_sorted( | |
sort_key="duration", | |
reverse=True, | |
key_max_value={"duration": hparams["avoid_if_longer_than"]}, | |
) | |
# when sorting do not shuffle in dataloader ! otherwise is pointless | |
hparams["dataloader_options"]["shuffle"] = False | |
elif hparams["sorting"] == "random": | |
pass | |
else: | |
raise NotImplementedError( | |
"sorting must be random, ascending or descending" | |
) | |
valid_data = sb.dataio.dataset.DynamicItemDataset.from_csv( | |
csv_path=hparams["valid_csv"], replacements={"data_root": data_folder}, | |
) | |
# We also sort the validation data so it is faster to validate | |
valid_data = valid_data.filtered_sorted(sort_key="duration") | |
test_datasets = {} | |
for csv_file in hparams["test_csv"]: | |
name = Path(csv_file).stem | |
test_datasets[name] = sb.dataio.dataset.DynamicItemDataset.from_csv( | |
csv_path=csv_file, replacements={"data_root": data_folder} | |
) | |
test_datasets[name] = test_datasets[name].filtered_sorted( | |
sort_key="duration" | |
) | |
datasets = [train_data, valid_data] + [i for k, i in test_datasets.items()] | |
# 2. Define audio pipeline: | |
def audio_pipeline(wav): | |
info = torchaudio.info(wav) | |
sig = sb.dataio.dataio.read_audio(wav) | |
if len(sig.shape)>1 : | |
sig = torch.mean(sig, dim=1) | |
resampled = torchaudio.transforms.Resample( | |
info.sample_rate, hparams["sample_rate"], | |
)(sig) | |
return resampled | |
sb.dataio.dataset.add_dynamic_item(datasets, audio_pipeline) | |
label_encoder = sb.dataio.encoder.CTCTextEncoder() | |
# 3. Define text pipeline: | |
def text_pipeline(wrd): | |
yield wrd | |
char_list = list(wrd) | |
yield char_list | |
tokens_list = label_encoder.encode_sequence(char_list) | |
yield tokens_list | |
tokens = torch.LongTensor(tokens_list) | |
yield tokens | |
sb.dataio.dataset.add_dynamic_item(datasets, text_pipeline) | |
lab_enc_file = os.path.join(hparams["save_folder"], "label_encoder.txt") | |
special_labels = { | |
"blank_label": hparams["blank_index"], | |
"unk_label": hparams["unk_index"] | |
} | |
label_encoder.load_or_create( | |
path=lab_enc_file, | |
from_didatasets=[train_data], | |
output_key="char_list", | |
special_labels=special_labels, | |
sequence_input=True, | |
) | |
# 4. Set output: | |
sb.dataio.dataset.set_output_keys( | |
datasets, ["id", "sig", "wrd", "char_list", "tokens"], | |
) | |
return train_data, valid_data,test_datasets, label_encoder | |
class ASR(sb.core.Brain): | |
def compute_forward(self, batch, stage): | |
"""Forward computations from the waveform batches to the output probabilities.""" | |
batch = batch.to(self.device) | |
wavs, wav_lens = batch.sig | |
wavs, wav_lens = wavs.to(self.device), wav_lens.to(self.device) | |
if stage == sb.Stage.TRAIN: | |
if hasattr(self.hparams, "augmentation"): | |
wavs = self.hparams.augmentation(wavs, wav_lens) | |
# Forward pass | |
feats = self.modules.wav2vec2(wavs, wav_lens) | |
x = self.modules.enc(feats) | |
logits = self.modules.ctc_lin(x) | |
p_ctc = self.hparams.log_softmax(logits) | |
return p_ctc, wav_lens | |
def custom_encode(self,wavs,wav_lens) : | |
wavs = wavs.to("cpu") | |
if(wav_lens is not None): wav_lens.to(self.device) | |
feats = self.modules.wav2vec2(wavs, wav_lens) | |
x = self.modules.enc(feats) | |
logits = self.modules.ctc_lin(x) | |
p_ctc = self.hparams.log_softmax(logits) | |
return feats,p_ctc | |
def compute_objectives(self, predictions, batch, stage): | |
"""Computes the loss (CTC) given predictions and targets.""" | |
p_ctc, wav_lens = predictions | |
ids = batch.id | |
tokens, tokens_lens = batch.tokens | |
loss = self.hparams.ctc_cost(p_ctc, tokens, wav_lens, tokens_lens) | |
if stage != sb.Stage.TRAIN: | |
predicted_tokens = sb.decoders.ctc_greedy_decode( | |
p_ctc, wav_lens, blank_id=self.hparams.blank_index | |
) | |
# Decode token terms to words | |
if self.hparams.use_language_modelling: | |
predicted_words = [] | |
for logs in p_ctc: | |
text = decoder.decode(logs.detach().cpu().numpy()) | |
predicted_words.append(text.split(" ")) | |
else: | |
predicted_words = [ | |
"".join(self.tokenizer.decode_ndim(utt_seq)).split(" ") | |
for utt_seq in predicted_tokens | |
] | |
# Convert indices to words | |
target_words = [wrd.split(" ") for wrd in batch.wrd] | |
self.wer_metric.append(ids, predicted_words, target_words) | |
self.cer_metric.append(ids, predicted_words, target_words) | |
return loss | |
def fit_batch(self, batch): | |
"""Train the parameters given a single batch in input""" | |
should_step = self.step % self.grad_accumulation_factor == 0 | |
# Managing automatic mixed precision | |
# TOFIX: CTC fine-tuning currently is unstable | |
# This is certainly due to CTC being done in fp16 instead of fp32 | |
if self.auto_mix_prec: | |
with torch.cuda.amp.autocast(): | |
with self.no_sync(): | |
outputs = self.compute_forward(batch, sb.Stage.TRAIN) | |
loss = self.compute_objectives(outputs, batch, sb.Stage.TRAIN) | |
with self.no_sync(not should_step): | |
self.scaler.scale( | |
loss / self.grad_accumulation_factor | |
).backward() | |
if should_step: | |
if not self.hparams.wav2vec2.freeze: | |
self.scaler.unscale_(self.wav2vec_optimizer) | |
self.scaler.unscale_(self.model_optimizer) | |
if self.check_gradients(loss): | |
if not self.hparams.wav2vec2.freeze: | |
if self.optimizer_step >= self.hparams.warmup_steps: | |
self.scaler.step(self.wav2vec_optimizer) | |
self.scaler.step(self.model_optimizer) | |
self.scaler.update() | |
self.zero_grad() | |
self.optimizer_step += 1 | |
else: | |
# This is mandatory because HF models have a weird behavior with DDP | |
# on the forward pass | |
with self.no_sync(): | |
outputs = self.compute_forward(batch, sb.Stage.TRAIN) | |
loss = self.compute_objectives(outputs, batch, sb.Stage.TRAIN) | |
with self.no_sync(not should_step): | |
(loss / self.grad_accumulation_factor).backward() | |
if should_step: | |
if self.check_gradients(loss): | |
if not self.hparams.wav2vec2.freeze: | |
if self.optimizer_step >= self.hparams.warmup_steps: | |
self.wav2vec_optimizer.step() | |
self.model_optimizer.step() | |
self.zero_grad() | |
self.optimizer_step += 1 | |
self.on_fit_batch_end(batch, outputs, loss, should_step) | |
return loss.detach().cpu() | |
def evaluate_batch(self, batch, stage): | |
"""Computations needed for validation/test batches""" | |
predictions = self.compute_forward(batch, stage=stage) | |
with torch.no_grad(): | |
loss = self.compute_objectives(predictions, batch, stage=stage) | |
return loss.detach() | |
def on_stage_start(self, stage, epoch): | |
"""Gets called at the beginning of each epoch""" | |
if stage != sb.Stage.TRAIN: | |
self.cer_metric = self.hparams.cer_computer() | |
self.wer_metric = self.hparams.error_rate_computer() | |
def on_stage_end(self, stage, stage_loss, epoch): | |
"""Gets called at the end of an epoch.""" | |
# Compute/store important stats | |
stage_stats = {"loss": stage_loss} | |
if stage == sb.Stage.TRAIN: | |
self.train_stats = stage_stats | |
else: | |
stage_stats["CER"] = self.cer_metric.summarize("error_rate") | |
stage_stats["WER"] = self.wer_metric.summarize("error_rate") | |
# Perform end-of-iteration things, like annealing, logging, etc. | |
if stage == sb.Stage.VALID: | |
old_lr_model, new_lr_model = self.hparams.lr_annealing_model( | |
stage_stats["loss"] | |
) | |
old_lr_wav2vec, new_lr_wav2vec = self.hparams.lr_annealing_wav2vec( | |
stage_stats["loss"] | |
) | |
sb.nnet.schedulers.update_learning_rate( | |
self.model_optimizer, new_lr_model | |
) | |
if not self.hparams.wav2vec2.freeze: | |
sb.nnet.schedulers.update_learning_rate( | |
self.wav2vec_optimizer, new_lr_wav2vec | |
) | |
self.hparams.train_logger.log_stats( | |
stats_meta={ | |
"epoch": epoch, | |
"lr_model": old_lr_model, | |
"lr_wav2vec": old_lr_wav2vec, | |
}, | |
train_stats=self.train_stats, | |
valid_stats=stage_stats, | |
) | |
self.checkpointer.save_and_keep_only( | |
meta={"WER": stage_stats["WER"]}, min_keys=["WER"], | |
) | |
elif stage == sb.Stage.TEST: | |
self.hparams.train_logger.log_stats( | |
stats_meta={"Epoch loaded": self.hparams.epoch_counter.current}, | |
test_stats=stage_stats, | |
) | |
with open(self.hparams.wer_file, "w") as w: | |
self.wer_metric.write_stats(w) | |
def init_optimizers(self): | |
"Initializes the wav2vec2 optimizer and model optimizer" | |
# If the wav2vec encoder is unfrozen, we create the optimizer | |
if not self.hparams.wav2vec2.freeze: | |
self.wav2vec_optimizer = self.hparams.wav2vec_opt_class( | |
self.modules.wav2vec2.parameters() | |
) | |
if self.checkpointer is not None: | |
self.checkpointer.add_recoverable( | |
"wav2vec_opt", self.wav2vec_optimizer | |
) | |
self.model_optimizer = self.hparams.model_opt_class( | |
self.hparams.model.parameters() | |
) | |
if self.checkpointer is not None: | |
self.checkpointer.add_recoverable("modelopt", self.model_optimizer) | |
def zero_grad(self, set_to_none=False): | |
if not self.hparams.wav2vec2.freeze: | |
self.wav2vec_optimizer.zero_grad(set_to_none) | |
self.model_optimizer.zero_grad(set_to_none) | |
from speechbrain.pretrained import EncoderASR,EncoderDecoderASR | |
french_asr_model = EncoderASR.from_hparams(source="asr-wav2vec2-commonvoice-fr", savedir="pretrained_models/asr-wav2vec2-commonvoice-fr") | |
french_asr_model.to("cpu") | |
cvhparams_file, cvrun_opts, cvoverrides = sb.parse_arguments(["EnglishCV/train_en_with_wav2vec.yaml"]) | |
with open(cvhparams_file) as cvfin: | |
cvhparams = load_hyperpyyaml(cvfin, cvoverrides) | |
cvrun_opts["device"]="cpu" | |
english_asr_model = ASRCV( | |
modules=cvhparams["modules"], | |
hparams=cvhparams, | |
run_opts=cvrun_opts, | |
checkpointer=cvhparams["checkpointer"], | |
) | |
english_asr_model.modules.to("cpu") | |
english_asr_model.device="cpu" | |
english_asr_model.checkpointer.recover_if_possible(device="cpu") | |
run_opts["device"]="cpu" | |
print("moving to tunisian model") | |
asr_brain = ASR( | |
modules=hparams["modules"], | |
hparams=hparams, | |
run_opts=run_opts, | |
checkpointer=hparams["checkpointer"], | |
) | |
asr_brain.modules.to("cpu") | |
asr_brain.checkpointer.recover_if_possible(device="cpu") | |
asr_brain.modules.eval() | |
english_asr_model.modules.eval() | |
french_asr_model.mods.eval() | |
asr_brain.modules.to("cpu") | |
# Commented out IPython magic to ensure Python compatibility. | |
# %ls | |
#UTILS FUNCTIOJNS | |
def get_size_dimensions(arr): | |
size_dimensions = [] | |
while isinstance(arr, list): | |
size_dimensions.append(len(arr)) | |
arr = arr[0] | |
return size_dimensions | |
def scale_array(batch,n): | |
scaled_batch = [] | |
for array in batch: | |
if(n < len(array)): raise ValueError("Cannot scale Array down") | |
repeat = round(n/len(array))+1 | |
scaled_length_array= [] | |
for i in array: | |
for j in range(repeat) : | |
if(len(scaled_length_array) == n): break | |
scaled_length_array.append(i) | |
scaled_batch.append(scaled_length_array) | |
return torch.tensor(scaled_batch) | |
def load_paths(wavs_path): | |
waveforms = [] | |
for path in wavs_path : | |
waveform, _ = torchaudio.load(path) | |
waveforms.append(waveform.squeeze(0)) | |
# normalize array length to the bigger arrays by pading with 0's | |
padded_arrays = pad_sequence(waveforms, batch_first=True) | |
return torch.tensor(padded_arrays) | |
device = 'cpu' | |
verbose = 0 | |
#FLOW LEVEL FUNCTIONS | |
def merge_strategy(embeddings1, embeddings2, embeddings3,post1, post2,post3): | |
post1 = post1.to(device) | |
post2 = post2.to(device) | |
post3 = post3.to(device) | |
embeddings1 = embeddings1.to(device) | |
embeddings2 = embeddings2.to(device) | |
embeddings3 = embeddings3.to(device) | |
posteriograms_merged = torch.cat((post1,post2,post3),dim=2) | |
embeddings_merged = torch.cat((embeddings1,embeddings2,embeddings3),dim=2) | |
if(verbose !=0): | |
print('MERGED POST ',posteriograms_merged.shape) | |
print('MERGED emb ',embeddings_merged.shape) | |
return torch.cat((posteriograms_merged,embeddings_merged),dim=2).to(device) | |
def decode(model,wavs,wav_lens): | |
with torch.no_grad(): | |
wav_lens = wav_lens.to(model.device) | |
encoder_out = model.encode_batch(wavs, wav_lens) | |
predictions = model.decoding_function(encoder_out, wav_lens) | |
return predictions | |
def middle_layer(batch, lens): | |
tn_embeddings, tn_posteriogram = asr_brain.custom_encode(batch,None) | |
fr_embeddings = french_asr_model.mods.encoder.wav2vec2(batch) | |
fr_posteriogram =french_asr_model.encode_batch(batch,lens) | |
en_embeddings = english_asr_model.modules.wav2vec2(batch, lens) | |
x = english_asr_model.modules.enc(en_embeddings) | |
en_posteriogram = english_asr_model.modules.ctc_lin(x) | |
#scores, en_posteriogram = english_asr_model.mods.decoder(en_embeddings ,lens) | |
if(verbose !=0): | |
print('[EMBEDDINGS] FR:',fr_embeddings.shape, "EN:",en_embeddings.shape, "TN:", tn_embeddings.shape) | |
print('[POSTERIOGRAM] FR:',fr_posteriogram.shape, "EN:",en_posteriogram.shape,"TN:",tn_posteriogram.shape) | |
bilangual_sample = merge_strategy(fr_embeddings,en_embeddings,tn_embeddings,fr_posteriogram,en_posteriogram,tn_posteriogram) | |
return bilangual_sample | |
class Mixer(sb.core.Brain): | |
def compute_forward(self, batch, stage): | |
"""Forward computations from the waveform batches to the output probabilities.""" | |
wavs, wav_lens = batch.sig | |
wavs, wav_lens = wavs.to(self.device), wav_lens.to(self.device) | |
if stage == sb.Stage.TRAIN: | |
if hasattr(self.hparams, "augmentation"): | |
wavs = self.hparams.augmentation(wavs, wav_lens) | |
multi_langual_feats = middle_layer(wavs, wav_lens) | |
multi_langual_feats= multi_langual_feats.to(device) | |
feats, _ = self.modules.enc(multi_langual_feats) | |
logits = self.modules.ctc_lin(feats) | |
p_ctc = self.hparams.log_softmax(logits) | |
if stage!= sb.Stage.TRAIN: | |
p_tokens = sb.decoders.ctc_greedy_decode( | |
p_ctc, wav_lens, blank_id=self.hparams.blank_index | |
) | |
else : | |
p_tokens = None | |
return p_ctc, wav_lens, p_tokens | |
def treat_wav(self,sig): | |
multi_langual_feats = middle_layer(sig.to("cpu"), torch.tensor([1]).to("cpu")) | |
multi_langual_feats= multi_langual_feats.to(device) | |
feats, _ = self.modules.enc(multi_langual_feats) | |
logits = self.modules.ctc_lin(feats) | |
p_ctc = self.hparams.log_softmax(logits) | |
predicted_words =[] | |
for logs in p_ctc: | |
text = decoder.decode(logs.detach().cpu().numpy()) | |
predicted_words.append(text.split(" ")) | |
return " ".join(predicted_words[0]) | |
def compute_objectives(self, predictions, batch, stage): | |
"""Computes the loss (CTC) given predictions and targets.""" | |
p_ctc, wav_lens , predicted_tokens= predictions | |
ids = batch.id | |
tokens, tokens_lens = batch.tokens | |
loss = self.hparams.ctc_cost(p_ctc, tokens, wav_lens, tokens_lens) | |
if stage == sb.Stage.VALID: | |
predicted_words = [ | |
"".join(self.tokenizer.decode_ndim(utt_seq)).split(" ") | |
for utt_seq in predicted_tokens | |
] | |
target_words = [wrd.split(" ") for wrd in batch.wrd] | |
self.wer_metric.append(ids, predicted_words, target_words) | |
self.cer_metric.append(ids, predicted_words, target_words) | |
if stage ==sb.Stage.TEST : | |
if self.hparams.language_modelling: | |
predicted_words = [] | |
for logs in p_ctc: | |
text = decoder.decode(logs.detach().cpu().numpy()) | |
predicted_words.append(text.split(" ")) | |
else : | |
predicted_words = [ | |
"".join(self.tokenizer.decode_ndim(utt_seq)).split(" ") | |
for utt_seq in predicted_tokens | |
] | |
target_words = [wrd.split(" ") for wrd in batch.wrd] | |
self.wer_metric.append(ids, predicted_words, target_words) | |
self.cer_metric.append(ids, predicted_words, target_words) | |
return loss | |
def fit_batch(self, batch): | |
"""Train the parameters given a single batch in input""" | |
should_step = self.step % self.grad_accumulation_factor == 0 | |
# Managing automatic mixed precision | |
# TOFIX: CTC fine-tuning currently is unstable | |
# This is certainly due to CTC being done in fp16 instead of fp32 | |
if self.auto_mix_prec: | |
with torch.cuda.amp.autocast(): | |
with self.no_sync(): | |
outputs = self.compute_forward(batch, sb.Stage.TRAIN) | |
loss = self.compute_objectives(outputs, batch, sb.Stage.TRAIN) | |
with self.no_sync(not should_step): | |
self.scaler.scale( | |
loss / self.grad_accumulation_factor | |
).backward() | |
if should_step: | |
self.scaler.unscale_(self.model_optimizer) | |
if self.check_gradients(loss): | |
self.scaler.step(self.model_optimizer) | |
self.scaler.update() | |
self.zero_grad() | |
self.optimizer_step += 1 | |
else: | |
# This is mandatory because HF models have a weird behavior with DDP | |
# on the forward pass | |
with self.no_sync(): | |
outputs = self.compute_forward(batch, sb.Stage.TRAIN) | |
loss = self.compute_objectives(outputs, batch, sb.Stage.TRAIN) | |
with self.no_sync(not should_step): | |
(loss / self.grad_accumulation_factor).backward() | |
if should_step: | |
if self.check_gradients(loss): | |
self.model_optimizer.step() | |
self.zero_grad() | |
self.optimizer_step += 1 | |
self.on_fit_batch_end(batch, outputs, loss, should_step) | |
return loss.detach().cpu() | |
def evaluate_batch(self, batch, stage): | |
"""Computations needed for validation/test batches""" | |
predictions = self.compute_forward(batch, stage=stage) | |
with torch.no_grad(): | |
loss = self.compute_objectives(predictions, batch, stage=stage) | |
return loss.detach() | |
def on_stage_start(self, stage, epoch): | |
"""Gets called at the beginning of each epoch""" | |
if stage != sb.Stage.TRAIN: | |
self.cer_metric = self.hparams.cer_computer() | |
self.wer_metric = self.hparams.error_rate_computer() | |
def on_stage_end(self, stage, stage_loss, epoch): | |
"""Gets called at the end of an epoch.""" | |
# Compute/store important stats | |
stage_stats = {"loss": stage_loss} | |
if stage == sb.Stage.TRAIN: | |
self.train_stats = stage_stats | |
else: | |
stage_stats["CER"] = self.cer_metric.summarize("error_rate") | |
stage_stats["WER"] = self.wer_metric.summarize("error_rate") | |
# Perform end-of-iteration things, like annealing, logging, etc. | |
if stage == sb.Stage.VALID: | |
old_lr_model, new_lr_model = self.hparams.lr_annealing_model( | |
stage_stats["loss"] | |
) | |
sb.nnet.schedulers.update_learning_rate( | |
self.model_optimizer, new_lr_model | |
) | |
self.hparams.train_logger.log_stats( | |
stats_meta={ | |
"epoch": epoch, | |
"lr_model": old_lr_model, | |
}, | |
train_stats=self.train_stats, | |
valid_stats=stage_stats, | |
) | |
self.checkpointer.save_and_keep_only( | |
meta={"WER": stage_stats["WER"]}, min_keys=["WER"], | |
) | |
elif stage == sb.Stage.TEST: | |
self.hparams.train_logger.log_stats( | |
stats_meta={"Epoch loaded": self.hparams.epoch_counter.current}, | |
test_stats=stage_stats, | |
) | |
with open(self.hparams.wer_file, "w") as w: | |
self.wer_metric.write_stats(w) | |
def init_optimizers(self): | |
self.model_optimizer = self.hparams.model_opt_class( | |
self.hparams.model.parameters() | |
) | |
if self.checkpointer is not None: | |
self.checkpointer.add_recoverable("modelopt", self.model_optimizer) | |
def zero_grad(self, set_to_none=False): | |
self.model_optimizer.zero_grad(set_to_none) | |
hparams_file, run_opts, overrides = sb.parse_arguments(["cs.yaml"]) | |
# If distributed_launch=True then | |
# create ddp_group with the right communication protocol | |
sb.utils.distributed.ddp_init_group(run_opts) | |
with open(hparams_file) as fin: | |
hparams = load_hyperpyyaml(fin, overrides) | |
# Create experiment directory | |
sb.create_experiment_directory( | |
experiment_directory=hparams["output_folder"], | |
hyperparams_to_save=hparams_file, | |
overrides=overrides, | |
) | |
def read_labels_file(labels_file): | |
with open(labels_file, "r",encoding="utf-8") as lf: | |
lines = lf.read().splitlines() | |
division = "===" | |
numbers = {} | |
for line in lines : | |
if division in line : | |
break | |
string, number = line.split("=>") | |
number = int(number) | |
string = string[1:-2] | |
numbers[number] = string | |
return [numbers[x] for x in range(len(numbers))] | |
label_encoder = sb.dataio.encoder.CTCTextEncoder() | |
lab_enc_file = os.path.join(hparams["save_folder"], "label_encoder.txt") | |
special_labels = { | |
"blank_label": hparams["blank_index"], | |
"unk_label": hparams["unk_index"] | |
} | |
label_encoder.load_or_create( | |
path=lab_enc_file, | |
from_didatasets=[[]], | |
output_key="char_list", | |
special_labels=special_labels, | |
sequence_input=True, | |
) | |
labels = read_labels_file(os.path.join(hparams["save_folder"], "label_encoder.txt")) | |
labels = [""] + labels[1:-1] + ["1"] | |
if hparams["language_modelling"]: | |
decoder = build_ctcdecoder( | |
labels, | |
kenlm_model_path=hparams["ngram_lm_path"], # either .arpa or .bin file | |
alpha=0.5, # tuned on a val set | |
beta=1, # tuned on a val set | |
) | |
description = """This is a speechbrain-based Automatic Speech Recognition (ASR) model for Tunisian arabic. It outputs code-switched Tunisian transcriptions written in Arabic and Latin characters. It handles Tunisian Arabic, English and French outputs. | |
Code-switching is notoriously hard to handle for speech recognition models, the main errors you man encounter using this model are spelling/language identification errors due to code-switching. We may work on improving this in further models. However if you do not need code-switching in your transcripts, you would better use the non-code switched model, available in another space from the same author. (https://huggingface.co/spaces/SalahZa/Tunisian-Speech-Recognition) | |
Run is done on CPU to keep it free in this space. This leads to quite long running times on long sequences. If for your project or research, you want to transcribe long sequences, you would better use the model directly from its page, some instructions for inference on a test set have been provided there. (https://huggingface.co/SalahZa/Code_Switched_Tunisian_Speech_Recognition). If you need help, feel free to drop an email here : zaiemsalah@gmail.com | |
Authors : | |
* [Salah Zaiem](https://fr.linkedin.com/in/salah-zaiem) | |
* [Ahmed Amine Ben Aballah](https://www.linkedin.com/in/aabenz/) | |
* [Ata Kaboudi](https://www.linkedin.com/in/ata-kaboudi-63365b1a8) | |
* [Amir Kanoun](https://tn.linkedin.com/in/ahmed-amir-kanoun) | |
More in-depth details and insights are available in a released preprint. Please find the paper [here](https://arxiv.org/abs/2309.11327). | |
If you use or refer to this model, please cite : | |
``` | |
@misc{abdallah2023leveraging, | |
title={Leveraging Data Collection and Unsupervised Learning for Code-switched Tunisian Arabic Automatic Speech Recognition}, | |
author={Ahmed Amine Ben Abdallah and Ata Kabboudi and Amir Kanoun and Salah Zaiem}, | |
year={2023}, | |
eprint={2309.11327}, | |
archivePrefix={arXiv}, | |
primaryClass={eess.AS} | |
} | |
""" | |
title = "Code-Switched Tunisian Speech Recognition" | |
run_opts["device"]="cpu" | |
mixer = Mixer( | |
modules=hparams["modules"], | |
hparams=hparams, | |
run_opts=run_opts, | |
checkpointer=hparams["checkpointer"], | |
) | |
mixer.tokenizer = label_encoder | |
mixer.device = "cpu" | |
mixer.checkpointer.recover_if_possible(device="cpu") | |
mixer.modules.eval() | |
device = "cpu" | |
mixer.device= "cpu" | |
mixer.modules.to("cpu") | |
from enum import Enum, auto | |
class Stage(Enum): | |
TRAIN = auto() | |
VALID = auto() | |
TEST = auto() | |
asr_brain.on_evaluate_start() | |
asr_brain.modules.eval() | |
import gradio as gr | |
def treat_wav_file(file_mic,file_upload ,asr=mixer, device="cpu") : | |
if (file_mic is not None) and (file_upload is not None): | |
warn_output = "WARNING: You've uploaded an audio file and used the microphone. The recorded file from the microphone will be used and the uploaded file will be discarded.\n" | |
wav = file_mic | |
elif (file_mic is None) and (file_upload is None): | |
return "ERROR: You have to either use the microphone or upload an audio file" | |
elif file_mic is not None: | |
wav = file_mic | |
else: | |
wav = file_upload | |
info = torchaudio.info(wav) | |
sr = info.sample_rate | |
sig = sb.dataio.dataio.read_audio(wav) | |
if len(sig.shape)>1 : | |
sig = torch.mean(sig, dim=1) | |
sig = torch.unsqueeze(sig, 0) | |
tensor_wav = sig.to(device) | |
resampled = torchaudio.functional.resample( tensor_wav, sr, 16000) | |
sentence = asr.treat_wav(resampled) | |
return sentence | |
gr.Interface( | |
fn=treat_wav_file, | |
title = title, | |
description = description, | |
inputs=[gr.Audio(source="microphone", type='filepath', label = "record", optional = True), | |
gr.Audio(source="upload", type='filepath', label="filein", optional=True)] | |
,outputs="text").launch() | |