DOCS/openwakeword/automatic_model_training_simple_tfr.py

# -*- coding: utf-8 -*-
"""Automatic_model_training_simple_TFR.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1EM2MgEL9cJlMLy6mMYwie8Njo6jl4h8X
"""

from google.colab import drive
drive.mount('/content/drive')

"""## Cleanup Options (for training a new wake word)

Use the cell below to clear previous training outputs if you're starting a new model or encountering issues. You can choose to delete generated models, configuration files, and test audio.
"""

# @title Clear Previous Training Outputs
import os
import shutil

clear_generated_models = False #@param {type:"boolean", label:"Clear models for the CURRENT wake word from Google Drive"}
clear_config_file = True #@param {type:"boolean", label:"Clear my_model.yaml configuration file"}
clear_test_audio = True #@param {type:"boolean", label:"Clear test_generation.wav audio file"}

if 'drive_output_dir' in locals() and 'target_word' in locals() and clear_generated_models:
    # Construct the model-specific output directory based on the current target_word
    current_model_name = target_word.replace(" ", "_")
    current_model_specific_output_dir = os.path.join(drive_output_dir, current_model_name)

    if os.path.exists(current_model_specific_output_dir):
        print(f"Deleting specific model directory for '{current_model_name}': {current_model_specific_output_dir}")
        shutil.rmtree(current_model_specific_output_dir)
    else:
        print(f"Model directory for '{current_model_name}' not found: {current_model_specific_output_dir}")

if clear_config_file:
    if os.path.exists('my_model.yaml'):
        print("Deleting my_model.yaml")
        os.remove('my_model.yaml')
    else:
        print("my_model.yaml not found.")

if clear_test_audio:
    if os.path.exists('test_generation.wav'):
        print("Deleting test_generation.wav")
        os.remove('test_generation.wav')
    else:
        print("test_generation.wav not found.")

print("Cleanup complete. If you changed fundamental settings or encountered persistent issues, consider restarting the runtime via 'Runtime > Restart runtime...'.")

import os

# Create a directory in Google Drive to store the models
drive_output_dir = '/content/drive/My Drive/Colab Notebooks/Wakeword/openwakeword'
os.makedirs(drive_output_dir, exist_ok=True)
print(f"Models will be saved in: {drive_output_dir}")

"""## Training your own openWakeWord models

***NOTE (2026-04-11):*** If you get error messages saying that the "runtime needs to be restarted" because packages have already been imported, select restart and run the notebook from the top again.

**Quick-start:** If you just want to train a basic custom model for openWakeWord!

Follow the instructions for Step 1 below. Each time you change the wake word, click the play icon to the left of the title to generate a sample and make sure it sounds correct. The first time it takes a few minutes but subsequent runs will be quick.

Once you're satisfied with the pronounciation, go to the "Runtime" dropdown menu in the upper left of the page, and select "run all". Keep the tab open but feel free to do something else. After ~1 hour, your custom model will be ready and will automatically be downloaded to your computer!

If you are a Home Assistant user with the openWakeWord add-on, follow the instructions [here](https://github.com/home-assistant/addons/blob/master/openwakeword/DOCS.md#custom-wake-word-models) to install and enable your custom model.

---

If you are interested in learning more about the custom model training process (and increasing the accuracy of your custom models), read through each step in this notebook and try experimenting with different training parameters. If you have any questions or problems, feel free to start a discussion at the openWakeWord [repo](https://github.com/dscripka/openWakeWord/discussions).
"""

# @title  { display-mode: "form" }
# @markdown # 1. Test Example Training Clip Generation
# @markdown Since openWakeWord models are trained on synthetic examples of your
# @markdown target wake word, it's a good idea to make sure that the examples
# @markdown sound correct. Type in your target wake word below, and run the
# @markdown cell to listen to it.
# @markdown
# @markdown Here are some tips that can help get the wake word to sound right:

