# Install the dependencies
!pip install --upgrade --quiet pip
!pip install --upgrade --quiet datasets[audio] transformers accelerate evaluate jiwer tensorboard gradio

import gc
import re
from dataclasses import dataclass
from typing import Any

import evaluate
import jiwer
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch
from datasets import Audio, Dataset, load_dataset
from huggingface_hub import notebook_login
from IPython.display import Audio as AudioPlayer, display
from sklearn.model_selection import StratifiedGroupKFold
from transformers import (
    EarlyStoppingCallback,
    Seq2SeqTrainer,
    Seq2SeqTrainingArguments,
    WhisperForConditionalGeneration,
    WhisperProcessor,
    set_seed,
)
from transformers.models.whisper.english_normalizer import EnglishTextNormalizer

# Verify that the notebook is connected to a GPU
if torch.cuda.is_available():
    print(f"Connected to GPU: {torch.cuda.get_device_name(0)}")
else:
    print("Not connected to a GPU. In Colab, select Runtime > Change runtime type > T4 GPU.")

Connected to GPU: Tesla T4

# Log in to the Hugging Face Hub so checkpoints can be pushed during training.
# Paste an access token with write permission when prompted; never hard-code it in the notebook.
notebook_login()

# Configuration shared by the rest of the notebook
MODEL_ID = "openai/whisper-tiny"
DATASET_ID = "DynamicSuperb/AAVESpeechRecognition_CORAAL"
LIBRISPEECH_ID = "openslr/librispeech_asr"
HUB_MODEL_ID = "melvinrajendran/whisper-tiny-coraal"
OUTPUT_DIR = "./whisper-tiny-coraal"
SAMPLING_RATE = 16_000  # Whisper expects 16 kHz audio
SEED = 11

N_FOLDS = 5  # Number of speaker-level cross-validation folds
HUB_FOLD = 0  # The fold whose fine-tuned model is pushed to the Hugging Face Hub
N_LIBRISPEECH = 100  # Number of standard English clips used to check for regressions
N_BOOTSTRAP = 10_000  # Number of resamples for the bootstrap confidence intervals

# Seed Python, NumPy, and PyTorch so that runs are reproducible
set_seed(SEED)

# Load the CORAAL dataset as a DataFrame
coraal_df = load_dataset(DATASET_ID, split="test").to_pandas()

# Each file name encodes its speaker. For example, "utt_651_ROC_se0_ag2_f_04_1" is a clip of speaker
# "ROC_se0_ag2_f_04": site ROC, socioeconomic group 0, age group 2, female, speaker number 04.
speaker_pattern = (
    r"^utt_\d+_(?P<speaker>(?P<site>[A-Z]+)_se(?P<socioeconomic_group>\d)"
    r"_ag(?P<age_group>\d)_(?P<gender>[mf])_\d+)_\d+$"
)
speaker_info = coraal_df["file"].str.extract(speaker_pattern)
assert speaker_info.notna().all().all(), "Some file names do not follow the CORAAL naming scheme"
coraal_df = pd.concat([coraal_df, speaker_info], axis=1)

# CORAAL recording sites
SITE_NAMES = {
    "ATL": "Atlanta, GA",
    "DCA": "Washington, DC (1968)",
    "DCB": "Washington, DC (2016)",
    "DTA": "Detroit, MI",
    "LES": "Lower East Side, NY",
    "PRV": "Princeville, NC",
    "ROC": "Rochester, NY",
    "VLD": "Valdosta, GA",
}
coraal_df["site_name"] = coraal_df["site"].map(SITE_NAMES)

print(f"{len(coraal_df)} clips from {coraal_df['speaker'].nunique()} speakers at {coraal_df['site'].nunique()} sites")

800 clips from 169 speakers at 8 sites

# Assign every speaker to one of N_FOLDS test folds. Grouping by speaker keeps all clips of a speaker in the
# same fold, and stratifying by site keeps each fold's mix of sites similar to the full dataset.
outer_cv = StratifiedGroupKFold(n_splits=N_FOLDS, shuffle=True, random_state=SEED)
coraal_df["fold"] = -1
for fold, (_, test_idx) in enumerate(outer_cv.split(coraal_df, y=coraal_df["site"], groups=coraal_df["speaker"])):
    coraal_df.loc[test_idx, "fold"] = fold

# Check the size of each fold
print(coraal_df.groupby("fold").agg(clips=("file", "size"), speakers=("speaker", "nunique")))

      clips  speakers
fold                 
0       131        34
1       210        35
2       152        34
3       154        33
4       153        33

def split_fold(fold):
    """Return the training, validation, and test row indices for one cross-validation fold.

    The given fold is the test set. The speakers in the remaining folds are split 7:1 into training and
    validation sets (again by speaker and stratified by site), giving roughly a 70/10/20 split overall.
    """
    test_mask = coraal_df["fold"] == fold
    train_val_df = coraal_df[~test_mask]

    inner_cv = StratifiedGroupKFold(n_splits=8, shuffle=True, random_state=SEED)
    train_pos, val_pos = next(inner_cv.split(train_val_df, y=train_val_df["site"], groups=train_val_df["speaker"]))

    train_idx = train_val_df.index[train_pos].to_numpy()
    val_idx = train_val_df.index[val_pos].to_numpy()
    test_idx = coraal_df.index[test_mask].to_numpy()

    return train_idx, val_idx, test_idx

