put onnx to tflite conversion behind flag in training code

Bump python version to 3.10+, switch tflite to new ai-edge-litert library

.gitattributes

@@ -1 +0,0 @@
-*.onnx filter=lfs diff=lfs merge=lfs -text

.github/workflows/build_and_publish_to_pypi.yml

@@ -3,6 +3,9 @@ name: Publish Python distributions to PyPI
 on: 
   push:
   workflow_dispatch:
+  create:
+    tags:
+      - "*"
 
 jobs:
   build-n-publish:
@@ -10,12 +13,10 @@ jobs:
     runs-on: ubuntu-latest
     steps:
     - uses: actions/checkout@master
-      with:
-        lfs: true
-    - name: Set up Python 3.8
+    - name: Set up Python 3.10
       uses: actions/setup-python@v3
       with:
-        python-version: "3.8"
+        python-version: "3.10"
     - name: Install pypa/build
       run: >-
         python -m

.github/workflows/tests.yml

@@ -8,18 +8,17 @@ on:
     branches: [ "main" ]
   pull_request:
     branches: [ "main" ]
+  workflow_dispatch:
 
 jobs:
   unit_tests_linux:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.8"]
+        python-version: ["3.10"]
 
     steps:
     - uses: actions/checkout@v3
-      with:
-        lfs: true
     - name: Set up Python ${{ matrix.python-version }}
       uses: actions/setup-python@v3
       with:
@@ -29,7 +28,6 @@ jobs:
         sudo apt-get install libspeexdsp-dev
         python -m pip install --upgrade pip
         pip install -e .[test]
-        pip install https://github.com/dscripka/openWakeWord/releases/download/v0.1.1/speexdsp_ns-0.1.2-cp38-cp38-linux_x86_64.whl
     - name: Test with pytest
       run: |
         pytest
@@ -38,12 +36,10 @@ jobs:
     runs-on: windows-latest
     strategy:
       matrix:
-        python-version: ["3.8"]
+        python-version: ["3.10"]
 
     steps:
     - uses: actions/checkout@v3
-      with:
-        lfs: true
     - name: Set up Python ${{ matrix.python-version }}
       uses: actions/setup-python@v3
       with:

.gitignore

@@ -127,3 +127,18 @@ dmypy.json
 
 # Pyre type checker
 .pyre/
+
+# Most notebooks
+notebooks/
+!notebooks/automatic_model_training.ipynb
+!notebooks/converting_google_speech_embedding_model.ipynb
+!notebooks/performance_metrics.ipynb
+!notebooks/training_models.ipynb
+!training_tutorial_data
+
+# Most example files
+examples/
+!examples/audio/activation.wav
+
+# archive files
+archive/

CHANGELOG.md