# @markdown - If your wake word isn't being pronounced in the way
# @markdown you want, try spelling out the sounds phonetically with underscores
# @markdown separating each part.
# @markdown For example: "hey siri" --> "hey_seer_e".

# @markdown - Spell out numbers ("2" --> "two")

# @markdown - Avoid all punctuation except for "?" and "!", and remove unicode characters

import os
import sys
from IPython.display import Audio

if not os.path.exists("./piper-sample-generator"):
    !git clone https://github.com/rhasspy/piper-sample-generator
    !wget -O piper-sample-generator/models/en_US-libritts_r-medium.pt 'https://github.com/rhasspy/piper-sample-generator/releases/download/v2.0.0/en_US-libritts_r-medium.pt'
    !cd piper-sample-generator && git checkout 213d4d5

    # Install system dependencies
    !pip install piper-tts piper-phonemize-cross
    !pip install webrtcvad
    # !pip install 'torch<=2.5' torchvision torchaudio
    !pip install torch==2.5.0 torchvision==0.20.0 torchaudio==2.5.0 --index-url https://download.pytorch.org/whl/cu121

target_word = 'hey_viva' # @param {type:"string"}

if "piper-sample-generator/" not in sys.path:
    sys.path.append("piper-sample-generator/")
from generate_samples import generate_samples

def text_to_speech(text):
    generate_samples(text = text,
                max_samples=1,
                length_scales=[1.1],
                noise_scales=[0.7], noise_scale_ws = [0.7],
                output_dir = './', batch_size=1, auto_reduce_batch_size=True,
                file_names=["test_generation.wav"]
                )

text_to_speech(target_word)

from IPython.display import Audio
Audio("test_generation.wav", autoplay=True)

# @title  { display-mode: "form" }
# @markdown # 2. Download Data
# @markdown Training custom models requires downloading a wide variety of data
# @markdown that will help make the model perform well in real-world scenarios.
# @markdown This example notebook will download small samples of background noise,
# @markdown music, and Room Impulse Responses (to add echo). This will still produce
# @markdown a custom model that performs well, but if you are interested in adding even more,
# @markdown feel free to extend this notebook to download the full datasets and even add
# @markdown your own!
# @markdown
# @markdown Downloading this example data will usually take about 15 minutes.

# @markdown **Important note!** The data downloaded here has a mixture of different
# @markdown licenses and usage restrictions. As such, any custom models trained with this
# @markdown data should be considered as appropriate for **non-commercial** personal use only.

# ## Install all dependencies
# !pip install datasets
# !pip install scipy
# !pip install tqdm

import locale
def getpreferredencoding(do_setlocale = True):
    return "UTF-8"
locale.getpreferredencoding = getpreferredencoding

import os
import shutil # Import shutil for rmtree

# install openwakeword (full installation to support training)
# Ensure a clean install by removing existing directory before cloning
if os.path.exists("./openwakeword"):
    shutil.rmtree("./openwakeword")
# Added a shell-level removal for robustness in case python's rmtree fails
!rm -rf ./openwakeword
!git clone https://github.com/dscripka/openwakeword
!pip install -e ./openwakeword --no-deps
# !cd openwakeword

# install other dependencies
!pip install mutagen==1.47.0
!pip install torchinfo==1.8.0
!pip install torchmetrics==1.2.0
!pip install speechbrain==0.5.14
!pip install audiomentations==0.33.0
!pip install torch-audiomentations==0.11.0
!pip install acoustics==0.2.6
# !pip uninstall tensorflow -y
# !pip install tensorflow-cpu==2.8.1
# !pip install protobuf==3.20.3
!pip install onnx_tf==1.10.0 # Added onnx_tf
!pip install onnxruntime==1.22.1 ai_edge_litert==1.4.0 onnxsim
!pip install onnx2tf
!pip install onnx==1.19.1
# !pip install ai_edge_litert==1.2.0
!pip install onnx_graphsurgeon
!pip install sng4onnx
!pip install pronouncing==0.2.0
!pip install datasets==2.14.6
!pip install deep-phonemizer==0.0.19