# Convert the audio and transcripts into a Dataset with the same row order as coraal_df
coraal = Dataset.from_pandas(coraal_df[["audio", "label"]], preserve_index=False)

# Decode the raw audio bytes into waveforms sampled at 16 kHz
coraal = coraal.cast_column("audio", Audio(sampling_rate=SAMPLING_RATE))

# Inspect the first example
example = coraal[0]
duration = len(example["audio"]["array"]) / example["audio"]["sampling_rate"]

print(f"File:          {coraal_df.loc[0, 'file']}")
print(f"Speaker:       {coraal_df.loc[0, 'speaker']} ({coraal_df.loc[0, 'site_name']})")
print(f"Transcript:    {example['label']}")
print(f"Duration:      {duration:.2f} s")
print(f"Sampling rate: {example['audio']['sampling_rate']} Hz")
display(AudioPlayer(example["audio"]["array"], rate=example["audio"]["sampling_rate"]))

File:          utt_0_ATL_se0_ag1_m_04_1
Speaker:       ATL_se0_ag1_m_04 (Atlanta, GA)
Transcript:    NOTHING LIKE THAT GO OVER MY HEAD CAUSE YOU KNOW WELL KILL OURSELVES TRYNA WORRY ABOUT
Duration:      3.12 s
Sampling rate: 16000 Hz

# Stream the first N_LIBRISPEECH clips of the LibriSpeech test set instead of downloading all of it.
# Audio decoding is deferred so that the streamed clips can be collected into a regular Dataset.
librispeech_stream = load_dataset(LIBRISPEECH_ID, "clean", split="test", streaming=True)
librispeech_stream = librispeech_stream.cast_column("audio", Audio(decode=False)).take(N_LIBRISPEECH)

librispeech = Dataset.from_list(list(librispeech_stream)).select_columns(["audio", "text"])
librispeech = librispeech.rename_column("text", "label").cast_column("audio", Audio(sampling_rate=SAMPLING_RATE))

print(f"Loaded {len(librispeech)} LibriSpeech clips")

README.md:   0%|          | 0.00/11.0k [00:00<?, ?B/s]

Resolving data files:   0%|          | 0/48 [00:00<?, ?it/s]

Resolving data files:   0%|          | 0/48 [00:00<?, ?it/s]

Loaded 100 LibriSpeech clips

# Load the feature extractor and tokenizer, configured for English transcription
processor = WhisperProcessor.from_pretrained(MODEL_ID, language="English", task="transcribe")

def normalize_text(text):
    """Lowercase text, drop apostrophes, and replace other punctuation with spaces.

    This mirrors the CORAAL transcription convention (e.g. "didn't" becomes "didnt"),
    so that WER is computed on words alone.
    """
    text = text.lower()
    text = re.sub(r"['’`]", "", text)
    text = re.sub(r"[^a-z0-9\s]", " ", text)
    return " ".join(text.split())

# Whisper's English normalizer standardizes numbers ("ten" -> "10"), spellings, and common reductions
# ("wanna" -> "want to"), and removes fillers such as "um" and "uh"
whisper_normalizer = EnglishTextNormalizer(processor.tokenizer.english_spelling_normalizer)

# Informal spellings in CORAAL transcripts that Whisper's normalizer does not cover, mapped to standard forms
INFORMAL_SPELLINGS = {
    "cause": "because",
    "cuz": "because",
    "tryna": "trying to",
    "yall": "you all",
    "kinda": "kind of",
    "sorta": "sort of",
    "outta": "out of",
    "lemme": "let me",
    "gimme": "give me",
    "everyday": "every day",
    "ninth": "9th",  # Whisper's normalizer converts every other ordinal to digits, but not "ninth"
}


def lenient_normalize(text):
    """Normalize text so that differences in transcription style are not counted as errors.

    Applies normalize_text(), expands informal spellings, then applies Whisper's English normalizer.
    """
    words = normalize_text(text).split()
    text = " ".join(INFORMAL_SPELLINGS.get(word, word) for word in words)
    return " ".join(whisper_normalizer(text).split())

def prepare_dataset(batch):
    """Convert one example into Whisper model inputs and target labels."""
    audio = batch["audio"]

    # Compute the log-Mel spectrogram input features from the 1-dimensional audio array
    batch["input_features"] = processor.feature_extractor(
        audio["array"], sampling_rate=audio["sampling_rate"]
    ).input_features[0]

    # Encode the normalized transcript into label IDs
    batch["labels"] = processor.tokenizer(normalize_text(batch["label"])).input_ids

    return batch

# Convert every clip into model inputs and labels. Rows keep the order of coraal_df,
# so the indices returned by split_fold() select the matching examples.
coraal_features = coraal.map(prepare_dataset, remove_columns=coraal.column_names, num_proc=2)
librispeech_features = librispeech.map(prepare_dataset, remove_columns=librispeech.column_names, num_proc=2)

# Store the normalized reference transcripts for scoring
coraal_df["reference"] = coraal_df["label"].map(normalize_text)
librispeech_references = [normalize_text(text) for text in librispeech["label"]]

Map (num_proc=2):   0%|          | 0/800 [00:00<?, ? examples/s]

Map (num_proc=2):   0%|          | 0/100 [00:00<?, ? examples/s]

