Cómo entrenar modelos de reconocimiento de voz con Whisper

Paso 1: Configurar el entorno y los datos de entrenamiento

El primer paso para completar nuestro objetivo es configurar correctamente el entorno de trabajo. Desde OpenWebinars te recomendamos utilizar Google Colab para seguir los pasos que te presentamos a continuación. Además, para el almacenamiento y despliegue del modelo que vamos a entrenar utilizaremos Hugging Face Hub, la plataforma ofrecida por Hugging Face para almacenar y gestionar modelos de inteligencia artificial.

Instalación de Whisper y dependencias necesarias

El primer paso en todos los proyectos que desarrollamos con Python es actualizar pip (el instalador de paquetes Python más usado) e instalar los paquetes necesarios.

!pip install --upgrade -q pip
!pip install --upgrade -q datasets[audio] transformers accelerate evaluate jiwer tensorboard gradio

Una vez instalados los paquetes necesarios procederemos a importar todas las librerías necesarias.

from dataclasses import dataclass
from typing import Any, Dict, List, Union

import evaluate
import gradio as gr
import torch

from datasets import Audio, DatasetDict, load_dataset
from huggingface_hub import notebook_login
from transformers import (
    Seq2SeqTrainer,
    Seq2SeqTrainingArguments,
    WhisperFeatureExtractor,
    WhisperForConditionalGeneration,
    WhisperProcessor,
    WhisperTokenizer,
    pipeline,
)

El último paso previo a recopilar los datos y entrenar Whisper es hacer login en Hugging Face Hub para posteriormente poder almacenar y disponer de nuestro modelo de manera sencilla. Para ello, en primer lugar, ejecutaremos la siguiente instrucción.

notebook_login()

Una vez ejecutada nos encontraremos un mensaje similar al mostrado en la siguiente imagen.

Aquí tendremos que introducir un token de acceso a Hugging Face Hub que podemos encontrar aquí.

Recolección de datos de voz

El siguiente paso es elegir el conjunto de datos con el que vamos a trabajar. En esta ocasión utilizaremos el dataset Common Voice 11. Concretamente elegiremos el subconjunto de datos del idioma hindi. El hindi es una lengua indioeuropea oficial en la India.

Hemos seleccionado esta lengua ya que es una de las lenguas con una menor cantidad de datos. Este hecho es importante ya que el entrenamiento de Whisper es un proceso computacionalmente extremadamente caro. Por este motivo para este ejercicio utilizaremos únicamente 50 muestras para el entrenamiento y 25 muestras para la validación.

dataset_name = "mozilla-foundation/common_voice_11_0"
dataset = DatasetDict()

dataset["train"] = load_dataset(dataset_name, "hi", split="train+validation", trust_remote_code=True).select(range(50))
dataset["test"] = load_dataset(dataset_name, "hi", split="test", trust_remote_code=True).select(range(25))
print(dataset)

Tal y como podemos ver en la siguiente imagen al descargar el conjunto de datos podemos observar un total de 11 features, las cuales podemos apreciar en la siguiente imagen.

Por lo tanto, procederemos a eliminar todas las columnas innecesarias. Además, procesaremos la columna “audio” para que todas las muestran presenten un sampling rate de 16 000 muestras por segundo (o hercios).

dataset = dataset.remove_columns(["accent", "age", "client_id", "down_votes", "gender", "locale", "path", "segment", "up_votes"])
dataset = dataset.cast_column("audio", Audio(sampling_rate=16000))

print(dataset)

Preprocesamiento del audio

Una vez tenemos la columna “audio” representada con el sampling rate correcto procederemos a preprocesar el audio. En primer lugar definiremos los distintos modelos que utilizaremos para preprocesar las muestras disponibles.

model_name = "openai/whisper-small"
language = "hi"

preprocessor = WhisperProcessor.from_pretrained(model_name, language=language, task="transcribe")
feature_extractor = WhisperFeatureExtractor.from_pretrained(model_name)
tokenizer = WhisperTokenizer.from_pretrained(model_name, language=language, task="transcribe")
model = WhisperForConditionalGeneration.from_pretrained(model_name)