# Download required models (workaround for Colab)
os.makedirs("./openwakeword/openwakeword/resources/models", exist_ok=True)
!wget https://github.com/dscripka/openWakeWord/releases/download/v0.5.1/embedding_model.onnx -O ./openwakeword/openwakeword/resources/models/embedding_model.onnx
!wget https://github.com/dscripka/openWakeWord/releases/download/v0.5.1/embedding_model.tflite -O ./openwakeword/openwakeword/resources/models/embedding_model.tflite
!wget https://github.com/dscripka/openWakeWord/releases/download/v0.5.1/melspectrogram.onnx -O ./openwakeword/openwakeword/resources/models/melspectrogram.onnx
!wget https://github.com/dscripka/openWakeWord/releases/download/v0.5.1/melspectrogram.tflite -O ./openwakeword/openwakeword/resources/models/melspectrogram.tflite

# Imports
import sys

if "piper-sample-generator/" not in sys.path:
    sys.path.append("piper-sample-generator/")
from generate_samples import generate_samples

import numpy as np
import torch
import sys
from pathlib import Path
import uuid
import yaml
import datasets
import scipy
from tqdm import tqdm

## Download all data

## Download MIR RIR data (takes about ~2 minutes)
output_dir = "./mit_rirs"
if not os.path.exists(output_dir):
    os.mkdir(output_dir)
    !git lfs install
    !git clone https://huggingface.co/datasets/davidscripka/MIT_environmental_impulse_responses
    rir_dataset = datasets.Dataset.from_dict({"audio": [str(i) for i in Path("./MIT_environmental_impulse_responses/16khz").glob("*.wav")]}).cast_column("audio", datasets.Audio())
    # Save clips to 16-bit PCM wav files
    for row in tqdm(rir_dataset):
        name = row['audio']['path'].split('/')[-1]
        scipy.io.wavfile.write(os.path.join(output_dir, name), 16000, (row['audio']['array']*32767).astype(np.int16))

## Download noise and background audio (takes about ~3 minutes)

# Audioset Dataset (https://research.google.com/audioset/index.html)
# The previous download link for 'bal_train09.tar' resulted in a 404 Not Found error.
# To proceed, this section is commented out. The model will be trained without this
# part of the AudioSet dataset, which might affect its robustness in noisy environments.
# For full-scale training, it's recommended to find an updated or alternative source for
# a diverse background noise dataset.

# if not os.path.exists("audioset"):
#     os.mkdir("audioset")

#     fname = "bal_train09.tar"
#     out_dir = f"audioset/{fname}"
#     link = "https://huggingface.co/datasets/agkphysics/AudioSet/resolve/main/data/" + fname
#     !wget -O {out_dir} {link}
#     !cd audioset && tar -xvf bal_train09.tar

#     output_dir = "./audioset_16k"
#     if not os.path.exists(output_dir):
#         os.mkdir(output_dir)

#     # Save clips to 16-bit PCM wav files
#     audioset_dataset = datasets.Dataset.from_dict({"audio": [str(i) for i in Path("audioset/audio").glob("**/*.flac")]})
#     audioset_dataset = audioset_dataset.cast_column("audio", datasets.Audio(sampling_rate=16000))
#     for row in tqdm(audioset_dataset):
#         name = row['audio']['path'].split('/')[-1].replace(".flac", ".wav")
#         scipy.io.wavfile.write(os.path.join(output_dir, name), 16000, (row['audio']['array']*32767).astype(np.int16))

# Free Music Archive dataset
# https://github.com/mdeff/fma