def load_model():
    """Load the pre-trained Whisper model, configured for English transcription and SpecAugment."""
    model = WhisperForConditionalGeneration.from_pretrained(MODEL_ID)

    # Transcribe in English. The language and task are set on the generation config, so the
    # legacy forced decoder IDs are cleared to avoid conflicting with them.
    model.generation_config.language = "english"
    model.generation_config.task = "transcribe"
    model.generation_config.forced_decoder_ids = None
    model.config.forced_decoder_ids = None

    # Enable SpecAugment, which randomly masks short spans of time in the spectrogram during
    # training. This regularizes the model on the small training set.
    model.config.apply_spec_augment = True
    model.config.mask_time_prob = 0.05
    model.config.mask_feature_prob = 0.0

    return model

@dataclass
class DataCollatorSpeechSeq2SeqWithPadding:
    """Pad audio features and label sequences into a batch for the Whisper model.

    Inputs and labels have different lengths and need different padding methods,
    so they are padded separately.
    """

    processor: Any
    decoder_start_token_id: int

    def __call__(self, features: list[dict[str, Any]]) -> dict[str, torch.Tensor]:
        # Pad the log-Mel spectrograms and return them as PyTorch tensors
        input_features = [{"input_features": feature["input_features"]} for feature in features]
        batch = self.processor.feature_extractor.pad(input_features, return_tensors="pt")

        # Pad the tokenized label sequences to the longest sequence in the batch
        label_features = [{"input_ids": feature["labels"]} for feature in features]
        labels_batch = self.processor.tokenizer.pad(label_features, return_tensors="pt")

        # Replace the padding with -100 so that it is ignored by the loss
        labels = labels_batch["input_ids"].masked_fill(labels_batch.attention_mask.ne(1), -100)

        # The model prepends the start-of-transcript token itself, so remove it from the labels if present
        if (labels[:, 0] == self.decoder_start_token_id).all().cpu().item():
            labels = labels[:, 1:]

        batch["labels"] = labels

        return batch

# Use the Word Error Rate (WER) metric for checkpoint selection
metric = evaluate.load("wer")

Downloading builder script: 0.00B [00:00, ?B/s]

def decode_predictions(pred):
    """Decode generated and label token IDs into normalized prediction and reference strings."""
    # Replace -100 (used as padding when batches are concatenated) with the padding token ID
    pad_token_id = processor.tokenizer.pad_token_id
    pred_ids = np.where(pred.predictions != -100, pred.predictions, pad_token_id)
    label_ids = np.where(pred.label_ids != -100, pred.label_ids, pad_token_id)

    pred_str = processor.tokenizer.batch_decode(pred_ids, skip_special_tokens=True)
    label_str = processor.tokenizer.batch_decode(label_ids, skip_special_tokens=True)

    return [normalize_text(text) for text in pred_str], [normalize_text(text) for text in label_str]


def compute_metrics(pred):
    """Compute the WER (as a percentage) of the model predictions."""
    pred_str, label_str = decode_predictions(pred)

    # WER is undefined for empty references, so skip any examples without words
    pairs = [(pred, ref) for pred, ref in zip(pred_str, label_str) if ref]
    predictions, references = zip(*pairs)

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

    return {"wer": wer}

def make_trainer(model, output_dir, eval_dataset, train_dataset=None, push_to_hub=False):
    """Create a Seq2SeqTrainer with the training configuration shared by every fold."""
    training_args = Seq2SeqTrainingArguments(
        output_dir=output_dir,
        seed=SEED,

        # Optimization
        per_device_train_batch_size=16,
        gradient_accumulation_steps=1,  # Increase by 2x for every 2x decrease in batch size
        learning_rate=1e-5,
        warmup_steps=50,
        max_steps=500,
        gradient_checkpointing=True,
        fp16=True,

        # Evaluation with generation; transcripts are short, so cap the generated length
        eval_strategy="steps",
        eval_steps=25,
        per_device_eval_batch_size=8,
        predict_with_generate=True,
        generation_max_length=128,
        label_names=["labels"],

        # Checkpointing: keep the checkpoint with the lowest validation WER
        save_steps=25,
        save_total_limit=2,
        load_best_model_at_end=True,
        metric_for_best_model="wer",
        greater_is_better=False,

        # Logging and Hugging Face Hub
        logging_steps=25,
        report_to=["tensorboard"],
        push_to_hub=push_to_hub,
        hub_model_id=HUB_MODEL_ID,
    )

    # Stop early if the validation WER does not improve for 4 consecutive evaluations
    return Seq2SeqTrainer(
        args=training_args,
        model=model,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        data_collator=DataCollatorSpeechSeq2SeqWithPadding(
            processor=processor,
            decoder_start_token_id=model.config.decoder_start_token_id,
        ),
        compute_metrics=compute_metrics,
        processing_class=processor,
        callbacks=[EarlyStoppingCallback(early_stopping_patience=4)],
    )

# Define the model card metadata for the model pushed to the Hugging Face Hub
hub_kwargs = {
    "dataset_tags": DATASET_ID,
    "dataset": "AAVE Speech Recognition CORAAL",
    "dataset_args": f"speaker-level {N_FOLDS}-fold cross-validation, fold {HUB_FOLD}",
    "model_name": "Whisper Tiny - CORAAL",
    "finetuned_from": MODEL_ID,
    "tasks": "automatic-speech-recognition",
}