model.generation_config.language = language
model.generation_config.task = "transcribe"

En el anterior bloque de código hemos definidos los siguientes cuatro modelos:

• Feature_extractor: es una instancia de WhisperFeatureExtractor que se encarga de dividir los audios en segmentos de 30 segundos y posteriormiente convierte estos audios en espectogramas.
• Tokenizer: es una instancia de WhisperTokenizer encargada de convertir los vectores numéricos generados por el modelo en secuencias de texto.
• Preprocessor: es una instancia de WhisperProcessor capaz de realizar las tareas del feature_extractor y el tokenizer.
• Model: es una instancia de WhisperForConditionalGeneration, es decir, la una instancia de Whisper, el modelo que vamos a entrenar.

Una vez tenemos definidos los diferentes modelos que vamos a utilizar en el preprocesamiento del audio vamos a definir la función encargada de realizar la tarea y la vamos a aplicar sobre todos los datos recuperados.

def preprocess_dataset(batch):
    audio = batch["audio"]

    batch["input_features"] = feature_extractor(audio["array"], sampling_rate=audio["sampling_rate"]).input_features[0]

    batch["labels"] = tokenizer(batch["sentence"]).input_ids
    return batch

dataset = dataset.map(preprocess_dataset, remove_columns=dataset.column_names["train"])

Si quieres conocer más detalles de Whisper te recomendamos nuestra publicación Whisper: Introducción al reconocimiento de voz.

Paso 2: Configurar el modelo Whisper para el entrenamiento

Una vez que tenemos preprocesado el conjunto de datos vamos a adaptarlo para poder entrenar nuestro modelo con él.

Preparación de los datos

Para preparar los datos vamos a definir una clase encargada de recuperar datos del conjunto de entrenamiento y procesarlos para que las secuencias de audio correspondan con las secuencias de texto transcrito.

@dataclass
class DataCollatorSpeechSeq2SeqWithPadding:
    processor: Any
    decoder_start_token_id: int

    def __call__(self, features):

        input_features = [{"input_features": feature["input_features"]} for feature in features]
        batch = self.processor.feature_extractor.pad(input_features, return_tensors="pt")

        label_features = [{"input_ids": feature["labels"]} for feature in features]

        labels_batch = self.processor.tokenizer.pad(label_features, return_tensors="pt")

        labels = labels_batch["input_ids"].masked_fill(labels_batch.attention_mask.ne(1), -100)


        if (labels[:, 0] == self.decoder_start_token_id).all().cpu().item():
            labels = labels[:, 1:]

        batch["labels"] = labels

        return batch

data_collator = DataCollatorSpeechSeq2SeqWithPadding(
    processor=preprocessor,
    decoder_start_token_id=model.config.decoder_start_token_id,
)

Ajuste de parámetros de entrenamiento

El siguiente paso es definir los parámetros de entrenamiento que vamos a utilizar. Estos se pueden observar en el siguiente bloque de código.

training_args = Seq2SeqTrainingArguments(
    output_dir=f"./{model_name}-hi",
    per_device_train_batch_size=32,
    gradient_accumulation_steps=1,
    learning_rate=1e-5,
    warmup_steps=500,
    num_train_epochs=1,
    gradient_checkpointing=True,
    fp16=True,
    evaluation_strategy="epoch",
    save_strategy="epoch",
    per_device_eval_batch_size=8,
    predict_with_generate=True,
    generation_max_length=225,
    logging_steps=500,
    report_to=["tensorboard"],
    load_best_model_at_end=True,
    metric_for_best_model="wer",
    greater_is_better=False,
    push_to_hub=True,
)

De entre todos los parámetros queremos destacar los dos siguientes:

• num_train_epochs: el número de epochs que vamos a entrenar nuestro modelo. Este parámetro es el que afecta en mayor medida al tiempo de entrenamiento. Con la finalidad de reducir el tiempo de entrenamiento lo hemos fijado al mínimo valor posible, pero para obtener los mejores resultados deberíamos fijarlo a 20. Si decidimos fijarlo a 20 epochs el entrenamiento tardará unas 8 horas utilizando una tarjeta gráfica.

• push_to_hub: con este parámetro indicamos que queremos guardar automáticamente el progreso de nuestro entrenamiento en Hugging Face Hub.

Paso 3: Entrenar el modelo

El siguiente paso a realizar es el propio entrenamiento del modelo.

Monitorear el progreso

Antes de nada definiremos la función que monitorizará el proceso de entrenamiento y nos indicará cómo de eficaz está siendo este.

def compute_metrics(pred):
    pred_ids = pred.predictions
    label_ids = pred.label_ids

    label_ids[label_ids == -100] = tokenizer.pad_token_id

    # we do not want to group tokens when computing the metrics
    pred_str = tokenizer.batch_decode(pred_ids, skip_special_tokens=True)
    label_str = tokenizer.batch_decode(label_ids, skip_special_tokens=True)

    wer = 100 * metric.compute(predictions=pred_str, references=label_str)

    return {"wer": wer}

Inicio del proceso de entrenamiento

Una vez definida la función de calcula el valor de las métricas definimos una instancia de entrenamiento indicando la configuración de entrenamiento previamente definida, qué conjuntos de datos se van a utilizar, qué función se encargará de seleccionar los datos, cómo vamos a calcular las métricas y qué modelo es el encargado de preprocesar los datos.

trainer = Seq2SeqTrainer(
    args=training_args,
    model=model,
    train_dataset=dataset["train"],
    eval_dataset=dataset["test"],
    data_collator=data_collator,
    compute_metrics=compute_metrics,
    tokenizer=preprocessor.feature_extractor,
)

El siguiente paso es definir la métrica de entrenamiento y comenzar el entrenamiento.

metric = evaluate.load("wer")

trainer.train()

En nuestro curso dominando el reconocimiento de voz con Whisper podrás encontrar mucha más información sobre el proceso de entrenamiento.

Finalización y almacenamiento del modelo

Una vez el proceso de entrenamiento ha finalizado debemos identificar nuestro modelo y almacenarlo en Hugging Face Hub.

kwargs = {
    "dataset_tags": dataset_name,
    "dataset": f"{dataset_name} {language}",  # a 'pretty' name for the training dataset
    "dataset_args": "config: {language}, split: test",
    "language": "{language}",
    "model_name": "{model_name} - OpenWebinars",  # a 'pretty' name for our model
    "finetuned_from": "{model_name}",
    "tasks": "automatic-speech-recognition",
}

trainer.push_to_hub(\**kwargs)

Paso 4: Evaluar y ajustar el modelo entrenado

Evaluar el modelo para comprobar que el rendimiento de este es el adecuado es tan importante como el proceso de entrenamiento.

Evaluación del modelo con datos de prueba

Uno de los motivos por los que hemos decidido utilizar el objeto Trainer para entrenar Whisper es que podemos realizar la evaluación del modelo a la vez que se entrena. Esto lo hemos hecho al utilizar el parámetro evaluation_strategy en la variable training_args y en la inicialización del objeto Trainer el parámetro compute_metrics.

Métricas de evaluación

Existen muchas métricas de evaluación que se pueden utilizar para evaluar el rendimiento de Whisper, desde OpenWebinars queremos destacar las siguientes:

• Character Error Rate (CER): mide qué porcentaje de carácteres se han predicho incorrectamente.
• Word Error Rate (WER): mide qué porcentaje de palabras se han predicho incorrectamente.
• BLEU (Bilingual Evaluation Understudy): evalúa el nivel de solape existente entre las palabras generadas por Whisper y la transcripción correcta.

Paso 5: Implementar el modelo en aplicaciones

Una vez entrenado y evaluado el modelo es la hora de implementar el modelo en una aplicación real y utilizar el modelo en un entorno productivo.