output_dir = "./fma"
if not os.path.exists(output_dir):
    os.mkdir(output_dir)
    fma_dataset = datasets.load_dataset("rudraml/fma", name="small", split="train", streaming=True)
    fma_dataset = iter(fma_dataset.cast_column("audio", datasets.Audio(sampling_rate=16000)))

    # Save clips to 16-bit PCM wav files
    n_hours = 1  # use only 1 hour of clips for this example notebook, recommend increasing for full-scale training
    for i in tqdm(range(n_hours*3600//30)):  # this works because the FMA dataset is all 30 second clips
        row = next(fma_dataset)
        name = row['audio']['path'].split('/')[-1].replace(".mp3", ".wav")
        scipy.io.wavfile.write(os.path.join(output_dir, name), 16000, (row['audio']['array']*32767).astype(np.int16))
        i += 1
        if i == n_hours*3600//30:
            break

# Download pre-computed openWakeWord features for training and validation

# training set (~2,000 hours from the ACAV100M Dataset)
# See https://huggingface.co/datasets/davidscripka/openwakeword_features for more information
if not os.path.exists("./openwakeword_features_ACAV100M_2000_hrs_16bit.npy"):
    !wget https://huggingface.co/datasets/davidscripka/openwakeword_features/resolve/main/openwakeword_features_ACAV100M_2000_hrs_16bit.npy

# validation set for false positive rate estimation (~11 hours)
if not os.path.exists("validation_set_features.npy"):
    !wget https://huggingface.co/datasets/davidscripka/openwakeword_features/resolve/main/validation_set_features.npy

# @title  { display-mode: "form" }
# @markdown # 3. Train the Model
# @markdown Now that you have verified your target wake word and downloaded the data,
# @markdown the last step is to adjust the training paramaters (or keep
# @markdown the defaults below) and start the training!

# @markdown Each paramater controls a different aspect of training:
# @markdown - `number_of_examples` controls how many examples of your wakeword
# @markdown are generated. The default (1,000) usually produces a good model,
# @markdown but between 30,000 and 50,000 is often the best.

# @markdown - `number_of_training_steps` controls how long to train the model.
# @markdown Similar to the number of examples, the default (10,000) usually works well
# @markdown but training longer usually helps.

# @markdown - `false_activation_penalty` controls how strongly false activations
# @markdown are penalized during the training process. Higher values can make the model
# @markdown much less likely to activate when it shouldn't, but may also cause it
# @markdown to not activate when the wake word isn't spoken clearly and there is
# @markdown background noise.

# @markdown With the default values shown below,
# @markdown this takes about 30 - 60 minutes total on the normal CPU Colab runtime.
# @markdown If you want to train on more examples or train for longer,
# @markdown try changing the runtime type to a GPU to significantly speedup
# @markdown the example generating and model training.

# @markdown When the model finishes training, you can navigate to the `my_custom_model` folder
# @markdown in the file browser on the left (click on the folder icon), and download
# @markdown the [your target wake word].onnx or  <your target wake word>.tflite files.
# @markdown You can then use these as you would any other openWakeWord model!

# Load default YAML config file for training
import yaml
import os # Import os for path manipulation
import shutil # Import shutil for file operations

config = yaml.load(open("openwakeword/examples/custom_model.yml", 'r').read(), yaml.Loader)

# Modify values in the config and save a new version
number_of_examples = 5000 # @param {type:"slider", min:100, max:50000, step:50}
number_of_training_steps = 12500  # @param {type:"slider", min:0, max:50000, step:100}
false_activation_penalty = 1500  # @param {type:"slider", min:100, max:5000, step:50}
config["target_phrase"] = [target_word]
config["model_name"] = config["target_phrase"][0].replace(" ", "_")
config["n_samples"] = number_of_examples
config["n_samples_val"] = max(500, number_of_examples//10)
config["steps"] = number_of_training_steps
config["target_accuracy"] = 0.5
config["target_recall"] = 0.25

# Define the final model output directory in Google Drive
# This path will contain only the final .onnx and .tflite models
final_drive_output_dir = os.path.join(drive_output_dir, config['model_name'])
os.makedirs(final_drive_output_dir, exist_ok=True)

# Define a local temporary directory for all intermediate training artifacts
# This will improve I/O performance during clip generation and training
local_intermediate_dir = './openwakeword_training_temp'
os.makedirs(local_intermediate_dir, exist_ok=True)

config["output_dir"] = local_intermediate_dir # Set training output to the local temporary directory
config["max_negative_weight"] = false_activation_penalty

config["background_paths"] = ['./fma']  # Removed './audioset_16k' as it was not downloaded
config["false_positive_validation_data_path"] = "validation_set_features.npy"
config["feature_data_files"] = {"ACAV100M_sample": "openwakeword_features_ACAV100M_2000_hrs_16bit.npy"}

with open('my_model.yaml', 'w') as file:
    documents = yaml.dump(config, file)

# Generate clips (outputs to local_intermediate_dir)
!{sys.executable} openwakeword/openwakeword/train.py --training_config my_model.yaml --generate_clips

# Step 2: Augment the generated clips (outputs to local_intermediate_dir)
!{sys.executable} openwakeword/openwakeword/train.py --training_config my_model.yaml --augment_clips

# Step 3: Train model (outputs initial ONNX to local_intermediate_dir)
!{sys.executable} openwakeword/openwakeword/train.py --training_config my_model.yaml --train_model

# Copy the final ONNX model from the local temporary directory to Google Drive
local_onnx_model_path = os.path.join(local_intermediate_dir, f"{config['model_name']}.onnx")
final_drive_onnx_model_path = os.path.join(final_drive_output_dir, f"{config['model_name']}.onnx")
shutil.copy(local_onnx_model_path, final_drive_onnx_model_path)
print(f"Copied final ONNX model to Google Drive: {final_drive_onnx_model_path}")

# Convert ONNX model to tflite using `onnx2tf` library
# This conversion will now operate on the ONNX model in Google Drive and output the TFLite model to Google Drive
onnx_model_path_for_conversion = final_drive_onnx_model_path # Use the ONNX model from Drive for conversion
drive_tflite_name1 = f"{final_drive_output_dir}/{config['model_name']}_float32.tflite"
drive_tflite_name2 = f"{final_drive_output_dir}/{config['model_name']}.tflite"

!onnx2tf -i "{onnx_model_path_for_conversion}" -o "{final_drive_output_dir}" -kat onnx____Flatten_0
!mv "{drive_tflite_name1}" "{drive_tflite_name2}"

print(f"ONNX model converted to TFLite and saved to: {drive_tflite_name2}")

# Automatically download the trained model files from Google Drive
from google.colab import files # Ensure files is imported

files.download(final_drive_onnx_model_path)
files.download(drive_tflite_name2)

"""## Convert ONNX to TFLite

This section converts the trained ONNX model to TFLite format and saves it to your Google Drive.
"""

!pip install onnx_tf==1.10.0 # Added onnx_tf

import os
from google.colab import files

# Ensure drive_output_dir and config['model_name'] are available
# (These variables are typically defined in previous cells)

# Define the model-specific output directory
# This cell should be run after cell qgaKWIY6WlJ1 for config to be defined
model_specific_output_dir = os.path.join(drive_output_dir, config['model_name'])

# Construct the paths for the ONNX and TFLite models
onnx_model_path = f"{model_specific_output_dir}/{config['model_name']}.onnx"
name1, name2 = f"{model_specific_output_dir}/{config['model_name']}_float32.tflite", f"{model_specific_output_dir}/{config['model_name']}.tflite"

# Convert ONNX model to tflite using `onnx2tf` library
# Paths are quoted to handle spaces in directory names
!onnx2tf -i "{onnx_model_path}" -o "{model_specific_output_dir}" -kat onnx____Flatten_0

# Rename the generated float32 TFLite model to the desired name
!mv "{name1}" "{name2}"

print(f"ONNX model converted to TFLite and saved to: {name2}")

# Automatically download the trained model files for convenience
files.download(onnx_model_path)
files.download(name2)