# Transcribe every CORAAL clip and the LibriSpeech clips with the pre-trained model (zero-shot baseline).
# The baseline is not trained, so all 800 clips are valid test data for it.
baseline_trainer = make_trainer(load_model(), output_dir=f"{OUTPUT_DIR}/zero-shot", eval_dataset=coraal_features)
coraal_df["zero_shot"] = decode_predictions(baseline_trainer.predict(coraal_features))[0]
librispeech_zero_shot = decode_predictions(baseline_trainer.predict(librispeech_features))[0]

# Free GPU memory before training
del baseline_trainer
gc.collect()
torch.cuda.empty_cache()

Loading weights:   0%|          | 0/167 [00:00<?, ?it/s]

[transformers] Passing `generation_config` together with generation-related arguments=({'begin_suppress_tokens', 'suppress_tokens'}) is deprecated and will be removed in future versions. Please pass either a `generation_config` object OR all generation parameters explicitly, but not both.
[transformers] The attention mask is not set with a batched input, and cannot be inferred from input because pad token is same as eos token. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.

[transformers] Ignoring clean_up_tokenization_spaces=True for BPE tokenizer WhisperTokenizer. The clean_up_tokenization post-processing step is designed for WordPiece tokenizers and is destructive for BPE (it strips spaces before punctuation). Set clean_up_tokenization_spaces=False to suppress this warning, or set clean_up_tokenization_spaces_for_bpe_even_though_it_will_corrupt_output=True to force cleanup anyway.

# Train one model per fold, and transcribe that fold's held-out speakers with its best checkpoint
coraal_df["fine_tuned"] = ""
fold_histories = {}
librispeech_fine_tuned = {}

for fold in range(N_FOLDS):
    print(f"\n===== Fold {fold + 1} of {N_FOLDS} =====")
    train_idx, val_idx, test_idx = split_fold(fold)
    set_seed(SEED)

    trainer = make_trainer(
        load_model(),
        output_dir=f"{OUTPUT_DIR}/fold-{fold}",
        train_dataset=coraal_features.select(train_idx),
        eval_dataset=coraal_features.select(val_idx),
        push_to_hub=(fold == HUB_FOLD),
    )
    trainer.train()

    # Transcribe the held-out speakers and the standard English clips
    coraal_df.loc[test_idx, "fine_tuned"] = decode_predictions(trainer.predict(coraal_features.select(test_idx)))[0]
    librispeech_fine_tuned[fold] = decode_predictions(trainer.predict(librispeech_features))[0]

    # Keep the training logs and the selected checkpoint for plotting
    fold_histories[fold] = {
        "log_history": trainer.state.log_history,
        "best_checkpoint": trainer.state.best_model_checkpoint,
    }

    if fold == HUB_FOLD:
        # The splits were built with pandas, so they have no Hub config name. Mark them as pandas-built so
        # that Transformers does not write an empty config into the model card, which the Hub rejects.
        trainer.eval_dataset.info.builder_name = "pandas"

        # Push the best model and its model card to the Hugging Face Hub
        trainer.push_to_hub(**hub_kwargs)

    # Free GPU memory before the next fold
    del trainer
    gc.collect()
    torch.cuda.empty_cache()

# Save the predictions so that the analysis can be redone without retraining
coraal_df.drop(columns="audio").to_csv(f"{OUTPUT_DIR}/cv_predictions.csv", index=False)

===== Fold 1 of 5 =====

Loading weights:   0%|          | 0/167 [00:00<?, ?it/s]

Writing model shards:   0%|          | 0/1 [00:00<?, ?it/s]

Writing model shards:   0%|          | 0/1 [00:00<?, ?it/s]

Writing model shards:   0%|          | 0/1 [00:00<?, ?it/s]

Writing model shards:   0%|          | 0/1 [00:00<?, ?it/s]

Writing model shards:   0%|          | 0/1 [00:00<?, ?it/s]

Writing model shards:   0%|          | 0/1 [00:00<?, ?it/s]

Writing model shards:   0%|          | 0/1 [00:00<?, ?it/s]

Writing model shards:   0%|          | 0/1 [00:00<?, ?it/s]

def word_errors(reference, hypothesis):
    """Align one hypothesis with its reference.

    Returns the number of word errors, and for each reference word whether it was transcribed correctly.
    """
    ref_words, hyp_words = reference.split(), hypothesis.split()

    # jiwer cannot align empty strings: every reference word is a deletion and every hypothesis word an insertion
    if not ref_words or not hyp_words:
        return len(ref_words) + len(hyp_words), [False] * len(ref_words)

    output = jiwer.process_words(reference, hypothesis)
    hits = [False] * len(ref_words)
    for chunk in output.alignments[0]:
        if chunk.type == "equal":
            hits[chunk.ref_start_idx:chunk.ref_end_idx] = [True] * (chunk.ref_end_idx - chunk.ref_start_idx)

    return output.substitutions + output.deletions + output.insertions, hits


def corpus_wer(df, system, scoring="strict"):
    """Compute the corpus-level WER (as a percentage) of one system on the clips in df."""
    return 100 * df[f"{system}_errors_{scoring}"].sum() / df[f"words_{scoring}"].sum()