Implementación en tiempo real

Una de las principales ventajas de Hugging Face Hub es la facilidad de utilizar de manera productiva los modelos entrenados. Para ello únicamente necesitamos crear una pipeline de Hugging Face con el nombre del modelo.

pipe = pipeline(model="{model_name} - OpenWebinars")

def transcribe(audio):
    text = pipe(audio)["text"]
    return text

Y una vez que disponemos de la pipeline la podemos utilizar. El siguiente bloque de código es un ejemplo de despliegue real.

iface = gr.Interface(
    fn=transcribe,
    inputs=gr.Audio(source="microphone", type="filepath"),
    outputs="text",
    title="Whisper Hindi",
    description="Despliegue de Whisper.",
)

iface.launch()

Integración con backend

Mediante las pipelines de Hugging Face podemos integrar fácilmente nuestros modelos en los frameworks más utilizados para desarrollar backends. Algunos de los backends más comunes son:

• Djanjo: es un framework de desarrollo web de alto nivel que permite realizar desarrollos extremadamente rápidos.
• Flask: este framework de desarrollo web tiene como objetivo minimizar el número de líneas necesarias para desarrollar una aplicación.
• FastAPI: es un framework_ de desarrollo web centrado en servidores HTTP que hagan uso de los modelos de OpenAI.

Desafíos comunes durante el proceso

En las secciones anteriores hemos utilizado diversas herramientas que facilitan el proceso de entrenamiento de Whisper, pero el proceso no está carente de desafíos. Es el momento de analizar los desafíos más comunes.

Problemas con la calidad del audio

Todos los modelos de machine learning y deep learning presentan el mismo desafío, y Whisper no iba a ser menos. Estos modelos aprenden a partir de datos y muchas veces es difícil encontrar datos de calidad. Los problemas más comunes que afectan al tipo de datos que utiliza Whisper son los siguientes:

• Variabilidad en las frecuencias de muestreo: cada dispositivo de grabación utiliza frecuencias de muestreo diferentes, sin embargo para utilizar Whisper debemos utilizar audios con una frecuencia de muestro de 16 000 hercios.
• Volumen y nivel de intensidad inconsistentes: en la mayor parte de los conjuntos de datos existen audios en los que el volumen varía en gran medida a lo largo de la grabación. Además, el volumen medio de los audios puede variar drásticamente afectando al rendimiento del modelo.
• Compresión de audio y pérdida de calidad: muchas veces a la hora de almacenar información se comprimen los datos sacrificando calidad de audio a cambio de reducir el peso de los audios.

Si quieres conocer en mayor profundidad los retos comunes a todas las técnicas de inteligencia artificial te recomendamos nuestra ruta formativa de especialización en inteligencia artificial.

Limitaciones de hardware

El reconocimiento y transcripción de voz es una tarea computacionalmente costosa que requiere de hardware especializado como son las GPU (tarjetas gráficas) o las TPU (Tensor Processing Units).

Sin embargo, acceder a dichas herramientas especializadas suele tener un coste monetario. La necesidad de herramientas de cómputo y el coste de estas hace que el entrenamiento de Whisper con grandes conjuntos de datos suela ser una tarea inabordable.

Conclusiones

Whisper es el modelo más popular utilizado para reconocer y transcribir audios. Fue desarrollado por OpenAI en 2022 y liberado a través de una licencia MIT. Gracias a esto podemos adaptar el modelo a nuestras necesidades y utilizarlo libremente tanto en nuestros proyectos personales como en proyectos laborales.

Whisper representa el presente y el futuro del reconocimiento y transcripción de voz. Decenas de científicos, desarrolladores y emprendedores crean a diario versiones mejoradas de este modelo que compiten con las mejoras ocasionales desarrolladas por OpenAI.

Estos hechos combinados con todas las formas diferentes de acceder a Whisper permiten afirmar sin temor a equivocaciones que Whisper seguirá siendo el modelo rey en su campo al menos los próximos cinco años.