@@ -0,0 +1,29 @@
+# Change Log
+
+## v0.6.0 - 2023/06/15
+
+### Added
+
+* Various bug fixes, and some new functionality in `model.py` to control repeated detections
+
+### Changed
+
+* Models are no longer included in the PyPi package, and must be downloaded separately
+
+### Removed
+
+## v0.5.0 - 2023/06/15
+
+### Added
+
+* A new wakeword model, "hey rhasspy"
+* Added support for tflite versions of the melspectrogram model, embedding model, and pre-trained wakeword models
+* Added an inference framework argument to allow users to select either ONNX or tflite as the inference framework
+* The `detect_from_microphone.py` example now supports additional arguments and has improved console formatting
+
+### Changed
+
+* Made tflite the default inference framework for linux platforms due to improved efficiency, with windows still using ONNX as the default given the lack of pre-built Windows WHLs for the tflite runtime (https://pypi.org/project/tflite/)
+* Adjusted the default provider arguments for onnx models to avoid warnings (https://github.com/dscripka/openWakeWord/issues/27)
+
+### Removed

MANIFEST.in

@@ -1 +0,0 @@
-recursive-include openwakeword *.onnx

README.md

@@ -4,13 +4,32 @@
 
 openWakeWord is an open-source wakeword library that can be used to create voice-enabled applications and interfaces. It includes pre-trained models for common words & phrases that work well in real-world environments.
 
+**Quick Links**
+- [Installation](#installation)
+- [Training New Models](#training-new-models)
+- [FAQ](#faq)
+
+# Updates
+
+**2024/02/11**
+- v0.6.0 of openWakeWord released. See the [releases](https://github.com/dscripka/openWakeWord/releases) for a full descriptions of new features and changes.
+
+**2023/11/09**
+- Added example scripts under `examples/web` that demonstrate streaming audio from a web application into openWakeWord.
+
+**2023/10/11**
+- Significant improvements to the process of [training new models](#training-new-models), including an example Google Colab notebook demonstrating how to train a basic wake word model in <1 hour.
+
+**2023/06/15**
+- v0.5.0 of openWakeWord released. See the [releases](https://github.com/dscripka/openWakeWord/releases) for a full descriptions of new features and changes.
+
 # Demo
 
-You can try an online demo of the included pre-trained models via HuggingFace Spaces [right here!](https://huggingface.co/spaces/davidscripka/openWakeWord).
+You can try an online demo of the included pre-trained models via HuggingFace Spaces [right here](https://huggingface.co/spaces/davidscripka/openWakeWord)!
 
 Note that real-time detection of a microphone stream can occasionally behave strangely in Spaces. For the most reliable testing, perform a local installation as described below.
 
-# Installation & Usage
+# Installation
 
 Installing openWakeWord is simple and has minimal dependencies:
 
@@ -18,6 +37,8 @@ Installing openWakeWord is simple and has minimal dependencies:
 pip install openwakeword
 ```
 
+On Linux systems, both the [onnxruntime](https://pypi.org/project/onnxruntime/) package and [tflite-runtime](https://pypi.org/project/tflite-runtime/) packages will be installed as dependencies since both inference frameworks are supported. On Windows, only onnxruntime is installed due to a lack of support for modern versions of tflite.
+
 To (optionally) use [Speex](https://www.speex.org/) noise suppression on Linux systems to improve performance in noisy environments, install the Speex dependencies and then the pre-built Python package (see the assets [here](https://github.com/dscripka/openWakeWord/releases/tag/v0.1.1) for all .whl versions), adjusting for your python version and system architecture as needed.
 
 ```
@@ -27,16 +48,22 @@ pip install https://github.com/dscripka/openWakeWord/releases/download/v0.1.1/sp
 
 Many thanks to [TeaPoly](https://github.com/TeaPoly/speexdsp-ns-python) for their Python wrapper of the Speex noise suppression libraries.
 
-For quick local testing, clone this repository and use the included [example script](examples/detect_from_microphone.py) to try streaming detection from a local microphone. **Important note!** The model files are stored in this repo using [git-lfs](https://git-lfs.com/); make sure it is installed on your system and if needed use `git-lfs fetch --all` to make sure the the models download correctly.
+# Usage
+
+For quick local testing, clone this repository and use the included [example script](examples/detect_from_microphone.py) to try streaming detection from a local microphone. You can individually download pre-trained models from current and past [releases](https://github.com/dscripka/openWakeWord/releases/), or you can download them using Python (see below).
 
 Adding openWakeWord to your own Python code requires just a few lines:
 
 ```python
+import openwakeword
 from openwakeword.model import Model
 
-# Instantiate the model
+# One-time download of all pre-trained models (or only select models)
+openwakeword.utils.download_models()
+
+# Instantiate the model(s)
 model = Model(
-    wakeword_model_paths=["path/to/model.onnx"],  # can also leave this argument empty to load all of the included pre-trained models
+    wakeword_models=["path/to/model.tflite"],  # can also leave this argument empty to load all of the included pre-trained models
 )
 
 # Get audio data containing 16-bit 16khz PCM audio data from a file, microphone, network stream, etc.
@@ -48,6 +75,27 @@ frame = my_function_to_get_audio_frame()
 prediction = model.predict(frame)
 ```
 
+Additionally, openWakeWord provides other useful utility functions. For example:
+
+```python
+# Get predictions for individual WAV files (16-bit 16khz PCM)
+from openwakeword.model import Model
+
+model = Model()
+model.predict_clip("path/to/wav/file")
+
+# Get predictions for a large number of files using multiprocessing
+from openwakeword.utils import bulk_predict
+
+bulk_predict(
+    file_paths = ["path/to/wav/file/1", "path/to/wav/file/2"],
+    wakeword_models = ["hey jarvis"],
+    ncpu=2
+)
+```
+
+See `openwakeword/utils.py` and `openwakeword/model.py` for the full specification of class methods and utility functions.
+
 # Recommendations for Usage
 
 ## Noise Suppression and Voice Activity Detection (VAD)
@@ -91,10 +139,11 @@ The table below lists each model, examples of the word/phrases it is trained to
 | alexa | "alexa"| [docs](docs/models/alexa.md) |
 | hey mycroft | "hey mycroft" | [docs](docs/models/hey_mycroft.md) |
 | hey jarvis | "hey jarvis" | [docs](docs/models/hey_jarvis.md) |
+| hey rhasspy | "hey rhasspy" | TBD
 | current weather | "what's the weather" | [docs](docs/models/weather.md) |
 | timers | "set a 10 minute timer" | [docs](docs/models/timers.md) |
 
-Based on the methods discussed in [performance testing](#performance-and-evaluation), each included model aims to meet the target performance criteria of <5% false-reject rates and <0.5/hour false-accept rates with appropriate threshold tuning. These levels are subjective, but hopefully are below the annoyance threshold where the average user becomes frustrated with a system that often misses intended activations and/or causes disruption by activating too frequently at undesired times. For example, at these performance levels a user could expect to have the model process continuous mixed content audio of several hours with at most a few false activations, and have a failed intended activation in only 1/20 attempts (and a failed retry in only 1/400 attempts). 
+Based on the methods discussed in [performance testing](#performance-and-evaluation), each included model aims to meet the target performance criteria of <5% false-reject rates and <0.5/hour false-accept rates with appropriate threshold tuning. These levels are subjective, but hopefully are below the annoyance threshold where the average user becomes frustrated with a system that often misses intended activations and/or causes disruption by activating too frequently at undesired times. For example, at these performance levels a user could expect to have the model process continuous mixed content audio of several hours with at most a few false activations, and have a failed intended activation in only 1/20 attempts (and a failed retry in only 1/400 attempts).
 
 If you have a new wake word or phrase that you would like to see included in the next release, please open an issue, and we'll do a best to train a model! The focus of these requests and future release will be on words and phrases that have broad general usage versus highly specific application.
 
@@ -170,7 +219,15 @@ While the models are trained with background noise to increase robustness, in so
 
 # Training New Models
 
-Training new models is conceptually simple, and the entire process is demonstrated in a [tutorial notebook](notebooks/training_models.ipynb).
+openWakeWord includes an automated utility that greatly simplifies the process of training custom models. This can be used in two ways:
+
+1) A simple [Google Colab](https://colab.research.google.com/drive/1q1oe2zOyZp7UsB3jJiQ1IFn8z5YfjwEb?usp=sharing) notebook with an easy to use interface and simple end-to-end process. This allows anyone to produce a custom model very quickly (<1 hour) and doesn't require any development experience, but the performance of the model may be low in some deployment scenarios.
+
+2) A more detailed [notebook](notebooks/automatic_model_training.ipynb) that describes the training process in more details, and enables more customization. This can produce high quality models, but requires more development experience.
+
+For a collection of models trained using the notebooks above by the Home Assistant Community (and with much gratitude to @fwartner), see the excellent repository [here](https://github.com/fwartner/home-assistant-wakewords-collection).
+
+For users interested in understanding the fundamental concepts behind model training there is a more detailed, educational [tutorial notebook](notebooks/training_models.ipynb) also available. However, this specific notebook is not intended for training production models, and the automated process above is recommended for that purpose.
 
 Fundamentally, a new model requires two data generation and collection steps:
 
@@ -184,6 +241,38 @@ Currently, openWakeWord only supports English, primarily because the pre-trained
 
 Future release road maps may have non-english support. In particular, [Mycroft.AIs Mimic 3](https://github.com/MycroftAI/mimic3-voices) TTS engine may work well to help extend some support to other languages.
 
+# FAQ
+
+**Is there a Docker implementation for openWakeWord?**
+- While there isn't an official Docker implementation, [@dalehumby](https://github.com/dalehumby) [has created one](https://github.com/dalehumby/openWakeWord-rhasspy) that works very well!
+
+**Can openWakeWord be run in a browser with javascript?**
+- While the ONNX runtime [does support javascript](https://onnxruntime.ai/docs/get-started/with-javascript.html), much of the other functionality required for openWakeWord models would need to be ported. This is not currently on the roadmap, but please open an issue/start a discussion if this feature is of particular interest.
+- As a potential work-around for some applications, the example scripts in `examples/web` demonstrate how audio can be captured in a browser and streaming via websockets into openWakeWord running in a Python backend server.
+- Other potential options could include projects like `pyodide` (see [here](https://github.com/pyodide/pyodide/issues/4220)) for a related issue.
+
+**Is there a C++ version of openWakeWord?**
+- While the ONNX runtime [also has a C++ API](https://onnxruntime.ai/docs/get-started/with-cpp.html), there isn't an official C++ implementation of the full openWakeWord library. However, [@synesthesiam](https://github.com/synesthesiam) has created a [C++ version of openWakeWord](https://github.com/rhasspy/openWakeWord-cpp) with basic functionality implemented.
+
+**Is openWakeWord suitable for edge devices and microcontrollers?**
+- openWakeWord is generally small and efficient, but likely not enough to be suitable for deployment on very low power edge devices. For example, some experimentation by other openWakeWord users & contributors indicates that it may still take several seconds to process a single 80 ms frame on an [ESP32-S3](https://www.espressif.com/en/products/socs/esp32-s3) with quantized openWakeWord models. Instead, I would recommend the excellent [microWakeWord](https://github.com/kahrendt/microWakeWord) library from @kahrendt. It uses a similar synthetic-only training data approach and can produce high quality models that are efficient enough to run on very low power edge devices.
+
+**Why are there three separate models instead of just one?**
+- Separating the models was an intentional choice to provide flexibility and optimize the efficiency of the end-to-end prediction process. For example, with separate melspectrogram, embedding, and prediction models, each one can operate on different size inputs of audio to optimize overall latency and share computations between models. It certainly is possible to make a combined model with all of the steps integrated, though, if that was a requirement of a particular use case.
+
+**I still get a large number of false activations when I use the pre-trained models, how can I reduce these?**
+- First, review the [recommendations for usage](#recommendations-for-usage) and ensure that these options do not improve overall system accuracy. Second, experiment with [custom verifier models](#user-specific-models), if possible. If neither of these approaches are helping, please open an issue with details of the deployment environment and the types of false activations that you are experiencing. We certainly appreciate feedback & requests on how to improve the base pre-trained models!
+
+# Acknowledgements
+
+I am very grateful for the encouraging and positive response from the open-source community since the release of openWakeWord in January 2023. In particular, I want to acknowledge and thank the following individuals and groups for their feedback, collaboration, and development support:
+
+- [synesthesiam](https://github.com/synesthesiam)
+- [SecretSauceAI](https://github.com/secretsauceai)
+- [OpenVoiceOS](https://github.com/OpenVoiceOS)
+- [Nabu Casa](https://github.com/NabuCasa)
+- [Home Assistant](https://github.com/home-assistant)
+
 # License
 
-All of the code in openWakeWord is licensed under the **Apache 2.0** license. All of the included pre-trained models are licensed under the [Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International](https://creativecommons.org/licenses/by-nc-sa/4.0/) license due to the inclusion of datasets with unknown or restrictive licensing as part of the training data. If you are interested in pre-trained models with more permissive licensing, please raise an issue and we will try to add them to a future release.
+All of the code in this repository is licensed under the **Apache 2.0** license. All of the included pre-trained models are licensed under the [Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International](https://creativecommons.org/licenses/by-nc-sa/4.0/) license due to the inclusion of datasets with unknown or restrictive licensing as part of the training data. If you are interested in pre-trained models with more permissive licensing, please raise an issue and we will try to add them to a future release.

examples/capture_activations.py

@@ -68,10 +68,26 @@ parser.add_argument(
     default=False,
     required=False
 )
+parser=argparse.ArgumentParser()
 parser.add_argument(
-    "--model",
-    help="The model to use for openWakeWord, leave blank to use all available models",
+    "--chunk_size",
+    help="How much audio (in number of 16khz samples) to predict on at once",
+    type=int,
+    default=1280,
+    required=False
+)
+parser.add_argument(
+    "--model_path",
+    help="The path of a specific model to load",
     type=str,
+    default="",
+    required=False
+)
+parser.add_argument(
+    "--inference_framework",
+    help="The inference framework to use (either 'onnx' or 'tflite'",
+    type=str,
+    default='tflite',
     required=False
 )
 parser.add_argument(
@@ -87,25 +103,26 @@ args=parser.parse_args()
 FORMAT = pyaudio.paInt16
 CHANNELS = 1
 RATE = 16000
-CHUNK = 1280
+CHUNK = args.chunk_size
 audio = pyaudio.PyAudio()
 mic_stream = audio.open(format=FORMAT, channels=CHANNELS, rate=RATE, input=True, frames_per_buffer=CHUNK)
 
 # Load pre-trained openwakeword models
-if args.model:
+if args.model_path:
     model_paths = openwakeword.get_pretrained_model_paths()
     for path in model_paths:
-        if args.model in path:
+        if args.model_path in path:
             model_path = path
             
     if model_path:
         owwModel = Model(
-            wakeword_model_paths=[model_path],
+            wakeword_models=[model_path],
             enable_speex_noise_suppression=args.noise_suppression,
-            vad_threshold = args.vad_threshold
-            )
+            vad_threshold = args.vad_threshold,
+            inference_framework=args.inference_framework
+        )
     else: 
-        print(f'Could not find model \"{args.model}\"')
+        print(f'Could not find model \"{args.model_path}\"')
         exit()
 else:
     owwModel = Model(

examples/custom_model.yml

@@ -0,0 +1,101 @@
+## Configuration file to be used with `train.py` to create custom wake word/phrase models
+
+# The name of the model (will be used when creating directoires and when saving the final .onnx and .tflite files)
+model_name: "my_model"
+
+# The target word/phrase to be detected by the model. Adding multiple unique words/phrases will
+# still only train a binary model detection model, but it will activate on any one of the provided words/phrases.
+target_phrase:
+  - "hey jarvis"
+
+# Specific phrases that you do *not* want the model to activate on, outside of those generated automatically via phoneme overlap
+# This can be a good way to reduce false positives if you notice that, in practice, certain words or phrases are problematic
+custom_negative_phrases: []
+
+# The total number of positive samples to generate for training (minimum of 20,000 recommended, often 100,000+ is best)
+n_samples: 10000
+
+# The total number of positive samples to generate for validation and early stopping of model training
+n_samples_val: 2000
+
+# The batch size to use with Piper TTS when generating synthetic training data
+tts_batch_size: 50
+
+# The batch size to use when performing data augmentation on generated clips prior to training
+# It's recommended that this not be too large to ensure that there is enough variety in the augmentation
+augmentation_batch_size: 16
+
+# The path to a fork of the piper-sample-generator repository for TTS (https://github.com/dscripka/piper-sample-generator)
+piper_sample_generator_path: "./piper-sample-generator"
+
+# The output directory for the generated synthetic clips, openwakeword features, and trained models
+# Sub-directories will be automatically created for train and test clips for both positive and negative examples
+output_dir: "./my_custom_model"
+
+# The directories containing Room Impulse Response recordings
+rir_paths:
+  - "./mit_rirs"
+
+# The directories containing background audio files to mix with training data
+background_paths:
+  - "./background_clips"
+
+# The duplication rate for the background audio clips listed above (1 or higher). Can be useful as a way to oversample
+# a particular type of background noise more relevant to a given deployment environment. Values apply in the same
+# order as the background_paths list above. Only useful when multiple directories are provided above.
+background_paths_duplication_rate:
+  - 1
+
+# The location of pre-computed openwakeword features for false-positive validation data
+# If you do not have deployment environment validation data, a good general purpose dataset with
+# a reasonable mix with ~11 hours of speech, noise, and music is available here: https://huggingface.co/datasets/davidscripka/openwakeword_features
+false_positive_validation_data_path: "./validation_set_features.npy"
+
+# The number of times to apply augmentations to the generated training data
+# Values greater than 1 reuse each generation that many times, producing overall unique
+# clips for training due to the randomness intrinsic to the augmentation despite using
+# the same original synthetic generation. Can be a useful way to increase model robustness
+# without having to generate extremely large numbers of synthetic examples.
+augmentation_rounds: 1
+
+# Paths to pre-computed openwakeword features for positive and negative data. Each file must be a saved
+# .npy array (see the example notebook on manually training new models for details on how to create these).
+# There is no limit on the number of files but training speed will decrease as more
+# data will need to be read from disk for each additional file.
+# Also, there is a custom dataloader that uses memory-mapping with loading data, so the total size
+# of the files is not limited by the amount of available system memory (though this will result
+# in decreased training throughput depending on the speed of the underlying storage device). A fast
+# NVME SSD is recommended for optimal performance.
+
+feature_data_files:
+  "ACAV100M_sample": "./openwakeword_features_ACAV100M_2000_hrs_16bit.npy"
+
+# Define the number of examples from each data file per batch. Note that the key names here
+# must correspond to those define in the `feature_data_files` dictionary above (except for
+# the `positive` and `adversarial_negative` keys, which are automatically defined). The sum
+# of the values for each key define the total batch size for training. Initial testing indicates
+# that batch sizes of 1024-4096 work well in practice.
+
+batch_n_per_class:
+  "ACAV100M_sample": 1024
+  "adversarial_negative": 50
+  "positive": 50
+
+# Define the type of size of the openwakeword model to train. Increasing the layer size
+# may result in a more capable model, at the cost of decreased inference speed. The default
+# value (32) seems to work well in practice for most wake words/phrases.
+
+model_type: "dnn"
+layer_size: 32
+
+# Define training parameters. The values below are recommended defaults for most applications,
+# but unique deployment environments will likely require testing to determine which values
+# are the most appropriate.
+
+# The maximum number of steps to train the model
+steps: 50000
+
+# The maximum negative weight and target false positives per hour, used to control the auto training process
+# The target false positive rate may not be achieved, and adjusting the maximum negative weight may be necessary
+max_negative_weight: 1500
+target_false_positives_per_hour: 0.2

examples/detect_from_microphone.py

@@ -22,10 +22,24 @@ import argparse
 parser=argparse.ArgumentParser()
 parser.add_argument(
     "--chunk_size",
-    help="How much audio (in samples) to predict on at once",
+    help="How much audio (in number of samples) to predict on at once",
     type=int,
     default=1280,
-    required=True
+    required=False
+)
+parser.add_argument(
+    "--model_path",
+    help="The path of a specific model to load",
+    type=str,
+    default="",
+    required=False
+)
+parser.add_argument(
+    "--inference_framework",
+    help="The inference framework to use (either 'onnx' or 'tflite'",
+    type=str,
+    default='tflite',
+    required=False
 )
 
 args=parser.parse_args()
@@ -39,7 +53,12 @@ audio = pyaudio.PyAudio()
 mic_stream = audio.open(format=FORMAT, channels=CHANNELS, rate=RATE, input=True, frames_per_buffer=CHUNK)
 
 # Load pre-trained openwakeword models
-owwModel = Model()
+if args.model_path != "":
+    owwModel = Model(wakeword_models=[args.model_path], inference_framework=args.inference_framework)
+else:
+    owwModel = Model(inference_framework=args.inference_framework)
+
+n_models = len(owwModel.models.keys())
 
 # Run capture loop continuosly, checking for wakewords
 if __name__ == "__main__":
@@ -48,7 +67,7 @@ if __name__ == "__main__":
     print("#"*100)
     print("Listening for wakewords...")
     print("#"*100)
-    print("\n"*13)
+    print("\n"*(n_models*3))
 
     while True:
         # Get audio
@@ -73,5 +92,5 @@ if __name__ == "__main__":
             """
 
         # Print results table
-        print("\033[F"*14)
+        print("\033[F"*(4*n_models+1))
         print(output_string_header, "                             ", end='\r')

examples/web/README.md

@@ -0,0 +1,21 @@
+# Examples
+
+This folder contains examples of using openWakeWord with web applications.
+
+## Websocket Streaming
+
+As openWakeWord does not have a native Javascript port, using it within a web browswer is best accomplished with websocket streaming of the audio data from the browser to a simple Python application. To install the requirements for this example:
+
+```
+pip install aiohttp
+pip install resampy
+```
+
+The `streaming_client.html` page shows a simple implementation of audio capture and streamimng from a microphone and streaming in a browser, and the `streaming_server.py` file is the corresponding websocket server that passes the audio into openWakeWord.
+
+To run the example, execute `python streaming_server.py` (add the `--help` argument to see options) and navigate to `localhost:9000` in your browser.
+
+Note that this example is illustrative only, and integration of this approach with other web applications may have different requirements. In particular, some key considerations:
+
+- This example captures PCM audio from the web browser and streams full 16-bit integer representations of ~250 ms audio chunks over the websocket connection. In practice, bandwidth efficient streams of compressed audio may be more suitable for some applications.
+- The browser captures audio at the native sampling rate of the capture device, which can require re-sampling prior to passing the audio data to openWakeWord. This example uses the `resampy` library which has a good balance between performance and quality, but other resampling approaches that optimize different aspects may be more suitable for some applications.

examples/web/streaming_client.html

@@ -0,0 +1,197 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+  <meta charset="UTF-8">
+  <meta name="viewport" content="width=device-width, initial-scale=1.0">
+  <title>Websocket Microphone Streaming</title>
+  <style>
+    body {
+      text-align: center;
+      font-family: 'Roboto', sans-serif;
+    }
+    #startButton {
+      padding: 15px 30px;
+      font-size: 18px;
+      background-color: #03A9F4;
+      border: none;
+      border-radius: 4px;
+      color: white;
+      cursor: pointer;
+      outline: none;
+      transition: background-color 0.3s;
+    }
+    #startButton.listening {
+      background-color: #4CAF50;
+    }
+    table {
+      margin: 20px auto;
+      border-collapse: collapse;
+      width: 60%;
+    }
+    th, td {
+      border: 1px solid #E0E0E0;
+      padding: 10px;
+      text-align: left;
+    }
+    th {
+      background-color: #F5F5F5;
+    }
+
+    @keyframes fadeOut {
+      from {
+        opacity: 1;
+      }
+      to {
+        opacity: 0;
+      }
+    }
+
+    .detected-animation {
+      animation: fadeOut 2s forwards;
+    }
+  </style>
+</head>
+<body>
+  <h1>Streaming Audio to openWakeWord Using Websockets</h1>
+  <button id="startButton">Start Listening</button>
+
+  <table>
+    <tr>
+      <th>Wakeword</th>
+      <th>Detected</th>
+    </tr>
+    <tr>
+      <td></td>
+      <td></td>
+    </tr>
+  </table>
+
+  <script>
+  // Create websocket connection
+  const ws = new WebSocket('ws://localhost:9000/ws');
+
+  // When the websocket connection is open
+  ws.onopen = function() {
+    console.log('WebSocket connection is open');
+  };
+
+  // Get responses from websocket and display information
+  ws.onmessage = (event) => {
+    console.log(event.data);
+    const model_payload = JSON.parse(event.data);
+    if ("loaded_models" in model_payload) {
+      // Add loaded models to the rows of the first column in the table, inserting rows as needed
+      const table = document.querySelector('table');
+      const rows = table.querySelectorAll('tr');
+      for (let i = 1; i < model_payload.loaded_models.length + 1; i++) {
+        if (i < rows.length) {
+          const row = rows[i];
+          const cell = row.querySelectorAll('td')[0];
+          cell.textContent = model_payload.loaded_models[i - 1];
+        } else {
+          // Insert extra rows if needed, both column 1 and 2
+          const row = table.insertRow();
+          const cell1 = row.insertCell();
+          const cell2 = row.insertCell();
+          cell1.textContent = model_payload.loaded_models[i - 1];
+          cell2.textContent = '';
+        }
+      }
+
+    }
+
+    if ("activations" in model_payload) {
+      // Add detected wakeword to the rows of the second column in the table
+      const table = document.querySelector('table');
+      const rows = table.querySelectorAll('tr');
+      for (let i = 1; i < rows.length; i++) {
+        // Check for the model name in the first column and add "Detected!" to the second column if they match
+        if (model_payload.activations.includes(rows[i].querySelectorAll('td')[0].textContent)) {
+          const cell = rows[i].querySelectorAll('td')[1];
+          cell.textContent = "Detected!";
+          cell.classList.add('detected-animation'); // animate fade out
+      
+          // Remove the CSS class after the fade out animation ends to reset the state
+          cell.addEventListener('animationend', () => {
+            cell.textContent = '';
+            cell.classList.remove('detected-animation');
+          }, { once: true });
+        }
+      }
+    }
+  };
+
+  // Create microphone capture stream for 16-bit PCM audio data
+  // Code based on the excellent tutorial by Ragy Morkas: https://medium.com/@ragymorkos/gettineg-monochannel-16-bit-signed-integer-pcm-audio-samples-from-the-microphone-in-the-browser-8d4abf81164d
+  navigator.getUserMedia = navigator.getUserMedia || 
+                         navigator.webkitGetUserMedia || 
+                         navigator.mozGetUserMedia || 
+                         navigator.msGetUserMedia;
+ 
+  let audioStream;
+  let audioContext;
+  let recorder;
+  let volume;
+  let sampleRate;
+
+  if (navigator.getUserMedia) {
+    navigator.getUserMedia({audio: true}, function(stream) {
+      audioStream = stream;
+
+      // creates the an instance of audioContext
+      const context = window.AudioContext || window.webkitAudioContext;
+      audioContext = new context();
+      
+      // retrieve the current sample rate of microphone the browser is using and send to Python server
+      sampleRate = audioContext.sampleRate;
+      
+      // creates a gain node
+      volume = audioContext.createGain();
+      
+      // creates an audio node from the microphone incoming stream
+      const audioInput = audioContext.createMediaStreamSource(audioStream);
+      
+      // connect the stream to the gain node
+      audioInput.connect(volume);
+      
+      const bufferSize = 4096;
+      recorder = (audioContext.createScriptProcessor || 
+                  audioContext.createJavaScriptNode).call(audioContext, 
+                                                          bufferSize, 
+                                                          1, 
+                                                          1);
+
+      recorder.onaudioprocess = function(event) {
+        const samples = event.inputBuffer.getChannelData(0);
+        const PCM16iSamples = samples.map(sample => {
+          let val = Math.floor(32767 * sample);
+          return Math.min(32767, Math.max(-32768, val));
+        });
+
+        // Push audio to websocket
+        const int16Array = new Int16Array(PCM16iSamples);
+        const blob = new Blob([int16Array], { type: 'application/octet-stream' });
+        ws.send(blob);
+      };
+
+    }, function(error) {
+      alert('Error capturing audio.');
+    });
+  } else {
+    alert('getUserMedia not supported in this browser.');
+  }
+
+  // start recording
+  const startButton = document.getElementById('startButton');
+  startButton.addEventListener('click', function() {
+    if (!startButton.classList.contains('listening')) {
+      volume.connect(recorder);
+      recorder.connect(audioContext.destination);
+      ws.send(sampleRate);
+      startButton.classList.add('listening');
+      startButton.textContent = 'Listening...';
+    }
+  });
+  </script>
+</body>
+</html>

examples/web/streaming_server.py

@@ -0,0 +1,112 @@
+# Copyright 2023 David Scripka. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+#######################################################################################
+
+# This example scripts runs openWakeWord in a simple web server receiving audio
+# from a web page using websockets.
+
+#######################################################################################
+
+# Imports
+import aiohttp
+from aiohttp import web
+import numpy as np
+from openwakeword import Model
+import resampy
+import argparse
+import json
+
+# Define websocket handler
+async def websocket_handler(request):
+    ws = web.WebSocketResponse()
+    await ws.prepare(request)
+
+    # Send loaded models
+    await ws.send_str(json.dumps({"loaded_models": list(owwModel.models.keys())}))
+
+    # Start listening for websocket messages
+    async for msg in ws:
+        # Get the sample rate of the microphone from the browser
+        if msg.type == aiohttp.WSMsgType.TEXT:
+            sample_rate = int(msg.data)
+        elif msg.type == aiohttp.WSMsgType.ERROR:
+            print(f"WebSocket error: {ws.exception()}")
+        else:
+            # Get audio data from websocket
+            audio_bytes = msg.data
+
+            # Add extra bytes of silence if needed
+            if len(msg.data) % 2 == 1:
+                audio_bytes += (b'\x00')
+
+            # Convert audio to correct format and sample rate
+            data = np.frombuffer(audio_bytes, dtype=np.int16)
+            if sample_rate != 16000:
+                data = resampy.resample(data, sample_rate, 16000)
+
+            # Get openWakeWord predictions and set to browser client
+            predictions = owwModel.predict(data)
+
+            activations = []
+            for key in predictions:
+                if predictions[key] >= 0.5:
+                    activations.append(key)
+
+            if activations != []:
+                await ws.send_str(json.dumps({"activations": activations}))
+
+    return ws
+
+# Define static file handler
+async def static_file_handler(request):
+    return web.FileResponse('./streaming_client.html')
+
+app = web.Application()
+app.add_routes([web.get('/ws', websocket_handler), web.get('/', static_file_handler)])
+
+if __name__ == '__main__':
+    # Parse CLI arguments
+    parser=argparse.ArgumentParser()
+    parser.add_argument(
+        "--chunk_size",
+        help="How much audio (in number of samples) to predict on at once",
+        type=int,
+        default=1280,
+        required=False
+    )
+    parser.add_argument(
+        "--model_path",
+        help="The path of a specific model to load",
+        type=str,
+        default="",
+        required=False
+    )
+    parser.add_argument(
+        "--inference_framework",
+        help="The inference framework to use (either 'onnx' or 'tflite'",
+        type=str,
+        default='tflite',
+        required=False
+    )
+    args=parser.parse_args()
+
+    # Load openWakeWord models
+    if args.model_path != "":
+        owwModel = Model(wakeword_models=[args.model_path], inference_framework=args.inference_framework)
+    else:
+        owwModel = Model(inference_framework=args.inference_framework)
+
+    # Start webapp
+    web.run_app(app, host='localhost', port=9000)

notebooks/.gitignore

@@ -0,0 +1 @@
+cv11_test_clips

notebooks/automatic_model_training.ipynb

@@ -0,0 +1,504 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "c1eab0b3",
+   "metadata": {
+    "id": "c1eab0b3"
+   },
+   "source": [
+    "# Introduction"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "882058c5",
+   "metadata": {
+    "id": "882058c5"
+   },
+   "source": [
+    "This notebook demonstrates how to train custom openWakeWord models using pre-defined datasets and an automated process for dataset generation and training. While not guaranteed to always produce the best performing model, the methods shown in this notebook often produce baseline models with releatively strong performance.\n",
+    "\n",
+    "Manual data preparation and model training (e.g., see the [training models](training_models.ipynb) notebook) remains an option for when full control over the model development process is needed.\n",
+    "\n",
+    "At a high level, the automatic training process takes advantages of several techniques to try and produce a good model, including:\n",
+    "\n",
+    "- Early-stopping and checkpoint averaging (similar to [stochastic weight averaging](https://arxiv.org/abs/1803.05407)) to search for the best models found during training, according to the validation data\n",
+    "- Variable learning rates with cosine decay and multiple cycles\n",
+    "- Adaptive batch construction to focus on only high-loss examples when the model begins to converge, combined with gradient accumulation to ensure that batch sizes are still large enough for stable training\n",
+    "- Cycical weight schedules for negative examples to help the model reduce false-positive rates\n",
+    "\n",
+    "See the contents of the `train.py` file for more details."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e08d031b",
+   "metadata": {
+    "id": "e08d031b"
+   },
+   "source": [
+    "# Environment Setup"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "aee78c37",
+   "metadata": {
+    "id": "aee78c37"
+   },
+   "source": [
+    "To begin, we'll need to install the requirements for training custom models. In particular, a relatively recent version of Pytorch and custom fork of the [piper-sample-generator](https://github.com/dscripka/piper-sample-generator) library for generating synthetic examples for the custom model.\n",
+    "\n",
+    "**Important Note!** Currently, automated model training is only supported on linux systems due to the requirements of the text to speech library used for synthetic sample generation (Piper). It may be possible to use Piper on Windows/Mac systems, but that has not (yet) been tested."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4b1227eb",
+   "metadata": {
+    "id": "4b1227eb"
+   },
+   "outputs": [],
+   "source": [
+    "## Environment setup\n",
+    "\n",
+    "# install piper-sample-generator (currently only supports linux systems)\n",
+    "!git clone https://github.com/rhasspy/piper-sample-generator\n",
+    "!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'\n",
+    "!pip install piper-phonemize\n",
+    "!pip install webrtcvad\n",
+    "\n",
+    "# install openwakeword (full installation to support training)\n",
+    "!git clone https://github.com/dscripka/openwakeword\n",
+    "!pip install -e ./openwakeword\n",
+    "!cd openwakeword\n",
+    "\n",
+    "# install other dependencies\n",
+    "!pip install mutagen==1.47.0\n",
+    "!pip install torchinfo==1.8.0\n",
+    "!pip install torchmetrics==1.2.0\n",
+    "!pip install speechbrain==0.5.14\n",
+    "!pip install audiomentations==0.33.0\n",
+    "!pip install torch-audiomentations==0.11.0\n",
+    "!pip install acoustics==0.2.6\n",
+    "!pip install tensorflow-cpu==2.8.1\n",
+    "!pip install tensorflow_probability==0.16.0\n",
+    "!pip install onnx_tf==1.10.0\n",
+    "!pip install pronouncing==0.2.0\n",
+    "!pip install datasets==2.14.6\n",
+    "!pip install deep-phonemizer==0.0.19\n",
+    "\n",
+    "# Download required models (workaround for Colab)\n",
+    "import os\n",
+    "os.makedirs(\"./openwakeword/openwakeword/resources/models\")\n",
+    "!wget https://github.com/dscripka/openWakeWord/releases/download/v0.5.1/embedding_model.onnx -O ./openwakeword/openwakeword/resources/models/embedding_model.onnx\n",
+    "!wget https://github.com/dscripka/openWakeWord/releases/download/v0.5.1/embedding_model.tflite -O ./openwakeword/openwakeword/resources/models/embedding_model.tflite\n",
+    "!wget https://github.com/dscripka/openWakeWord/releases/download/v0.5.1/melspectrogram.onnx -O ./openwakeword/openwakeword/resources/models/melspectrogram.onnx\n",
+    "!wget https://github.com/dscripka/openWakeWord/releases/download/v0.5.1/melspectrogram.tflite -O ./openwakeword/openwakeword/resources/models/melspectrogram.tflite\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d4c1056e",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-09-04T13:42:01.183840Z",
+     "start_time": "2023-09-04T13:41:59.752153Z"
+    },
+    "id": "d4c1056e"
+   },
+   "outputs": [],
+   "source": [
+    "# Imports\n",
+    "\n",
+    "import os\n",
+    "import numpy as np\n",
+    "import torch\n",
+    "import sys\n",
+    "from pathlib import Path\n",
+    "import uuid\n",
+    "import yaml\n",
+    "import datasets\n",
+    "import scipy\n",
+    "from tqdm import tqdm\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e9d7a05a",
+   "metadata": {
+    "id": "e9d7a05a"
+   },
+   "source": [
+    "# Download Data"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c52f75cc",
+   "metadata": {
+    "id": "c52f75cc"
+   },
+   "source": [
+    "When training new openWakeWord models using the automated procedure, four specific types of data are required:\n",
+    "\n",
+    "1) Synthetic examples of the target word/phrase generated with text-to-speech models\n",
+    "\n",
+    "2) Synthetic examples of adversarial words/phrases generated with text-to-speech models\n",
+    "\n",
+    "3) Room impulse reponses and noise/background audio data to augment the synthetic examples and make them more realistic\n",
+    "\n",
+    "4) Generic \"negative\" audio data that is very unlikely to contain examples of the target word/phrase in the context where the model should detect it. This data can be the original audio data, or precomputed openWakeWord features ready for model training.\n",
+    "\n",
+    "5) Validation data to use for early-stopping when training the model.\n",
+    "\n",
+    "For the purposes of this notebook, all five of these sources will either be generated manually or can be obtained from HuggingFace thanks to their excellent `datasets` library and extremely generous hosting policy. Also note that while only a portion of some datasets are downloaded, for the best possible performance it is recommended to download the entire dataset and keep a local copy for future training runs."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d25a93b1",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-09-04T01:07:17.746749Z",
+     "start_time": "2023-09-04T01:07:17.740846Z"
+    },
+    "id": "d25a93b1"
+   },
+   "outputs": [],
+   "source": [
+    "# Download room impulse responses collected by MIT\n",
+    "# https://mcdermottlab.mit.edu/Reverb/IR_Survey.html\n",
+    "\n",
+    "output_dir = \"./mit_rirs\"\n",
+    "if not os.path.exists(output_dir):\n",
+    "    os.mkdir(output_dir)\n",
+    "rir_dataset = datasets.load_dataset(\"davidscripka/MIT_environmental_impulse_responses\", split=\"train\", streaming=True)\n",
+    "\n",
+    "# Save clips to 16-bit PCM wav files\n",
+    "for row in tqdm(rir_dataset):\n",
+    "    name = row['audio']['path'].split('/')[-1]\n",
+    "    scipy.io.wavfile.write(os.path.join(output_dir, name), 16000, (row['audio']['array']*32767).astype(np.int16))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2c0e178b",
+   "metadata": {
+    "id": "2c0e178b"
+   },
+   "outputs": [],
+   "source": [
+    "## Download noise and background audio\n",
+    "\n",
+    "# Audioset Dataset (https://research.google.com/audioset/dataset/index.html)\n",
+    "# Download one part of the audioset .tar files, extract, and convert to 16khz\n",
+    "# For full-scale training, it's recommended to download the entire dataset from\n",
+    "# https://huggingface.co/datasets/agkphysics/AudioSet, and\n",
+    "# even potentially combine it with other background noise datasets (e.g., FSD50k, Freesound, etc.)\n",
+    "\n",
+    "if not os.path.exists(\"audioset\"):\n",
+    "    os.mkdir(\"audioset\")\n",
+    "\n",
+    "fname = \"bal_train09.tar\"\n",
+    "out_dir = f\"audioset/{fname}\"\n",
+    "link = \"https://huggingface.co/datasets/agkphysics/AudioSet/resolve/main/data/\" + fname\n",
+    "!wget -O {out_dir} {link}\n",
+    "!cd audioset && tar -xvf bal_train09.tar\n",
+    "\n",
+    "output_dir = \"./audioset_16k\"\n",
+    "if not os.path.exists(output_dir):\n",
+    "    os.mkdir(output_dir)\n",
+    "\n",
+    "# Convert audioset files to 16khz sample rate\n",
+    "audioset_dataset = datasets.Dataset.from_dict({\"audio\": [str(i) for i in Path(\"audioset/audio\").glob(\"**/*.flac\")]})\n",
+    "audioset_dataset = audioset_dataset.cast_column(\"audio\", datasets.Audio(sampling_rate=16000))\n",
+    "for row in tqdm(audioset_dataset):\n",
+    "    name = row['audio']['path'].split('/')[-1].replace(\".flac\", \".wav\")\n",
+    "    scipy.io.wavfile.write(os.path.join(output_dir, name), 16000, (row['audio']['array']*32767).astype(np.int16))\n",
+    "\n",
+    "# Free Music Archive dataset (https://github.com/mdeff/fma)\n",
+    "output_dir = \"./fma\"\n",
+    "if not os.path.exists(output_dir):\n",
+    "    os.mkdir(output_dir)\n",
+    "fma_dataset = datasets.load_dataset(\"rudraml/fma\", name=\"small\", split=\"train\", streaming=True)\n",
+    "fma_dataset = iter(fma_dataset.cast_column(\"audio\", datasets.Audio(sampling_rate=16000)))\n",
+    "\n",
+    "n_hours = 1  # use only 1 hour of clips for this example notebook, recommend increasing for full-scale training\n",
+    "for i in tqdm(range(n_hours*3600//30)):  # this works because the FMA dataset is all 30 second clips\n",
+    "    row = next(fma_dataset)\n",
+    "    name = row['audio']['path'].split('/')[-1].replace(\".mp3\", \".wav\")\n",
+    "    scipy.io.wavfile.write(os.path.join(output_dir, name), 16000, (row['audio']['array']*32767).astype(np.int16))\n",
+    "    i += 1\n",
+    "    if i == n_hours*3600//30:\n",
+    "        break\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d01ec467",
+   "metadata": {
+    "id": "d01ec467"
+   },
+   "outputs": [],
+   "source": [
+    "# Download pre-computed openWakeWord features for training and validation\n",
+    "\n",
+    "# training set (~2,000 hours from the ACAV100M Dataset)\n",
+    "# See https://huggingface.co/datasets/davidscripka/openwakeword_features for more information\n",
+    "!wget https://huggingface.co/datasets/davidscripka/openwakeword_features/resolve/main/openwakeword_features_ACAV100M_2000_hrs_16bit.npy\n",
+    "\n",
+    "# validation set for false positive rate estimation (~11 hours)\n",
+    "!wget https://huggingface.co/datasets/davidscripka/openwakeword_features/resolve/main/validation_set_features.npy"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "cfe82647",
+   "metadata": {
+    "id": "cfe82647"
+   },
+   "source": [
+    "# Define Training Configuration"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b2e71329",
+   "metadata": {
+    "id": "b2e71329"
+   },
+   "source": [
+    "For automated model training openWakeWord uses a specially designed training script and a [YAML](https://yaml.org/) configuration file that defines all of the information required for training a new wake word/phrase detection model.\n",
+    "\n",
+    "It is strongly recommended that you review [the example config file](../examples/custom_model.yml), as each value is fully documented there. For the purposes of this notebook, we'll read in the YAML file to modify certain configuration parameters before saving a new YAML file for training our example model. Specifically:\n",
+    "\n",
+    "- We'll train a detection model for the phrase \"hey sebastian\"\n",
+    "- We'll only generate 5,000 positive and negative examples (to save on time for this example)\n",
+    "- We'll only generate 1,000 validation positive and negative examples for early stopping (again to save time)\n",
+    "- The model will only be trained for 10,000 steps (larger datasets will benefit from longer training)\n",
+    "- We'll reduce the target metrics to account for the small dataset size and limited training.\n",
+    "\n",
+    "On the topic of target metrics, there are *not* specific guidelines about what these metrics should be in practice, and you will need to conduct testing in your target deployment environment to establish good thresholds. However, from very limited testing the default values in the config file (accuracy >= 0.7, recall >= 0.5, false-positive rate <= 0.2 per hour) seem to produce models with reasonable performance.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fb0b6e4f",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-09-04T18:11:33.893397Z",
+     "start_time": "2023-09-04T18:11:33.878938Z"
+    },
+    "id": "fb0b6e4f"
+   },
+   "outputs": [],
+   "source": [
+    "# Load default YAML config file for training\n",
+    "config = yaml.load(open(\"openwakeword/examples/custom_model.yml\", 'r').read(), yaml.Loader)\n",
+    "config"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "482cf2d0",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-09-04T15:07:00.859210Z",
+     "start_time": "2023-09-04T15:07:00.841472Z"
+    },
+    "id": "482cf2d0"
+   },
+   "outputs": [],
+   "source": [
+    "# Modify values in the config and save a new version\n",
+    "\n",
+    "config[\"target_phrase\"] = [\"hey sebastian\"]\n",
+    "config[\"model_name\"] = config[\"target_phrase\"][0].replace(\" \", \"_\")\n",
+    "config[\"n_samples\"] = 1000\n",
+    "config[\"n_samples_val\"] = 1000\n",
+    "config[\"steps\"] = 10000\n",
+    "config[\"target_accuracy\"] = 0.6\n",
+    "config[\"target_recall\"] = 0.25\n",
+    "\n",
+    "config[\"background_paths\"] = ['./audioset_16k', './fma']  # multiple background datasets are supported\n",
+    "config[\"false_positive_validation_data_path\"] = \"validation_set_features.npy\"\n",
+    "config[\"feature_data_files\"] = {\"ACAV100M_sample\": \"openwakeword_features_ACAV100M_2000_hrs_16bit.npy\"}\n",
+    "\n",
+    "with open('my_model.yaml', 'w') as file:\n",
+    "    documents = yaml.dump(config, file)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "aa6b2ab0",
+   "metadata": {
+    "id": "aa6b2ab0"
+   },
+   "source": [
+    "# Train the Model"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a51202c0",
+   "metadata": {
+    "id": "a51202c0"
+   },
+   "source": [
+    "With the data downloaded and training configuration set, we can now start training the model. We'll do this in parts to better illustrate the sequence, but you can also execute every step at once for a fully automated process."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f01531fa",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-09-04T13:50:08.803326Z",
+     "start_time": "2023-09-04T13:50:06.790241Z"
+    },
+    "id": "f01531fa"
+   },
+   "outputs": [],
+   "source": [
+    "# Step 1: Generate synthetic clips\n",
+    "# For the number of clips we are using, this should take ~10 minutes on a free Google Colab instance with a T4 GPU\n",
+    "# If generation fails, you can simply run this command again as it will continue generating until the\n",
+    "# number of files meets the targets specified in the config file\n",
+    "\n",
+    "!{sys.executable} openwakeword/openwakeword/train.py --training_config my_model.yaml --generate_clips"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "afeedae4",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-09-04T13:56:08.781018Z",
+     "start_time": "2023-09-04T13:55:40.203515Z"
+    },
+    "id": "afeedae4"
+   },
+   "outputs": [],
+   "source": [
+    "# Step 2: Augment the generated clips\n",
+    "\n",
+    "!{sys.executable} openwakeword/openwakeword/train.py --training_config my_model.yaml --augment_clips"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9ad81ea0",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-09-04T15:11:14.742260Z",
+     "start_time": "2023-09-04T15:07:03.755159Z"
+    },
+    "id": "9ad81ea0"
+   },
+   "outputs": [],
+   "source": [
+    "# Step 3: Train model\n",
+    "\n",
+    "!{sys.executable} openwakeword/openwakeword/train.py --training_config my_model.yaml --train_model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "JSKWWLalnYzR",
+   "metadata": {
+    "id": "JSKWWLalnYzR"
+   },
+   "outputs": [],
+   "source": [
+    "# Step 4 (Optional): On Google Colab, sometimes the .tflite model isn't saved correctly\n",
+    "# If so, run this cell to retry\n",
+    "\n",
+    "# Manually save to tflite as this doesn't work right in colab\n",
+    "def convert_onnx_to_tflite(onnx_model_path, output_path):\n",
+    "    \"\"\"Converts an ONNX version of an openwakeword model to the Tensorflow tflite format.\"\"\"\n",
+    "    # imports\n",
+    "    import onnx\n",
+    "    import logging\n",
+    "    import tempfile\n",
+    "    from onnx_tf.backend import prepare\n",
+    "    import tensorflow as tf\n",
+    "\n",
+    "    # Convert to tflite from onnx model\n",
+    "    onnx_model = onnx.load(onnx_model_path)\n",
+    "    tf_rep = prepare(onnx_model, device=\"CPU\")\n",
+    "    with tempfile.TemporaryDirectory() as tmp_dir:\n",
+    "        tf_rep.export_graph(os.path.join(tmp_dir, \"tf_model\"))\n",
+    "        converter = tf.lite.TFLiteConverter.from_saved_model(os.path.join(tmp_dir, \"tf_model\"))\n",
+    "        tflite_model = converter.convert()\n",
+    "\n",
+    "        logging.info(f\"####\\nSaving tflite mode to '{output_path}'\")\n",
+    "        with open(output_path, 'wb') as f:\n",
+    "            f.write(tflite_model)\n",
+    "\n",
+    "    return None\n",
+    "\n",
+    "convert_onnx_to_tflite(f\"my_custom_model/{config['model_name']}.onnx\", f\"my_custom_model/{config['model_name']}.tflite\")\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f9OyUW3ltOSs",
+   "metadata": {
+    "id": "f9OyUW3ltOSs"
+   },
+   "source": [
+    "After the model finishes training, the auto training script will automatically convert it to ONNX and tflite versions, saving them as `my_custom_model/<model_name>.onnx/tflite` in the present working directory, where `<model_name>` is defined in the YAML training config file. Either version can be used as normal with `openwakeword`. I recommend testing them with the [`detect_from_microphone.py`](https://github.com/dscripka/openWakeWord/blob/main/examples/detect_from_microphone.py) example script to see how the model performs!"
+   ]
+  }
+ ],
+ "metadata": {
+  "colab": {
+   "provenance": []
+  },
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.12"
+  },
+  "toc": {
+   "base_numbering": 1,
+   "nav_menu": {},
+   "number_sections": true,
+   "sideBar": true,
+   "skip_h1_title": false,
+   "title_cell": "Table of Contents",
+   "title_sidebar": "Contents",
+   "toc_cell": false,
+   "toc_position": {},
+   "toc_section_display": true,
+   "toc_window_display": false
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

openwakeword/__init__.py

@@ -5,21 +5,48 @@ from openwakeword.custom_verifier_model import train_custom_verifier
 
 __all__ = ['Model', 'VAD', 'train_custom_verifier']
 
-models = {
+FEATURE_MODELS = {
+    "embedding": {
+        "model_path": os.path.join(os.path.dirname(os.path.abspath(__file__)), "resources/models/embedding_model.tflite"),
+        "download_url": "https://github.com/dscripka/openWakeWord/releases/download/v0.5.1/embedding_model.tflite"
+    },
+    "melspectrogram": {
+        "model_path": os.path.join(os.path.dirname(os.path.abspath(__file__)), "resources/models/melspectrogram.tflite"),
+        "download_url": "https://github.com/dscripka/openWakeWord/releases/download/v0.5.1/melspectrogram.tflite"
+    }
+}
+
+VAD_MODELS = {
+    "silero_vad": {
+        "model_path": os.path.join(os.path.dirname(os.path.abspath(__file__)), "resources/models/silero_vad.onnx"),
+        "download_url": "https://github.com/dscripka/openWakeWord/releases/download/v0.5.1/silero_vad.onnx"
+    }
+}
+
+MODELS = {
     "alexa": {
-        "model_path": os.path.join(os.path.dirname(os.path.abspath(__file__)), "resources/models/alexa_v0.1.onnx")
+        "model_path": os.path.join(os.path.dirname(os.path.abspath(__file__)), "resources/models/alexa_v0.1.tflite"),
+        "download_url": "https://github.com/dscripka/openWakeWord/releases/download/v0.5.1/alexa_v0.1.tflite"
     },
     "hey_mycroft": {
-        "model_path": os.path.join(os.path.dirname(os.path.abspath(__file__)), "resources/models/hey_mycroft_v0.1.onnx")
+        "model_path": os.path.join(os.path.dirname(os.path.abspath(__file__)), "resources/models/hey_mycroft_v0.1.tflite"),
+        "download_url": "https://github.com/dscripka/openWakeWord/releases/download/v0.5.1/hey_mycroft_v0.1.tflite"
     },
     "hey_jarvis": {
-        "model_path": os.path.join(os.path.dirname(os.path.abspath(__file__)), "resources/models/hey_jarvis_v0.1.onnx")
+        "model_path": os.path.join(os.path.dirname(os.path.abspath(__file__)), "resources/models/hey_jarvis_v0.1.tflite"),
+        "download_url": "https://github.com/dscripka/openWakeWord/releases/download/v0.5.1/hey_jarvis_v0.1.tflite"
+    },
+    "hey_rhasspy": {
+        "model_path": os.path.join(os.path.dirname(os.path.abspath(__file__)), "resources/models/hey_rhasspy_v0.1.tflite"),
+        "download_url": "https://github.com/dscripka/openWakeWord/releases/download/v0.5.1/hey_rhasspy_v0.1.tflite"
     },
     "timer": {
-        "model_path": os.path.join(os.path.dirname(os.path.abspath(__file__)), "resources/models/timer_v0.1.onnx")
+        "model_path": os.path.join(os.path.dirname(os.path.abspath(__file__)), "resources/models/timer_v0.1.tflite"),
+        "download_url": "https://github.com/dscripka/openWakeWord/releases/download/v0.5.1/timer_v0.1.tflite"
     },
     "weather": {
-        "model_path": os.path.join(os.path.dirname(os.path.abspath(__file__)), "resources/models/weather_v0.1.onnx")
+        "model_path": os.path.join(os.path.dirname(os.path.abspath(__file__)), "resources/models/weather_v0.1.tflite"),
+        "download_url": "https://github.com/dscripka/openWakeWord/releases/download/v0.5.1/weather_v0.1.tflite"
     }
 }
 
@@ -35,5 +62,8 @@ model_class_mappings = {
 }
 
 
-def get_pretrained_model_paths():
-    return [models[i]["model_path"] for i in models.keys()]
+def get_pretrained_model_paths(inference_framework="tflite"):
+    if inference_framework == "tflite":
+        return [MODELS[i]["model_path"] for i in MODELS.keys()]
+    elif inference_framework == "onnx":
+        return [MODELS[i]["model_path"].replace(".tflite", ".onnx") for i in MODELS.keys()]

openwakeword/custom_verifier_model.py

@@ -13,17 +13,19 @@
 # limitations under the License.
 
 # Imports
-import os
-from tqdm import tqdm
 import collections
-import openwakeword
+import os
+import pickle
+from typing import List, Union
+
 import numpy as np
 import scipy
-import pickle
-
 from sklearn.linear_model import LogisticRegression
 from sklearn.pipeline import make_pipeline
 from sklearn.preprocessing import FunctionTransformer, StandardScaler
+from tqdm import tqdm
+
+import openwakeword
 
 
 # Define functions to prepare data for speaker dependent verifier model
@@ -60,7 +62,7 @@ def get_reference_clip_features(
     # Get predictions
     for _ in range(N):
         # Load clip
-        if type(reference_clip) == str:
+        if isinstance(reference_clip, str):
             sr, dat = scipy.io.wavfile.read(reference_clip)
         else:
             dat = reference_clip
@@ -112,8 +114,8 @@ def train_verifier_model(features: np.ndarray, labels: np.ndarray):
 
 
 def train_custom_verifier(
-        positive_reference_clips: str,
-        negative_reference_clips: str,
+        positive_reference_clips: List[Union[str, os.PathLike]],
+        negative_reference_clips: List[Union[str, os.PathLike]],
         output_path: str,
         model_name: str,
         **kwargs
@@ -123,11 +125,11 @@ def train_custom_verifier(
     from a single user.
 
     Args:
-        positive_reference_clips (str): The path to a directory containing single-channel 16khz, 16-bit WAV files
+        positive_reference_clips (List[Union[str, os.PathLike]]): The path(s) to single-channel 16khz, 16-bit WAV files
                                         of the target wake word/phrase.
-        negative_reference_clips (str): The path to a directory containing single-channel 16khz, 16-bit WAV files
+        negative_reference_clips (List[Union[str, os.PathLike]]): The path(s) to single-channel 16khz, 16-bit WAV files
                                         of miscellaneous speech not containing the target wake word/phrase.
-        output_path (str): The location to save the trained verifier model (as a scikit-learn .joblib file)
+        output_path (str): The location to save the trained verifier model (as a Python pickle file (.pkl))
         model_name (str): The name or path of the trained openWakeWord model that the verifier model will be
                           based on. If only a name, it must be one of the pre-trained models included in the
                           openWakeWord release.
@@ -139,10 +141,10 @@ def train_custom_verifier(
     # Load target openWakeWord model
     if os.path.exists(model_name):
         oww = openwakeword.Model(
-            wakeword_model_paths=[model_name],
+            wakeword_models=[model_name],
             **kwargs
         )
-        model_name = model_name.split(os.path.sep)[-1][0:-5]
+        model_name = os.path.splitext(model_name)[0].split(os.path.sep)[-1]
     else:
         oww = openwakeword.Model(**kwargs)
 
@@ -152,9 +154,9 @@ def train_custom_verifier(
          for i in tqdm(positive_reference_clips, desc="Processing positive reference clips")]
     )
     if positive_features.shape[0] == 0:
-        raise ValueError("The positive features were created! Make sure that"
+        raise ValueError("The positive features were not created! Make sure that"
                          " the positive reference clips contain the appropriate audio"
-                         " for the desired model")
+                         " for the desired model.")
 
     # Get features from negative reference clips
     negative_features = np.vstack(
@@ -171,4 +173,5 @@ def train_custom_verifier(
 
     # Save logistic regression model to specified output location
     print("Done!")
-    pickle.dump(lr_model, open(output_path, "wb"))
+    with open(output_path, "wb") as f:
+        pickle.dump(lr_model, f)

openwakeword/data.py

@@ -15,13 +15,19 @@
 # imports
 from multiprocessing.pool import ThreadPool
 import os
+import re
+import logging
 from functools import partial
 from pathlib import Path
 import random
 from tqdm import tqdm
 from typing import List, Tuple
 import numpy as np
+import itertools
+import pronouncing
 import torch
+import audiomentations
+import torch_audiomentations
 from numpy.lib.format import open_memmap
 from speechbrain.dataio.dataio import read_audio
 from speechbrain.processing.signal_processing import reverberate
@@ -111,7 +117,7 @@ def load_audio_clips(files, clip_size=32000):
 # Convert clips with sox
 def _convert_clip(input_file, output_file, backend="ffmpeg"):
     if backend == "sox":
-        cmd = f"sox \"{input_file}\" -G -r 16000 -c 1 \"{output_file}\""
+        cmd = f"sox \"{input_file}\" -G -r 16000 -c 1 -b 16 \"{output_file}\""
     elif backend == "ffmpeg":
         cmd = f"ffmpeg -y -i \"{input_file}\" -ar 16000 \"{output_file}\""
     os.system(cmd)
@@ -445,7 +451,7 @@ def mix_clips_batch(
         # Apply volume augmentation
         if volume_augmentation:
             volume_levels = np.random.uniform(0.02, 1.0, mixed_clips_batch.shape[0])
-            mixed_clips_batch = (volume_levels/mixed_clips_batch.max(axis=1)[0])[..., None]*mixed_clips_batch
+            mixed_clips_batch = (volume_levels/mixed_clips_batch.max(dim=1)[0])[..., None]*mixed_clips_batch
         else:
             # Normalize clips only if max value is outside of [-1, 1]
             abs_max, _ = torch.max(
@@ -457,7 +463,7 @@ def mix_clips_batch(
         mixed_clips_batch = (mixed_clips_batch.numpy()*32767).astype(np.int16)
 
         # Remove any clips that are silent (happens rarely when mixing/reverberating)
-        error_index = np.where(mixed_clips_batch.max(axis=1) != 0)[0]
+        error_index = torch.from_numpy(np.where(mixed_clips_batch.max(dim=1) != 0)[0])
         mixed_clips_batch = mixed_clips_batch[error_index]
         labels_batch = labels_batch[error_index]
         sequence_labels_batch = sequence_labels_batch[error_index]
@@ -548,6 +554,181 @@ def apply_reverb(x, rir_files):
     return reverbed.numpy()
 
 
+# Alternate data augmentation method using audiomentations library (https://pypi.org/project/audiomentations/)
+def augment_clips(
+        clip_paths: List[str],
+        total_length: int,
+        sr: int = 16000,
+        batch_size: int = 128,
+        augmentation_probabilities: dict = {
+            "SevenBandParametricEQ": 0.25,
+            "TanhDistortion": 0.25,
+            "PitchShift": 0.25,
+            "BandStopFilter": 0.25,
+            "AddColoredNoise": 0.25,
+            "AddBackgroundNoise": 0.75,
+            "Gain": 1.0,
+            "RIR": 0.5
+        },
+        background_clip_paths: List[str] = [],
+        RIR_paths: List[str] = []
+        ):
+    """
+    Applies audio augmentations to the specified audio clips, returning a generator that applies
+    the augmentations in batches to support very large quantities of input audio files.
+
+    The augmentations (and probabilities) are chosen from experience based on training openWakeWord models, as well
+    as for the efficiency of the augmentation. The individual probabilities of each augmentation may be adjusted
+    with the "augmentation_probabilities" argument.
+
+    Args:
+        clip_paths (List[str]) = The input audio files (as paths) to augment. Note that these should be shorter
+                                 than the "total_length" argument, else they will be truncated.
+        total_length (int): The total length of audio files (in samples) after augmentation. All input clips
+                            will be left-padded with silence to reach this size, with between 0 and 200 ms
+                            of other audio after the end of the original input clip.
+        sr (int): The sample size of the input audio files
+        batch_size (int): The number of audio files to augment at once.
+        augmentation_probabilities (dict): The individual probabilities of each augmentation. If all probabilities
+                                           are zero, the input audio files will simply be padded with silence. THe
+                                           default values are:
+
+                                            {
+                                                "SevenBandParametricEQ": 0.25,
+                                                "TanhDistortion": 0.25,
+                                                "PitchShift": 0.25,
+                                                "BandStopFilter": 0.25,
+                                                "AddColoredNoise": 0.25,
+                                                "AddBackgroundNoise": 0.75,
+                                                "Gain": 1.0,
+                                                "RIR": 0.5
+                                            }
+
+        background_clip_paths (List[str]) = The paths to background audio files to mix with the input files
+        RIR_paths (List[str]) = The paths to room impulse response functions (RIRs) to convolve with the input files,
+                                producing a version of the input clip with different acoustic characteristics.
+
+    Returns:
+        ndarray: A batch of augmented audio clips of size (batch_size, total_length)
+    """
+    # Define augmentations
+
+    # First pass augmentations that can't be done as a batch
+    augment1 = audiomentations.Compose([
+        audiomentations.SevenBandParametricEQ(min_gain_db=-6, max_gain_db=6, p=augmentation_probabilities["SevenBandParametricEQ"]),
+        audiomentations.TanhDistortion(
+            min_distortion=0.0001,
+            max_distortion=0.10,
+            p=augmentation_probabilities["TanhDistortion"]
+        ),
+    ])
+
+    # Augmentations that can be done as a batch
+    if background_clip_paths != []:
+        augment2 = torch_audiomentations.Compose([
+            torch_audiomentations.PitchShift(
+                min_transpose_semitones=-3,
+                max_transpose_semitones=3,
+                p=augmentation_probabilities["PitchShift"],
+                sample_rate=16000,
+                mode="per_batch"
+            ),
+            torch_audiomentations.BandStopFilter(p=augmentation_probabilities["BandStopFilter"], mode="per_batch"),
+            torch_audiomentations.AddColoredNoise(
+                min_snr_in_db=10, max_snr_in_db=30,
+                min_f_decay=-1, max_f_decay=2, p=augmentation_probabilities["AddColoredNoise"],
+                mode="per_batch"
+            ),
+            torch_audiomentations.AddBackgroundNoise(
+                p=augmentation_probabilities["AddBackgroundNoise"],
+                background_paths=background_clip_paths,
+                min_snr_in_db=-10,
+                max_snr_in_db=15,
+                mode="per_batch"
+            ),
+            torch_audiomentations.Gain(max_gain_in_db=0, p=augmentation_probabilities["Gain"]),
+        ])
+    else:
+        augment2 = torch_audiomentations.Compose([
+            torch_audiomentations.PitchShift(
+                min_transpose_semitones=-3,
+                max_transpose_semitones=3,
+                p=augmentation_probabilities["PitchShift"],
+                sample_rate=16000,
+                mode="per_batch"
+            ),
+            torch_audiomentations.BandStopFilter(p=augmentation_probabilities["BandStopFilter"], mode="per_batch"),
+            torch_audiomentations.AddColoredNoise(
+                min_snr_in_db=10, max_snr_in_db=30,
+                min_f_decay=-1, max_f_decay=2, p=augmentation_probabilities["AddColoredNoise"],
+                mode="per_batch"
+            ),
+            torch_audiomentations.Gain(max_gain_in_db=0, p=augmentation_probabilities["Gain"]),
+        ])
+
+    # Iterate through all clips and augment them
+    for i in range(0, len(clip_paths), batch_size):
+        batch = clip_paths[i:i+batch_size]
+        augmented_clips = []
+        for clip in batch:
+            clip_data, clip_sr = torchaudio.load(clip)
+            clip_data = clip_data[0]
+            if clip_data.shape[0] > total_length:
+                clip_data = clip_data[0:total_length]
+
+            if clip_sr != sr:
+                raise ValueError("Error! Clip does not have the correct sample rate!")
+
+            clip_data = create_fixed_size_clip(clip_data, total_length, clip_sr)
+
+            # Do first pass augmentations
+            augmented_clips.append(torch.from_numpy(augment1(samples=clip_data, sample_rate=sr)))
+
+        # Do second pass augmentations
+        device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+        augmented_batch = augment2(samples=torch.vstack(augmented_clips).unsqueeze(dim=1).to(device), sample_rate=sr).squeeze(axis=1)
+
+        # Do reverberation
+        if augmentation_probabilities["RIR"] >= np.random.random() and RIR_paths != []:
+            rir_waveform, sr = torchaudio.load(random.choice(RIR_paths))
+            augmented_batch = reverberate(augmented_batch.cpu(), rir_waveform, rescale_amp="avg")
+
+        # yield batch of 16-bit PCM audio data
+        yield (augmented_batch.cpu().numpy()*32767).astype(np.int16)
+
+
+def create_fixed_size_clip(x, n_samples, sr=16000, start=None, end_jitter=.200):
+    """
+    Create a fixed-length clip of the specified size by padding an input clip with zeros
+    Optionally specify the start/end position of the input clip, or let it be chosen randomly.
+
+    Args:
+        x (ndarray): The input audio to pad to a fixed size
+        n_samples (int): The total number of samples for the fixed length clip
+        sr (int): The sample rate of the audio
+        start (int): The start position of the clip in the fixed length output, in samples (default: None)
+        end_jitter (float): The time (in seconds) from the end of the fixed length output
+                            that the input clip should end, if `start` is None.
+
+    Returns:
+        ndarray: A new array of audio data of the specified length
+    """
+    dat = np.zeros(n_samples)
+    end_jitter = int(np.random.uniform(0, end_jitter)*sr)
+    if start is None:
+        start = max(0, n_samples - (int(len(x))+end_jitter))
+
+    if len(x) > n_samples:
+        if np.random.random() >= 0.5:
+            dat = x[0:n_samples].numpy()
+        else:
+            dat = x[-n_samples:].numpy()
+    else:
+        dat[start:start+len(x)] = x
+
+    return dat
+
+
 # Load batches of data from mmaped numpy arrays
 class mmap_batch_generator:
     """
@@ -622,8 +803,8 @@ class mmap_batch_generator:
             self.n_per_class = {}
             for lbl, shape in self.shapes.items():
                 dummy_data = np.random.random((10, self.shapes[lbl][1], self.shapes[lbl][2]))
-                if self.data_transform_funcs.get(lbl, None):
-                    scale_factor = self.data_transform_funcs.get(lbl, None)(dummy_data).shape[0]/10
+                if (transform_func := self.data_transform_funcs.get(lbl, None)):
+                    scale_factor = transform_func(dummy_data).shape[0]/10
 
                 ratio = self.shapes[lbl][0]/sum([i[0] for i in self.shapes.values()])
                 self.n_per_class[lbl] = max(1, int(int(batch_size*ratio)/scale_factor))
@@ -645,7 +826,6 @@ class mmap_batch_generator:
                 # Restart at zeroth index if an array reaches the end
                 if self.data_counter[label] >= self.shapes[label][0]:
                     self.data_counter[label] = 0
-                    # self.data[label] = np.load(self.data_files[label], mmap_mode='r')
 
                 # Get data from mmaped file
                 x = self.data[label][self.data_counter[label]:self.data_counter[label]+n]
@@ -697,7 +877,7 @@ def trim_mmap(mmap_path):
     mmap_file2 = open_memmap(output_file2, mode='w+', dtype=np.float32,
                              shape=(N_new, mmap_file1.shape[1], mmap_file1.shape[2]))
 
-    for i in tqdm(range(0, mmap_file1.shape[0], 1024), total=mmap_file1.shape[0]//1024):
+    for i in tqdm(range(0, mmap_file1.shape[0], 1024), total=mmap_file1.shape[0]//1024, desc="Trimming empty rows"):
         if i + 1024 > N_new:
             mmap_file2[i:N_new] = mmap_file1[i:N_new].copy()
             mmap_file2.flush()
@@ -710,3 +890,126 @@ def trim_mmap(mmap_path):
 
     # Rename new mmap file to match original
     os.rename(output_file2, mmap_path)
+
+
+# Generate words that sound similar ("adversarial") to the input phrase using phoneme overlap
+def generate_adversarial_texts(input_text: str, N: int, include_partial_phrase: float = 0, include_input_words: float = 0):
+    """
+    Generate adversarial words and phrases based on phoneme overlap.
+    Currently only works for english texts.
+    Note that homophones are excluded, as this wouldn't actually be an adversarial example for the input text.
+
+    Args:
+        input_text (str): The target text for adversarial phrases
+        N (int): The total number of adversarial texts to return. Uses sampling,
+                 so not all possible combinations will be included and some duplicates
+                 may be present.
+        include_partial_phrase (float): The probability of returning a number of words less than the input
+                                        text (but always between 1 and the number of input words)
+        include_input_words (float): The probability of including individual input words in the adversarial
+                                     texts when the input text consists of multiple words. For example,
+                                     if the `input_text` was "ok google", then setting this value > 0.0
+                                     will allow for adversarial texts like "ok noodle", versus the word "ok"
+                                     never being present in the adversarial texts.
+
+    Returns:
+        list: A list of strings corresponding to words and phrases that are phonetically similar (but not identical)
+              to the input text.
+    """
+    # Get phonemes for english vowels (CMUDICT labels)
+    vowel_phones = ["AA", "AE", "AH", "AO", "AW", "AX", "AXR", "AY", "EH", "ER", "EY", "IH", "IX", "IY", "OW", "OY", "UH", "UW", "UX"]
+
+    word_phones = []
+    input_text_phones = [pronouncing.phones_for_word(i) for i in input_text.split()]
+
+    # Download phonemizer model for OOV words, if needed
+    if [] in input_text_phones:
+        phonemizer_mdl_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "resources", "en_us_cmudict_forward.pt")
+        if not os.path.exists(os.path.join(os.path.dirname(os.path.abspath(__file__)), "resources")):
+            os.mkdir(os.path.join(os.path.dirname(os.path.abspath(__file__)), "resources"))
+        if not os.path.exists(phonemizer_mdl_path):
+            logging.warning("Downloading phonemizer model from DeepPhonemizer library...")
+            import requests
+            file_url = "https://public-asai-dl-models.s3.eu-central-1.amazonaws.com/DeepPhonemizer/en_us_cmudict_forward.pt"
+            r = requests.get(file_url, stream=True)
+            with open(phonemizer_mdl_path, "wb") as f:
+                for chunk in r.iter_content(chunk_size=2048):
+                    if chunk:
+                        f.write(chunk)
+
+        # Create phonemizer object
+        from dp.phonemizer import Phonemizer
+        phonemizer = Phonemizer.from_checkpoint(phonemizer_mdl_path)
+
+    for phones, word in zip(input_text_phones, input_text.split()):
+        if phones != []:
+            word_phones.extend(phones)
+        elif phones == []:
+            logging.warning(f"The word '{word}' was not found in the pronunciation dictionary! "
+                            "Using the DeepPhonemizer library to predict the phonemes.")
+            phones = phonemizer(word, lang='en_us')
+            logging.warning(f"Phones for '{word}': {phones}")
+            word_phones.append(re.sub(r"[\]|\[]", "", re.sub(r"\]\[", " ", phones)))
+        elif isinstance(phones[0], list):
+            logging.warning(f"There are multiple pronunciations for the word '{word}'.")
+            word_phones.append(phones[0])
+
+    # add all possible lexical stresses to vowels
+    word_phones = [re.sub('|'.join(vowel_phones), lambda x: str(x.group(0)) + '[0|1|2]', re.sub(r'\d+', '', i)) for i in word_phones]
+
+    adversarial_phrases = []
+    for phones, word in zip(word_phones, input_text.split()):
+        query_exps = []
+        phones = phones.split()
+        adversarial_words = []
+        if len(phones) <= 2:
+            query_exps.append(" ".join(phones))
+        else:
+            query_exps.extend(phoneme_replacement(phones, max_replace=max(0, len(phones)-2), replace_char="(.){1,3}"))
+
+        for query in query_exps:
+            matches = pronouncing.search(query)
+            matches_phones = [pronouncing.phones_for_word(i)[0] for i in matches]
+            allowed_matches = [i for i, j in zip(matches, matches_phones) if j != phones]
+            adversarial_words.extend([i for i in allowed_matches if word.lower() != i])
+
+        if adversarial_words != []:
+            adversarial_phrases.append(adversarial_words)
+
+    # Build combinations for final output
+    adversarial_texts = []
+    for i in range(N):
+        txts = []
+        for j, k in zip(adversarial_phrases, input_text.split()):
+            if np.random.random() > (1 - include_input_words):
+                txts.append(k)
+            else:
+                txts.append(np.random.choice(j))
+
+        if include_partial_phrase is not None and len(input_text.split()) > 1 and np.random.random() <= include_partial_phrase:
+            n_words = np.random.randint(1, len(input_text.split())+1)
+            adversarial_texts.append(" ".join(np.random.choice(txts, size=n_words, replace=False)))
+        else:
+            adversarial_texts.append(" ".join(txts))
+
+    # Remove any exact matches to input phrase
+    adversarial_texts = [i for i in adversarial_texts if i != input_text]
+
+    return adversarial_texts
+
+
+def phoneme_replacement(input_chars, max_replace, replace_char='"(.){1,3}"'):
+    results = []
+    chars = list(input_chars)
+
+    # iterate over the number of characters to replace (1 to max_replace)
+    for r in range(1, max_replace+1):
+        # get all combinations for a fixed r
+        comb = itertools.combinations(range(len(chars)), r)
+        for indices in comb:
+            chars_copy = chars.copy()
+            for i in indices:
+                chars_copy[i] = replace_char
+            results.append(' '.join(chars_copy))
+
+    return results

openwakeword/model.py

@@ -14,12 +14,13 @@
 
 # Imports
 import numpy as np
-import onnxruntime as ort
 import openwakeword
-from openwakeword.utils import AudioFeatures
+from openwakeword.utils import AudioFeatures, re_arg
 
 import wave
 import os
+import logging
+import functools
 import pickle
 from collections import deque, defaultdict
 from functools import partial
@@ -33,23 +34,26 @@ class Model():
     The main model class for openWakeWord. Creates a model object with the shared audio pre-processer
     and for arbitrarily many custom wake word/wake phrase models.
     """
+    @re_arg({"wakeword_model_paths": "wakeword_models"})  # temporary handling of keyword argument change
     def __init__(
             self,
-            wakeword_model_paths: List[str] = [],
+            wakeword_models: List[str] = [],
             class_mapping_dicts: List[dict] = [],
             enable_speex_noise_suppression: bool = False,
             vad_threshold: float = 0,
             custom_verifier_models: dict = {},
             custom_verifier_threshold: float = 0.1,
+            inference_framework: str = "tflite",
             **kwargs
             ):
         """Initialize the openWakeWord model object.
 
         Args:
-            wakeword_model_paths (List[str]): A list of paths of ONNX models to load into the openWakeWord model object.
-                                              If not provided, will load all of the pre-trained models.
+            wakeword_models (List[str]): A list of paths of ONNX/tflite models to load into the openWakeWord model object.
+                                              If not provided, will load all of the pre-trained models. Alternatively,
+                                              just the names of pre-trained models can be provided to select a subset of models.
             class_mapping_dicts (List[dict]): A list of dictionaries with integer to string class mappings for
-                                              each model in the `wakeword_model_paths` arguments
+                                              each model in the `wakeword_models` arguments
                                               (e.g., {"0": "class_1", "1": "class_2"})
             enable_speex_noise_suppression (bool): Whether to use the noise suppresion from the SpeexDSP
                                                    library to pre-process all incoming audio. May increase
@@ -63,49 +67,119 @@ class Model():
                                    with VAD scores above the threshold will be returned. The default value (0),
                                    disables voice activity detection entirely.
             custom_verifier_models (dict): A dictionary of paths to custom verifier models, where
-                                           the keys are the model names (corresponding to the openwakeword.models
+                                           the keys are the model names (corresponding to the openwakeword.MODELS
                                            attribute) and the values are the filepaths of the
                                            custom verifier models.
             custom_verifier_threshold (float): The score threshold to use a custom verifier model. If the score
                                                from a model for a given frame is greater than this value, the
                                                associated custom verifier model will also predict on that frame, and
                                                the verifier score will be returned.
+            inference_framework (str): The inference framework to use when for model prediction. Options are
+                                       "tflite" or "onnx". The default is "tflite" as this results in better
+                                       efficiency on common platforms (x86, ARM64), but in some deployment
+                                       scenarios ONNX models may be preferable.
             kwargs (dict): Any other keyword arguments to pass the the preprocessor instance
         """
-
-        # Initialize the ONNX models and store them
-        sessionOptions = ort.SessionOptions()
-        sessionOptions.inter_op_num_threads = 1
-        sessionOptions.intra_op_num_threads = 1
-
         # Get model paths for pre-trained models if user doesn't provide models to load
-        if wakeword_model_paths == []:
-            wakeword_model_paths = openwakeword.get_pretrained_model_paths()
-            wakeword_model_names = list(openwakeword.models.keys())
-        else:
-            wakeword_model_names = [os.path.basename(i[0:-5]) for i in wakeword_model_paths]
+        pretrained_model_paths = openwakeword.get_pretrained_model_paths(inference_framework)
+        wakeword_model_names = []
+        if wakeword_models == []:
+            wakeword_models = pretrained_model_paths
+            wakeword_model_names = list(openwakeword.MODELS.keys())
+        elif len(wakeword_models) >= 1:
+            for ndx, i in enumerate(wakeword_models):
+                if os.path.exists(i):
+                    wakeword_model_names.append(os.path.splitext(os.path.basename(i))[0])
+                else:
+                    # Find pre-trained path by modelname
+                    matching_model = [j for j in pretrained_model_paths if i.replace(" ", "_") in j.split(os.path.sep)[-1]]
+                    if matching_model == []:
+                        raise ValueError("Could not find pretrained model for model name '{}'".format(i))
+                    else:
+                        wakeword_models[ndx] = matching_model[0]
+                        wakeword_model_names.append(i)
 
         # Create attributes to store models and metadata
         self.models = {}
         self.model_inputs = {}
         self.model_outputs = {}
+        self.model_prediction_function = {}
         self.class_mapping = {}
-        self.model_input_names = {}
         self.custom_verifier_models = {}
         self.custom_verifier_threshold = custom_verifier_threshold
-        for mdl_path, mdl_name in zip(wakeword_model_paths, wakeword_model_names):
+
+        # Do imports for  inference framework
+        if inference_framework == "tflite":
+            try:
+                import ai_edge_litert.interpreter as tflite
+
+                def tflite_predict(tflite_interpreter, input_index, output_index, x):
+                    tflite_interpreter.set_tensor(input_index, x)
+                    tflite_interpreter.invoke()
+                    return tflite_interpreter.get_tensor(output_index)[None, ]
+
+            except ImportError:
+                logging.warning("Tried to import the tflite runtime, but it was not found. "
+                                "Trying to switching to onnxruntime instead, if appropriate models are available.")
+                if wakeword_models != [] and all(['.onnx' in i for i in wakeword_models]):
+                    inference_framework = "onnx"
+                elif wakeword_models != [] and all([os.path.exists(i.replace('.tflite', '.onnx')) for i in wakeword_models]):
+                    inference_framework = "onnx"
+                    wakeword_models = [i.replace('.tflite', '.onnx') for i in wakeword_models]
+                else:
+                    raise ValueError("Tried to import the LiteRT runtime for provided LiteRT models, but it was not found. "
+                                     "Please install it using `pip install ai-edge-litert`")
+
+        if inference_framework == "onnx":
+            try:
+                import onnxruntime as ort
+
+                def onnx_predict(onnx_model, x):
+                    return onnx_model.run(None, {onnx_model.get_inputs()[0].name: x})
+
+            except ImportError:
+                raise ValueError("Tried to import onnxruntime, but it was not found. Please install it using `pip install onnxruntime`")
+
+        for mdl_path, mdl_name in zip(wakeword_models, wakeword_model_names):
             # Load openwakeword models
-            self.models[mdl_name] = ort.InferenceSession(mdl_path, sess_options=sessionOptions,
-                                                         providers=["CPUExecutionProvider"])
-            self.model_inputs[mdl_name] = self.models[mdl_name].get_inputs()[0].shape[1]
-            self.model_outputs[mdl_name] = self.models[mdl_name].get_outputs()[0].shape[1]
-            if class_mapping_dicts and class_mapping_dicts[wakeword_model_paths.index(mdl_path)].get(mdl_name, None):
-                self.class_mapping[mdl_name] = class_mapping_dicts[wakeword_model_paths.index(mdl_path)]
+            if inference_framework == "onnx":
+                if ".tflite" in mdl_path:
+                    raise ValueError("The onnx inference framework is selected, but tflite models were provided!")
+
+                sessionOptions = ort.SessionOptions()
+                sessionOptions.inter_op_num_threads = 1
+                sessionOptions.intra_op_num_threads = 1
+
+                self.models[mdl_name] = ort.InferenceSession(mdl_path, sess_options=sessionOptions,
+                                                             providers=["CPUExecutionProvider"])
+
+                self.model_inputs[mdl_name] = self.models[mdl_name].get_inputs()[0].shape[1]
+                self.model_outputs[mdl_name] = self.models[mdl_name].get_outputs()[0].shape[1]
+                pred_function = functools.partial(onnx_predict, self.models[mdl_name])
+                self.model_prediction_function[mdl_name] = pred_function
+
+            if inference_framework == "tflite":
+                if ".onnx" in mdl_path:
+                    raise ValueError("The tflite inference framework is selected, but onnx models were provided!")
+
+                self.models[mdl_name] = tflite.Interpreter(model_path=mdl_path, num_threads=1)
+                self.models[mdl_name].allocate_tensors()
+
+                self.model_inputs[mdl_name] = self.models[mdl_name].get_input_details()[0]['shape'][1]
+                self.model_outputs[mdl_name] = self.models[mdl_name].get_output_details()[0]['shape'][1]
+
+                tflite_input_index = self.models[mdl_name].get_input_details()[0]['index']
+                tflite_output_index = self.models[mdl_name].get_output_details()[0]['index']
+
+                pred_function = functools.partial(tflite_predict, self.models[mdl_name], tflite_input_index, tflite_output_index)
+                self.model_prediction_function[mdl_name] = pred_function
+
+            if class_mapping_dicts and class_mapping_dicts[wakeword_models.index(mdl_path)].get(mdl_name, None):
+                self.class_mapping[mdl_name] = class_mapping_dicts[wakeword_models.index(mdl_path)]
             elif openwakeword.model_class_mappings.get(mdl_name, None):
                 self.class_mapping[mdl_name] = openwakeword.model_class_mappings[mdl_name]
             else:
                 self.class_mapping[mdl_name] = {str(i): str(i) for i in range(0, self.model_outputs[mdl_name])}
-            self.model_input_names[mdl_name] = self.models[mdl_name].get_inputs()[0].name
 
             # Load custom verifier models
             if isinstance(custom_verifier_models, dict):
@@ -136,7 +210,7 @@ class Model():
             self.vad = openwakeword.VAD()
 
         # Create AudioFeatures object
-        self.preprocessor = AudioFeatures(**kwargs)
+        self.preprocessor = AudioFeatures(inference_framework=inference_framework, **kwargs)
 
     def get_parent_model_from_label(self, label):
         """Gets the parent model associated with a given prediction label"""
@@ -150,25 +224,32 @@ class Model():
         return parent_model
 
     def reset(self):
-        """Reset the prediction buffer"""
+        """Reset the prediction and audio feature buffers. Useful for re-initializing the model, though may not be efficient
+        when called too frequently."""
         self.prediction_buffer = defaultdict(partial(deque, maxlen=30))
+        self.preprocessor.reset()
 
-    def predict(self, x: np.ndarray, patience: dict = {}, threshold: dict = {}, timing: bool = False):
+    def predict(self, x: np.ndarray, patience: dict = {},
+                threshold: dict = {}, debounce_time: float = 0.0, timing: bool = False):
         """Predict with all of the wakeword models on the input audio frames
 
         Args:
-            x (Union[ndarray]): The input audio data to predict on with the models. Should be multiples of 80 ms
+            x (ndarray): The input audio data to predict on with the models. Ideally should be multiples of 80 ms
                                 (1280 samples), with longer lengths reducing overall CPU usage
-                                but decreasing detection latency.
+                                but decreasing detection latency. Input audio with durations greater than or less
+                                than 80 ms is also supported, though this will add a detection delay of up to 80 ms
+                                as the appropriate number of samples are accumulated.
             patience (dict): How many consecutive frames (of 1280 samples or 80 ms) above the threshold that must
                              be observed before the current frame will be returned as non-zero.
                              Must be provided as an a dictionary where the keys are the
                              model names and the values are the number of frames. Can reduce false-positive
                              detections at the cost of a lower true-positive rate.
                              By default, this behavior is disabled.
-            threshold (dict): The threshold values to use when the `patience` behavior is enabled.
+            threshold (dict): The threshold values to use when the `patience` or `debounce_time` behavior is enabled.
                               Must be provided as an a dictionary where the keys are the
                               model names and the values are the thresholds.
+            debounce_time (float): The time (in seconds) to wait before returning another non-zero prediction
+                                   after a non-zero prediction. Can preven multiple detections of the same wake-word.
             timing (bool): Whether to return timing information of the models. Can be useful to debug and
                            assess how efficiently models are running on the current hardware.
 
@@ -177,6 +258,9 @@ class Model():
                   wake-word/wake-phrase detected. If the `timing` argument is true, returns a
                   tuple of dicts containing model predictions and timing information, respectively.
         """
+        # Check input data type
+        if not isinstance(x, np.ndarray):
+            raise ValueError(f"The input audio data (x) must by a Numpy array, instead received an object of type {type(x)}.")
 
         # Setup timing dict
         if timing:
@@ -186,9 +270,9 @@ class Model():
 
         # Get audio features (optionally with Speex noise suppression)
         if self.speex_ns:
-            self.preprocessor(self._suppress_noise_with_speex(x))
+            n_prepared_samples = self.preprocessor(self._suppress_noise_with_speex(x))
         else:
-            self.preprocessor(x)
+            n_prepared_samples = self.preprocessor(x)
 
         if timing:
             timing_dict["models"]["preprocessor"] = time.time() - feature_start
@@ -196,30 +280,35 @@ class Model():
         # Get predictions from model(s)
         predictions = {}
         for mdl in self.models.keys():
-            input_name = self.model_input_names[mdl]
-
             if timing:
                 model_start = time.time()
 
             # Run model to get predictions
-            if len(x) > 1280:
+            if n_prepared_samples > 1280:
                 group_predictions = []
-                for i in np.arange(len(x)//1280-1, -1, -1):
+                for i in np.arange(n_prepared_samples//1280-1, -1, -1):
                     group_predictions.extend(
-                        self.models[mdl].run(
-                            None,
-                            {input_name: self.preprocessor.get_features(
+                        self.model_prediction_function[mdl](
+                            self.preprocessor.get_features(
                                     self.model_inputs[mdl],
                                     start_ndx=-self.model_inputs[mdl] - i
-                            )}
+                            )
                         )
                     )
                 prediction = np.array(group_predictions).max(axis=0)[None, ]
-            else:
-                prediction = self.models[mdl].run(
-                                        None,
-                                        {input_name: self.preprocessor.get_features(self.model_inputs[mdl])}
-                                    )
+            elif n_prepared_samples == 1280:
+                prediction = self.model_prediction_function[mdl](
+                    self.preprocessor.get_features(self.model_inputs[mdl])
+                )
+            elif n_prepared_samples < 1280:  # get previous prediction if there aren't enough samples
+                if self.model_outputs[mdl] == 1:
+                    if len(self.prediction_buffer[mdl]) > 0:
+                        prediction = [[[self.prediction_buffer[mdl][-1]]]]
+                    else:
+                        prediction = [[[0]]]
+                elif self.model_outputs[mdl] != 1:
+                    n_classes = max([int(i) for i in self.class_mapping[mdl].keys()])
+                    prediction = [[[0]*(n_classes+1)]]
 
             if self.model_outputs[mdl] == 1:
                 predictions[mdl] = prediction[0][0][0]
@@ -238,27 +327,40 @@ class Model():
                             )[0][-1]
                             predictions[cls] = verifier_prediction
 
-            # Update prediction buffer, and zero predictions for first 5 frames during model initialization
+            # Zero predictions for first 5 frames during model initialization
             for cls in predictions.keys():
                 if len(self.prediction_buffer[cls]) < 5:
                     predictions[cls] = 0.0
-                self.prediction_buffer[cls].append(predictions[cls])
 
             # Get timing information
             if timing:
                 timing_dict["models"][mdl] = time.time() - model_start
 
         # Update scores based on thresholds or patience arguments
-        if patience != {}:
+        if patience != {} or debounce_time > 0:
             if threshold == {}:
                 raise ValueError("Error! When using the `patience` argument, threshold "
                                  "values must be provided via the `threshold` argument!")
+            if patience != {} and debounce_time > 0:
+                raise ValueError("Error! The `patience` and `debounce_time` arguments cannot be used together!")
             for mdl in predictions.keys():
                 parent_model = self.get_parent_model_from_label(mdl)
-                if parent_model in patience.keys():
-                    scores = np.array(self.prediction_buffer[mdl])[-patience[parent_model]:]
-                    if (scores >= threshold[parent_model]).sum() < patience[parent_model]:
-                        predictions[mdl] = 0.0
+                if predictions[mdl] != 0.0:
+                    if parent_model in patience.keys():
+                        scores = np.array(self.prediction_buffer[mdl])[-patience[parent_model]:]
+                        if (scores >= threshold[parent_model]).sum() < patience[parent_model]:
+                            predictions[mdl] = 0.0
+                    elif debounce_time > 0:
+                        if parent_model in threshold.keys():
+                            n_frames = int(np.ceil(debounce_time/(n_prepared_samples/16000)))
+                            recent_predictions = np.array(self.prediction_buffer[mdl])[-n_frames:]
+                            if predictions[mdl] >= threshold[parent_model] and \
+                               (recent_predictions >= threshold[parent_model]).sum() > 0:
+                                predictions[mdl] = 0.0
+
+        # Update prediction buffer
+        for mdl in predictions.keys():
+            self.prediction_buffer[mdl].append(predictions[mdl])
 
         # (optionally) get voice activity detection scores and update model scores
         if self.vad_threshold > 0:

openwakeword/resources/models/alexa_v0.1.onnx

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:6ff566a01d12670e8d9e3c59da32651db1575d17272a601b7f8a39283dfbae3e
-size 854246

openwakeword/resources/models/embedding_model.onnx

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:ba754db3cd768a524c655ea90655ee5e6055a43b8dfd29366a11e93716ae9e51
-size 1328103

openwakeword/resources/models/hey_jarvis_v0.1.onnx

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:94a13cfe60075b132f6a472e7e462e8123ee70861bc3fb58434a73712ee0d2cb
-size 1271370

openwakeword/resources/models/hey_marvin_v0.1.onnx

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:b6d4b794ddf2e1d6f29e9f45848e24858e2edd0d810b14e0c1c70dda9a1fcbf0
-size 857691

openwakeword/resources/models/hey_mycroft_v0.1.onnx

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:785bdf5655863ae47553b23793aa108c7b0152d4823f7869b41f2d2d765912fc
-size 503850

openwakeword/resources/models/melspectrogram.onnx

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:ba2b0e0f8b7b875369a2c89cb13360ff53bac436f2895cced9f479fa65eb176f
-size 1087958

openwakeword/resources/models/silero_vad.onnx

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:a35ebf52fd3ce5f1469b2a36158dba761bc47b973ea3382b3186ca15b1f5af28
-size 1807522

openwakeword/resources/models/timer_v0.1.onnx

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:371e44535470a29248b3b8f1bbbbaf2525c86417fd8f75c67fcf02ae0b9626df
-size 1742475

openwakeword/resources/models/weather_v0.1.onnx

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:8441da8e746899e8d969528d5bad5651cdd563079c05962788f77753041f60e7
-size 1149158

openwakeword/train.py

@@ -0,0 +1,910 @@
+import torch
+from torch import optim, nn
+import torchinfo
+import torchmetrics
+import copy
+import os
+import sys
+import tempfile
+import uuid
+import numpy as np
+import scipy
+import collections
+import argparse
+import logging
+from tqdm import tqdm
+import yaml
+from pathlib import Path
+import openwakeword
+from openwakeword.data import generate_adversarial_texts, augment_clips, mmap_batch_generator
+from openwakeword.utils import compute_features_from_generator
+from openwakeword.utils import AudioFeatures
+
+
+# Base model class for an openwakeword model
+class Model(nn.Module):
+    def __init__(self, n_classes=1, input_shape=(16, 96), model_type="dnn",
+                 layer_dim=128, n_blocks=1, seconds_per_example=None):
+        super().__init__()
+
+        # Store inputs as attributes
+        self.n_classes = n_classes
+        self.input_shape = input_shape
+        self.seconds_per_example = seconds_per_example
+        self.device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+        self.best_models = []
+        self.best_model_scores = []
+        self.best_val_fp = 1000
+        self.best_val_accuracy = 0
+        self.best_val_recall = 0
+        self.best_train_recall = 0
+
+        # Define model (currently on fully-connected network supported)
+        if model_type == "dnn":
+            # self.model = nn.Sequential(
+            #     nn.Flatten(),
+            #     nn.Linear(input_shape[0]*input_shape[1], layer_dim),
+            #     nn.LayerNorm(layer_dim),
+            #     nn.ReLU(),
+            #     nn.Linear(layer_dim, layer_dim),
+            #     nn.LayerNorm(layer_dim),
+            #     nn.ReLU(),
+            #     nn.Linear(layer_dim, n_classes),
+            #     nn.Sigmoid() if n_classes == 1 else nn.ReLU(),
+            # )
+
+            class FCNBlock(nn.Module):
+                def __init__(self, layer_dim):
+                    super().__init__()
+                    self.fcn_layer = nn.Linear(layer_dim, layer_dim)
+                    self.relu = nn.ReLU()
+                    self.layer_norm = nn.LayerNorm(layer_dim)
+
+                def forward(self, x):
+                    return self.relu(self.layer_norm(self.fcn_layer(x)))
+
+            class Net(nn.Module):
+                def __init__(self, input_shape, layer_dim, n_blocks=1, n_classes=1):
+                    super().__init__()
+                    self.flatten = nn.Flatten()
+                    self.layer1 = nn.Linear(input_shape[0]*input_shape[1], layer_dim)
+                    self.relu1 = nn.ReLU()
+                    self.layernorm1 = nn.LayerNorm(layer_dim)
+                    self.blocks = nn.ModuleList([FCNBlock(layer_dim) for i in range(n_blocks)])
+                    self.last_layer = nn.Linear(layer_dim, n_classes)
+                    self.last_act = nn.Sigmoid() if n_classes == 1 else nn.ReLU()
+
+                def forward(self, x):
+                    x = self.relu1(self.layernorm1(self.layer1(self.flatten(x))))
+                    for block in self.blocks:
+                        x = block(x)
+                    x = self.last_act(self.last_layer(x))
+                    return x
+            self.model = Net(input_shape, layer_dim, n_blocks=n_blocks, n_classes=n_classes)
+        elif model_type == "rnn":
+            class Net(nn.Module):
+                def __init__(self, input_shape, n_classes=1):
+                    super().__init__()
+                    self.layer1 = nn.LSTM(input_shape[-1], 64, num_layers=2, bidirectional=True,
+                                          batch_first=True, dropout=0.0)
+                    self.layer2 = nn.Linear(64*2, n_classes)
+                    self.layer3 = nn.Sigmoid() if n_classes == 1 else nn.ReLU()
+
+                def forward(self, x):
+                    out, h = self.layer1(x)
+                    return self.layer3(self.layer2(out[:, -1]))
+            self.model = Net(input_shape, n_classes)
+
+        # Define metrics
+        if n_classes == 1:
+            self.fp = lambda pred, y: (y-pred <= -0.5).sum()
+            self.recall = torchmetrics.Recall(task='binary')
+            self.accuracy = torchmetrics.Accuracy(task='binary')
+        else:
+            def multiclass_fp(p, y, threshold=0.5):
+                probs = torch.nn.functional.softmax(p, dim=1)
+                neg_ndcs = y == 0
+                fp = (probs[neg_ndcs].argmax(axis=1) != 0 & (probs[neg_ndcs].max(axis=1)[0] > threshold)).sum()
+                return fp
+
+            def positive_class_recall(p, y, negative_class_label=0, threshold=0.5):
+                probs = torch.nn.functional.softmax(p, dim=1)
+                pos_ndcs = y != 0
+                rcll = (probs[pos_ndcs].argmax(axis=1) > 0
+                        & (probs[pos_ndcs].max(axis=1)[0] >= threshold)).sum()/pos_ndcs.sum()
+                return rcll
+
+            def positive_class_accuracy(p, y, negative_class_label=0):
+                probs = torch.nn.functional.softmax(p, dim=1)
+                pos_preds = probs.argmax(axis=1) != negative_class_label
+                acc = (probs[pos_preds].argmax(axis=1) == y[pos_preds]).sum()/pos_preds.sum()
+                return acc
+
+            self.fp = multiclass_fp
+            self.acc = positive_class_accuracy
+            self.recall = positive_class_recall
+
+        self.n_fp = 0
+        self.val_fp = 0
+
+        # Define logging dict (in-memory)
+        self.history = collections.defaultdict(list)
+
+        # Define optimizer and loss
+        self.loss = torch.nn.functional.binary_cross_entropy if n_classes == 1 else nn.functional.cross_entropy
+        self.optimizer = optim.Adam(self.model.parameters(), lr=0.0001)
+
+    def save_model(self, output_path):
+        """
+        Saves the weights of a trained Pytorch model
+        """
+        if self.n_classes == 1:
+            torch.save(self.model, output_path)
+
+    def export_to_onnx(self, output_path, class_mapping=""):
+        obj = self
+        # Make simple model for export based on model structure
+        if self.n_classes == 1:
+            # Save ONNX model
+            torch.onnx.export(self.model.to("cpu"), torch.rand(self.input_shape)[None, ], output_path,
+                              output_names=[class_mapping])
+
+        elif self.n_classes >= 1:
+            class M(nn.Module):
+                def __init__(self):
+                    super().__init__()
+
+                    # Define model
+                    self.model = obj.model.to("cpu")
+
+                def forward(self, x):
+                    return torch.nn.functional.softmax(self.model(x), dim=1)
+
+            # Save ONNX model
+            torch.onnx.export(M(), torch.rand(self.input_shape)[None, ], output_path,
+                              output_names=[class_mapping])
+
+    def lr_warmup_cosine_decay(self,
+                               global_step,
+                               warmup_steps=0,
+                               hold=0,
+                               total_steps=0,
+                               start_lr=0.0,
+                               target_lr=1e-3
+                               ):
+        # Cosine decay
+        learning_rate = 0.5 * target_lr * (1 + np.cos(np.pi * (global_step - warmup_steps - hold)
+                                           / float(total_steps - warmup_steps - hold)))
+
+        # Target LR * progress of warmup (=1 at the final warmup step)
+        warmup_lr = target_lr * (global_step / warmup_steps)
+
+        # Choose between `warmup_lr`, `target_lr` and `learning_rate` based on whether
+        # `global_step < warmup_steps` and we're still holding.
+        # i.e. warm up if we're still warming up and use cosine decayed lr otherwise
+        if hold > 0:
+            learning_rate = np.where(global_step > warmup_steps + hold,
+                                     learning_rate, target_lr)
+
+        learning_rate = np.where(global_step < warmup_steps, warmup_lr, learning_rate)
+        return learning_rate
+
+    def forward(self, x):
+        return self.model(x)
+
+    def summary(self):
+        return torchinfo.summary(self.model, input_size=(1,) + self.input_shape, device='cpu')
+
+    def average_models(self, models=None):
+        """Averages the weights of the provided models together to make a new model"""
+
+        if models is None:
+            models = self.best_models
+
+        # Clone a model from the list as the base for the averaged model
+        averaged_model = copy.deepcopy(models[0])
+        averaged_model_dict = averaged_model.state_dict()
+
+        # Initialize a running total of the weights
+        for key in averaged_model_dict:
+            averaged_model_dict[key] *= 0  # set to 0
+
+        for model in models:
+            model_dict = model.state_dict()
+            for key, value in model_dict.items():
+                averaged_model_dict[key] += value
+
+        for key in averaged_model_dict:
+            averaged_model_dict[key] /= len(models)
+
+        # Load the averaged weights into the model
+        averaged_model.load_state_dict(averaged_model_dict)
+
+        return averaged_model
+
+    def _select_best_model(self, false_positive_validate_data, val_set_hrs=11.3, max_fp_per_hour=0.5, min_recall=0.20):
+        """
+        Select the top model based on the false positive rate on the validation data
+
+        Args:
+            false_positive_validate_data (torch.DataLoader): A dataloader with validation data
+            n (int): The number of models to select
+
+        Returns:
+            list: A list of the top n models
+        """
+        # Get false positive rates for each model
+        false_positive_rates = [0]*len(self.best_models)
+        for batch in false_positive_validate_data:
+            x_val, y_val = batch[0].to(self.device), batch[1].to(self.device)
+            for mdl_ndx, model in tqdm(enumerate(self.best_models), total=len(self.best_models),
+                                       desc="Find best checkpoints by false positive rate"):
+                with torch.no_grad():
+                    val_ps = model(x_val)
+                    false_positive_rates[mdl_ndx] = false_positive_rates[mdl_ndx] + self.fp(val_ps, y_val[..., None]).detach().cpu().numpy()
+        false_positive_rates = [fp/val_set_hrs for fp in false_positive_rates]
+
+        candidate_model_ndx = [ndx for ndx, fp in enumerate(false_positive_rates) if fp <= max_fp_per_hour]
+        candidate_model_recall = [self.best_model_scores[ndx]["val_recall"] for ndx in candidate_model_ndx]
+        if max(candidate_model_recall) <= min_recall:
+            logging.warning(f"No models with recall >= {min_recall} found!")
+            return None
+        else:
+            best_model = self.best_models[candidate_model_ndx[np.argmax(candidate_model_recall)]]
+            best_model_training_step = self.best_model_scores[candidate_model_ndx[np.argmax(candidate_model_recall)]]["training_step_ndx"]
+            logging.info(f"Best model from training step {best_model_training_step} out of {len(candidate_model_ndx)}"
+                         f"models has recall of {np.max(candidate_model_recall)} and false positive rate of"
+                         f" {false_positive_rates[candidate_model_ndx[np.argmax(candidate_model_recall)]]}")
+
+        return best_model
+
+    def auto_train(self, X_train, X_val, false_positive_val_data, steps=50000, max_negative_weight=1000,
+                   target_fp_per_hour=0.2):
+        """A sequence of training steps that produce relatively strong models
+        automatically, based on validation data and performance targets provided.
+        After training merges the best checkpoints and returns a single model.
+        """
+
+        # Get false positive validation data duration
+        val_set_hrs = 11.3
+
+        # Sequence 1
+        logging.info("#"*50 + "\nStarting training sequence 1...\n" + "#"*50)
+        lr = 0.0001
+        weights = np.linspace(1, max_negative_weight, int(steps)).tolist()
+        val_steps = np.linspace(steps-int(steps*0.25), steps, 20).astype(np.int64)
+        self.train_model(
+                    X=X_train,
+                    X_val=X_val,
+                    false_positive_val_data=false_positive_val_data,
+                    max_steps=steps,
+                    negative_weight_schedule=weights,
+                    val_steps=val_steps, warmup_steps=steps//5,
+                    hold_steps=steps//3, lr=lr, val_set_hrs=val_set_hrs)
+
+        # Sequence 2
+        logging.info("#"*50 + "\nStarting training sequence 2...\n" + "#"*50)
+        lr = lr/10
+        steps = steps/10
+
+        # Adjust weights as needed based on false positive per hour performance from first sequence
+        if self.best_val_fp > target_fp_per_hour:
+            max_negative_weight = max_negative_weight*2
+            logging.info("Increasing weight on negative examples to reduce false positives...")
+
+        weights = np.linspace(1, max_negative_weight, int(steps)).tolist()
+        val_steps = np.linspace(1, steps, 20).astype(np.int16)
+        self.train_model(
+                    X=X_train,
+                    X_val=X_val,
+                    false_positive_val_data=false_positive_val_data,
+                    max_steps=steps,
+                    negative_weight_schedule=weights,
+                    val_steps=val_steps, warmup_steps=steps//5,
+                    hold_steps=steps//3, lr=lr, val_set_hrs=val_set_hrs)
+
+        # Sequence 3
+        logging.info("#"*50 + "\nStarting training sequence 3...\n" + "#"*50)
+        lr = lr/10
+
+        # Adjust weights as needed based on false positive per hour performance from second sequence
+        if self.best_val_fp > target_fp_per_hour:
+            max_negative_weight = max_negative_weight*2
+            logging.info("Increasing weight on negative examples to reduce false positives...")
+
+        weights = np.linspace(1, max_negative_weight, int(steps)).tolist()
+        val_steps = np.linspace(1, steps, 20).astype(np.int16)
+        self.train_model(
+                    X=X_train,
+                    X_val=X_val,
+                    false_positive_val_data=false_positive_val_data,
+                    max_steps=steps,
+                    negative_weight_schedule=weights,
+                    val_steps=val_steps, warmup_steps=steps//5,
+                    hold_steps=steps//3, lr=lr, val_set_hrs=val_set_hrs)
+
+        # Merge best models
+        logging.info("Merging checkpoints above the 90th percentile into single model...")
+        accuracy_percentile = np.percentile(self.history["val_accuracy"], 90)
+        recall_percentile = np.percentile(self.history["val_recall"], 90)
+        fp_percentile = np.percentile(self.history["val_fp_per_hr"], 10)
+
+        # Get models above the 90th percentile
+        models = []
+        for model, score in zip(self.best_models, self.best_model_scores):
+            if score["val_accuracy"] >= accuracy_percentile and \
+                    score["val_recall"] >= recall_percentile and \
+                    score["val_fp_per_hr"] <= fp_percentile:
+                models.append(model)
+
+        if len(models) > 0:
+            combined_model = self.average_models(models=models)
+        else:
+            combined_model = self.model
+
+        # Report validation metrics for combined model
+        with torch.no_grad():
+            for batch in X_val:
+                x, y = batch[0].to(self.device), batch[1].to(self.device)
+                val_ps = combined_model(x)
+
+            combined_model_recall = self.recall(val_ps, y[..., None]).detach().cpu().numpy()
+            combined_model_accuracy = self.accuracy(val_ps, y[..., None].to(torch.int64)).detach().cpu().numpy()
+
+            combined_model_fp = 0
+            for batch in false_positive_val_data:
+                x_val, y_val = batch[0].to(self.device), batch[1].to(self.device)
+                val_ps = combined_model(x_val)
+                combined_model_fp += self.fp(val_ps, y_val[..., None])
+
+            combined_model_fp_per_hr = (combined_model_fp/val_set_hrs).detach().cpu().numpy()
+
+        logging.info(f"\n################\nFinal Model Accuracy: {combined_model_accuracy}"
+                     f"\nFinal Model Recall: {combined_model_recall}\nFinal Model False Positives per Hour: {combined_model_fp_per_hr}"
+                     "\n################\n")
+
+        return combined_model
+
+    def predict_on_features(self, features, model=None):
+        """
+        Predict on Tensors of openWakeWord features corresponding to single audio clips
+
+        Args:
+            features (torch.Tensor): A Tensor of openWakeWord features with shape (batch, features)
+            model (torch.nn.Module): A Pytorch model to use for prediction (default None, which will use self.model)
+
+        Returns:
+            torch.Tensor: An array of predictions of shape (batch, prediction), where 0 is negative and 1 is positive
+        """
+        if len(features) < 3:
+            features = features[None, ]
+
+        features = features.to(self.device)
+        predictions = []
+        for x in tqdm(features, desc="Predicting on clips"):
+            x = x[None, ]
+            batch = []
+            for i in range(0, x.shape[1]-16, 1):  # step size of 1 (80 ms)
+                batch.append(x[:, i:i+16, :])
+            batch = torch.vstack(batch)
+            if model is None:
+                preds = self.model(batch)
+            else:
+                preds = model(batch)
+            predictions.append(preds.detach().cpu().numpy()[None, ])
+
+        return np.vstack(predictions)
+
+    def predict_on_clips(self, clips, model=None):
+        """
+        Predict on Tensors of 16-bit 16 khz audio data
+
+        Args:
+            clips (np.ndarray): A Numpy array of audio clips with shape (batch, samples)
+            model (torch.nn.Module): A Pytorch model to use for prediction (default None, which will use self.model)
+
+        Returns:
+            np.ndarray: An array of predictions of shape (batch, prediction), where 0 is negative and 1 is positive
+        """
+
+        # Get features from clips
+        F = AudioFeatures(device='cpu', ncpu=4)
+        features = F.embed_clips(clips, batch_size=16)
+
+        # Predict on features
+        preds = self.predict_on_features(torch.from_numpy(features), model=model)
+
+        return preds
+
+    def export_model(self, model, model_name, output_dir):
+        """Saves the trained openwakeword model to both onnx and tflite formats"""
+
+        if self.n_classes != 1:
+            raise ValueError("Exporting models to both onnx and tflite with more than one class is currently not supported! "
+                             "Use the `export_to_onnx` function instead.")
+
+        # Save ONNX model
+        logging.info(f"####\nSaving ONNX mode as '{os.path.join(output_dir, model_name + '.onnx')}'")
+        model_to_save = copy.deepcopy(model)
+        torch.onnx.export(model_to_save.to("cpu"), torch.rand(self.input_shape)[None, ],
+                          os.path.join(output_dir, model_name + ".onnx"), opset_version=13)
+
+        return None
+
+    def train_model(self, X, max_steps, warmup_steps, hold_steps, X_val=None,
+                    false_positive_val_data=None, positive_test_clips=None,
+                    negative_weight_schedule=[1],
+                    val_steps=[250], lr=0.0001, val_set_hrs=1):
+        # Move models and main class to target device
+        self.to(self.device)
+        self.model.to(self.device)
+
+        # Train model
+        accumulation_steps = 1
+        accumulated_samples = 0
+        accumulated_predictions = torch.Tensor([]).to(self.device)
+        accumulated_labels = torch.Tensor([]).to(self.device)
+        for step_ndx, data in tqdm(enumerate(X, 0), total=max_steps, desc="Training"):
+            # get the inputs; data is a list of [inputs, labels]
+            x, y = data[0].to(self.device), data[1].to(self.device)
+            y_ = y[..., None].to(torch.float32)
+
+            # Update learning rates
+            for g in self.optimizer.param_groups:
+                g['lr'] = self.lr_warmup_cosine_decay(step_ndx, warmup_steps=warmup_steps, hold=hold_steps,
+                                                      total_steps=max_steps, target_lr=lr)
+
+            # zero the parameter gradients
+            self.optimizer.zero_grad()
+
+            # Get predictions for batch
+            predictions = self.model(x)
+
+            # Construct batch with only samples that have high loss
+            neg_high_loss = predictions[(y == 0) & (predictions.squeeze() >= 0.001)]  # thresholds were chosen arbitrarily but work well
+            pos_high_loss = predictions[(y == 1) & (predictions.squeeze() < 0.999)]
+            y = torch.cat((y[(y == 0) & (predictions.squeeze() >= 0.001)], y[(y == 1) & (predictions.squeeze() < 0.999)]))
+            y_ = y[..., None].to(torch.float32)
+            predictions = torch.cat((neg_high_loss, pos_high_loss))
+
+            # Set weights for batch
+            if len(negative_weight_schedule) == 1:
+                w = torch.ones(y.shape[0])*negative_weight_schedule[0]
+                pos_ndcs = y == 1
+                w[pos_ndcs] = 1
+                w = w[..., None]
+            else:
+                if self.n_classes == 1:
+                    w = torch.ones(y.shape[0])*negative_weight_schedule[step_ndx]
+                    pos_ndcs = y == 1
+                    w[pos_ndcs] = 1
+                    w = w[..., None]
+
+            if predictions.shape[0] != 0:
+                # Do backpropagation, with gradient accumulation if the batch-size after selecting high loss examples is too small
+                loss = self.loss(predictions, y_ if self.n_classes == 1 else y, w.to(self.device))
+                loss = loss/accumulation_steps
+                accumulated_samples += predictions.shape[0]
+
+                if predictions.shape[0] >= 128:
+                    accumulated_predictions = predictions
+                    accumulated_labels = y_
+                if accumulated_samples < 128:
+                    accumulation_steps += 1
+                    accumulated_predictions = torch.cat((accumulated_predictions, predictions))
+                    accumulated_labels = torch.cat((accumulated_labels, y_))
+                else:
+                    loss.backward()
+                    self.optimizer.step()
+                    accumulation_steps = 1
+                    accumulated_samples = 0
+
+                    self.history["loss"].append(loss.detach().cpu().numpy())
+
+                    # Compute training metrics and log them
+                    fp = self.fp(accumulated_predictions, accumulated_labels if self.n_classes == 1 else y)
+                    self.n_fp += fp
+                    self.history["recall"].append(self.recall(accumulated_predictions, accumulated_labels).detach().cpu().numpy())
+
+                    accumulated_predictions = torch.Tensor([]).to(self.device)
+                    accumulated_labels = torch.Tensor([]).to(self.device)
+
+            # Run validation and log validation metrics
+            if step_ndx in val_steps and step_ndx > 1 and false_positive_val_data is not None:
+                # Get false positives per hour with false positive data
+                val_fp = 0
+                for val_step_ndx, data in enumerate(false_positive_val_data):
+                    with torch.no_grad():
+                        x_val, y_val = data[0].to(self.device), data[1].to(self.device)
+                        val_predictions = self.model(x_val)
+                        val_fp += self.fp(val_predictions, y_val[..., None])
+                val_fp_per_hr = (val_fp/val_set_hrs).detach().cpu().numpy()
+                self.history["val_fp_per_hr"].append(val_fp_per_hr)
+
+            # Get recall on test clips
+            if step_ndx in val_steps and step_ndx > 1 and positive_test_clips is not None:
+                tp = 0
+                fn = 0
+                for val_step_ndx, data in enumerate(positive_test_clips):
+                    with torch.no_grad():
+                        x_val = data[0].to(self.device)
+                        batch = []
+                        for i in range(0, x_val.shape[1]-16, 1):
+                            batch.append(x_val[:, i:i+16, :])
+                        batch = torch.vstack(batch)
+                        preds = self.model(batch)
+                        if any(preds >= 0.5):
+                            tp += 1
+                        else:
+                            fn += 1
+                self.history["positive_test_clips_recall"].append(tp/(tp + fn))
+
+            if step_ndx in val_steps and step_ndx > 1 and X_val is not None:
+                # Get metrics for balanced test examples of positive and negative clips
+                for val_step_ndx, data in enumerate(X_val):
+                    with torch.no_grad():
+                        x_val, y_val = data[0].to(self.device), data[1].to(self.device)
+                        val_predictions = self.model(x_val)
+                        val_recall = self.recall(val_predictions, y_val[..., None]).detach().cpu().numpy()
+                        val_acc = self.accuracy(val_predictions, y_val[..., None].to(torch.int64))
+                        val_fp = self.fp(val_predictions, y_val[..., None])
+                self.history["val_accuracy"].append(val_acc.detach().cpu().numpy())
+                self.history["val_recall"].append(val_recall)
+                self.history["val_n_fp"].append(val_fp.detach().cpu().numpy())
+
+            # Save models with a validation score above/below the 90th percentile
+            # of the validation scores up to that point
+            if step_ndx in val_steps and step_ndx > 1:
+                if self.history["val_n_fp"][-1] <= np.percentile(self.history["val_n_fp"], 50) and \
+                   self.history["val_recall"][-1] >= np.percentile(self.history["val_recall"], 5):
+                    # logging.info("Saving checkpoint with metrics >= to targets!")
+                    self.best_models.append(copy.deepcopy(self.model))
+                    self.best_model_scores.append({"training_step_ndx": step_ndx, "val_n_fp": self.history["val_n_fp"][-1],
+                                                   "val_recall": self.history["val_recall"][-1],
+                                                   "val_accuracy": self.history["val_accuracy"][-1],
+                                                   "val_fp_per_hr": self.history.get("val_fp_per_hr", [0])[-1]})
+                    self.best_val_recall = self.history["val_recall"][-1]
+                    self.best_val_accuracy = self.history["val_accuracy"][-1]
+
+            if step_ndx == max_steps-1:
+                break
+
+
+# Separate function to convert onnx models to tflite format
+def convert_onnx_to_tflite(onnx_model_path, output_path):
+    """Converts an ONNX version of an openwakeword model to the Tensorflow tflite format."""
+    # imports
+    import onnx
+    from onnx_tf.backend import prepare
+    import tensorflow as tf
+
+    # Convert to tflite from onnx model
+    onnx_model = onnx.load(onnx_model_path)
+    tf_rep = prepare(onnx_model, device="CPU")
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        tf_rep.export_graph(os.path.join(tmp_dir, "tf_model"))
+        converter = tf.lite.TFLiteConverter.from_saved_model(os.path.join(tmp_dir, "tf_model"))
+        tflite_model = converter.convert()
+
+        logging.info(f"####\nSaving tflite mode to '{output_path}'")
+        with open(output_path, 'wb') as f:
+            f.write(tflite_model)
+
+    return None
+
+
+if __name__ == '__main__':
+    # Get training config file
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--training_config",
+        help="The path to the training config file (required)",
+        type=str,
+        required=True
+    )
+    parser.add_argument(
+        "--generate_clips",
+        help="Execute the synthetic data generation process",
+        action="store_true",
+        default="False",
+        required=False
+    )
+    parser.add_argument(
+        "--augment_clips",
+        help="Execute the synthetic data augmentation process",
+        action="store_true",
+        default="False",
+        required=False
+    )
+    parser.add_argument(
+        "--overwrite",
+        help="Overwrite existing openwakeword features when the --augment_clips flag is used",
+        action="store_true",
+        default="False",
+        required=False
+    )
+    parser.add_argument(
+        "--train_model",
+        help="Execute the model training process",
+        action="store_true",
+        default="False",
+        required=False
+    )
+    parser.add_argument(
+        "--convert_to_tflite",
+        help="Convert the trained ONNX model to TFLite format",
+        action="store_true",
+        default="False",
+        required=False
+    )
+
+    args = parser.parse_args()
+    config = yaml.load(open(args.training_config, 'r').read(), yaml.Loader)
+
+    # imports Piper for synthetic sample generation
+    sys.path.insert(0, os.path.abspath(config["piper_sample_generator_path"]))
+    from generate_samples import generate_samples
+
+    # Define output locations
+    config["output_dir"] = os.path.abspath(config["output_dir"])
+    if not os.path.exists(config["output_dir"]):
+        os.mkdir(config["output_dir"])
+    if not os.path.exists(os.path.join(config["output_dir"], config["model_name"])):
+        os.mkdir(os.path.join(config["output_dir"], config["model_name"]))
+
+    positive_train_output_dir = os.path.join(config["output_dir"], config["model_name"], "positive_train")
+    positive_test_output_dir = os.path.join(config["output_dir"], config["model_name"], "positive_test")
+    negative_train_output_dir = os.path.join(config["output_dir"], config["model_name"], "negative_train")
+    negative_test_output_dir = os.path.join(config["output_dir"], config["model_name"], "negative_test")
+    feature_save_dir = os.path.join(config["output_dir"], config["model_name"])
+
+    # Get paths for impulse response and background audio files
+    rir_paths = [i.path for j in config["rir_paths"] for i in os.scandir(j)]
+    background_paths = []
+    if len(config["background_paths_duplication_rate"]) != len(config["background_paths"]):
+        config["background_paths_duplication_rate"] = [1]*len(config["background_paths"])
+    for background_path, duplication_rate in zip(config["background_paths"], config["background_paths_duplication_rate"]):
+        background_paths.extend([i.path for i in os.scandir(background_path)]*duplication_rate)
+
+    if args.generate_clips is True:
+        # Generate positive clips for training
+        logging.info("#"*50 + "\nGenerating positive clips for training\n" + "#"*50)
+        if not os.path.exists(positive_train_output_dir):
+            os.mkdir(positive_train_output_dir)
+        n_current_samples = len(os.listdir(positive_train_output_dir))
+        if n_current_samples <= 0.95*config["n_samples"]:
+            generate_samples(
+                text=config["target_phrase"], max_samples=config["n_samples"]-n_current_samples,
+                batch_size=config["tts_batch_size"],
+                noise_scales=[0.98], noise_scale_ws=[0.98], length_scales=[0.75, 1.0, 1.25],
+                output_dir=positive_train_output_dir, auto_reduce_batch_size=True,
+                file_names=[uuid.uuid4().hex + ".wav" for i in range(config["n_samples"])]
+            )
+            torch.cuda.empty_cache()
+        else:
+            logging.warning(f"Skipping generation of positive clips for training, as ~{config['n_samples']} already exist")
+
+        # Generate positive clips for testing
+        logging.info("#"*50 + "\nGenerating positive clips for testing\n" + "#"*50)
+        if not os.path.exists(positive_test_output_dir):
+            os.mkdir(positive_test_output_dir)
+        n_current_samples = len(os.listdir(positive_test_output_dir))
+        if n_current_samples <= 0.95*config["n_samples_val"]:
+            generate_samples(text=config["target_phrase"], max_samples=config["n_samples_val"]-n_current_samples,
+                             batch_size=config["tts_batch_size"],
+                             noise_scales=[1.0], noise_scale_ws=[1.0], length_scales=[0.75, 1.0, 1.25],
+                             output_dir=positive_test_output_dir, auto_reduce_batch_size=True)
+            torch.cuda.empty_cache()
+        else:
+            logging.warning(f"Skipping generation of positive clips testing, as ~{config['n_samples_val']} already exist")
+
+        # Generate adversarial negative clips for training
+        logging.info("#"*50 + "\nGenerating negative clips for training\n" + "#"*50)
+        if not os.path.exists(negative_train_output_dir):
+            os.mkdir(negative_train_output_dir)
+        n_current_samples = len(os.listdir(negative_train_output_dir))
+        if n_current_samples <= 0.95*config["n_samples"]:
+            adversarial_texts = config["custom_negative_phrases"]
+            for target_phrase in config["target_phrase"]:
+                adversarial_texts.extend(generate_adversarial_texts(
+                    input_text=target_phrase,
+                    N=config["n_samples"]//len(config["target_phrase"]),
+                    include_partial_phrase=1.0,
+                    include_input_words=0.2))
+            generate_samples(text=adversarial_texts, max_samples=config["n_samples"]-n_current_samples,
+                             batch_size=config["tts_batch_size"]//7,
+                             noise_scales=[0.98], noise_scale_ws=[0.98], length_scales=[0.75, 1.0, 1.25],
+                             output_dir=negative_train_output_dir, auto_reduce_batch_size=True,
+                             file_names=[uuid.uuid4().hex + ".wav" for i in range(config["n_samples"])]
+                             )
+            torch.cuda.empty_cache()
+        else:
+            logging.warning(f"Skipping generation of negative clips for training, as ~{config['n_samples']} already exist")
+
+        # Generate adversarial negative clips for testing
+        logging.info("#"*50 + "\nGenerating negative clips for testing\n" + "#"*50)
+        if not os.path.exists(negative_test_output_dir):
+            os.mkdir(negative_test_output_dir)
+        n_current_samples = len(os.listdir(negative_test_output_dir))
+        if n_current_samples <= 0.95*config["n_samples_val"]:
+            adversarial_texts = config["custom_negative_phrases"]
+            for target_phrase in config["target_phrase"]:
+                adversarial_texts.extend(generate_adversarial_texts(
+                    input_text=target_phrase,
+                    N=config["n_samples_val"]//len(config["target_phrase"]),
+                    include_partial_phrase=1.0,
+                    include_input_words=0.2))
+            generate_samples(text=adversarial_texts, max_samples=config["n_samples_val"]-n_current_samples,
+                             batch_size=config["tts_batch_size"]//7,
+                             noise_scales=[1.0], noise_scale_ws=[1.0], length_scales=[0.75, 1.0, 1.25],
+                             output_dir=negative_test_output_dir, auto_reduce_batch_size=True)
+            torch.cuda.empty_cache()
+        else:
+            logging.warning(f"Skipping generation of negative clips for testing, as ~{config['n_samples_val']} already exist")
+
+    # Set the total length of the training clips based on the ~median generated clip duration, rounding to the nearest 1000 samples
+    # and setting to 32000 when the median + 750 ms is close to that, as it's a good default value
+    n = 50  # sample size
+    positive_clips = [str(i) for i in Path(positive_test_output_dir).glob("*.wav")]
+    duration_in_samples = []
+    for i in range(n):
+        sr, dat = scipy.io.wavfile.read(positive_clips[np.random.randint(0, len(positive_clips))])
+        duration_in_samples.append(len(dat))
+
+    config["total_length"] = int(round(np.median(duration_in_samples)/1000)*1000) + 12000  # add 750 ms to clip duration as buffer
+    if config["total_length"] < 32000:
+        config["total_length"] = 32000  # set a minimum of 32000 samples (2 seconds)
+    elif abs(config["total_length"] - 32000) <= 4000:
+        config["total_length"] = 32000
+
+    # Do Data Augmentation
+    if args.augment_clips is True:
+        if not os.path.exists(os.path.join(feature_save_dir, "positive_features_train.npy")) or args.overwrite is True:
+            positive_clips_train = [str(i) for i in Path(positive_train_output_dir).glob("*.wav")]*config["augmentation_rounds"]
+            positive_clips_train_generator = augment_clips(positive_clips_train, total_length=config["total_length"],
+                                                           batch_size=config["augmentation_batch_size"],
+                                                           background_clip_paths=background_paths,
+                                                           RIR_paths=rir_paths)
+
+            positive_clips_test = [str(i) for i in Path(positive_test_output_dir).glob("*.wav")]*config["augmentation_rounds"]
+            positive_clips_test_generator = augment_clips(positive_clips_test, total_length=config["total_length"],
+                                                          batch_size=config["augmentation_batch_size"],
+                                                          background_clip_paths=background_paths,
+                                                          RIR_paths=rir_paths)
+
+            negative_clips_train = [str(i) for i in Path(negative_train_output_dir).glob("*.wav")]*config["augmentation_rounds"]
+            negative_clips_train_generator = augment_clips(negative_clips_train, total_length=config["total_length"],
+                                                           batch_size=config["augmentation_batch_size"],
+                                                           background_clip_paths=background_paths,
+                                                           RIR_paths=rir_paths)
+
+            negative_clips_test = [str(i) for i in Path(negative_test_output_dir).glob("*.wav")]*config["augmentation_rounds"]
+            negative_clips_test_generator = augment_clips(negative_clips_test, total_length=config["total_length"],
+                                                          batch_size=config["augmentation_batch_size"],
+                                                          background_clip_paths=background_paths,
+                                                          RIR_paths=rir_paths)
+
+            # Compute features and save to disk via memmapped arrays
+            logging.info("#"*50 + "\nComputing openwakeword features for generated samples\n" + "#"*50)
+            n_cpus = os.cpu_count()
+            if n_cpus is None:
+                n_cpus = 1
+            else:
+                n_cpus = n_cpus//2
+            compute_features_from_generator(positive_clips_train_generator, n_total=len(os.listdir(positive_train_output_dir)),
+                                            clip_duration=config["total_length"],
+                                            output_file=os.path.join(feature_save_dir, "positive_features_train.npy"),
+                                            device="gpu" if torch.cuda.is_available() else "cpu",
+                                            ncpu=n_cpus if not torch.cuda.is_available() else 1)
+
+            compute_features_from_generator(negative_clips_train_generator, n_total=len(os.listdir(negative_train_output_dir)),
+                                            clip_duration=config["total_length"],
+                                            output_file=os.path.join(feature_save_dir, "negative_features_train.npy"),
+                                            device="gpu" if torch.cuda.is_available() else "cpu",
+                                            ncpu=n_cpus if not torch.cuda.is_available() else 1)
+
+            compute_features_from_generator(positive_clips_test_generator, n_total=len(os.listdir(positive_test_output_dir)),
+                                            clip_duration=config["total_length"],
+                                            output_file=os.path.join(feature_save_dir, "positive_features_test.npy"),
+                                            device="gpu" if torch.cuda.is_available() else "cpu",
+                                            ncpu=n_cpus if not torch.cuda.is_available() else 1)
+
+            compute_features_from_generator(negative_clips_test_generator, n_total=len(os.listdir(negative_test_output_dir)),
+                                            clip_duration=config["total_length"],
+                                            output_file=os.path.join(feature_save_dir, "negative_features_test.npy"),
+                                            device="gpu" if torch.cuda.is_available() else "cpu",
+                                            ncpu=n_cpus if not torch.cuda.is_available() else 1)
+        else:
+            logging.warning("Openwakeword features already exist, skipping data augmentation and feature generation")
+
+    # Create openwakeword model
+    if args.train_model is True:
+        F = openwakeword.utils.AudioFeatures(device='cpu')
+        input_shape = np.load(os.path.join(feature_save_dir, "positive_features_test.npy")).shape[1:]
+
+        oww = Model(n_classes=1, input_shape=input_shape, model_type=config["model_type"],
+                    layer_dim=config["layer_size"], seconds_per_example=1280*input_shape[0]/16000)
+
+        # Create data transform function for batch generation to handle differ clip lengths (todo: write tests for this)
+        def f(x, n=input_shape[0]):
+            """Simple transformation function to ensure negative data is the appropriate shape for the model size"""
+            if n > x.shape[1] or n < x.shape[1]:
+                x = np.vstack(x)
+                new_batch = np.array([x[i:i+n, :] for i in range(0, x.shape[0]-n, n)])
+            else:
+                return x
+            return new_batch
+
+        # Create label transforms as needed for model (currently only supports binary classification models)
+        data_transforms = {key: f for key in config["feature_data_files"].keys()}
+        label_transforms = {}
+        for key in ["positive"] + list(config["feature_data_files"].keys()) + ["adversarial_negative"]:
+            if key == "positive":
+                label_transforms[key] = lambda x: [1 for i in x]
+            else:
+                label_transforms[key] = lambda x: [0 for i in x]
+
+        # Add generated positive and adversarial negative clips to the feature data files dictionary
+        config["feature_data_files"]['positive'] = os.path.join(feature_save_dir, "positive_features_train.npy")
+        config["feature_data_files"]['adversarial_negative'] = os.path.join(feature_save_dir, "negative_features_train.npy")
+
+        # Make PyTorch data loaders for training and validation data
+        batch_generator = mmap_batch_generator(
+            config["feature_data_files"],
+            n_per_class=config["batch_n_per_class"],
+            data_transform_funcs=data_transforms,
+            label_transform_funcs=label_transforms
+        )
+
+        class IterDataset(torch.utils.data.IterableDataset):
+            def __init__(self, generator):
+                self.generator = generator
+
+            def __iter__(self):
+                return self.generator
+
+        n_cpus = os.cpu_count()
+        if n_cpus is None:
+            n_cpus = 1
+        else:
+            n_cpus = n_cpus//2
+        X_train = torch.utils.data.DataLoader(IterDataset(batch_generator),
+                                              batch_size=None, num_workers=n_cpus, prefetch_factor=16)
+
+        X_val_fp = np.load(config["false_positive_validation_data_path"])
+        X_val_fp = np.array([X_val_fp[i:i+input_shape[0]] for i in range(0, X_val_fp.shape[0]-input_shape[0], 1)])  # reshape to match model
+        X_val_fp_labels = np.zeros(X_val_fp.shape[0]).astype(np.float32)
+        X_val_fp = torch.utils.data.DataLoader(
+            torch.utils.data.TensorDataset(torch.from_numpy(X_val_fp), torch.from_numpy(X_val_fp_labels)),
+            batch_size=len(X_val_fp_labels)
+        )
+
+        X_val_pos = np.load(os.path.join(feature_save_dir, "positive_features_test.npy"))
+        X_val_neg = np.load(os.path.join(feature_save_dir, "negative_features_test.npy"))
+        labels = np.hstack((np.ones(X_val_pos.shape[0]), np.zeros(X_val_neg.shape[0]))).astype(np.float32)
+
+        X_val = torch.utils.data.DataLoader(
+            torch.utils.data.TensorDataset(
+                torch.from_numpy(np.vstack((X_val_pos, X_val_neg))),
+                torch.from_numpy(labels)
+                ),
+            batch_size=len(labels)
+        )
+
+        # Run auto training
+        best_model = oww.auto_train(
+            X_train=X_train,
+            X_val=X_val,
+            false_positive_val_data=X_val_fp,
+            steps=config["steps"],
+            max_negative_weight=config["max_negative_weight"],
+            target_fp_per_hour=config["target_false_positives_per_hour"],
+        )
+
+        # Export the trained model to onnx
+        oww.export_model(model=best_model, model_name=config["model_name"], output_dir=config["output_dir"])
+
+        # Convert the model from onnx to tflite format
+        if args.convert_to_tflite:
+            convert_onnx_to_tflite(os.path.join(config["output_dir"], config["model_name"] + ".onnx"),
+                                   os.path.join(config["output_dir"], config["model_name"] + ".tflite"))

openwakeword/utils.py

@@ -14,15 +14,18 @@
 
 # Imports
 import os
-import onnxruntime as ort
 import numpy as np
 import pathlib
 from collections import deque
 from multiprocessing.pool import ThreadPool
 from multiprocessing import Process, Queue
 import time
+import logging
+from tqdm import tqdm
 import openwakeword
+from numpy.lib.format import open_memmap
 from typing import Union, List, Callable, Deque
+import requests
 
 
 # Base class for computing audio features using Google's speech_embedding
@@ -33,44 +36,147 @@ class AudioFeatures():
     `speech_embedding` features.
     """
     def __init__(self,
-                 melspec_onnx_model_path: str = os.path.join(
-                                                            pathlib.Path(__file__).parent.resolve(),
-                                                            "resources", "models", "melspectrogram.onnx"
-                 ),
-                 embedding_onnx_model_path: str = os.path.join(
-                                                                pathlib.Path(__file__).parent.resolve(),
-                                                                "resources", "models", "embedding_model.onnx"
-                 ),
+                 melspec_model_path: str = "",
+                 embedding_model_path: str = "",
                  sr: int = 16000,
-                 ncpu: int = 1
+                 ncpu: int = 1,
+                 inference_framework: str = "onnx",
+                 device: str = 'cpu'
                  ):
         """
         Initialize the AudioFeatures object.
 
         Args:
-            melspec_onnx_model_path (str): The path to the ONNX model for computing melspectograms from audio data
-            embedding_onnx_model_path (str): The path to the ONNX model for Google's `speech_embedding` model
+            melspec_model_path (str): The path to the model for computing melspectograms from audio data
+            embedding_model_path (str): The path to the model for Google's `speech_embedding` model
             sr (int): The sample rate of the audio (default: 16000 khz)
             ncpu (int): The number of CPUs to use when computing melspectrograms and audio features (default: 1)
+            inference_framework (str): The inference framework to use when for model prediction. Options are
+                                       "tflite" or "onnx". The default is "tflite" as this results in better
+                                       efficiency on common platforms (x86, ARM64), but in some deployment
+                                       scenarios ONNX models may be preferable.
+            device (str): The device to use when running the models, either "cpu" or "gpu" (default is "cpu".)
+                          Note that depending on the inference framework selected and system configuration,
+                          this setting may not have an effect. For example, to use a GPU with the ONNX
+                          framework the appropriate onnxruntime package must be installed.
         """
-        # Initialize the ONNX models
-        sessionOptions = ort.SessionOptions()
-        sessionOptions.inter_op_num_threads = ncpu
-        sessionOptions.intra_op_num_threads = ncpu
-        self.melspec_model = ort.InferenceSession(melspec_onnx_model_path, sess_options=sessionOptions,
-                                                  providers=["CUDAExecutionProvider", "CPUExecutionProvider"])
-        self.embedding_model = ort.InferenceSession(embedding_onnx_model_path, sess_options=sessionOptions,
-                                                    providers=["CUDAExecutionProvider", "CPUExecutionProvider"])
-        self.onnx_execution_provider = self.melspec_model.get_providers()[0]
+        # Initialize the models with the appropriate framework
+        if inference_framework == "onnx":
+            try:
+                import onnxruntime as ort
+            except ImportError:
+                raise ValueError("Tried to import onnxruntime, but it was not found. Please install it using `pip install onnxruntime`")
 
-        # Create databuffers
+            if melspec_model_path == "":
+                melspec_model_path = os.path.join(pathlib.Path(__file__).parent.resolve(), "resources", "models", "melspectrogram.onnx")
+            if embedding_model_path == "":
+                embedding_model_path = os.path.join(pathlib.Path(__file__).parent.resolve(), "resources", "models", "embedding_model.onnx")
+
+            if ".tflite" in melspec_model_path or ".tflite" in embedding_model_path:
+                raise ValueError("The onnx inference framework is selected, but tflite models were provided!")
+
+            # Initialize ONNX options
+            sessionOptions = ort.SessionOptions()
+            sessionOptions.inter_op_num_threads = ncpu
+            sessionOptions.intra_op_num_threads = ncpu
+
+            # Melspectrogram model
+            self.melspec_model = ort.InferenceSession(melspec_model_path, sess_options=sessionOptions,
+                                                      providers=["CUDAExecutionProvider"] if device == "gpu" else ["CPUExecutionProvider"])
+            self.onnx_execution_provider = self.melspec_model.get_providers()[0]
+            self.melspec_model_predict = lambda x: self.melspec_model.run(None, {'input': x})
+
+            # Audio embedding model
+            self.embedding_model = ort.InferenceSession(embedding_model_path, sess_options=sessionOptions,
+                                                        providers=["CUDAExecutionProvider"] if device == "gpu"
+                                                        else ["CPUExecutionProvider"])
+            self.embedding_model_predict = lambda x: self.embedding_model.run(None, {'input_1': x})[0].squeeze()
+
+        elif inference_framework == "tflite":
+            try:
+                import ai_edge_litert.interpreter as tflite
+            except ImportError:
+                raise ValueError("Tried to import the LiteRT runtime, but it was not found."
+                                 "Please install it using `pip install ai-edge-litert`")
+
+            if melspec_model_path == "":
+                melspec_model_path = os.path.join(pathlib.Path(__file__).parent.resolve(),
+                                                  "resources", "models", "melspectrogram.tflite")
+            if embedding_model_path == "":
+                embedding_model_path = os.path.join(pathlib.Path(__file__).parent.resolve(),
+                                                    "resources", "models", "embedding_model.tflite")
+
+            if ".onnx" in melspec_model_path or ".onnx" in embedding_model_path:
+                raise ValueError("The tflite inference framework is selected, but onnx models were provided!")
+
+            # Melspectrogram model
+            self.melspec_model = tflite.Interpreter(model_path=melspec_model_path, num_threads=ncpu)
+            self.melspec_model.resize_tensor_input(0, [1, 1280], strict=True)  # initialize with fixed input size
+            self.melspec_model.allocate_tensors()
+
+            melspec_input_index = self.melspec_model.get_input_details()[0]['index']
+            melspec_output_index = self.melspec_model.get_output_details()[0]['index']
+
+            self._tflite_current_melspec_input_size = 1280
+
+            def tflite_melspec_predict(x):
+                if x.shape[1] != 1280:
+                    self.melspec_model.resize_tensor_input(0, [1, x.shape[1]], strict=True)  # initialize with fixed input size
+                    self.melspec_model.allocate_tensors()
+                    self._tflite_current_melspec_input_size = x.shape[1]
+                elif self._tflite_current_melspec_input_size != 1280:
+                    self.melspec_model.resize_tensor_input(0, [1, 1280], strict=True)  # initialize with fixed input size
+                    self.melspec_model.allocate_tensors()
+                    self._tflite_current_melspec_input_size = 1280
+
+                self.melspec_model.set_tensor(melspec_input_index, x)
+                self.melspec_model.invoke()
+                return self.melspec_model.get_tensor(melspec_output_index)
+
+            self.melspec_model_predict = tflite_melspec_predict
+
+            # Audio embedding model
+            self.embedding_model = tflite.Interpreter(model_path=embedding_model_path, num_threads=ncpu)
+            self.embedding_model.allocate_tensors()
+
+            embedding_input_index = self.embedding_model.get_input_details()[0]['index']
+            embedding_output_index = self.embedding_model.get_output_details()[0]['index']
+
+            self._tflite_current_embedding_batch_size = 1
+
+            def tflite_embedding_predict(x):
+                if x.shape[0] != 1:
+                    self.embedding_model.resize_tensor_input(0, [x.shape[0], 76, 32, 1], strict=True)  # initialize with fixed input size
+                    self.embedding_model.allocate_tensors()
+                    self._tflite_current_embedding_batch_size = x.shape[0]
+                elif self._tflite_current_embedding_batch_size != 1:
+                    self.embedding_model.resize_tensor_input(0, [1, 76, 32, 1], strict=True)  # initialize with fixed input size
+                    self.embedding_model.allocate_tensors()
+                    self._tflite_current_embedding_batch_size = x.shape[0]
+
+                self.embedding_model.set_tensor(embedding_input_index, x)
+                self.embedding_model.invoke()
+                return self.embedding_model.get_tensor(embedding_output_index).squeeze()
+
+            self.embedding_model_predict = tflite_embedding_predict
+
+        # Create databuffers with empty/random data
         self.raw_data_buffer: Deque = deque(maxlen=sr*10)
         self.melspectrogram_buffer = np.ones((76, 32))  # n_frames x num_features
         self.melspectrogram_max_len = 10*97  # 97 is the number of frames in 1 second of 16hz audio
         self.accumulated_samples = 0  # the samples added to the buffer since the audio preprocessor was last called
-        self.feature_buffer = self._get_embeddings(np.zeros(160000).astype(np.int16))  # fill with blank data to start
+        self.raw_data_remainder = np.empty(0)
+        self.feature_buffer = self._get_embeddings(np.random.randint(-1000, 1000, 16000*4).astype(np.int16))
         self.feature_buffer_max_len = 120  # ~10 seconds of feature buffer history
 
+    def reset(self):
+        """Reset the internal buffers"""
+        self.raw_data_buffer.clear()
+        self.melspectrogram_buffer = np.ones((76, 32))
+        self.accumulated_samples = 0
+        self.raw_data_remainder = np.empty(0)
+        self.feature_buffer = self._get_embeddings(np.random.randint(-1000, 1000, 16000*4).astype(np.int16))
+
     def _get_melspectrogram(self, x: Union[np.ndarray, List], melspec_transform: Callable = lambda x: x/10 + 2):
         """
         Function to compute the mel-spectrogram of the provided audio samples.
@@ -93,7 +199,7 @@ class AudioFeatures():
         x = x.astype(np.float32) if x.dtype != np.float32 else x
 
         # Get melspectrogram
-        outputs = self.melspec_model.run(None, {'input': x})
+        outputs = self.melspec_model_predict(x)
         spec = np.squeeze(outputs[0])
 
         # Arbitrary transform of melspectrogram
@@ -113,7 +219,7 @@ class AudioFeatures():
         """
         if melspec.shape[0] != 1:
             melspec = melspec[None, ]
-        embedding = self.embedding_model.run(None, {'input_1': melspec})[0].squeeze()
+        embedding = self.embedding_model_predict(melspec)
         return embedding
 
     def _get_embeddings(self, x: np.ndarray, window_size: int = 76, step_size: int = 8, **kwargs):
@@ -126,7 +232,7 @@ class AudioFeatures():
                 windows.append(window)
 
         batch = np.expand_dims(np.array(windows), axis=-1).astype(np.float32)
-        embedding = self.embedding_model.run(None, {'input_1': batch})[0].squeeze()
+        embedding = self.embedding_model_predict(batch)
         return embedding
 
     def get_embedding_shape(self, audio_length: float, sr: int = 16000):
@@ -171,8 +277,9 @@ class AudioFeatures():
                 result = self._get_melspectrogram(batch)
 
             elif pool:
+                chunksize = batch.shape[0]//ncpu if batch.shape[0] >= ncpu else 1
                 result = np.array(pool.map(self._get_melspectrogram,
-                                           batch, chunksize=batch.shape[0]//ncpu))
+                                           batch, chunksize=chunksize))
 
             melspecs[i:i+batch_size, :, :] = result.squeeze()
 
@@ -229,11 +336,12 @@ class AudioFeatures():
             if len(batch) >= batch_size or ndx+1 == x.shape[0]:
                 batch = np.array(batch).astype(np.float32)
                 if "CUDA" in self.onnx_execution_provider:
-                    result = self.embedding_model.run(None, {'input_1': batch})[0].squeeze()
+                    result = self.embedding_model_predict(batch)
 
                 elif pool:
+                    chunksize = batch.shape[0]//ncpu if batch.shape[0] >= ncpu else 1
                     result = np.array(pool.map(self._get_embeddings_from_melspec,
-                                      batch, chunksize=batch.shape[0]//ncpu))
+                                      batch, chunksize=chunksize))
 
                 for j, ndx2 in zip(range(0, result.shape[0], n_frames), ndcs):
                     embeddings[ndx2, :, :] = result[j:j+n_frames]
@@ -282,6 +390,9 @@ class AudioFeatures():
         clip is calculated. It's unclear if this difference is significant and will impact model performance.
         In particular padding with 0 or very small values seems to demonstrate the differences well.
         """
+        if len(self.raw_data_buffer) < 400:
+            raise ValueError("The number of input frames must be at least 400 samples @ 16khz (25 ms)!")
+
         self.melspectrogram_buffer = np.vstack(
             (self.melspectrogram_buffer, self._get_melspectrogram(list(self.raw_data_buffer)[-n_samples-160*3:]))
         )
@@ -293,21 +404,33 @@ class AudioFeatures():
         """
         Adds raw audio data to the input buffer
         """
-        if len(x) < 400:
-            raise ValueError("The number of input frames must be at least 400 samples @ 16khz (25 ms)!")
         self.raw_data_buffer.extend(x.tolist() if isinstance(x, np.ndarray) else x)
 
     def _streaming_features(self, x):
-        # if len(x) != 1280:
-        #     raise ValueError("You must provide input samples in frames of 1280 samples @ 1600khz."
-        #                      f"Received a frame of {len(x)} samples.")
+        # Add raw audio data to buffer, temporarily storing extra frames if not an even number of 80 ms chunks
+        processed_samples = 0
 
-        # Add raw audio data to buffer
-        self._buffer_raw_data(x)
-        self.accumulated_samples += len(x)
+        if self.raw_data_remainder.shape[0] != 0:
+            x = np.concatenate((self.raw_data_remainder, x))
+            self.raw_data_remainder = np.empty(0)
 
-        # Only calculate melspectrogram every ~0.5 seconds to significantly increase efficiency
-        if self.accumulated_samples >= 1280:
+        if self.accumulated_samples + x.shape[0] >= 1280:
+            remainder = (self.accumulated_samples + x.shape[0]) % 1280
+            if remainder != 0:
+                x_even_chunks = x[0:-remainder]
+                self._buffer_raw_data(x_even_chunks)
+                self.accumulated_samples += len(x_even_chunks)
+                self.raw_data_remainder = x[-remainder:]
+            elif remainder == 0:
+                self._buffer_raw_data(x)
+                self.accumulated_samples += x.shape[0]
+                self.raw_data_remainder = np.empty(0)
+        else:
+            self.accumulated_samples += x.shape[0]
+            self._buffer_raw_data(x)
+
+        # Only calculate melspectrogram once minimum samples are accumulated
+        if self.accumulated_samples >= 1280 and self.accumulated_samples % 1280 == 0:
             self._streaming_melspectrogram(self.accumulated_samples)
 
             # Calculate new audio embeddings/features based on update melspectrograms
@@ -317,14 +440,17 @@ class AudioFeatures():
                 x = self.melspectrogram_buffer[-76 + ndx:ndx].astype(np.float32)[None, :, :, None]
                 if x.shape[1] == 76:
                     self.feature_buffer = np.vstack((self.feature_buffer,
-                                                    self.embedding_model.run(None, {'input_1': x})[0].squeeze()))
+                                                    self.embedding_model_predict(x)))
 
             # Reset raw data buffer counter
+            processed_samples = self.accumulated_samples
             self.accumulated_samples = 0
 
         if self.feature_buffer.shape[0] > self.feature_buffer_max_len:
             self.feature_buffer = self.feature_buffer[-self.feature_buffer_max_len:, :]
 
+        return processed_samples if processed_samples != 0 else self.accumulated_samples
+
     def get_features(self, n_feature_frames: int = 16, start_ndx: int = -1):
         if start_ndx != -1:
             end_ndx = start_ndx + int(n_feature_frames) \
@@ -334,15 +460,16 @@ class AudioFeatures():
             return self.feature_buffer[int(-1*n_feature_frames):, :][None, ].astype(np.float32)
 
     def __call__(self, x):
-        self._streaming_features(x)
+        return self._streaming_features(x)
 
 
 # Bulk prediction function
 def bulk_predict(
                  file_paths: List[str],
-                 wakeword_model_paths: List[str],
+                 wakeword_models: List[str],
                  prediction_function: str = 'predict_clip',
                  ncpu: int = 1,
+                 inference_framework: str = "tflite",
                  **kwargs
                  ):
     """
@@ -350,10 +477,14 @@ def bulk_predict(
 
     Args:
         input_paths (List[str]): The list of input file to predict
-        wakeword_model_path (List[str])): The paths to the wakeword ONNX model files
+        wakeword_models (List[str])): The paths to the wakeword model files
         prediction_function (str): The name of the method used to predict on the input audio files
                                    (default is the `predict_clip` method)
         ncpu (int): How many processes to create (up to max of available CPUs)
+        inference_framework (str): The inference framework to use when for model prediction. Options are
+                                    "tflite" or "onnx". The default is "tflite" as this results in better
+                                    efficiency on common platforms (x86, ARM64), but in some deployment
+                                    scenarios ONNX models may be preferable.
         kwargs (dict): Any other keyword arguments to pass to the model initialization or
                        specified prediction function
 
@@ -376,7 +507,8 @@ def bulk_predict(
         filtered_kwargs = {key: value for key, value in kwargs.items()
                            if key in openwakeword.Model.__init__.__code__.co_varnames}
         oww = openwakeword.Model(
-            wakeword_model_paths=wakeword_model_paths,
+            wakeword_models=wakeword_models,
+            inference_framework=inference_framework,
             **filtered_kwargs
         )
         mdls.append(oww)
@@ -405,3 +537,152 @@ def bulk_predict(
 
     # Consolidate results and return
     return {list(i.keys())[0]: list(i.values())[0] for i in results}
+
+
+def compute_features_from_generator(generator, n_total, clip_duration, output_file, device="cpu", ncpu=1):
+    """
+    Computes audio features from a generator that produces Numpy arrays of shape (batch_size, samples)
+    containing 16-bit PCM audio data.
+
+    Args:
+        generator (Generator): The generator that process the arrays of audio data
+        n_total (int): The total number of rows (audio clips) that the generator will produce.
+                       Ideally this is precise, but it can be approximate as well as the output
+                       .npy file will be automatically trimmed to remove empty values.
+        clip_duration (float): The duration (in samples) of the audio produced by the generator
+        output_file (str): The output file (.npy) containing the audio features. Note that this file
+                           will be written to using memmap arrays, so it can be substantially larger
+                           than the available system memory.
+        device (str): The device ("cpu" or "gpu") to use for computing features.
+        ncpu (int): The number of cores to use when process the audio features (if computing on CPU)
+
+    Returns:
+        None
+    """
+    # Function specific imports
+    from openwakeword.data import trim_mmap
+
+    # Create audio features object
+    F = AudioFeatures(device=device)
+
+    # Determine the output shape and create output file
+    n_feature_cols = F.get_embedding_shape(clip_duration/16000)
+    output_shape = (n_total, n_feature_cols[0], n_feature_cols[1])
+    fp = open_memmap(output_file, mode='w+', dtype=np.float32, shape=output_shape)
+
+    # Get batch size by pulling one value from the generator and store features
+    row_counter = 0
+    audio_data = next(generator)
+    batch_size = audio_data.shape[0]
+
+    if batch_size > n_total:
+        raise ValueError(f"The value of 'n_total' ({n_total}) is less than the batch size ({batch_size})."
+                         " Please increase 'n_total' to be >= batch size.")
+
+    features = F.embed_clips(audio_data, batch_size=batch_size)
+    fp[row_counter:row_counter+features.shape[0], :, :] = features
+    row_counter += features.shape[0]
+    fp.flush()
+
+    # Compute features and add data to output file
+    for audio_data in tqdm(generator, total=n_total//batch_size, desc="Computing features"):
+        if row_counter >= n_total:
+            break
+
+        features = F.embed_clips(audio_data, batch_size=batch_size, ncpu=ncpu)
+        if row_counter + features.shape[0] > n_total:
+            features = features[0:n_total-row_counter]
+
+        fp[row_counter:row_counter+features.shape[0], :, :] = features
+        row_counter += features.shape[0]
+        fp.flush()
+
+    # Trip empty rows from the mmapped array
+    trim_mmap(output_file)
+
+
+# Function to download files from a URL with a progress bar
+def download_file(url, target_directory, file_size=None):
+    """A simple function to download a file from a URL with a progress bar using only the requests library"""
+    local_filename = url.split('/')[-1]
+
+    with requests.get(url, stream=True) as r:
+        if file_size is not None:
+            progress_bar = tqdm(total=file_size, unit='iB', unit_scale=True, desc=f"{local_filename}")
+        else:
+            total_size = int(r.headers.get('content-length', 0))
+            progress_bar = tqdm(total=total_size, unit='iB', unit_scale=True, desc=f"{local_filename}")
+
+        with open(os.path.join(target_directory, local_filename), 'wb') as f:
+            for chunk in r.iter_content(chunk_size=8192):
+                f.write(chunk)
+                progress_bar.update(len(chunk))
+
+    progress_bar.close()
+
+
+# Function to download models from GitHub release assets
+def download_models(
+        model_names: List[str] = [],
+        target_directory: str = os.path.join(pathlib.Path(__file__).parent.resolve(), "resources", "models")
+        ):
+    """
+    Download the specified models from the release assets in the openWakeWord GitHub repository.
+    Uses the official urls in the MODELS dictionary in openwakeword/__init__.py.
+
+    Args:
+        model_names (List[str]): The names of the models to download (e.g., hey_jarvis_v0.1). Both ONNX and
+                                 tflite models will be downloaded. If not provided (the default),
+                                 the latest versions of all models will be downloaded.
+        target_directory (str): The directory to save the models to. Defaults to the install location
+                                of openWakeWord (i.e., the `resources/models` directory).
+    Returns:
+        None
+    """
+    if not isinstance(model_names, list):
+        raise ValueError("The model_names argument must be a list of strings")
+
+    # Always download melspectrogram and embedding models, if they don't already exist
+    if not os.path.exists(target_directory):
+        os.makedirs(target_directory)
+    for feature_model in openwakeword.FEATURE_MODELS.values():
+        if not os.path.exists(os.path.join(target_directory, feature_model["download_url"].split("/")[-1])):
+            download_file(feature_model["download_url"], target_directory)
+            download_file(feature_model["download_url"].replace(".tflite", ".onnx"), target_directory)
+
+    # Always download VAD models, if they don't already exist
+    for vad_model in openwakeword.VAD_MODELS.values():
+        if not os.path.exists(os.path.join(target_directory, vad_model["download_url"].split("/")[-1])):
+            download_file(vad_model["download_url"], target_directory)
+
+    # Get all model urls
+    official_model_urls = [i["download_url"] for i in openwakeword.MODELS.values()]
+    official_model_names = [i["download_url"].split("/")[-1] for i in openwakeword.MODELS.values()]
+
+    if model_names != []:
+        for model_name in model_names:
+            url = [i for i, j in zip(official_model_urls, official_model_names) if model_name in j]
+            if url != []:
+                if not os.path.exists(os.path.join(target_directory, url[0].split("/")[-1])):
+                    download_file(url[0], target_directory)
+                    download_file(url[0].replace(".tflite", ".onnx"), target_directory)
+    else:
+        for official_model_url in official_model_urls:
+            if not os.path.exists(os.path.join(target_directory, official_model_url.split("/")[-1])):
+                download_file(official_model_url, target_directory)
+                download_file(official_model_url.replace(".tflite", ".onnx"), target_directory)
+
+
+# Handle deprecated arguments and naming (thanks to https://stackoverflow.com/a/74564394)
+def re_arg(kwarg_map):
+    def decorator(func):
+        def wrapped(*args, **kwargs):
+            new_kwargs = {}
+            for k, v in kwargs.items():
+                if k in kwarg_map:
+                    logging.warning(f"DEPRECATION: keyword argument '{k}' is no longer valid and "
+                                    f"will be removed in future releases. Use '{kwarg_map[k]}' instead.")
+                new_kwargs[kwarg_map.get(k, k)] = v
+            return func(*args, **new_kwargs)
+        return wrapped
+    return decorator

openwakeword/vad.py

@@ -63,18 +63,20 @@ class VAD():
                     "resources",
                     "models",
                     "silero_vad.onnx"
-                 )
+                 ),
+                 n_threads: int = 1
                  ):
         """Initialize the VAD model object.
 
             Args:
                 model_path (str): The path to the Silero VAD ONNX model.
+                n_threads (int): The number of threads to use for the VAD model.
         """
 
         # Initialize the ONNX model
         sessionOptions = ort.SessionOptions()
-        sessionOptions.inter_op_num_threads = 1
-        sessionOptions.intra_op_num_threads = 1
+        sessionOptions.inter_op_num_threads = n_threads
+        sessionOptions.intra_op_num_threads = n_threads
         self.model = ort.InferenceSession(model_path, sess_options=sessionOptions,
                                           providers=["CPUExecutionProvider"])
 

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [tool.pytest.ini_options]
 addopts = "--cov=openwakeword --cov-report term-missing --flake8 --mypy --mypy-ignore-missing-imports"
-flake8-max-line-length = "120"
+flake8-max-line-length = "140"
 testpaths = [
     "tests",
     "openwakeword"
@@ -12,18 +12,19 @@ testpaths = [
 
 [project]
 name = "openwakeword"
-version = "0.4.0"
+version = "0.6.0"
 authors = [
   { name="David Scripka", email="david.scripka@gmail.com" },
 ]
 description = "An open-source audio wake word (or phrase) detection framework with a focus on performance and simplicity"
 readme = "README.md"
-requires-python = ">=3.7"
+requires-python = ">=3.10"
 classifiers = [
     "Programming Language :: Python :: 3",
     "License :: OSI Approved :: Apache Software License",
     "Operating System :: OS Independent",
 ]
+dynamic = ["dependencies", "optional-dependencies"]
 
 [project.urls]
 "Homepage" = "https://github.com/dscripka/openWakeWord"

setup.py

@@ -4,19 +4,20 @@ import setuptools
 with open("README.md", "r", encoding="utf-8") as fh:
     long_description = fh.read()
 
+
 # Build extras_requires based on platform
 def build_additional_requires():
-    py_version = platform.python_version()[0:3].replace('.', "")
-    if platform.system() == "Linux" and platform.machine() == "x86_64":
-        additional_requires=[
-            f"speexdsp_ns @ https://github.com/dscripka/openWakeWord/releases/download/v0.1.1/speexdsp_ns-0.1.2-cp{py_version}-cp{py_version}-linux_x86_64.whl",
-        ]
-    elif platform.system() == "Linux" and platform.machine() == "aarch64":
-        additional_requires=[
-            f"speexdsp_ns @ https://github.com/dscripka/openWakeWord/releases/download/v0.1.1/speexdsp_ns-0.1.2-cp{py_version}-cp{py_version}-linux_aarch64.whl",
-        ],
-    elif platform.system() == "Windows" and platform.machine() == "x86_64":
-        additional_requires=[
+    # py_version = platform.python_version()[0:3].replace('.', "")
+    # if platform.system() == "Linux" and platform.machine() == "x86_64":
+    #     additional_requires=[
+    #         f"speexdsp_ns @ https://github.com/dscripka/openWakeWord/releases/download/v0.1.1/speexdsp_ns-0.1.2-cp{py_version}-cp{py_version}-linux_x86_64.whl",
+    #     ]
+    # elif platform.system() == "Linux" and platform.machine() == "aarch64":
+    #     additional_requires=[
+    #         f"speexdsp_ns @ https://github.com/dscripka/openWakeWord/releases/download/v0.1.1/speexdsp_ns-0.1.2-cp{py_version}-cp{py_version}-linux_aarch64.whl",
+    #     ],
+    if platform.system() == "Windows" and platform.machine() == "x86_64":
+        additional_requires = [
             'PyAudioWPatch'
         ]
     else:
@@ -24,27 +25,56 @@ def build_additional_requires():
 
     return additional_requires
 
+
 setuptools.setup(
     name="openwakeword",
-    version="0.4.0",
-    install_requires=['onnxruntime>=1.10.0,<2', 'tqdm>=4.0,<5.0', 'scipy>=1.3,<2', 'scikit-learn>=1,<2'],
+    version="0.6.0",
+    install_requires=[
+        'onnxruntime>=1.10.0,<2',
+        'ai-edge-litert>=2.0.2,<3; platform_system == "Linux" or platform_system == "Darwin"',
+        'speexdsp-ns>=0.1.2,<1; platform_system == "Linux"',
+        'tqdm>=4.0,<5.0',
+        'scipy>=1.3,<2',
+        'scikit-learn>=1,<2',
+        'requests>=2.0,<3',
+    ],
     extras_require={
         'test': [
                     'pytest>=7.2.0,<8',
                     'pytest-cov>=2.10.1,<3',
                     'pytest-flake8>=1.1.1,<2',
-                    'flake8>=4.0,<4.1',
-                    'pytest-mypy>=0.10.0,<1'
+                    'flake8>=5.0,<7.1',
+                    'pytest-mypy>=0.10.0,<1',
+                    'types-requests',
+                    'types-PyYAML',
+                    'mock>=5.1,<6',
+                    'types-mock>=5.1,<6',
+                    'types-requests>=2.0,<3'
                 ],
         'full': [
                     'mutagen>=1.46.0,<2',
-                    'speechbrain>=0.5.13,<1',
+                    'torch>=1.13.1,<3',
+                    'torchaudio>=0.13.1,<1',
+                    'torchinfo>=1.8.0,<2',
+                    'torchmetrics>=0.11.4,<1',
+                    'speechbrain>=0.5.14,<1',
+                    'audiomentations>=0.30.0,<1',
+                    'torch-audiomentations>=0.11.0,<1',
+                    'tqdm>=4.64.0,<5',
                     'pytest>=7.2.0,<8',
                     'pytest-cov>=2.10.1,<3',
                     'pytest-flake8>=1.1.1,<2',
                     'pytest-mypy>=0.10.0,<1',
-                    'plotext>=5.2.7,<6',
-                    'sounddevice>=0.4.1,<1'
+                    'acoustics>=0.2.6,<1',
+                    'pyyaml>=6.0,<7',
+                    'tensorflow-cpu==2.8.1',
+                    'tensorflow_probability==0.16.0',
+                    'protobuf>=3.20,<4',
+                    'onnx_tf==1.10.0',
+                    'onnx==1.14.0',
+                    'pronouncing>=0.2.0,<1',
+                    'datasets>=2.14.4,<3',
+                    'deep-phonemizer==0.0.19'
                 ]
     },
     author="David Scripka",
@@ -63,5 +93,5 @@ setuptools.setup(
     ],
     packages=setuptools.find_packages(),
     include_package_data=True,
-    python_requires=">=3.7",
+    python_requires=">=3.10",
 )

tests/test_custom_verifier_model.py

@@ -34,6 +34,9 @@ import scipy.io.wavfile
 import tempfile
 import pytest
 
+# Download models needed for tests
+openwakeword.utils.download_models(model_names=["alexa_v0.1", "hey_mycroft_v0.1"])
+
 
 # Tests
 class TestModels:
@@ -69,20 +72,21 @@ class TestModels:
                 positive_reference_clips=reference_clips,
                 negative_reference_clips=negative_clips,
                 output_path=os.path.join(tmp_dir, 'verifier_model.pkl'),
-                model_name=os.path.join("openwakeword", "resources", "models", "hey_mycroft_v0.1.onnx")
+                model_name=os.path.join("openwakeword", "resources", "models", "hey_mycroft_v0.1.tflite")
             )
 
             with pytest.raises(ValueError):
                 # Load model with verifier model incorrectly to catch ValueError
                 owwModel = openwakeword.Model(
-                    wakeword_model_paths=[os.path.join("openwakeword", "resources", "models", "hey_mycroft_v0.1.onnx")],
+                    wakeword_models=[os.path.join("openwakeword", "resources",
+                                     "models", "hey_mycroft_v0.1.tflite")],
                     custom_verifier_models={"bad_key": os.path.join(tmp_dir, "verifier_model.pkl")},
                     custom_verifier_threshold=0.3,
                 )
 
             # Load model with verifier model incorrectly to catch ValueError
             owwModel = openwakeword.Model(
-                wakeword_model_paths=[os.path.join("openwakeword", "resources", "models", "hey_mycroft_v0.1.onnx")],
+                wakeword_models=[os.path.join("openwakeword", "resources", "models", "hey_mycroft_v0.1.tflite")],
                 custom_verifier_models={"hey_mycroft_v0.1": os.path.join(tmp_dir, "verifier_model.pkl")},
                 custom_verifier_threshold=0.3,
             )

tests/test_models.py

@@ -29,33 +29,120 @@
 # Imports
 import openwakeword
 import os
+import sys
+import logging
 import numpy as np
 from pathlib import Path
 import collections
 import pytest
 import platform
+import pickle
+import tempfile
+import mock
+import wave
+
+# Download models needed for tests
+openwakeword.utils.download_models()
 
 
 # Tests
 class TestModels:
     def test_load_models_by_path(self):
         # Load model with defaults
-        owwModel = openwakeword.Model(wakeword_model_paths=[
+        owwModel = openwakeword.Model(wakeword_models=[
                                         os.path.join("openwakeword", "resources", "models", "alexa_v0.1.onnx")
-                                      ])
+                                      ], inference_framework="onnx")
+
+        # Prediction on random data
+        prediction = owwModel.predict(np.random.randint(-1000, 1000, 1280).astype(np.int16))
+        assert prediction["alexa_v0.1"] >= 0 and prediction["alexa_v0.1"] <= 1
+
+        owwModel = openwakeword.Model(wakeword_models=[
+                                        os.path.join("openwakeword", "resources", "models", "alexa_v0.1.tflite")
+                                      ], inference_framework="tflite")
+
+        # Prediction on random data
+        prediction = owwModel.predict(np.random.randint(-1000, 1000, 1280).astype(np.int16))
+        assert prediction["alexa_v0.1"] >= 0 and prediction["alexa_v0.1"] <= 1
+
+    def test_predict_with_different_frame_sizes(self):
+        # Test with binary model
+        owwModel1 = openwakeword.Model(wakeword_models=[
+                                        os.path.join("openwakeword", "resources", "models", "alexa_v0.1.onnx")
+                                      ], inference_framework="onnx")
+
+        owwModel2 = openwakeword.Model(wakeword_models=[
+                                        os.path.join("openwakeword", "resources", "models", "alexa_v0.1.onnx")
+                                      ], inference_framework="onnx")
+
+        # Prediction on random data with integer multiples of standard chunk size (1280 samples)
+        predictions1 = owwModel1.predict_clip(os.path.join("tests", "data", "alexa_test.wav"), chunk_size=1280)
+        predictions2 = owwModel2.predict_clip(os.path.join("tests", "data", "alexa_test.wav"), chunk_size=1280*2)
+        np.testing.assert_approx_equal(max([i['alexa_v0.1'] for i in predictions1]), max([i['alexa_v0.1'] for i in predictions2]), 5)
+
+        # Prediction on data with a chunk size not an integer multiple of 1280
+        predictions1 = owwModel1.predict_clip(os.path.join("tests", "data", "alexa_test.wav"), chunk_size=1024)
+        predictions2 = owwModel2.predict_clip(os.path.join("tests", "data", "alexa_test.wav"), chunk_size=1024*2)
+        np.testing.assert_approx_equal(max([i['alexa_v0.1'] for i in predictions1]), max([i['alexa_v0.1'] for i in predictions2]), 5)
+
+        # Test with multiclass model
+        owwModel1 = openwakeword.Model(wakeword_models=["timer"], inference_framework="onnx")
+        owwModel2 = openwakeword.Model(wakeword_models=["timer"], inference_framework="onnx")
+
+        # Prediction on random data with integer multiples of standard chunk size (1280 samples)
+        predictions1 = owwModel1.predict_clip(os.path.join("tests", "data", "alexa_test.wav"), chunk_size=1280)
+        predictions2 = owwModel2.predict_clip(os.path.join("tests", "data", "alexa_test.wav"), chunk_size=1280*2)
+        assert abs(max([i['1_minute_timer'] for i in predictions1]) - max([i['1_minute_timer'] for i in predictions2])) < 0.00001
+
+        # Prediction on data with a chunk size not an integer multiple of 1280
+        predictions1 = owwModel1.predict_clip(os.path.join("tests", "data", "alexa_test.wav"), chunk_size=1024)
+        predictions2 = owwModel2.predict_clip(os.path.join("tests", "data", "alexa_test.wav"), chunk_size=1024*2)
+        assert abs(max([i['1_minute_timer'] for i in predictions1]) - max([i['1_minute_timer'] for i in predictions2])) < 0.00001
+
+    def test_exception_handling_for_inference_framework(self):
+        with mock.patch.dict(sys.modules, {'onnxruntime': None}):
+            with pytest.raises(ValueError):
+                openwakeword.Model(wakeword_models=[
+                                                os.path.join("openwakeword", "resources", "models", "alexa_v0.1.onnx")
+                                            ], inference_framework="onnx")
+
+        with mock.patch.dict(sys.modules, {'tflite_runtime': None}):
+            openwakeword.Model(wakeword_models=[
+                                            os.path.join("openwakeword", "resources", "models", "alexa_v0.1.tflite")
+                                        ], inference_framework="tflite")
+
+    def test_predict_with_custom_verifier_model(self):
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            # Train custom verifier model with random data
+            verifier_model = openwakeword.custom_verifier_model.train_verifier_model(np.random.random((2, 1536)), np.array([0, 1]))
+            pickle.dump(verifier_model, open(os.path.join(tmp_dir, "test_verifier.pkl"), "wb"))
+
+            # Load model with verifier
+            owwModel = openwakeword.Model(
+                wakeword_models=[os.path.join("openwakeword", "resources", "models", "alexa_v0.1.onnx")],
+                inference_framework="onnx",
+                custom_verifier_models={"alexa_v0.1": os.path.join(tmp_dir, "test_verifier.pkl")},
+                custom_verifier_threshold=0.0
+            )
+
+            owwModel.predict(np.random.randint(-1000, 1000, 1280).astype(np.int16))
+
+    def test_load_pretrained_model_by_name(self):
+        # Load model with defaults
+        owwModel = openwakeword.Model(wakeword_models=["alexa", "hey mycroft"], inference_framework="onnx")
+
+        owwModel = openwakeword.Model(wakeword_models=["alexa", "hey mycroft"], inference_framework="tflite")
 
         # Prediction on random data
         owwModel.predict(np.random.randint(-1000, 1000, 1280).astype(np.int16))
 
-        # Prediction on random data with different chunk size
-        owwModel.predict(np.random.randint(-1000, 1000, 1280*2).astype(np.int16))
-
     def test_custom_model_label_mapping_dict(self):
         # Load model with model path
-        owwModel = openwakeword.Model(wakeword_model_paths=[
+        owwModel = openwakeword.Model(wakeword_models=[
                                         os.path.join("openwakeword", "resources", "models", "alexa_v0.1.onnx")
                                       ],
-                                      class_mapping_dicts=[{"alexa_v0.1": {"0": "positive"}}]
+                                      class_mapping_dicts=[{"alexa_v0.1": {"0": "positive"}}],
+                                      inference_framework="onnx"
                                       )
 
         # Prediction on random data
@@ -95,31 +182,79 @@ class TestModels:
             assert 1 == 1
         else:
             # Load model with defaults
-            owwModel = openwakeword.Model(enable_speex_noise_suppression=True)
+            try:
+                owwModel = openwakeword.Model(enable_speex_noise_suppression=True)
 
-            # Get clips for each model (assumes that test clips will have the model name in the filename)
-            test_dict = {}
-            for mdl_name in owwModel.models.keys():
-                all_clips = [str(i) for i in Path(os.path.join("tests", "data")).glob("*.wav")]
-                test_dict[mdl_name] = [i for i in all_clips if mdl_name in i]
+                # Get clips for each model (assumes that test clips will have the model name in the filename)
+                test_dict = {}
+                for mdl_name in owwModel.models.keys():
+                    all_clips = [str(i) for i in Path(os.path.join("tests", "data")).glob("*.wav")]
+                    test_dict[mdl_name] = [i for i in all_clips if mdl_name in i]
 
-            # Predict
-            for model, clips in test_dict.items():
-                for clip in clips:
-                    # Get predictions for reach frame in the clip
-                    predictions = owwModel.predict_clip(clip)
-                    owwModel.reset()  # reset after each clip to ensure independent results
+                # Predict
+                for model, clips in test_dict.items():
+                    for clip in clips:
+                        # Get predictions for reach frame in the clip
+                        predictions = owwModel.predict_clip(clip)
+                        owwModel.reset()  # reset after each clip to ensure independent results
 
-                    # Make predictions dictionary flatter
-                    predictions_flat = collections.defaultdict(list)
-                    [predictions_flat[key].append(i[key]) for i in predictions for key in i.keys()]
+                        # Make predictions dictionary flatter
+                        predictions_flat = collections.defaultdict(list)
+                        [predictions_flat[key].append(i[key]) for i in predictions for key in i.keys()]
 
-                # Check scores against default threshold (0.5)
-                for key in predictions_flat.keys():
-                    if key in clip:
-                        assert max(predictions_flat[key]) >= 0.5
-                    else:
-                        assert max(predictions_flat[key]) < 0.5
+                    # Check scores against default threshold (0.5)
+                    for key in predictions_flat.keys():
+                        if key in clip:
+                            assert max(predictions_flat[key]) >= 0.5
+                        else:
+                            assert max(predictions_flat[key]) < 0.5
+            except ImportError:
+                logging.warning("Attemped to test Speex noise cancelling functionality, but the 'speexdsp_ns' library was not installed!"
+                                " If you want these tests to be run, install this library as shown in the openwakeword documentation."
+                                )
+                assert 1 == 1
+
+    def test_models_with_debounce(self):
+        # Load model with defaults
+        owwModel = openwakeword.Model()
+
+        # Predict with chunks of 1280 with and without debounce
+        predictions = owwModel.predict_clip(os.path.join("tests", "data", "alexa_test.wav"),
+                                            debounce_time=0, threshold={"alexa_v0.1": 0.5})
+        scores = np.array([i['alexa'] for i in predictions])
+
+        predictions = owwModel.predict_clip(os.path.join("tests", "data", "alexa_test.wav"),
+                                            debounce_time=1.25, threshold={"alexa": 0.5})
+        scores_with_debounce = np.array([i['alexa'] for i in predictions])
+        print(scores, scores_with_debounce)
+        assert (scores >= 0.5).sum() > 1
+        assert (scores_with_debounce >= 0.5).sum() == 1
+
+    def test_model_reset(self):
+        # Load the model
+        owwModel = openwakeword.Model()
+
+        # Get test clip and load it
+        clip = os.path.join("tests", "data", "alexa_test.wav")
+        with wave.open(clip, mode='rb') as f:
+            data = np.frombuffer(f.readframes(f.getnframes()), dtype=np.int16)
+
+        # Predict frame by frame
+        for i in range(0, len(data), 1280):
+            prediction = owwModel.predict(data[i:i+1280])
+            if prediction['alexa'] > 0.5:
+                break
+
+        # Assert that next prediction is still > 0.5
+        prediction = owwModel.predict(data[i:i+1280])
+        assert prediction['alexa'] > 0.5
+
+        # Reset the model
+        owwModel.reset()
+
+        # Assert that next prediction is < 0.5
+        prediction = owwModel.predict(data[i:i+1280])
+        assert prediction['alexa'] < 0.5
 
     def test_models_with_vad(self):
         # Load model with defaults
@@ -186,9 +321,10 @@ class TestModels:
         owwModel.get_parent_model_from_label(target_model_name)
 
     def test_get_positive_prediction_frames(self):
-        owwModel = openwakeword.Model()
+        owwModel = openwakeword.Model(wakeword_models=[
+                                        os.path.join("openwakeword", "resources", "models", "alexa_v0.1.onnx")
+                                      ], inference_framework="onnx")
 
-        # Get a clip to use for the test
-        clip = [str(i) for i in Path(os.path.join("tests", "data")).glob("*.wav")][0]
+        clip = os.path.join("tests", "data", "alexa_test.wav")
         features = owwModel._get_positive_prediction_frames(clip)
         assert list(features.values())[0].shape[0] > 0