SYSTEMS = {"zero_shot": "Zero-shot", "fine_tuned": "Fine-tuned"}
SCORINGS = {"strict": normalize_text, "lenient": lenient_normalize}

# Count the reference words and each system's word errors for every clip under both scorings
for scoring, normalize in SCORINGS.items():
    references = coraal_df["reference"].map(normalize)
    coraal_df[f"words_{scoring}"] = references.str.split().str.len()
    for system in SYSTEMS:
        hypotheses = coraal_df[system].map(normalize)
        coraal_df[f"{system}_errors_{scoring}"] = [word_errors(ref, hyp)[0] for ref, hyp in zip(references, hypotheses)]

# Compare the zero-shot and fine-tuned WER on each fold's held-out speakers
fold_table = pd.DataFrame([
    {
        "Fold": fold,
        "Speakers": group["speaker"].nunique(),
        "Clips": len(group),
        "Zero-shot WER": corpus_wer(group, "zero_shot"),
        "Fine-tuned WER": corpus_wer(group, "fine_tuned"),
    }
    for fold, group in coraal_df.groupby("fold")
]).set_index("Fold")
fold_table["Change"] = fold_table["Fine-tuned WER"] - fold_table["Zero-shot WER"]
display(fold_table.round(2))

# Summarize the variation across folds
for column in ["Zero-shot WER", "Fine-tuned WER", "Change"]:
    print(f"{column:>15}: {fold_table[column].mean():6.2f} ± {fold_table[column].std():.2f} (mean ± std across folds)")

  Zero-shot WER:  34.76 ± 1.68 (mean ± std across folds)
 Fine-tuned WER:  29.17 ± 3.55 (mean ± std across folds)
         Change:  -5.59 ± 2.95 (mean ± std across folds)

def bootstrap_wer(scoring):
    """Estimate 95% confidence intervals for the zero-shot WER, fine-tuned WER, and their difference.

    Speakers (not clips) are resampled with replacement, since clips from the same speaker are not independent.
    """
    columns = [f"zero_shot_errors_{scoring}", f"fine_tuned_errors_{scoring}", f"words_{scoring}"]
    speaker_totals = coraal_df.groupby("speaker")[columns].sum().to_numpy()

    # Each row of samples is one bootstrap resample of speaker indices
    rng = np.random.default_rng(SEED)
    samples = rng.integers(0, len(speaker_totals), size=(N_BOOTSTRAP, len(speaker_totals)))
    resampled = speaker_totals[samples].sum(axis=1)

    zero_shot = 100 * resampled[:, 0] / resampled[:, 2]
    fine_tuned = 100 * resampled[:, 1] / resampled[:, 2]
    change = fine_tuned - zero_shot

    def interval(values):
        low, high = np.percentile(values, [2.5, 97.5])
        return f"[{low:.2f}, {high:.2f}]"

    point_zero_shot = corpus_wer(coraal_df, "zero_shot", scoring)
    point_fine_tuned = corpus_wer(coraal_df, "fine_tuned", scoring)

    # One-sided p-value: the share of resamples in which fine-tuning did not lower the WER
    p_value = (np.sum(change >= 0) + 1) / (N_BOOTSTRAP + 1)

    return {
        "Scoring": scoring,
        "Zero-shot WER": f"{point_zero_shot:.2f} {interval(zero_shot)}",
        "Fine-tuned WER": f"{point_fine_tuned:.2f} {interval(fine_tuned)}",
        "Change": f"{point_fine_tuned - point_zero_shot:+.2f} {interval(change)}",
        "Relative change": f"{100 * (point_fine_tuned - point_zero_shot) / point_zero_shot:+.1f}%",
        "p-value": f"{p_value:.4f}",
    }


# WER on all 800 clips with 95% confidence intervals, under strict and lenient scoring
display(pd.DataFrame([bootstrap_wer(scoring) for scoring in SCORINGS]).set_index("Scoring"))

# Plot each fold's validation loss and WER during training
FOLD_COLORS = ["#2a78d6", "#eb6834", "#1baf7a", "#eda100", "#e87ba4"]  # One color per fold

fig, (loss_ax, wer_ax) = plt.subplots(1, 2, figsize=(14, 5))

for fold, history in fold_histories.items():
    eval_logs = [log for log in history["log_history"] if "eval_loss" in log]
    steps = [log["step"] for log in eval_logs]
    color = FOLD_COLORS[fold]

    loss_ax.plot(steps, [log["eval_loss"] for log in eval_logs], color=color, linewidth=2, label=f"Fold {fold}")
    wer_ax.plot(steps, [log["eval_wer"] for log in eval_logs], color=color, linewidth=2, label=f"Fold {fold}")

    # Mark the selected checkpoint, whose step is at the end of its directory name, e.g. "checkpoint-125"
    best_step = int(history["best_checkpoint"].rsplit("-", 1)[-1])
    best_wer = next(log["eval_wer"] for log in eval_logs if log["step"] == best_step)
    wer_ax.plot(best_step, best_wer, marker="o", markersize=9, color=color, markeredgecolor="white", markeredgewidth=2)

loss_ax.set_title("Validation Loss by Fold")
loss_ax.set_ylabel("Loss")
wer_ax.set_title("Validation WER by Fold (dots mark the selected checkpoints)")
wer_ax.set_ylabel("WER (%)")
wer_ax.legend(frameon=False)

for ax in (loss_ax, wer_ax):
    ax.set_xlabel("Step")
    ax.grid(alpha=0.3)
    ax.spines[["top", "right"]].set_visible(False)

plt.tight_layout()
plt.show()

def wer_by_group(column):
    """Compare the zero-shot and fine-tuned strict WER for each value of a speaker attribute."""
    table = pd.DataFrame([
        {
            column: value,
            "Speakers": group["speaker"].nunique(),
            "Clips": len(group),
            "Zero-shot WER": corpus_wer(group, "zero_shot"),
            "Fine-tuned WER": corpus_wer(group, "fine_tuned"),
        }
        for value, group in coraal_df.groupby(column)
    ]).set_index(column)
    table["Change"] = table["Fine-tuned WER"] - table["Zero-shot WER"]
    return table.round(2)


for column in ["site_name", "gender", "age_group", "socioeconomic_group"]:
    display(wer_by_group(column))

# Plot the WER of each site before and after fine-tuning
site_table = wer_by_group("site_name").sort_values("Zero-shot WER")
positions = np.arange(len(site_table))
bar_height = 0.38

fig, ax = plt.subplots(figsize=(10, 6))
ax.barh(positions + bar_height / 2, site_table["Zero-shot WER"], bar_height, color="#2a78d6", label="Zero-shot")
ax.barh(positions - bar_height / 2, site_table["Fine-tuned WER"], bar_height, color="#eb6834", label="Fine-tuned")

ax.set_yticks(positions, site_table.index)
ax.set_xlabel("WER (%)")
ax.set_title("WER by CORAAL Site")
ax.legend(frameon=False, loc="lower right")
ax.grid(axis="x", alpha=0.3)
ax.spines[["top", "right"]].set_visible(False)

plt.tight_layout()
plt.show()

PRONOUNS = {"i", "you", "he", "she", "we", "they", "it"}

# Each feature is a test of whether reference word i (in a list of words) is an occurrence of it
FEATURES = {
    "ain't": lambda words, i: words[i] == "aint",
    "y'all": lambda words, i: words[i] == "yall",
    "tryna": lambda words, i: words[i] == "tryna",
    "finna": lambda words, i: words[i] == "finna",
    "habitual be (pronoun + be)": lambda words, i: words[i] == "be" and i > 0 and words[i - 1] in PRONOUNS,
    "gonna (informal spelling)": lambda words, i: words[i] == "gonna",
    "wanna (informal spelling)": lambda words, i: words[i] == "wanna",
    "cause (informal spelling)": lambda words, i: words[i] == "cause",
}

# For every clip and system, record which reference words were transcribed correctly
hit_flags = {
    system: [word_errors(ref, hyp)[1] for ref, hyp in zip(coraal_df["reference"], coraal_df[system])]
    for system in SYSTEMS
}

rows = []
for feature, is_occurrence in FEATURES.items():
    # Find every (clip, word position) where the feature occurs in the reference
    occurrences = [
        (clip, i)
        for clip, reference in enumerate(coraal_df["reference"])
        for i in range(len(reference.split()))
        if is_occurrence(reference.split(), i)
    ]
    if not occurrences:
        continue

    row = {"Feature": feature, "Occurrences": len(occurrences)}
    for system, name in SYSTEMS.items():
        row[f"{name} recall (%)"] = 100 * np.mean([hit_flags[system][clip][i] for clip, i in occurrences])
    rows.append(row)

display(pd.DataFrame(rows).set_index("Feature").round(1))

def librispeech_wer(predictions):
    """Compute the corpus-level strict WER (as a percentage) on the LibriSpeech sample."""
    errors = sum(word_errors(ref, hyp)[0] for ref, hyp in zip(librispeech_references, predictions))
    words = sum(len(ref.split()) for ref in librispeech_references)
    return 100 * errors / words


librispeech_fold_wers = [librispeech_wer(predictions) for predictions in librispeech_fine_tuned.values()]

# Every fold model transcribes the same LibriSpeech clips, so the fine-tuned LibriSpeech WER is the mean ± std across
# fold models. CORAAL WER is computed over all clips, as in the bootstrap table.
display(pd.DataFrame(
    {
        "LibriSpeech WER": [
            f"{librispeech_wer(librispeech_zero_shot):.2f}",
            f"{np.mean(librispeech_fold_wers):.2f} ± {np.std(librispeech_fold_wers, ddof=1):.2f}",
        ],
        "CORAAL WER": [
            f"{corpus_wer(coraal_df, 'zero_shot'):.2f}",
            f"{corpus_wer(coraal_df, 'fine_tuned'):.2f}",
        ],
    },
    index=["Zero-shot", "Fine-tuned"],
))

# Compute the strict WER of each clip, leaving it undefined for clips without reference words
reference_words = coraal_df["words_strict"].replace(0, np.nan)
for system in SYSTEMS:
    coraal_df[f"{system}_wer"] = 100 * coraal_df[f"{system}_errors_strict"] / reference_words

columns = ["file", "reference", "zero_shot", "fine_tuned", "zero_shot_wer", "fine_tuned_wer"]

with pd.option_context("display.max_colwidth", None):
    # A random sample of clips
    display(coraal_df[columns].sample(10, random_state=SEED).round(1))

    # The clips that the fine-tuned models transcribed worst
    display(coraal_df[columns].nlargest(5, "fine_tuned_wer").round(1))

Step	Training Loss	Validation Loss	Wer
25	1.532963	1.216541	32.888592
50	1.041475	0.837900	29.486324
75	0.777600	0.749881	30.353569
100	0.605078	0.732287	28.485657
125	0.492425	0.727856	36.757839
150	0.452669	0.722165	36.290861
175	0.352071	0.725699	30.953969
200	0.315806	0.721394	29.419613

Step	Training Loss	Validation Loss	Wer
25	1.518719	1.055863	24.425059
50	1.063937	0.690619	22.759715
75	0.714119	0.620159	23.314830
100	0.554112	0.596866	23.473434
125	0.433761	0.598104	22.997621
150	0.386335	0.598007	23.552736

Step	Training Loss	Validation Loss	Wer
25	1.483469	1.309779	36.432927
50	1.073553	0.917662	35.137195
75	0.735161	0.819563	33.231707
100	0.569061	0.799030	36.814024
125	0.471097	0.785360	34.375000
150	0.422784	0.780064	34.222561
175	0.335217	0.773955	33.460366

Step	Training Loss	Validation Loss	Wer
25	1.495267	1.358871	43.197937
50	1.033694	0.970243	37.137331
75	0.722541	0.895377	38.168923
100	0.554089	0.869575	37.524178
125	0.438802	0.857510	37.008382
150	0.410482	0.879902	36.750484
175	0.344096	0.873185	45.003224
200	0.283520	0.871883	40.618956
225	0.253072	0.879229	44.294004
250	0.220570	0.888820	49.194068

Step	Training Loss	Validation Loss	Wer
25	1.525024	1.242497	31.806931
50	1.040351	0.882037	28.960396
75	0.764968	0.791002	28.960396
100	0.579931	0.769050	27.599010
125	0.449191	0.756157	28.589109
150	0.408312	0.749337	28.341584
175	0.345917	0.753084	28.836634
200	0.284342	0.751124	28.341584

Improving AAVE Speech-to-Text¶

Prepare Environment¶

Load Dataset¶

Load Standard English Comparison Set¶

Prepare Processor and Data¶

Training and Evaluation¶

Evaluate Performance¶

Results by Speaker Group¶

AAVE Feature Recall¶

Standard English¶

Example Predictions¶

Results & Notes¶

Overall¶

Training¶

Repetition Loops¶

Speaker Groups¶

AAVE Features¶

Standard English¶

Limitations¶

	Speakers	Clips	Zero-shot WER	Fine-tuned WER	Change
Fold
0	34	131	33.99	24.29	-9.70
1	35	210	37.09	30.28	-6.81
2	34	152	33.88	28.56	-5.31
3	33	154	35.89	34.15	-1.74
4	33	153	32.96	28.58	-4.38

	Zero-shot WER	Fine-tuned WER	Change	Relative change	p-value
Scoring
strict	34.92 [31.44, 38.56]	29.35 [26.14, 32.82]	-5.57 [-8.10, -3.03]	-16.0%	0.0001
lenient	33.51 [29.93, 37.20]	28.90 [25.66, 32.40]	-4.61 [-7.18, -2.05]	-13.8%	0.0007

	Speakers	Clips	Zero-shot WER	Fine-tuned WER	Change
site_name
Atlanta, GA	13	100	28.77	25.16	-3.62
Detroit, MI	30	100	37.48	38.89	1.41
Lower East Side, NY	10	100	25.65	23.37	-2.28
Princeville, NC	16	100	46.23	38.25	-7.98
Rochester, NY	14	100	22.56	18.34	-4.21
Valdosta, GA	12	100	34.80	26.33	-8.46
Washington, DC (1968)	38	100	40.95	32.78	-8.17
Washington, DC (2016)	36	100	43.18	31.73	-11.45

	Speakers	Clips	Zero-shot WER	Fine-tuned WER	Change
age_group
1	49	218	33.42	29.83	-3.59
2	39	194	32.48	27.14	-5.34
3	50	245	35.96	28.62	-7.34
4	31	143	38.83	32.91	-5.92

	Speakers	Clips	Zero-shot WER	Fine-tuned WER	Change
socioeconomic_group
0	65	500	31.57	26.28	-5.29
1	33	100	45.50	37.54	-7.95
2	32	91	39.61	33.64	-5.97
3	39	109	36.75	32.38	-4.37

	Occurrences	Zero-shot recall (%)	Fine-tuned recall (%)
Feature
ain't	17	35.3	29.4
y'all	9	22.2	66.7
tryna	3	0.0	0.0
habitual be (pronoun + be)	5	40.0	80.0
gonna (informal spelling)	50	46.0	80.0
wanna (informal spelling)	23	0.0	56.5
cause (informal spelling)	53	5.7	52.8

	file	reference	zero_shot	fine_tuned	zero_shot_wer	fine_tuned_wer
514	utt_563_PRV_se0_ag1_f_01_1	okay my grandmas like i said well she said well i put it out i was like	okay im gonna be like i said well she said well i put it out i was like	okay im gonna be like i said well she said well i put it out its like	17.6	29.4
168	utt_251_DCB_se1_ag2_f_02_1	so i enrolled in school i didnt know who was gonna watch my kids but it was just like	so ive run this school i didnt know who was gonna watch my kids but it was a sign	so i wrote this school i didnt know who was gonna watch my kids but it was a sign	26.3	21.1
144	utt_22_ATL_se0_ag1_f_02_1	and for them to say oh be careful of the sun dont look into the sun cause it might hurt your eyes or fuck with your eyes we are sun people	and for them to say oh be careful of the sun dont look into the sun because it might hurt your eye the fuck we are eyes we are our sun people	and for them to say oh be careful of the sun dont look into the sun cause they might hurt your eyes the fuck we are eyes we are our sun people	19.4	16.1
123	utt_210_DCB_se1_ag4_m_01_1	and she didnt say anything you know she didnt say watch your mouth or anything it was just like	and she didnt say anything you know see them say watch them out of anything it was just like	and she didnt say anything you know see them say watch them out anything it was just like	26.3	26.3
521	utt_56_ATL_se0_ag2_m_03_1	so we end up staying with him and his wife and his two daughters you feel me so	so were gonna stand with him in his wife and his two daughters your family	so we gonna stand with him and his wife and his two daughters you fave me so	50.0	22.2
376	utt_439_LES_se0_ag2_m_01_1	which allows me to be able to work more independently you know just with myself and my mother and just	which allows me to be able to work more independently you know just with myself and my mother and just	which allows me to be able to work more independently you know just for myself and my mother and just	0.0	5.0
798	utt_99_ATL_se0_ag2_f_01_1	for cartoons and stuff like that so theres a lot of things that im interested in its not just one so	for cartoons and stuff like that so theres a lot of things that i mentioned there and its not just one	for cartoons and stuff like that so theres a lot of things that i mentioned that and its not just one so	23.8	19.0
293	utt_364_DTA_se2_ag3_f_03_1	and you throw the ball up and its a certain way you have to serve the ball hit the ball when you serve it	and you throw the ball up and its the same way you have to serve the ball hit the ball when you serve it	and you throw the ball up and its certain where you have to serve the ball hit the ball when you serve it	8.3	8.3
597	utt_638_ROC_se0_ag1_f_01_1	mostly white people you know obviously we stand out its a group we went for his	mostly white people you know obviously we stand out as group we went for his	mostly white people you know obviously we stand out as group we went for his	12.5	12.5
576	utt_619_ROC_se0_ag1_m_03_1	they love playing outside so if anybody has other kids around you know they would definitely make friends	they love playing outside so if anybody has other kids around you know they would definitely make friends	they love playing outside so if anybody has other kids around you know they would definitely make friends	0.0	0.0

	file	reference	zero_shot	fine_tuned	zero_shot_wer	fine_tuned_wer
257	utt_331_DTA_se3_ag1_f_01_1	this i at uh automobile accident over there on my one of the mans car was all turned up	this is an automobile accident over there it might win mens car was all turned up	this uh all the well i said and over there were my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my my	47.4	652.6
499	utt_54_ATL_se0_ag1_m_01_1	cause you know you know what im saying you might it depend bruh some niggas be intimidated you cant walk around there intimidated like	because you know you know what im saying my eight of them bro sunday is being timid at the time you walk around and timid at the time	cause you know you know what i saying my editor peeing bro so i think it was me and ten of the dead he came and walked around and then ten of the dead just like	79.2	108.3
319	utt_388_DTA_se1_ag1_f_01_1	i dont do something like this i dont know i dont like fights all that much	i dont do something i dont know i dont like fighting against you	you know im doing some m and i im not doing some m and i dont know i dont like face like that	37.5	100.0
375	utt_438_LES_se0_ag4_m_01_1	an account executive on an important account at a major advertising agency in new york city	and the counter executive on the report and the counter major advertising agency in new york city	and the counter executive under the port and the counter manager at the time they were seen in the york city	50.0	100.0
12	utt_110_DCA_se3_ag1_f_06_1	theyll tell us one thing and hell tell us another then when we tell him what theyve told us he gets very upset	l t ellis one day in hill t ellis another day in hill and another day in hill and another day in hill and another day in hill and another day in hill and another day in hill and another day in hill and another day in hill and another day in hill and another day in hill and another day in hill and another day in hill and another day in hill and another day in hill and another day in hill and another day in hill and another day in hill and another day in hill and another day in hill	l t r s w m d and hilltiles and other then we re tow him and they started to get to bury himself	434.8	95.7

Scoring	Zero-shot WER	Fine-tuned WER	Change	p-value
Strict	34.92 [31.44, 38.56]	29.35 [26.14, 32.82]	−5.57 [−8.10, −3.03] (−16.0%)	0.0001
Lenient	33.51 [29.93, 37.20]	28.90 [25.66, 32.40]	−4.61 [−7.18, −2.05] (−13.8%)	0.0007

Feature	Occurrences	Zero-shot recall	Fine-tuned recall
ain't	17	35.3%	29.4%
y'all	9	22.2%	66.7%
tryna	3	0.0%	0.0%
habitual be	5	40.0%	80.0%
gonna (informal)	50	46.0%	80.0%
wanna (informal)	23	0.0%	56.5%
cause (informal)	53	5.7%	52.8%

	Speakers	Clips	Zero-shot WER	Fine-tuned WER	Change
gender
f	89	396	30.38	25.48	-4.91
m	80	404	39.38	33.16	-6.23