From e717f1ba67e24c898e54cac74183400cc4044b1a Mon Sep 17 00:00:00 2001 From: Google Health Date: Fri, 4 Apr 2025 17:04:11 +0000 Subject: [PATCH 01/15] huggingface notebook for pytorch PiperOrigin-RevId: 743983873 --- .dockerignore | 4 + ...uick_start_with_hugging_face_pytorch.ipynb | 320 +++++++++++++++ python/data_processing/audio_utils.py | 382 ++++++++++++++++++ 3 files changed, 706 insertions(+) create mode 100644 .dockerignore create mode 100644 notebooks/quick_start_with_hugging_face_pytorch.ipynb create mode 100644 python/data_processing/audio_utils.py diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..d8ee3b7 --- /dev/null +++ b/.dockerignore @@ -0,0 +1,4 @@ +# Exclude files that are not needed in the serving image from the Docker build +# context. + +/python/data_processing/audio_utils.py \ No newline at end of file diff --git a/notebooks/quick_start_with_hugging_face_pytorch.ipynb b/notebooks/quick_start_with_hugging_face_pytorch.ipynb new file mode 100644 index 0000000..d2ad1b9 --- /dev/null +++ b/notebooks/quick_start_with_hugging_face_pytorch.ipynb @@ -0,0 +1,320 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "id": "1fgVWTMK9SNz" + }, + "source": [ + "~~~\n", + "Copyright 2025 Google LLC\n", + "\n", + "Licensed under the Apache License, Version 2.0 (the \"License\");\n", + "you may not use this file except in compliance with the License.\n", + "You may obtain a copy of the License at\n", + "\n", + " https://www.apache.org/licenses/LICENSE-2.0\n", + "\n", + "Unless required by applicable law or agreed to in writing, software\n", + "distributed under the License is distributed on an \"AS IS\" BASIS,\n", + "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n", + "See the License for the specific language governing permissions and\n", + "limitations under the License.\n", + "~~~\n", + "\n", + "\n", + "# Quick start with Hugging Face (PyTorch model)\n", + "\n", + "\u003ctable\u003e\u003ctbody\u003e\u003ctr\u003e\n", + " \u003ctd style=\"text-align: center\"\u003e\n", + " \u003ca href=\"https://colab.research.google.com/github/google-health/hear/blob/master/notebooks/quick_start_with_hugging_face_pytorch.ipynb\"\u003e\n", + " \u003cimg alt=\"Google Colab logo\" src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" width=\"32px\"\u003e\u003cbr\u003e Run in Google Colab\n", + " \u003c/a\u003e\n", + " \u003c/td\u003e \n", + " \u003ctd style=\"text-align: center\"\u003e\n", + " \u003ca href=\"https://github.com/google-health/hear/blob/master/notebooks/quick_start_with_hugging_face_pytorch.ipynb\"\u003e\n", + " \u003cimg alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"\u003e\u003cbr\u003e View on GitHub\n", + " \u003c/a\u003e\n", + " \u003c/td\u003e\n", + " \u003ctd style=\"text-align: center\"\u003e\n", + " \u003ca href=\"https://huggingface.co/google/hear-pytorch\"\u003e\n", + " \u003cimg alt=\"HuggingFace logo\" src=\"https://huggingface.co/front/assets/huggingface_logo-noborder.svg\" width=\"32px\"\u003e\u003cbr\u003e View on HuggingFace\n", + " \u003c/a\u003e\n", + " \u003c/td\u003e\n", + "\u003c/tr\u003e\u003c/tbody\u003e\u003c/table\u003e\n", + "\n", + "This Colab notebook provides a basic usage example of the HeAR encoder that generates a machine learning representation (known as \"embeddings\") from health-related sounds (2-second audio clips sampled at 16kHz). These embeddings can be used to develop custom machine learning models for health acoustic use-cases with less data and compute compared to traditional model development methods.\n", + "\n", + " Learn more about embeddings and their benefits at [this page](https://developers.google.com/health-ai-developer-foundations/hear)." + ] + }, + { + "metadata": { + "id": "EvQJsuAZ50-Z" + }, + "cell_type": "markdown", + "source": [ + "## Install dependencies" + ] + }, + { + "metadata": { + "id": "qdrXvhfo6dMV" + }, + "cell_type": "code", + "source": [ + "! git clone https://github.com/Google-Health/hear.git\n", + "! pip install --upgrade --quiet transformers==4.50.3" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": { + "id": "_uZFXCSuqr1V" + }, + "source": [ + "## Authenticate with HuggingFace, skip if you have a HF_TOKEN secret" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "9XpCla68-Iol" + }, + "outputs": [], + "source": [ + "from huggingface_hub.utils import HfFolder\n", + "\n", + "if HfFolder.get_token() is None:\n", + " from huggingface_hub import notebook_login\n", + " notebook_login()" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "SWckxRXiqc3L" + }, + "source": [ + "## Load and play cough audio recording" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "z2vEauhtvTan" + }, + "outputs": [], + "source": [ + "SAMPLE_RATE = 16000 # Samples per second (Hz)\n", + "CLIP_DURATION = 2 # Duration of the audio clip in seconds\n", + "CLIP_LENGTH = SAMPLE_RATE * CLIP_DURATION # Total number of samples" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "j81KORIqloq-" + }, + "outputs": [], + "source": [ + "import numpy as np\n", + "from scipy.io import wavfile\n", + "from scipy import signal\n", + "from IPython.display import Audio, display\n", + "\n", + "\n", + "def resample_audio_and_convert_to_mono(\n", + " audio_array: np.ndarray,\n", + " sampling_rate: float,\n", + " new_sampling_rate: float = SAMPLE_RATE,\n", + " ) -\u003e np.ndarray:\n", + " \"\"\"\n", + " Resamples an audio array to 16kHz and converts it to mono if it has multiple channels.\n", + "\n", + " Args:\n", + " audio_array: A numpy array representing the audio data.\n", + " sampling_rate: The original sampling rate of the audio.\n", + " new_sampling_rate: Target sampling rate.\n", + "\n", + " Returns:\n", + " resampled_audio_mono: A numpy array representing the resampled mono audio at 16kHz.\n", + " \"\"\"\n", + " # Convert to mono if it's multi-channel\n", + " if audio_array.ndim \u003e 1:\n", + " audio_mono = np.mean(audio_array, axis=1)\n", + " else:\n", + " audio_mono = audio_array\n", + "\n", + " # Resample\n", + " original_sample_count = audio_mono.shape[0]\n", + " new_sample_count = int(round(original_sample_count * (new_sampling_rate / sampling_rate)))\n", + " resampled_audio_mono = signal.resample(audio_mono, new_sample_count)\n", + "\n", + " return resampled_audio_mono\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "h1Lw92M1b7Dd" + }, + "outputs": [], + "source": [ + "!wget -nc https://upload.wikimedia.org/wikipedia/commons/b/be/Woman_coughing_three_times.wav" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "5W7BMKk9cerB" + }, + "outputs": [], + "source": [ + "# Load file\n", + "with open('Woman_coughing_three_times.wav', 'rb') as f:\n", + " original_sampling_rate, audio_array = wavfile.read(f)\n", + "\n", + "print(f\"Sample Rate: {original_sampling_rate} Hz\")\n", + "print(f\"Data Shape: {audio_array.shape}\")\n", + "print(f\"Data Type: {audio_array.dtype}\")\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "1l-eFOh0mq1r" + }, + "outputs": [], + "source": [ + "audio_array = resample_audio_and_convert_to_mono(audio_array, original_sampling_rate, SAMPLE_RATE)\n", + "display(Audio(audio_array, rate=SAMPLE_RATE))" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "fJIdGmpFqhi8" + }, + "source": [ + "## Compute embeddings" + ] + }, + { + "metadata": { + "id": "MIRuwbfz2a-l" + }, + "cell_type": "code", + "source": [ + "from transformers import ViTConfig, ViTModel\n", + "\n", + "\n", + "# Load the model directly from Hugging Face Hub\n", + "configuration = ViTConfig(\n", + " image_size=(192, 128),\n", + " hidden_size=1024,\n", + " num_hidden_layers=24,\n", + " num_attention_heads=16,\n", + " intermediate_size=1024 * 4,\n", + " hidden_act=\"gelu\",\n", + " hidden_dropout_prob=0.0,\n", + " attention_probs_dropout_prob=0.0,\n", + " initializer_range=0.02,\n", + " layer_norm_eps=1e-12,\n", + " pooled_dim=512,\n", + " patch_size=16,\n", + " num_channels=1,\n", + " qkv_bias=True,\n", + " encoder_stride=16,\n", + " pooler_act='linear',\n", + " pooler_output_size=512,\n", + ")\n", + "loaded_model = ViTModel.from_pretrained(\n", + " \"google/hear-pytorch\",\n", + " config=configuration\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "Fhk4GKkwsAMj" + }, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt\n", + "import torch\n", + "import importlib\n", + "audio_utils = importlib.import_module(\n", + " \"hear.python.data_processing.audio_utils\"\n", + ")\n", + "\n", + "mel_pcen = audio_utils.mel_pcen\n", + "\n", + "\n", + "# This index corresponds to a cough and was determined by hand. In practice, you\n", + "# would need a detector.\n", + "START = 0\n", + "\n", + "# Add batch dimension\n", + "input_tensor = np.expand_dims(audio_array[START: START + CLIP_LENGTH], axis=0)\n", + "\n", + "# Call inference\n", + "infer = lambda audio_array: loaded_model.forward(\n", + " mel_pcen(audio_array), return_dict=True, output_hidden_states=True)\n", + "output = infer(torch.Tensor(input_tensor))\n", + "\n", + "# Extract the embedding vector\n", + "embedding_vector = np.asarray(output.pooler_output.detach()).flatten()\n", + "print(\"Size of embedding vector:\", len(embedding_vector))\n", + "\n", + "# Plot the embedding vector\n", + "plt.figure(figsize=(12, 4))\n", + "plt.plot(embedding_vector)\n", + "plt.title('Embedding Vector')\n", + "plt.xlabel('Index')\n", + "plt.ylabel('Value')\n", + "plt.grid(True)\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "PHTxQttKYNpa" + }, + "source": [ + "# Next steps\n", + "\n", + "Explore the other [notebooks](https://github.com/google-health/hear/blob/master/notebooks) to learn what else you can do with the model." + ] + } + ], + "metadata": { + "colab": { + "name": "quick_start_with_hugging_face_pytorch.ipynb", + "private_outputs": true, + "provenance": [] + }, + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + }, + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 0 +} diff --git a/python/data_processing/audio_utils.py b/python/data_processing/audio_utils.py new file mode 100644 index 0000000..5d07a8f --- /dev/null +++ b/python/data_processing/audio_utils.py @@ -0,0 +1,382 @@ +# +# Copyright 2025 Google LLC +# +# 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. + +"""Utility functions for audio processing.""" + +import math +from typing import Callable + +import torch +import torch.nn.functional as F + + +def _enclosing_power_of_two(value: int) -> int: + """Calculates the smallest power of 2 greater than or equal to `value`.""" + return int(2**math.ceil(math.log2(value))) if value > 0 else 1 + + +def _compute_stft( + signals: torch.Tensor, + frame_length: int, + frame_step: int, + fft_length: int | None = None, + window_fn: Callable[[int], torch.Tensor] | None = torch.hann_window, + pad_end: bool = True, +) -> torch.Tensor: + """Computes the Short-time Fourier Transform of `signals`. + + Args: + signals: A `[..., samples]` `torch.Tensor` of real-valued signals. + frame_length: The window length in samples. + frame_step: The number of samples to step. + fft_length: The size of the FFT to apply. If not provided, uses the smallest + power of 2 enclosing `frame_length`. + window_fn: A callable that takes a window length and returns a + `[window_length]` `torch.Tensor` of samples. If set to `None`, no + windowing is used. + pad_end: Whether to pad the end of `signals` with zeros when the provided + frame length and step produces a frame that lies partially past its end. + + Returns: + A `[..., frames, fft_unique_bins]` `torch.Tensor` of `complex64` + STFT values where `fft_unique_bins` is `fft_length // 2 + 1` (the unique + components of the FFT). + + Raises: + ValueError: If `signals` is not at least rank 1, `frame_length` is + not scalar, or `frame_step` is not scalar. + """ + if signals.ndim < 1: + raise ValueError( + f'Input signals must have rank at least 1, got rank {signals.ndim}' + ) + if not isinstance(frame_length, int): + raise ValueError( + f'frame_length must be an integer scalar, got type {type(frame_length)}' + ) + if not isinstance(frame_step, int): + raise ValueError( + f'frame_step must be an integer scalar, got type {type(frame_step)}' + ) + + if fft_length is None: + fft_length = _enclosing_power_of_two(frame_length) + elif not isinstance(fft_length, int): + raise ValueError( + 'fft_length must be an integer scalar or None, got type' + f' {type(fft_length)}' + ) + + if pad_end: + n_frames = ( + math.ceil(signals.shape[-1] / frame_step) + if signals.shape[-1] > 0 + else 0 + ) + padded_length = ( + max(0, (n_frames - 1) * frame_step + frame_length) + if n_frames > 0 + else frame_length + ) + padding_needed = max( + 0, padded_length - signals.shape[-1] + ) + if padding_needed > 0: + signals = F.pad(signals, (0, padding_needed)) + + framed_signals = signals.unfold(-1, frame_length, frame_step) + + if framed_signals.shape[-2] == 0: + return torch.empty( + *signals.shape[:-1], + 0, + fft_length // 2 + 1, + dtype=torch.complex64, + device=signals.device, + ) + + # Optionally window the framed signals. + if window_fn is not None: + window = ( + window_fn(frame_length) + .to(framed_signals.device) + .to(framed_signals.dtype) + ) # Ensure window is on the same device and dtype + framed_signals = framed_signals * window + + # torch.fft.rfft produces the (fft_length/2 + 1) unique components of the + # FFT of the real windowed signals in framed_signals. + return torch.fft.rfft(framed_signals, n=fft_length, dim=-1) + + +def _ema( + inputs: torch.Tensor, + num_channels: int, + smooth_coef: float, + initial_state: torch.Tensor | None = None, +) -> torch.Tensor: + """Exponential Moving Average (EMA). + + Args: + inputs: A 3D tensor of shape (batch_size, timesteps, input_dim). + This is the input sequence to the EMA RNN. + num_channels: The number of channels/units in the EMA. + smooth_coef: The smoothing coefficient (alpha) for the EMA. + initial_state: (Optional) A 2D tensor of shape (batch_size, num_channels) + representing the initial state for the EMA. If None, it defaults to zeros. + + Returns: + A 3D tensor of shape (batch_size, timesteps, num_channels) representing + the EMA output sequence. + """ + batch_size, timesteps, _ = inputs.shape + + if initial_state is None: + ema_state = torch.zeros( + (batch_size, num_channels), dtype=torch.float32, device=inputs.device + ) + else: + ema_state = initial_state + + identity_kernel_gain = smooth_coef + identity_recurrent_gain = 1.0 - smooth_coef + + identity_kernel = ( + torch.eye(num_channels, dtype=torch.float32, device=inputs.device) + * identity_kernel_gain + ) + identity_recurrent_kernel = ( + torch.eye(num_channels, dtype=torch.float32, device=inputs.device) + * identity_recurrent_gain + ) + + output_sequence = [] + + start = initial_state is not None + if start: + output_sequence.append(ema_state) + for t in range(start, timesteps): + current_input = inputs[:, t, :] # Shape (batch_size, num_channels) + + # EMA update formula: + output = torch.matmul(current_input, identity_kernel) + torch.matmul( + ema_state, identity_recurrent_kernel + ) + + ema_state = output + output_sequence.append(output) + + # Shape (batch_size, timesteps, num_channels) + output_sequence = torch.stack(output_sequence, dim=1) + return output_sequence + + +def _pcen_function( + inputs: torch.Tensor, + num_channels: int = 128, + alpha: float = 0.8, + smooth_coef: float = 0.04, + delta: float = 2.0, + root: float = 2.0, + floor: float = 1e-8, +) -> torch.Tensor: + """Per-Channel Energy Normalization as a function. + + This applies a fixed normalization by an exponential moving + average smoother, and a compression. + See https://arxiv.org/abs/1607.05666 for more details. + + Args: + inputs: A `[..., num_channels]` `torch.Tensor` of input signals. + num_channels: Number of channels + alpha: Exponent of EMA smoother + smooth_coef: Smoothing coefficient of EMA + delta: Bias added before compression + root: One over exponent applied for compression (r in the paper) + floor: Offset added to EMA smoother + + Returns: + A `[..., num_channels]` `torch.Tensor` of PCEN values. + """ + alpha_param = torch.ones(num_channels) * alpha + delta_param = torch.ones(num_channels) * delta + root_param = torch.ones(num_channels) * root + + alpha_param = ( + torch.minimum(alpha_param, torch.ones_like(alpha_param)) + .to(inputs.device) + .to(inputs.dtype) + ) + root_param = ( + torch.maximum(root_param, torch.ones_like(root_param)) + .to(inputs.device) + .to(inputs.dtype) + ) + ema_smoother = _ema( + inputs, + num_channels=num_channels, + smooth_coef=smooth_coef, + # Handle cases where input is 2D or 3D + initial_state=inputs[:, 0] if inputs.ndim > 1 else None, + ) + + one_over_root = 1. / root_param + output = ((inputs / (floor + ema_smoother)**alpha_param + delta_param) + **one_over_root - delta_param**one_over_root) + return output + + +def _hertz_to_mel(frequencies_hertz: torch.Tensor) -> torch.Tensor: + """Scale filter frequencies to mel scale. + + https://en.wikipedia.org/wiki/Mel_scale + + Args: + frequencies_hertz: A float Tensor of frequencies in Hertz. + + Returns: + A Tensor of the same shape but in mel-scale. + """ + return 2595.0 * torch.log10(1.0 + frequencies_hertz / 700.0) + + +def _linear_to_mel_weight_matrix( + num_mel_bins: int = 128, + num_spectrogram_bins: int = 201, + sample_rate: float = 16000, + lower_edge_hertz: float = 0.0, + upper_edge_hertz: float = 8000.0, + dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """Returns a matrix to warp linear scale spectrograms to the mel scale. + + This function is a PyTorch version of the TensorFlow function with the same + name. See the TensorFlow documentation for detailed explanation and usage + https://www.tensorflow.org/api_docs/python/tf/signal/linear_to_mel_weight_matrix + + Args: + num_mel_bins: How many bands in the resulting mel spectrum. + num_spectrogram_bins: How many bins there are in the + source spectrogram data. + sample_rate: Samples per second of the input signal used to create the + spectrogram. + lower_edge_hertz: Lower bound on the frequencies to be included in the mel + spectrum. + upper_edge_hertz: The desired top edge of the highest frequency band. + dtype: The torch.dtype of the result matrix. Must be a floating point type. + + Returns: + A torch.Tensor of shape `[num_spectrogram_bins, num_mel_bins]`. + + Raises: + ValueError: If `num_mel_bins`/`num_spectrogram_bins`/`sample_rate` are not + positive, `lower_edge_hertz` is negative, frequency edges are incorrectly + ordered, `upper_edge_hertz` is larger than the Nyquist frequency. + """ + + if num_mel_bins <= 0: + raise ValueError(f'num_mel_bins must be positive. Got: {num_mel_bins}.') + if num_spectrogram_bins <= 0: + raise ValueError( + f'num_spectrogram_bins must be positive. Got: {num_spectrogram_bins}.' + ) + if sample_rate <= 0: + raise ValueError(f'sample_rate must be positive. Got: {sample_rate}.') + if lower_edge_hertz < 0.0: + raise ValueError( + f'lower_edge_hertz must be non-negative. Got: {lower_edge_hertz}.' + ) + if lower_edge_hertz >= upper_edge_hertz: + raise ValueError( + 'lower_edge_hertz must be smaller than upper_edge_hertz. Got: ' + f'lower_edge_hertz={lower_edge_hertz}, ' + f'upper_edge_hertz={upper_edge_hertz}.' + ) + if upper_edge_hertz > sample_rate / 2.0: + raise ValueError( + 'upper_edge_hertz must not be larger than the Nyquist frequency' + f'({sample_rate / 2.0}). Got: upper_edge_hertz={upper_edge_hertz}.' + ) + + sample_rate_tensor = torch.tensor(sample_rate, dtype=dtype) + lower_edge_hertz_tensor = torch.tensor(lower_edge_hertz, dtype=dtype) + upper_edge_hertz_tensor = torch.tensor(upper_edge_hertz, dtype=dtype) + zero = torch.tensor(0.0, dtype=dtype) + + # HTK excludes the spectrogram DC bin. + bands_to_zero = 1 + nyquist_hertz = sample_rate_tensor / 2.0 + linear_frequencies = torch.linspace( + zero, nyquist_hertz, num_spectrogram_bins, dtype=dtype)[bands_to_zero:] + spectrogram_bins_mel = _hertz_to_mel(linear_frequencies).unsqueeze(1) + + # Compute num_mel_bins triples of (lower_edge, center, upper_edge). + # The center of each band is the lower and upper edge of the adjacent bands. + # Accordingly, we divide [lower_edge_hertz, upper_edge_hertz] into + # num_mel_bins + 2 pieces. + band_edges_mel = torch.linspace( + _hertz_to_mel(lower_edge_hertz_tensor), + _hertz_to_mel(upper_edge_hertz_tensor), + num_mel_bins + 2, dtype=dtype) + # Create frames of size 3 with stride 1 + band_edges_mel = band_edges_mel.unfold(0, 3, 1) + + # Split the triples up and reshape them into [1, num_mel_bins] tensors. + lower_edge_mel = band_edges_mel[:, 0].unsqueeze(0) + center_mel = band_edges_mel[:, 1].unsqueeze(0) + upper_edge_mel = band_edges_mel[:, 2].unsqueeze(0) + + # Calculate lower and upper slopes for every spectrogram bin. + # Line segments are linear in the mel domain, not Hertz. + lower_slopes = (spectrogram_bins_mel - lower_edge_mel) / ( + center_mel - lower_edge_mel) + upper_slopes = (upper_edge_mel - spectrogram_bins_mel) / ( + upper_edge_mel - center_mel) + + # Intersect the line segments with each other and zero. + mel_weights_matrix = torch.maximum( + zero, torch.minimum(lower_slopes, upper_slopes)) + + # Re-add the zeroed lower bins we sliced out above. + return F.pad( + mel_weights_matrix, (0, 0, bands_to_zero, 0), mode='constant', value=0.0) + + +def mel_pcen( + x: torch.Tensor, +) -> torch.Tensor: + """Melspec followed by pcen.""" + x = x.float() + # Scale to -1, 1 range + x -= torch.min(x) + x = x / (torch.max(x) + 1e-8) + x = (x * 2) - 1 + + frame_length = 16 * 25 + frame_step = 160 + + stft = _compute_stft( + x, + frame_length=frame_length, + fft_length=frame_length, + frame_step=frame_step, + window_fn=torch.hann_window, + pad_end=True, + ) + spectrograms = torch.square(torch.abs(stft)) + + mel_transform = _linear_to_mel_weight_matrix() + mel_spectrograms = torch.matmul(spectrograms, mel_transform) + return _pcen_function(mel_spectrograms) From be4b61d2b25c73121b010346be913857b1cb7f5a Mon Sep 17 00:00:00 2001 From: tiffchen Date: Fri, 4 Apr 2025 17:32:52 +0000 Subject: [PATCH 02/15] Add notebooks README PiperOrigin-RevId: 743993650 --- notebooks/README.md | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) create mode 100644 notebooks/README.md diff --git a/notebooks/README.md b/notebooks/README.md new file mode 100644 index 0000000..7e6403f --- /dev/null +++ b/notebooks/README.md @@ -0,0 +1,19 @@ +# HeAR Notebooks + +* Quick start with Hugging Face - + Example of encoding a health-related audio clip into an + embedding vector by running the model locally from Hugging Face. There are + different versions of this notebook for the following model formats: + + * [TensorFlow](quick_start_with_hugging_face.ipynb) + * [PyTorch](quick_start_with_hugging_face_pytorch.ipynb) + +* [Quick start with Vertex Model Garden](quick_start_with_model_garden.ipynb) - + Example of serving the model on + [Vertex AI](https://cloud.google.com/vertex-ai/docs/predictions/overview) + and using Vertex AI APIs to encode health-related audio to + embeddings in online or batch workflows. + +* [Train a data efficient classifier](train_data_efficient_classifier.ipynb) - + Example of using the generated embeddings to train a custom classifier with + less data and compute. From d658034b0b7e1559ab2d38350197c0a5c8c8433f Mon Sep 17 00:00:00 2001 From: tiffchen Date: Fri, 4 Apr 2025 21:10:05 +0000 Subject: [PATCH 03/15] Remove pytorch notebook. PiperOrigin-RevId: 744065300 --- notebooks/README.md | 8 +- ...uick_start_with_hugging_face_pytorch.ipynb | 320 ------------------ 2 files changed, 2 insertions(+), 326 deletions(-) delete mode 100644 notebooks/quick_start_with_hugging_face_pytorch.ipynb diff --git a/notebooks/README.md b/notebooks/README.md index 7e6403f..f5a7220 100644 --- a/notebooks/README.md +++ b/notebooks/README.md @@ -1,12 +1,8 @@ # HeAR Notebooks -* Quick start with Hugging Face - +* [Quick start with Hugging Face](quick_start_with_hugging_face.ipynb) - Example of encoding a health-related audio clip into an - embedding vector by running the model locally from Hugging Face. There are - different versions of this notebook for the following model formats: - - * [TensorFlow](quick_start_with_hugging_face.ipynb) - * [PyTorch](quick_start_with_hugging_face_pytorch.ipynb) + embedding vector by running the model locally from Hugging Face. * [Quick start with Vertex Model Garden](quick_start_with_model_garden.ipynb) - Example of serving the model on diff --git a/notebooks/quick_start_with_hugging_face_pytorch.ipynb b/notebooks/quick_start_with_hugging_face_pytorch.ipynb deleted file mode 100644 index d2ad1b9..0000000 --- a/notebooks/quick_start_with_hugging_face_pytorch.ipynb +++ /dev/null @@ -1,320 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": { - "id": "1fgVWTMK9SNz" - }, - "source": [ - "~~~\n", - "Copyright 2025 Google LLC\n", - "\n", - "Licensed under the Apache License, Version 2.0 (the \"License\");\n", - "you may not use this file except in compliance with the License.\n", - "You may obtain a copy of the License at\n", - "\n", - " https://www.apache.org/licenses/LICENSE-2.0\n", - "\n", - "Unless required by applicable law or agreed to in writing, software\n", - "distributed under the License is distributed on an \"AS IS\" BASIS,\n", - "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n", - "See the License for the specific language governing permissions and\n", - "limitations under the License.\n", - "~~~\n", - "\n", - "\n", - "# Quick start with Hugging Face (PyTorch model)\n", - "\n", - "\u003ctable\u003e\u003ctbody\u003e\u003ctr\u003e\n", - " \u003ctd style=\"text-align: center\"\u003e\n", - " \u003ca href=\"https://colab.research.google.com/github/google-health/hear/blob/master/notebooks/quick_start_with_hugging_face_pytorch.ipynb\"\u003e\n", - " \u003cimg alt=\"Google Colab logo\" src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" width=\"32px\"\u003e\u003cbr\u003e Run in Google Colab\n", - " \u003c/a\u003e\n", - " \u003c/td\u003e \n", - " \u003ctd style=\"text-align: center\"\u003e\n", - " \u003ca href=\"https://github.com/google-health/hear/blob/master/notebooks/quick_start_with_hugging_face_pytorch.ipynb\"\u003e\n", - " \u003cimg alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"\u003e\u003cbr\u003e View on GitHub\n", - " \u003c/a\u003e\n", - " \u003c/td\u003e\n", - " \u003ctd style=\"text-align: center\"\u003e\n", - " \u003ca href=\"https://huggingface.co/google/hear-pytorch\"\u003e\n", - " \u003cimg alt=\"HuggingFace logo\" src=\"https://huggingface.co/front/assets/huggingface_logo-noborder.svg\" width=\"32px\"\u003e\u003cbr\u003e View on HuggingFace\n", - " \u003c/a\u003e\n", - " \u003c/td\u003e\n", - "\u003c/tr\u003e\u003c/tbody\u003e\u003c/table\u003e\n", - "\n", - "This Colab notebook provides a basic usage example of the HeAR encoder that generates a machine learning representation (known as \"embeddings\") from health-related sounds (2-second audio clips sampled at 16kHz). These embeddings can be used to develop custom machine learning models for health acoustic use-cases with less data and compute compared to traditional model development methods.\n", - "\n", - " Learn more about embeddings and their benefits at [this page](https://developers.google.com/health-ai-developer-foundations/hear)." - ] - }, - { - "metadata": { - "id": "EvQJsuAZ50-Z" - }, - "cell_type": "markdown", - "source": [ - "## Install dependencies" - ] - }, - { - "metadata": { - "id": "qdrXvhfo6dMV" - }, - "cell_type": "code", - "source": [ - "! git clone https://github.com/Google-Health/hear.git\n", - "! pip install --upgrade --quiet transformers==4.50.3" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "markdown", - "metadata": { - "id": "_uZFXCSuqr1V" - }, - "source": [ - "## Authenticate with HuggingFace, skip if you have a HF_TOKEN secret" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "9XpCla68-Iol" - }, - "outputs": [], - "source": [ - "from huggingface_hub.utils import HfFolder\n", - "\n", - "if HfFolder.get_token() is None:\n", - " from huggingface_hub import notebook_login\n", - " notebook_login()" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "SWckxRXiqc3L" - }, - "source": [ - "## Load and play cough audio recording" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "z2vEauhtvTan" - }, - "outputs": [], - "source": [ - "SAMPLE_RATE = 16000 # Samples per second (Hz)\n", - "CLIP_DURATION = 2 # Duration of the audio clip in seconds\n", - "CLIP_LENGTH = SAMPLE_RATE * CLIP_DURATION # Total number of samples" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "j81KORIqloq-" - }, - "outputs": [], - "source": [ - "import numpy as np\n", - "from scipy.io import wavfile\n", - "from scipy import signal\n", - "from IPython.display import Audio, display\n", - "\n", - "\n", - "def resample_audio_and_convert_to_mono(\n", - " audio_array: np.ndarray,\n", - " sampling_rate: float,\n", - " new_sampling_rate: float = SAMPLE_RATE,\n", - " ) -\u003e np.ndarray:\n", - " \"\"\"\n", - " Resamples an audio array to 16kHz and converts it to mono if it has multiple channels.\n", - "\n", - " Args:\n", - " audio_array: A numpy array representing the audio data.\n", - " sampling_rate: The original sampling rate of the audio.\n", - " new_sampling_rate: Target sampling rate.\n", - "\n", - " Returns:\n", - " resampled_audio_mono: A numpy array representing the resampled mono audio at 16kHz.\n", - " \"\"\"\n", - " # Convert to mono if it's multi-channel\n", - " if audio_array.ndim \u003e 1:\n", - " audio_mono = np.mean(audio_array, axis=1)\n", - " else:\n", - " audio_mono = audio_array\n", - "\n", - " # Resample\n", - " original_sample_count = audio_mono.shape[0]\n", - " new_sample_count = int(round(original_sample_count * (new_sampling_rate / sampling_rate)))\n", - " resampled_audio_mono = signal.resample(audio_mono, new_sample_count)\n", - "\n", - " return resampled_audio_mono\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "h1Lw92M1b7Dd" - }, - "outputs": [], - "source": [ - "!wget -nc https://upload.wikimedia.org/wikipedia/commons/b/be/Woman_coughing_three_times.wav" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "5W7BMKk9cerB" - }, - "outputs": [], - "source": [ - "# Load file\n", - "with open('Woman_coughing_three_times.wav', 'rb') as f:\n", - " original_sampling_rate, audio_array = wavfile.read(f)\n", - "\n", - "print(f\"Sample Rate: {original_sampling_rate} Hz\")\n", - "print(f\"Data Shape: {audio_array.shape}\")\n", - "print(f\"Data Type: {audio_array.dtype}\")\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "1l-eFOh0mq1r" - }, - "outputs": [], - "source": [ - "audio_array = resample_audio_and_convert_to_mono(audio_array, original_sampling_rate, SAMPLE_RATE)\n", - "display(Audio(audio_array, rate=SAMPLE_RATE))" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "fJIdGmpFqhi8" - }, - "source": [ - "## Compute embeddings" - ] - }, - { - "metadata": { - "id": "MIRuwbfz2a-l" - }, - "cell_type": "code", - "source": [ - "from transformers import ViTConfig, ViTModel\n", - "\n", - "\n", - "# Load the model directly from Hugging Face Hub\n", - "configuration = ViTConfig(\n", - " image_size=(192, 128),\n", - " hidden_size=1024,\n", - " num_hidden_layers=24,\n", - " num_attention_heads=16,\n", - " intermediate_size=1024 * 4,\n", - " hidden_act=\"gelu\",\n", - " hidden_dropout_prob=0.0,\n", - " attention_probs_dropout_prob=0.0,\n", - " initializer_range=0.02,\n", - " layer_norm_eps=1e-12,\n", - " pooled_dim=512,\n", - " patch_size=16,\n", - " num_channels=1,\n", - " qkv_bias=True,\n", - " encoder_stride=16,\n", - " pooler_act='linear',\n", - " pooler_output_size=512,\n", - ")\n", - "loaded_model = ViTModel.from_pretrained(\n", - " \"google/hear-pytorch\",\n", - " config=configuration\n", - ")" - ], - "outputs": [], - "execution_count": null - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "Fhk4GKkwsAMj" - }, - "outputs": [], - "source": [ - "import matplotlib.pyplot as plt\n", - "import torch\n", - "import importlib\n", - "audio_utils = importlib.import_module(\n", - " \"hear.python.data_processing.audio_utils\"\n", - ")\n", - "\n", - "mel_pcen = audio_utils.mel_pcen\n", - "\n", - "\n", - "# This index corresponds to a cough and was determined by hand. In practice, you\n", - "# would need a detector.\n", - "START = 0\n", - "\n", - "# Add batch dimension\n", - "input_tensor = np.expand_dims(audio_array[START: START + CLIP_LENGTH], axis=0)\n", - "\n", - "# Call inference\n", - "infer = lambda audio_array: loaded_model.forward(\n", - " mel_pcen(audio_array), return_dict=True, output_hidden_states=True)\n", - "output = infer(torch.Tensor(input_tensor))\n", - "\n", - "# Extract the embedding vector\n", - "embedding_vector = np.asarray(output.pooler_output.detach()).flatten()\n", - "print(\"Size of embedding vector:\", len(embedding_vector))\n", - "\n", - "# Plot the embedding vector\n", - "plt.figure(figsize=(12, 4))\n", - "plt.plot(embedding_vector)\n", - "plt.title('Embedding Vector')\n", - "plt.xlabel('Index')\n", - "plt.ylabel('Value')\n", - "plt.grid(True)\n", - "plt.show()" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "PHTxQttKYNpa" - }, - "source": [ - "# Next steps\n", - "\n", - "Explore the other [notebooks](https://github.com/google-health/hear/blob/master/notebooks) to learn what else you can do with the model." - ] - } - ], - "metadata": { - "colab": { - "name": "quick_start_with_hugging_face_pytorch.ipynb", - "private_outputs": true, - "provenance": [] - }, - "kernelspec": { - "display_name": "Python 3", - "name": "python3" - }, - "language_info": { - "name": "python" - } - }, - "nbformat": 4, - "nbformat_minor": 0 -} From f5e3e491f8aa60391a8524b32d29ace4195fafa0 Mon Sep 17 00:00:00 2001 From: Google Health Date: Mon, 7 Apr 2025 13:35:24 +0000 Subject: [PATCH 04/15] Refactor audio preprocessing PiperOrigin-RevId: 744702267 --- python/data_processing/audio_utils.py | 159 +++++++++++++++++++++++--- 1 file changed, 141 insertions(+), 18 deletions(-) diff --git a/python/data_processing/audio_utils.py b/python/data_processing/audio_utils.py index 5d07a8f..30677cd 100644 --- a/python/data_processing/audio_utils.py +++ b/python/data_processing/audio_utils.py @@ -18,13 +18,15 @@ import math from typing import Callable +import numpy as np +from scipy import signal import torch import torch.nn.functional as F def _enclosing_power_of_two(value: int) -> int: """Calculates the smallest power of 2 greater than or equal to `value`.""" - return int(2**math.ceil(math.log2(value))) if value > 0 else 1 + return int(2 ** math.ceil(math.log2(value))) if value > 0 else 1 def _compute_stft( @@ -90,9 +92,7 @@ def _compute_stft( if n_frames > 0 else frame_length ) - padding_needed = max( - 0, padded_length - signals.shape[-1] - ) + padding_needed = max(0, padded_length - signals.shape[-1]) if padding_needed > 0: signals = F.pad(signals, (0, padding_needed)) @@ -130,8 +130,8 @@ def _ema( """Exponential Moving Average (EMA). Args: - inputs: A 3D tensor of shape (batch_size, timesteps, input_dim). - This is the input sequence to the EMA RNN. + inputs: A 3D tensor of shape (batch_size, timesteps, input_dim). This is the + input sequence to the EMA RNN. num_channels: The number of channels/units in the EMA. smooth_coef: The smoothing coefficient (alpha) for the EMA. initial_state: (Optional) A 2D tensor of shape (batch_size, num_channels) @@ -232,9 +232,10 @@ def _pcen_function( initial_state=inputs[:, 0] if inputs.ndim > 1 else None, ) - one_over_root = 1. / root_param - output = ((inputs / (floor + ema_smoother)**alpha_param + delta_param) - **one_over_root - delta_param**one_over_root) + one_over_root = 1.0 / root_param + output = ( + inputs / (floor + ema_smoother) ** alpha_param + delta_param + ) ** one_over_root - delta_param**one_over_root return output @@ -268,8 +269,8 @@ def _linear_to_mel_weight_matrix( Args: num_mel_bins: How many bands in the resulting mel spectrum. - num_spectrogram_bins: How many bins there are in the - source spectrogram data. + num_spectrogram_bins: How many bins there are in the source spectrogram + data. sample_rate: Samples per second of the input signal used to create the spectrogram. lower_edge_hertz: Lower bound on the frequencies to be included in the mel @@ -319,7 +320,8 @@ def _linear_to_mel_weight_matrix( bands_to_zero = 1 nyquist_hertz = sample_rate_tensor / 2.0 linear_frequencies = torch.linspace( - zero, nyquist_hertz, num_spectrogram_bins, dtype=dtype)[bands_to_zero:] + zero, nyquist_hertz, num_spectrogram_bins, dtype=dtype + )[bands_to_zero:] spectrogram_bins_mel = _hertz_to_mel(linear_frequencies).unsqueeze(1) # Compute num_mel_bins triples of (lower_edge, center, upper_edge). @@ -329,7 +331,9 @@ def _linear_to_mel_weight_matrix( band_edges_mel = torch.linspace( _hertz_to_mel(lower_edge_hertz_tensor), _hertz_to_mel(upper_edge_hertz_tensor), - num_mel_bins + 2, dtype=dtype) + num_mel_bins + 2, + dtype=dtype, + ) # Create frames of size 3 with stride 1 band_edges_mel = band_edges_mel.unfold(0, 3, 1) @@ -341,20 +345,24 @@ def _linear_to_mel_weight_matrix( # Calculate lower and upper slopes for every spectrogram bin. # Line segments are linear in the mel domain, not Hertz. lower_slopes = (spectrogram_bins_mel - lower_edge_mel) / ( - center_mel - lower_edge_mel) + center_mel - lower_edge_mel + ) upper_slopes = (upper_edge_mel - spectrogram_bins_mel) / ( - upper_edge_mel - center_mel) + upper_edge_mel - center_mel + ) # Intersect the line segments with each other and zero. mel_weights_matrix = torch.maximum( - zero, torch.minimum(lower_slopes, upper_slopes)) + zero, torch.minimum(lower_slopes, upper_slopes) + ) # Re-add the zeroed lower bins we sliced out above. return F.pad( - mel_weights_matrix, (0, 0, bands_to_zero, 0), mode='constant', value=0.0) + mel_weights_matrix, (0, 0, bands_to_zero, 0), mode='constant', value=0.0 + ) -def mel_pcen( +def _mel_pcen( x: torch.Tensor, ) -> torch.Tensor: """Melspec followed by pcen.""" @@ -380,3 +388,118 @@ def mel_pcen( mel_transform = _linear_to_mel_weight_matrix() mel_spectrograms = torch.matmul(spectrograms, mel_transform) return _pcen_function(mel_spectrograms) + + +def _torch_resize_bilinear_tf_compat(images, size): + """PyTorch implementation of tf.image.resize. + + Internally matches the numerical output of: + ``` + tf.image.resize( + # TF input needs HWC/BHWC format + tf_permuted_input, + size, + method=tf.image.ResizeMethod.BILINEAR, + preserve_aspect_ratio=False, + antialias=False + ) + ``` + + Args: + images: Input tensor, shape [C, H, W] or [B, C, H, W]. Should be float32 + or convertible to float32 for numerical matching. + size: Target size as [new_height, new_width]. + + Returns: + Resized tensor, shape [C, new_H, new_W] or [B, C, new_H, new_W], dtype + torch.float32. + """ + original_dims = images.dim() + new_height, new_width = size + + if original_dims not in [3, 4]: + raise ValueError('Input tensor must be 3D [C, H, W] or 4D [B, C, H, W]') + + if len(size) != 2: + raise ValueError( + 'size must be a tuple or list of 2 integers (new_height, new_width)' + ) + + images = images.to(torch.float32) + + was_3d = False + if original_dims == 3: + images = images.unsqueeze(0) # Shape: [1, C, H, W] + was_3d = True + + resized_images_bchw = F.interpolate( + images, # Shape: [B, C, H, W] + size=(new_height, new_width), + mode='bilinear', + align_corners=False, + antialias=False, + ) + # Output shape: [B, C, new_H, new_W] + + # Remove batch dimension if original input was 3D (1CNHWC -> CNHWC) + if was_3d: + # Shape: [C, new_H, new_W] + resized_images_bchw = resized_images_bchw.squeeze(0) + + return resized_images_bchw + + +def preprocess_audio(audio: torch.Tensor) -> torch.Tensor: + """Preprocesses audio. + + Args: + audio: A `[..., samples]` `torch.Tensor` of real-valued signals. Represents + batched audio clips of duration 2s sampled at 16kHz. + + Returns: + A `[..., 1, 192, 128]` `torch.Tensor` of mel-pcen values. + """ + if audio.ndim != 2: + raise ValueError( + f'Input audio must have rank 2, got rank {audio.ndim}' + ) + if audio.shape[1] != 32000: + raise ValueError( + f'Input audio must have 32000 samples, got {audio.shape[1]}' + ) + spectrogram = _mel_pcen(audio) + # Add a channel dimension. Torch images have format [B, C, H, W]. + spectrogram = torch.unsqueeze(spectrogram, dim=1) + return _torch_resize_bilinear_tf_compat(spectrogram, size=(192, 128)) + + +def resample_audio_and_convert_to_mono( + audio_array: np.ndarray, + sampling_rate: float, + new_sampling_rate: float, +) -> np.ndarray: + """Resamples an audio array to 16kHz and converts it to mono. + + Args: + audio_array: A numpy array representing the audio data. + sampling_rate: The original sampling rate of the audio. + new_sampling_rate: Target sampling rate. + + Returns: + resampled_audio_mono: A numpy array representing the resampled mono audio at + 16kHz. + """ + # Convert to mono if it's multi-channel + if audio_array.ndim > 1: + audio_mono = np.mean(audio_array, axis=1) + else: + audio_mono = audio_array + + # Resample + original_sample_count = audio_mono.shape[0] + new_sample_count = int( + round(original_sample_count * (new_sampling_rate / sampling_rate)) + ) + resampled_audio_mono = signal.resample(audio_mono, new_sample_count) + + return resampled_audio_mono From 3141c98fb8e7a8818796e40549e7954f9f2c7763 Mon Sep 17 00:00:00 2001 From: Google Health Date: Mon, 7 Apr 2025 21:14:22 +0000 Subject: [PATCH 05/15] handle case where input is too short and can be padded PiperOrigin-RevId: 744846066 --- python/data_processing/audio_utils.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/python/data_processing/audio_utils.py b/python/data_processing/audio_utils.py index 30677cd..84dc350 100644 --- a/python/data_processing/audio_utils.py +++ b/python/data_processing/audio_utils.py @@ -463,7 +463,10 @@ def preprocess_audio(audio: torch.Tensor) -> torch.Tensor: raise ValueError( f'Input audio must have rank 2, got rank {audio.ndim}' ) - if audio.shape[1] != 32000: + if audio.shape[1] < 32000: + n = 32000 - audio.shape[1] + audio = torch.nn.functional.pad(audio, pad=(0, n), mode='constant', value=0) + elif audio.shape[1] > 32000: raise ValueError( f'Input audio must have 32000 samples, got {audio.shape[1]}' ) From 1c28ae0279b8335dc06b05225303d5d162c9828c Mon Sep 17 00:00:00 2001 From: Google Health Date: Thu, 10 Apr 2025 21:24:11 +0000 Subject: [PATCH 06/15] Notebook showing how to use HeAR embedding model with event detectors PiperOrigin-RevId: 746173637 --- notebooks/hear_event_detector_demo.ipynb | 638 +++++++++++++++++++++++ 1 file changed, 638 insertions(+) create mode 100644 notebooks/hear_event_detector_demo.ipynb diff --git a/notebooks/hear_event_detector_demo.ipynb b/notebooks/hear_event_detector_demo.ipynb new file mode 100644 index 0000000..227bc33 --- /dev/null +++ b/notebooks/hear_event_detector_demo.ipynb @@ -0,0 +1,638 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "id": "Mt64XBHqJiYo" + }, + "source": [ + "~~~markdown\n", + "Copyright 2025 Google LLC\n", + "\n", + "Licensed under the Apache License, Version 2.0 (the \"License\");\n", + "you may not use this file except in compliance with the License.\n", + "You may obtain a copy of the License at\n", + "\n", + " https://www.apache.org/licenses/LICENSE-2.0\n", + "\n", + "Unless required by applicable law or agreed to in writing, software\n", + "distributed under the License is distributed on an \"AS IS\" BASIS,\n", + "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n", + "See the License for the specific language governing permissions and\n", + "limitations under the License.\n", + "~~~\n", + "\n", + "# HeAR Event Detector Demo\n", + "\u003ctable\u003e\u003ctbody\u003e\u003ctr\u003e\n", + " \u003ctd style=\"text-align: center\"\u003e\n", + " \u003ca href=\"https://colab.research.google.com/github/google-health/hear/blob/master/notebooks/hear_event_detector_demo.ipynb\"\u003e\n", + " \u003cimg alt=\"Google Colab logo\" src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" width=\"32px\"\u003e\u003cbr\u003e Run in Google Colab\n", + " \u003c/a\u003e\n", + " \u003c/td\u003e \n", + " \u003ctd style=\"text-align: center\"\u003e\n", + " \u003ca href=\"https://github.com/google-health/hear/blob/master/notebooks/hear_event_detector_demo.ipynb\"\u003e\n", + " \u003cimg alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"\u003e\u003cbr\u003e View on GitHub\n", + " \u003c/a\u003e\n", + " \u003c/td\u003e\n", + " \u003ctd style=\"text-align: center\"\u003e\n", + " \u003ca href=\"https://huggingface.co/google/hear\"\u003e\n", + " \u003cimg alt=\"HuggingFace logo\" src=\"https://huggingface.co/front/assets/huggingface_logo-noborder.svg\" width=\"32px\"\u003e\u003cbr\u003e View on HuggingFace\n", + " \u003c/a\u003e\n", + " \u003c/td\u003e\n", + "\u003c/tr\u003e\u003c/tbody\u003e\u003c/table\u003e\n", + "\n", + "\n", + "This Colab notebook demonstrates using the HeAR (Health Acoustic Representations) model along with the included Health Event Detectors directly from Hugging Face, to identify audio clips with relevent health sounds such as coughing, breathing or sneezing, then create and utilize embeddings from this subset of health-related audio clips.\n", + "\n", + "This notebook is similar to `train_data_efficient_classifier.ipynb` and also uses the small [Wikimedia Commons](https://commons.wikimedia.org/wiki/Commons:Welcome) dataset of relevant health sounds. In this example the audio files are reduced to a smaller subset of clips using the event detector to identify clips containing interesting health sounds to embed with HeAR.\n", + "\n", + "\n", + "\n", + "#### This notebook demonstrates:\n", + "\n", + "1. Loading all supported Hugging Face Models (HeAR, Event Detector and Frontend).\n", + "\n", + "2. Detecting 2-second clips within the [Wikimedia Commons](https://commons.wikimedia.org/wiki/Commons:Welcome) dataset with high probability of containing one or more of the supported event detection labels, then generating HeAR embedddings for these clips.\n", + "\n", + "3. Finding the most similar audio files to a given query audio file based on the Cosine Similarity between the respective HeAR embeddings of the audio files.\n", + "\n", + "4. Optimizing the event detectors and frontend for low latency on-device usage using TFLite to support large scale feature event detection or feature generation" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "_uZFXCSuqr1V" + }, + "source": [ + "# Authenticate with HuggingFace, skip if you have a HF_TOKEN secret" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "l-5Tj0uqS3dI" + }, + "outputs": [], + "source": [ + "from huggingface_hub.utils import HfFolder\n", + "\n", + "if HfFolder.get_token() is None:\n", + " from huggingface_hub import notebook_login\n", + " notebook_login()" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "mubK5MmPQcUS" + }, + "source": [ + "# Clone HuggingFace repository snapshot\n", + "\n", + "This will store the HeAR and event detector models in local cache so they can be loaded later.\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "WjfuNYeZQJdc" + }, + "outputs": [], + "source": [ + "import numpy as np\n", + "import tensorflow as tf\n", + "\n", + "print(\"TensorFlow version:\", tf.__version__)\n", + "print(\"Keras version:\", tf.keras.__version__)\n", + "\n", + "from huggingface_hub import snapshot_download\n", + "hugging_face_repo = \"google/hear\"\n", + "local_snapshot_path = snapshot_download(repo_id=hugging_face_repo)\n", + "print(f\"Saved {hugging_face_repo} to {local_snapshot_path}\\n\")\n" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "mO-Z5BOtj3D1" + }, + "source": [ + "# Download Audio Data" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "NBZnJwzzj3D1" + }, + "source": [ + " Wiki Commons\n", + "https://commons.wikimedia.org/wiki/Category:Coughing_audio\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "tJ55XsJsj3D2" + }, + "outputs": [], + "source": [ + "# @title Download Public Domain Cough Examples to Notebook\n", + "import os\n", + "import subprocess\n", + "from urllib.parse import urlparse\n", + "\n", + "# More examples: https://commons.wikimedia.org/wiki/Category:Coughing_audio\n", + "wiki_cough_file_urls = [\n", + " 'https://upload.wikimedia.org/wikipedia/commons/c/cc/Man_coughing.ogg',\n", + " 'https://upload.wikimedia.org/wikipedia/commons/6/6a/Cough_1.ogg',\n", + " 'https://upload.wikimedia.org/wikipedia/commons/d/d9/Cough_2.ogg',\n", + " 'https://upload.wikimedia.org/wikipedia/commons/b/be/Woman_coughing_three_times.wav',\n", + " 'https://upload.wikimedia.org/wikipedia/commons/d/d0/Sneezing.ogg',\n", + " 'https://upload.wikimedia.org/wikipedia/commons/e/ef/Laughter_and_clearing_voice.ogg',\n", + " 'https://upload.wikimedia.org/wikipedia/commons/c/c6/Laughter.ogg',\n", + " 'https://upload.wikimedia.org/wikipedia/commons/1/1c/Knocking_on_wood_or_door.ogg',\n", + "]\n", + "\n", + "# Download the files.\n", + "files_map = {} # file name to file path map\n", + "for url in wiki_cough_file_urls:\n", + " filename = os.path.basename(urlparse(url).path)\n", + " print(f'Downloading {filename}...')\n", + " res = subprocess.run(['wget', '-nv', '-O', filename, url], capture_output=True, text=True)\n", + " if res.returncode != 0:\n", + " print(f\" Download failed. Return code: {res.returncode}\\nError: {res.stderr}\")\n", + " files_map[filename] = url\n", + "print(f'\\nLocal Files:\\n{os.listdir():}\\n')" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "K2Eyu4VIsbUB" + }, + "source": [ + "# Load Models and Run Inference\n", + "\n", + "### HeAR Model\n", + "\n", + "The HeAR model uses a powerful [ViT](https://huggingface.co/docs/transformers/en/model_doc/vit) backbone and generates rich 512 length embeddings from a 2 second single channel 16kHz audio clip. See the [HuggingFace Model Card](https://huggingface.co/google/hear) for more details.\n", + "\n", + "### Event Detector Models\n", + "\n", + "The event detector models use the efficient [MobileNet-V3](https://huggingface.co/docs/timm/en/models/mobilenet-v3) backbone paired with our custom TensorFlow spectrogram frontend.\n", + "\n", + "As with HeAR, these models expect a 2 second single channel 16kHz audio clip and output 8 detection probability scores for the following labels:\n", + "\n", + "```\n", + "['Cough', 'Snore', 'Baby Cough', 'Breathe', 'Sneeze','Throat Clear', 'Laugh', 'Speech']\n", + "```\n", + "\n", + "The event detector has two size variants which can be used interchangeably depending on the use-case.\n", + "\n", + "* `event_detector_small/` is based on [MobileNetV3Small](https://www.tensorflow.org/api_docs/python/tf/keras/applications/MobileNetV3Small) with approximately **1M** parameters (3.60 MB)\n", + "\n", + "* `event_detector_large/` is based on [MobileNetV3Large](https://www.tensorflow.org/api_docs/python/tf/keras/applications/MobileNetV3Large) with approximately **3M** parameters (11.46 MB)\n", + "\n", + "#### Spectrogram Frontend\n", + "\n", + "Our event detectors are fused with a custom, on-device optimized spectrogram frontend which efficiently converts 2 seconds of 16kHz audio into [PCEN](https://research.google/pubs/trainable-frontend-for-robust-and-far-field-keyword-spotting) scaled [Mel-spectrogram](https://huggingface.co/learn/audio-course/en/chapter1/audio_data#mel-spectrogram) features with 200 time steps and 48 Mel-frequency bins.\n", + "\n", + "* `spectrogram_frontend/` is based on the PCEN implementation from [LEAF](https://research.google/blog/leaf-a-learnable-frontend-for-audio-classification/) and has only **5k** non-trainable parameters (18.56 KB) which are frozen and not configurable.\n", + "\n", + "We provide a standalone version of this frontend so that features can be pre-computed for new event detector training applications. See **Extract Batch Frontend Spectrogram Features** section for usage examples.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "cQVAxkBHZ_JO" + }, + "outputs": [], + "source": [ + "# @title Load HeAR, Event Detector and Frontend Models\n", + "from huggingface_hub import from_pretrained_keras\n", + "\n", + "# Constants for all included models, each input should be 32,000 samples.\n", + "SAMPLE_RATE = 16000\n", + "CLIP_DURATION = 2\n", + "\n", + "# Select event detector variant\n", + "EVENT_DETECTOR = \"event_detector_small\" # @param [\"event_detector_large\", \"event_detector_small\"]\n", + "# Included event detectors are trained to doutput detection probabilities in this order.\n", + "LABEL_LIST = ['Cough', 'Snore', 'Baby Cough', 'Breathe', 'Sneeze', 'Throat Clear', 'Laugh', 'Speech']\n", + "\n", + "# HeAR Embedding Model\n", + "print(f\"\\nLoading HeAR model\")\n", + "hear_model = from_pretrained_keras(local_snapshot_path)\n", + "hear_infer = hear_model.signatures[\"serving_default\"]\n", + "\n", + "# Event detector models and frontend are nested in the \"event_detector/\" folder.\n", + "# Detector Frontend Model for efficiently computing spectrogram feature\n", + "frontend_path = os.path.join(\"event_detector/\", \"spectrogram_frontend\")\n", + "print(f\"\\nLoading frontend model from: {frontend_path}\")\n", + "frontend_model = from_pretrained_keras(\n", + " os.path.join(local_snapshot_path, frontend_path)\n", + ")\n", + "\n", + "# Detector Model based on size variant selection, frontend included\n", + "event_detector_path = os.path.join(\"event_detector/\", EVENT_DETECTOR)\n", + "print(f\"\\nLoading detector model from: {event_detector_path}\")\n", + "event_detector = from_pretrained_keras(\n", + " os.path.join(local_snapshot_path, event_detector_path)\n", + ")\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "33G4zKJHjGGc" + }, + "outputs": [], + "source": [ + "# @title Plot Helpers\n", + "import librosa\n", + "import matplotlib.pyplot as plt\n", + "import librosa.display\n", + "from IPython.display import Audio\n", + "import matplotlib.cm as cm\n", + "import warnings\n", + "\n", + "# Suppress the specific warning\n", + "warnings.filterwarnings(\"ignore\", category=UserWarning, module=\"soundfile\")\n", + "warnings.filterwarnings(\"ignore\", module=\"librosa\")\n", + "\n", + "def plot_waveform(sound, sr, title, figsize=(12, 4), color='blue', alpha=0.7):\n", + " \"\"\"Plots the waveform of the audio using librosa.display.\"\"\"\n", + " plt.figure(figsize=figsize)\n", + " librosa.display.waveshow(sound, sr=sr, color=color, alpha=alpha)\n", + " plt.title(f\"{title}\\nshape={sound.shape}, sr={sr}, dtype={sound.dtype}\")\n", + " plt.xlabel(\"Time (s)\")\n", + " plt.ylabel(\"Amplitude\")\n", + " plt.grid(True)\n", + " plt.tight_layout()\n", + " plt.show()\n", + "\n", + "\n", + "def plot_spectrogram(sound, sr, title, figsize=(12, 4), n_fft=2048, hop_length=256, n_mels=128, cmap='nipy_spectral'):\n", + " \"\"\"Plots the Mel spectrogram of the audio using librosa.\"\"\"\n", + " plt.figure(figsize=figsize)\n", + " mel_spectrogram = librosa.feature.melspectrogram(y=sound, sr=sr, n_fft=n_fft, hop_length=hop_length, n_mels=n_mels)\n", + " log_mel_spectrogram = librosa.power_to_db(mel_spectrogram, ref=np.max)\n", + " librosa.display.specshow(log_mel_spectrogram, sr=sr, hop_length=hop_length, x_axis='time', y_axis='mel', cmap=cmap)\n", + " plt.title(f\"{title} - Mel Spectrogram\")\n", + " plt.tight_layout()\n", + " plt.show()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "3dJKyHXznPgg" + }, + "outputs": [], + "source": [ + "# @title Event Detector Plot Helpers\n", + "\n", + "def plot_frontend_feature(\n", + " frontend_feature: np.ndarray,\n", + " title: str,\n", + " figsize: tuple[int, int] = (12, 4),\n", + " cmap: str = 'nipy_spectral',\n", + ") -\u003e None:\n", + " \"\"\"Plots the frontend spectrogram input feature.\n", + "\n", + " Args:\n", + " frontend_feature: The event detector frontend feature as a 2D NumPy array\n", + " with shape (number of time steps, number of frequency bins). The default\n", + " shape for the included event detectors is (200, 48), which represents a\n", + " 2-second audio clip with 48 frequency bins.\n", + " title: The title prefix of the plot.\n", + " figsize: Optional size of the figure.\n", + " cmap: Optional colormap to use.\n", + " \"\"\"\n", + " # Frontend features are typically rotated when fed into the model\n", + " # for spectrogram visualization it is more standard for x axis to be time\n", + " audio_spectrogram = np.rot90(frontend_feature)\n", + " plt.figure(figsize=figsize)\n", + " plt.imshow(audio_spectrogram, aspect='auto', cmap=cmap)\n", + " plt.title(f\"{title} - Frontend PCEN Mel Spectrogram\")\n", + " plt.tight_layout()\n", + " plt.show()\n", + "\n", + "def plot_detection_scores(\n", + " scores_batch: np.ndarray,\n", + " label_list: list[str],\n", + " title: str,\n", + " figsize: tuple[int, int] = (12, 4),\n", + " cmap: str = 'nipy_spectral',\n", + ") -\u003e None:\n", + " \"\"\"Plots per-label detection scores for batch sequentially with a consistent color scale.\n", + "\n", + " Args:\n", + " scores_batch: The event detection scores as a 2D NumPy array with shape\n", + " (number of clips, number of labels). Where number of labels should match\n", + " the length of the label_list, which is 8 for the included event detectors.\n", + " label_list: A list of labels representing the event detector classes.\n", + " title: The title prefix of the plot.\n", + " figsize: Optional size of the figure.\n", + " cmap: Optional colormap to use.\n", + " \"\"\"\n", + " plt.figure(figsize=figsize)\n", + " scores_img = np.transpose(scores_batch)\n", + " # Explicitly set the color limits for imshow\n", + " im = plt.imshow(scores_img, aspect='auto', cmap=cmap, vmin=0, vmax=1)\n", + " # Set up the 'y' label axis\n", + " plt.yticks(\n", + " np.arange(len(label_list)), [l.replace(' ', '\\n') for l in label_list]\n", + " )\n", + " # Add horizontal grid lines between labels\n", + " for i in range(1, scores_img.shape[0]):\n", + " plt.axhline(y=i - 0.5, color='gray', linestyle='--')\n", + " plt.grid(axis='y', which='major', color='white', alpha=0)\n", + " # Setup the 'x' time axis\n", + " n_clips = scores_img.shape[1]\n", + " plt.xticks(np.arange(n_clips), [f'Clip {i+1}' for i in range(n_clips)])\n", + " plt.xlabel(\"Time Step\")\n", + " # Add vertical grid lines between time steps\n", + " for j in range(1, n_clips):\n", + " plt.axvline(x=j - 0.5, color='gray', linestyle='--')\n", + " plt.title(f\"{title} - Sound Event Detections\")\n", + " # Add colorbar with a consistent scale from 0 to 1\n", + " plt.colorbar(im, ticks=[0, 0.2, 0.4, 0.6, 0.8, 1.0])\n", + " plt.tight_layout()\n", + " plt.show()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "KYbz904QXio2" + }, + "outputs": [], + "source": [ + "# @title Load Audio and Generate HeAR Embeddings\n", + "%%time\n", + "\n", + "# Audio display options\n", + "SHOW_WAVEFORM = False\n", + "SHOW_SPECTROGRAM = False\n", + "SHOW_PLAYER = True\n", + "SHOW_DETECTION_SCORES = True\n", + "COLORMAP = \"Blues\"\n", + "\n", + "# Keep clips with high detection scores for these respiratory labels\n", + "# then embed the clips using HeAR. Ignore clips with no detections\n", + "LABELS_TO_EMBED = ['Cough', 'Snore', 'Breathe', 'Sneeze']\n", + "# Assert that all labels to embed are actually present in the main label list\n", + "assert all(label in LABEL_LIST for label in LABELS_TO_EMBED)\n", + "\n", + "# Clips of length CLIP_DURATION seconds are extracted from the audio file\n", + "# using a sliding window. Adjecent clips are overlapped by CLIP_OVERLAP_PERCENT.\n", + "CLIP_OVERLAP_PERCENT = 10\n", + "\n", + "# Labels must have score above this threshold to be considered a detection\n", + "DETECTION_THRESHOLD = 0.9\n", + "\n", + "frame_length = int(CLIP_DURATION * SAMPLE_RATE)\n", + "frame_step = int(frame_length * (1 - CLIP_OVERLAP_PERCENT / 100))\n", + "hear_embeddings = {}\n", + "for file_key, file_url in files_map.items():\n", + " hear_embeddings[file_key] = {}\n", + " print(f\"\\nLoading file: {file_key} from {file_url}\")\n", + " audio, sample_rate = librosa.load(file_key, sr=SAMPLE_RATE, mono=True)\n", + "\n", + " # Display full audio file (optional).\n", + " if SHOW_WAVEFORM:\n", + " plot_waveform(audio, sample_rate, title=file_key, color='blue')\n", + " if SHOW_SPECTROGRAM:\n", + " plot_spectrogram(\n", + " audio, sample_rate, title=file_key, n_fft=2*1024, hop_length=64, n_mels=256, cmap=COLORMAP)\n", + " if SHOW_PLAYER:\n", + " display(Audio(data=audio, rate=sample_rate))\n", + "\n", + " # Segment an audio array into fixed length overlapping clips.\n", + " if len(audio) \u003c frame_length:\n", + " audio = np.pad(audio, (0, frame_length - len(audio)), mode='constant')\n", + " audio_clip_batch = tf.signal.frame(audio, frame_length, frame_step )\n", + " print(f\"Number of audio clips in batch: {len(audio_clip_batch)}.\")\n", + "\n", + " # Perform detector inference on the audio_clip_batch\n", + " # The model will generate the input feature, then infer the detection\n", + " print(f\"Running batched {EVENT_DETECTOR} model inference on audio clips.\")\n", + " detection_scores_batch = event_detector(audio_clip_batch)[\"scores\"].numpy()\n", + " print(\"Computed batch probability scores with shape:\",\n", + " f\"{detection_scores_batch.shape} from input audio clips batch with\",\n", + " f\"shape: {audio_clip_batch.shape}\"\n", + " )\n", + " hear_embeddings[file_key]['detections'] = detection_scores_batch\n", + "\n", + " if SHOW_DETECTION_SCORES:\n", + " plot_detection_scores(detection_scores_batch, LABEL_LIST, title=f'{file_key}: {EVENT_DETECTOR}', cmap=COLORMAP)\n", + "\n", + " # Filter clips for HeAR inference based on if in ANY detection scores for\n", + " # 'LABELS_TO_EMBED' are above the 'DETECTION_THRESHOLD'.\n", + " print(f\"Filtering clips based on detections for labels: {LABELS_TO_EMBED}\")\n", + " embed_hear_clips = []\n", + " for clip_i, scores in enumerate(detection_scores_batch):\n", + " for label_index, label in enumerate(LABEL_LIST):\n", + " if label in LABELS_TO_EMBED and scores[label_index] \u003e DETECTION_THRESHOLD:\n", + " embed_hear_clips.append(audio_clip_batch[clip_i])\n", + " break\n", + "\n", + " # Perform HeAR batch inference to extract the associated clip embedding.\n", + " # Only run inference on 'embed_hear_clips' which have have a high detection\n", + " # score for one of the 'LABELS_TO_EMBED'.\n", + " if len(embed_hear_clips):\n", + " print(f\"Computing HeAR embedding for batch of\",\n", + " f\"{len(embed_hear_clips)} selected clips.\")\n", + " hear_embedding_batch = hear_infer(x=np.asarray(embed_hear_clips))[ 'output_0'].numpy()\n", + " print(f\"Embedding batch shape: {hear_embedding_batch.shape},\",\n", + " f\"data type: {hear_embedding_batch.dtype}\")\n", + " else:\n", + " hear_embedding_batch = np.array([])\n", + " print(f\"None of the {len(audio_clip_batch)} clips in {file_key} have\",\n", + " f\"detections above the threshold: {DETECTION_THRESHOLD} for the\",\n", + " f\"labels: {LABELS_TO_EMBED}\")\n", + " hear_embeddings[file_key]['embeddings'] = hear_embedding_batch" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "PFLdq8WI3AJT" + }, + "outputs": [], + "source": [ + "# @title Use HeAR embeddings to find most similar file to query file\n", + "from scipy.spatial import distance\n", + "\n", + "# Set up query file and make sure if the query file_key exists in the dictionary and has embeddings.\n", + "query_file_key = 'Cough_1.ogg'\n", + "assert query_file_key in hear_embeddings and len(hear_embeddings[query_file_key]['embeddings'])\n", + "\n", + "# Get the average embedding for the query file and compare similarity to the average embedding for the other files.\n", + "query_embedding = np.mean(hear_embeddings[query_file_key]['embeddings'], axis=0)\n", + "similarities = {}\n", + "for file_key, model_outputs in hear_embeddings.items():\n", + " # Skip comparing file_key to itself or comparing to keys without HeAR embeddings.\n", + " if file_key == query_file_key or not len(model_outputs['embeddings']):\n", + " continue\n", + "\n", + " # Compute cosine similarity between the query file and the current file.\n", + " current_embedding = np.mean(model_outputs['embeddings'], axis=0)\n", + " similarities[file_key] = 1 - distance.cosine(query_embedding, current_embedding)\n", + "\n", + "# Find the top N most similar entries\n", + "N = 3\n", + "top_N_similar = dict(sorted(similarities.items(), key=lambda item: item[1], reverse=True)[:N])\n", + "print(f\"\\nTop {N} most similar entries to '{query_file_key}':\")\n", + "for key, similarity in top_N_similar.items():\n", + " print(f\" {key}: {similarity:.3f}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "q4LoPrUCs0F4" + }, + "source": [ + "# Extract Batch Frontend Spectrogram Features\n", + "\n", + "We provide a standalone, frozen frontend spectrogram model for efficient generation of [PCEN Mel-Spectrogram](https://research.google/pubs/trainable-frontend-for-robust-and-far-field-keyword-spotting/) features, which are used by the event detectors. This model is comprised of non-trainable TensorFlow operations, prioritizing portability, scalability, and optimization at the cost of configurability.\n", + "\n", + "The frontend can also be used as a standalone feature extractor for generating large amounts of training data for finetuning or retraining additional models as demonstrated below.\n", + "\n", + "The input must be 2-seconds and corresponding output feature will have shape `(200, 48)`. Typically, input clips will be segmented with some amount of overlap to avoid distorting sounds near the boundry." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "DLDc2OFAngcT" + }, + "outputs": [], + "source": [ + "# @title Plot Example Frontend Features\n", + "example_file_key = 'Sneezing.ogg'\n", + "audio, sample_rate = librosa.load(example_file_key, sr=SAMPLE_RATE, mono=True)\n", + "print(f\"Loaded {example_file_key} with duration {len(audio)/sample_rate:0.2f}s\")\n", + "\n", + "# Extract example clips with 50% overlap, then generate frontend feature batch.\n", + "frame_length = int(CLIP_DURATION * SAMPLE_RATE)\n", + "audio_clip_batch = tf.signal.frame(audio, frame_length, frame_length // 2 )\n", + "frontend_feature_batch = frontend_model(audio_clip_batch)\n", + "print(f\"Generated input feature batch from {example_file_key} with shape: {frontend_feature_batch.shape}\")\n", + "for clip_i, frontend_feature in enumerate(frontend_feature_batch):\n", + " plot_frontend_feature(frontend_feature.numpy(), title=f'Clip: {clip_i}', cmap=COLORMAP, figsize=(8, 3))" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "_9e_m3Wf0ClK" + }, + "source": [ + "# Convert Event Detectors to TFLite\n", + "\n", + "The included event detectors are based on [MobileNet-V3](https://huggingface.co/docs/timm/en/models/mobilenet-v3) and are designed to be run on-device with low latency and power requirements while also allowing for finetuning and quick training.\n", + "\n", + "Converting the event detector models to [LiteRT](https://ai.google.dev/edge/litert) (formally TFLite) allows for more optimal performance on specific hardware such as mobile devices or compute constrained servers.\n", + "\n", + "In order to save compute in realtime or large scale applications, these hardware optimized event detectors can be used as a gate for HeAR, only embedding clips detected to contain a health sound of interest (as shown above).\n", + "\n", + "The code below demonstrates converting the loaded event detector to a TensorFlow LiteRT model, there are many more options available to quantize or furthur optimize the converted models, see the [documentation](https://ai.google.dev/edge/litert/models/convert_tf) for more information." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "I5UimMm15FMy" + }, + "outputs": [], + "source": [ + "# @title TFLite Conversion\n", + "%%time\n", + "\n", + "def convert_to_tflite(\n", + " model: tf.keras.Model,\n", + " quantize: bool = False,\n", + ") -\u003e bytes:\n", + " \"\"\"Converts a SavedModel model to a TensorFlow Lite (TFLite) model.\n", + "\n", + " Args:\n", + " model: The model to convert.\n", + " quantize: If True, apply dynamic range quantization to optimize the model.\n", + "\n", + " Returns:\n", + " The raw byte representation of the converted TFLite model.\n", + " \"\"\"\n", + " converter = tf.lite.TFLiteConverter.from_keras_model(model)\n", + " converter.target_spec.supported_ops = [\n", + " tf.lite.OpsSet.TFLITE_BUILTINS,\n", + " tf.lite.OpsSet.SELECT_TF_OPS, # needed for frontend ops\n", + " ]\n", + " # See documentation for quantization options beyond dyanmic range quantization.\n", + " # https://ai.google.dev/edge/litert/models/post_training_quantization\n", + " if quantize:\n", + " converter.optimizations = [tf.lite.Optimize.DEFAULT]\n", + " return converter.convert()\n", + "\n", + "# Convert event detector to TFLite and save result.\n", + "event_detector_lite = convert_to_tflite(event_detector, quantize=False)\n", + "tflite_output_path ='event_detector.tflite'\n", + "with open(tflite_output_path, 'wb') as f:\n", + " f.write(event_detector_lite)\n", + "print(f\"Saved TFLite model to: {tflite_output_path}\")\n", + "\n", + "# Initalize TFLite Model and print I/O specification.\n", + "tflite_interp = tf.lite.Interpreter(model_content=event_detector_lite)\n", + "print(f\"Input details:\\n {tflite_interp.get_input_details()[0]}\")\n", + "print(f\"Output details: \\n {tflite_interp.get_output_details()[0]}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "PHTxQttKYNpa" + }, + "source": [ + "# Next steps\n", + "\n", + "Explore the other [notebooks](https://github.com/google-health/hear/blob/master/notebooks) to learn what else you can do with the model." + ] + } + ], + "metadata": { + "colab": { + "last_runtime": { + "build_target": "//experimental/health_foundation_models/colab:colab_deps", + "kind": "private" + }, + "name": "hear_event_detector_demo.ipynb", + "private_outputs": true, + "provenance": [] + }, + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + }, + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 0 +} From 71eaeefa0b9661c17e438c55d96a9ee04cba7fb3 Mon Sep 17 00:00:00 2001 From: Google Health Date: Thu, 10 Apr 2025 22:53:33 +0000 Subject: [PATCH 07/15] add HeAR pytorch notebook PiperOrigin-RevId: 746205713 --- notebooks/README.md | 8 +- ...uick_start_with_hugging_face_pytorch.ipynb | 286 ++++++++++++++++++ 2 files changed, 292 insertions(+), 2 deletions(-) create mode 100644 notebooks/quick_start_with_hugging_face_pytorch.ipynb diff --git a/notebooks/README.md b/notebooks/README.md index f5a7220..7e6403f 100644 --- a/notebooks/README.md +++ b/notebooks/README.md @@ -1,8 +1,12 @@ # HeAR Notebooks -* [Quick start with Hugging Face](quick_start_with_hugging_face.ipynb) - +* Quick start with Hugging Face - Example of encoding a health-related audio clip into an - embedding vector by running the model locally from Hugging Face. + embedding vector by running the model locally from Hugging Face. There are + different versions of this notebook for the following model formats: + + * [TensorFlow](quick_start_with_hugging_face.ipynb) + * [PyTorch](quick_start_with_hugging_face_pytorch.ipynb) * [Quick start with Vertex Model Garden](quick_start_with_model_garden.ipynb) - Example of serving the model on diff --git a/notebooks/quick_start_with_hugging_face_pytorch.ipynb b/notebooks/quick_start_with_hugging_face_pytorch.ipynb new file mode 100644 index 0000000..b5bd621 --- /dev/null +++ b/notebooks/quick_start_with_hugging_face_pytorch.ipynb @@ -0,0 +1,286 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "id": "1fgVWTMK9SNz" + }, + "source": [ + "~~~\n", + "Copyright 2025 Google LLC\n", + "\n", + "Licensed under the Apache License, Version 2.0 (the \"License\");\n", + "you may not use this file except in compliance with the License.\n", + "You may obtain a copy of the License at\n", + "\n", + " https://www.apache.org/licenses/LICENSE-2.0\n", + "\n", + "Unless required by applicable law or agreed to in writing, software\n", + "distributed under the License is distributed on an \"AS IS\" BASIS,\n", + "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n", + "See the License for the specific language governing permissions and\n", + "limitations under the License.\n", + "~~~\n", + "\n", + "\n", + "# Quick start with Hugging Face (PyTorch model)\n", + "\n", + "\u003ctable\u003e\u003ctbody\u003e\u003ctr\u003e\n", + " \u003ctd style=\"text-align: center\"\u003e\n", + " \u003ca href=\"https://colab.research.google.com/github/google-health/hear/blob/master/notebooks/quick_start_with_hugging_face_pytorch.ipynb\"\u003e\n", + " \u003cimg alt=\"Google Colab logo\" src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" width=\"32px\"\u003e\u003cbr\u003e Run in Google Colab\n", + " \u003c/a\u003e\n", + " \u003c/td\u003e \n", + " \u003ctd style=\"text-align: center\"\u003e\n", + " \u003ca href=\"https://github.com/google-health/hear/blob/master/notebooks/quick_start_with_hugging_face_pytorch.ipynb\"\u003e\n", + " \u003cimg alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"\u003e\u003cbr\u003e View on GitHub\n", + " \u003c/a\u003e\n", + " \u003c/td\u003e\n", + " \u003ctd style=\"text-align: center\"\u003e\n", + " \u003ca href=\"https://huggingface.co/google/hear-pytorch\"\u003e\n", + " \u003cimg alt=\"HuggingFace logo\" src=\"https://huggingface.co/front/assets/huggingface_logo-noborder.svg\" width=\"32px\"\u003e\u003cbr\u003e View on HuggingFace\n", + " \u003c/a\u003e\n", + " \u003c/td\u003e\n", + "\u003c/tr\u003e\u003c/tbody\u003e\u003c/table\u003e\n", + "\n", + "This Colab notebook provides a basic usage example of the HeAR encoder that generates a machine learning representation (known as \"embeddings\") from health-related sounds (2-second audio clips sampled at 16kHz). These embeddings can be used to develop custom machine learning models for health acoustic use-cases with less data and compute compared to traditional model development methods.\n", + "\n", + " Learn more about embeddings and their benefits at [this page](https://developers.google.com/health-ai-developer-foundations/hear)." + ] + }, + { + "metadata": { + "id": "EvQJsuAZ50-Z" + }, + "cell_type": "markdown", + "source": [ + "## Install dependencies" + ] + }, + { + "metadata": { + "id": "qdrXvhfo6dMV" + }, + "cell_type": "code", + "source": [ + "! git clone https://github.com/Google-Health/hear.git\n", + "! pip install --upgrade --quiet transformers==4.50.3" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": { + "id": "_uZFXCSuqr1V" + }, + "source": [ + "## Authenticate with HuggingFace, skip if you have a HF_TOKEN secret" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "9XpCla68-Iol" + }, + "outputs": [], + "source": [ + "from huggingface_hub.utils import HfFolder\n", + "\n", + "if HfFolder.get_token() is None:\n", + " from huggingface_hub import notebook_login\n", + " notebook_login()" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "SWckxRXiqc3L" + }, + "source": [ + "## Load and play cough audio recording" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "z2vEauhtvTan" + }, + "outputs": [], + "source": [ + "SAMPLE_RATE = 16000 # Samples per second (Hz)\n", + "CLIP_DURATION = 2 # Duration of the audio clip in seconds\n", + "CLIP_LENGTH = SAMPLE_RATE * CLIP_DURATION # Total number of samples" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "h1Lw92M1b7Dd" + }, + "outputs": [], + "source": [ + "!wget -nc https://upload.wikimedia.org/wikipedia/commons/b/be/Woman_coughing_three_times.wav" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "5W7BMKk9cerB" + }, + "outputs": [], + "source": [ + "from scipy.io import wavfile\n", + "\n", + "# Load file\n", + "with open('Woman_coughing_three_times.wav', 'rb') as f:\n", + " original_sampling_rate, audio_array = wavfile.read(f)\n", + "\n", + "print(f\"Sample Rate: {original_sampling_rate} Hz\")\n", + "print(f\"Data Shape: {audio_array.shape}\")\n", + "print(f\"Data Type: {audio_array.dtype}\")\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "1l-eFOh0mq1r" + }, + "outputs": [], + "source": [ + "from IPython.display import Audio, display\n", + "import importlib\n", + "audio_utils = importlib.import_module(\n", + " \"hear.python.data_processing.audio_utils\"\n", + ")\n", + "resample_audio_and_convert_to_mono = audio_utils.resample_audio_and_convert_to_mono\n", + "\n", + "\n", + "audio_array = resample_audio_and_convert_to_mono(\n", + " audio_array=audio_array, \n", + " sampling_rate=original_sampling_rate,\n", + " new_sampling_rate=SAMPLE_RATE,\n", + ")\n", + "display(Audio(audio_array, rate=SAMPLE_RATE))" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "fJIdGmpFqhi8" + }, + "source": [ + "## Compute embeddings" + ] + }, + { + "metadata": { + "id": "MIRuwbfz2a-l" + }, + "cell_type": "code", + "source": [ + "from transformers import ViTConfig, ViTModel\n", + "\n", + "\n", + "# Load the model directly from Hugging Face Hub\n", + "configuration = ViTConfig(\n", + " image_size=(192, 128),\n", + " hidden_size=1024,\n", + " num_hidden_layers=24,\n", + " num_attention_heads=16,\n", + " intermediate_size=1024 * 4,\n", + " hidden_act=\"gelu\",\n", + " hidden_dropout_prob=0.0,\n", + " attention_probs_dropout_prob=0.0,\n", + " initializer_range=0.02,\n", + " layer_norm_eps=1e-12,\n", + " pooled_dim=512,\n", + " patch_size=16,\n", + " num_channels=1,\n", + " qkv_bias=True,\n", + " encoder_stride=16,\n", + " pooler_act='linear',\n", + " pooler_output_size=512,\n", + ")\n", + "loaded_model = ViTModel.from_pretrained(\n", + " \"google/hear-pytorch\",\n", + " config=configuration\n", + ")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "Fhk4GKkwsAMj" + }, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "import torch\n", + "\n", + "preprocess_audio = audio_utils.preprocess_audio\n", + "\n", + "# This index corresponds to a cough and was determined by hand. In practice, you\n", + "# would need a detector.\n", + "START = 0\n", + "\n", + "# Add batch dimension\n", + "input_tensor = np.expand_dims(audio_array[START: START + CLIP_LENGTH], axis=0)\n", + "\n", + "# Call inference\n", + "infer = lambda audio_array: loaded_model.forward(\n", + " preprocess_audio(audio_array), return_dict=True, output_hidden_states=True)\n", + "output = infer(torch.Tensor(input_tensor))\n", + "\n", + "# Extract the embedding vector\n", + "embedding_vector = np.asarray(output.pooler_output.detach()).flatten()\n", + "print(\"Size of embedding vector:\", len(embedding_vector))\n", + "\n", + "# Plot the embedding vector\n", + "plt.figure(figsize=(12, 4))\n", + "plt.plot(embedding_vector)\n", + "plt.title('Embedding Vector')\n", + "plt.xlabel('Index')\n", + "plt.ylabel('Value')\n", + "plt.grid(True)\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "PHTxQttKYNpa" + }, + "source": [ + "# Next steps\n", + "\n", + "Explore the other [notebooks](https://github.com/google-health/hear/blob/master/notebooks) to learn what else you can do with the model." + ] + } + ], + "metadata": { + "colab": { + "name": "quick_start_with_hugging_face_pytorch.ipynb", + "private_outputs": true, + "provenance": [] + }, + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + }, + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 0 +} From 92ca9fa84e14f4d28f9ef533ae22d32c2f443bfe Mon Sep 17 00:00:00 2001 From: Google Health Date: Thu, 10 Apr 2025 22:55:41 +0000 Subject: [PATCH 08/15] Update HeAR PyTorch config so that it closely matches internal Jax implementation PiperOrigin-RevId: 746206516 --- notebooks/quick_start_with_hugging_face_pytorch.ipynb | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/notebooks/quick_start_with_hugging_face_pytorch.ipynb b/notebooks/quick_start_with_hugging_face_pytorch.ipynb index b5bd621..1475491 100644 --- a/notebooks/quick_start_with_hugging_face_pytorch.ipynb +++ b/notebooks/quick_start_with_hugging_face_pytorch.ipynb @@ -194,11 +194,11 @@ " num_hidden_layers=24,\n", " num_attention_heads=16,\n", " intermediate_size=1024 * 4,\n", - " hidden_act=\"gelu\",\n", + " hidden_act=\"gelu_fast\",\n", " hidden_dropout_prob=0.0,\n", " attention_probs_dropout_prob=0.0,\n", " initializer_range=0.02,\n", - " layer_norm_eps=1e-12,\n", + " layer_norm_eps=1e-6,\n", " pooled_dim=512,\n", " patch_size=16,\n", " num_channels=1,\n", From 3cba3ba843350a9166227cf56e2a4a58399640d3 Mon Sep 17 00:00:00 2001 From: Google Health Date: Mon, 21 Apr 2025 17:21:35 +0000 Subject: [PATCH 09/15] enable audio preprocessing to run on accelerators PiperOrigin-RevId: 749837178 --- python/data_processing/audio_utils.py | 20 ++++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/python/data_processing/audio_utils.py b/python/data_processing/audio_utils.py index 84dc350..7b15211 100644 --- a/python/data_processing/audio_utils.py +++ b/python/data_processing/audio_utils.py @@ -214,6 +214,7 @@ def _pcen_function( delta_param = torch.ones(num_channels) * delta root_param = torch.ones(num_channels) * root + delta_param = delta_param.to(inputs.device).to(inputs.dtype) alpha_param = ( torch.minimum(alpha_param, torch.ones_like(alpha_param)) .to(inputs.device) @@ -230,7 +231,7 @@ def _pcen_function( smooth_coef=smooth_coef, # Handle cases where input is 2D or 3D initial_state=inputs[:, 0] if inputs.ndim > 1 else None, - ) + ).to(inputs.device) one_over_root = 1.0 / root_param output = ( @@ -254,6 +255,7 @@ def _hertz_to_mel(frequencies_hertz: torch.Tensor) -> torch.Tensor: def _linear_to_mel_weight_matrix( + device: torch.device, num_mel_bins: int = 128, num_spectrogram_bins: int = 201, sample_rate: float = 16000, @@ -268,6 +270,7 @@ def _linear_to_mel_weight_matrix( https://www.tensorflow.org/api_docs/python/tf/signal/linear_to_mel_weight_matrix Args: + device: The device to use for the result matrix. num_mel_bins: How many bands in the resulting mel spectrum. num_spectrogram_bins: How many bins there are in the source spectrogram data. @@ -312,15 +315,19 @@ def _linear_to_mel_weight_matrix( ) sample_rate_tensor = torch.tensor(sample_rate, dtype=dtype) - lower_edge_hertz_tensor = torch.tensor(lower_edge_hertz, dtype=dtype) - upper_edge_hertz_tensor = torch.tensor(upper_edge_hertz, dtype=dtype) - zero = torch.tensor(0.0, dtype=dtype) + lower_edge_hertz_tensor = torch.tensor( + lower_edge_hertz, dtype=dtype, device=device + ) + upper_edge_hertz_tensor = torch.tensor( + upper_edge_hertz, dtype=dtype, device=device + ) + zero = torch.tensor(0.0, dtype=dtype, device=device) # HTK excludes the spectrogram DC bin. bands_to_zero = 1 nyquist_hertz = sample_rate_tensor / 2.0 linear_frequencies = torch.linspace( - zero, nyquist_hertz, num_spectrogram_bins, dtype=dtype + zero, nyquist_hertz, num_spectrogram_bins, dtype=dtype, device=device )[bands_to_zero:] spectrogram_bins_mel = _hertz_to_mel(linear_frequencies).unsqueeze(1) @@ -333,6 +340,7 @@ def _linear_to_mel_weight_matrix( _hertz_to_mel(upper_edge_hertz_tensor), num_mel_bins + 2, dtype=dtype, + device=device, ) # Create frames of size 3 with stride 1 band_edges_mel = band_edges_mel.unfold(0, 3, 1) @@ -385,7 +393,7 @@ def _mel_pcen( ) spectrograms = torch.square(torch.abs(stft)) - mel_transform = _linear_to_mel_weight_matrix() + mel_transform = _linear_to_mel_weight_matrix(x.device) mel_spectrograms = torch.matmul(spectrograms, mel_transform) return _pcen_function(mel_spectrograms) From 6651d145d0d2222347423dada3da6de3fbf88931 Mon Sep 17 00:00:00 2001 From: Google Health Date: Wed, 23 Apr 2025 02:57:14 +0000 Subject: [PATCH 10/15] Add event detector demo notebook to README PiperOrigin-RevId: 750420516 --- notebooks/README.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/notebooks/README.md b/notebooks/README.md index 7e6403f..40bab9f 100644 --- a/notebooks/README.md +++ b/notebooks/README.md @@ -17,3 +17,7 @@ * [Train a data efficient classifier](train_data_efficient_classifier.ipynb) - Example of using the generated embeddings to train a custom classifier with less data and compute. + +* [Health event detector demo](hear_event_detector_demo.ipynb) - Example of + using HeAR with the health event detector models to first identify and then + generate embeddings for health-related sounds like coughing and breathing. \ No newline at end of file From 3ebc8454a0953d609c4d7a0a1c49293e6067c2d9 Mon Sep 17 00:00:00 2001 From: bramsterling Date: Tue, 6 May 2025 10:49:10 -0700 Subject: [PATCH 11/15] Refactor inline_prediction_executor to inject ModelRunner implementation. This prepares to add support for alternative model server types requiring their own ModelRunner adaptors. PiperOrigin-RevId: 755434582 --- python/serving/server_gunicorn.py | 4 +++- .../inline_prediction_executor.py | 5 ++-- .../inline_prediction_executor_test.py | 23 ++++++++++++------- 3 files changed, 21 insertions(+), 11 deletions(-) diff --git a/python/serving/server_gunicorn.py b/python/serving/server_gunicorn.py index de070df..1190e2c 100644 --- a/python/serving/server_gunicorn.py +++ b/python/serving/server_gunicorn.py @@ -28,6 +28,7 @@ from serving.serving_framework import inline_prediction_executor from serving.serving_framework import server_gunicorn +from serving.serving_framework import server_model_runner from serving import predictor @@ -52,7 +53,8 @@ def main(argv: Sequence[str]) -> None: logging.info('Launching gunicorn application.') server_gunicorn.PredictionApplication( inline_prediction_executor.InlinePredictionExecutor( - predictor_instance.predict + predictor_instance.predict, + server_model_runner.ServerModelRunner ), health_check=health_checker, options=options, diff --git a/python/serving/serving_framework/inline_prediction_executor.py b/python/serving/serving_framework/inline_prediction_executor.py index a3e25d2..0be41bd 100644 --- a/python/serving/serving_framework/inline_prediction_executor.py +++ b/python/serving/serving_framework/inline_prediction_executor.py @@ -30,7 +30,6 @@ from serving.serving_framework import model_runner from serving.serving_framework import server_gunicorn -from serving.serving_framework import server_model_runner class InlinePredictionExecutor(server_gunicorn.PredictionExecutor): @@ -49,9 +48,11 @@ def __init__( predictor: Callable[ [dict[str, Any], model_runner.ModelRunner], dict[str, Any] ], + model_runner_source: Callable[[], model_runner.ModelRunner] ): self._predictor = predictor self._model_runner = None + self._model_runner_source = model_runner_source @override def start(self) -> None: @@ -61,7 +62,7 @@ def start(self) -> None: which needs to be done post-fork. """ # Safer to instantiate the RPC stub post-fork. - self._model_runner = server_model_runner.ServerModelRunner() + self._model_runner = self._model_runner_source() def predict(self, input_json: dict[str, Any]) -> dict[str, Any]: """Executes the given request payload.""" diff --git a/python/serving/serving_framework/inline_prediction_executor_test.py b/python/serving/serving_framework/inline_prediction_executor_test.py index af21676..70b3b09 100644 --- a/python/serving/serving_framework/inline_prediction_executor_test.py +++ b/python/serving/serving_framework/inline_prediction_executor_test.py @@ -23,28 +23,35 @@ class InlinePredictionExecutorTest(absltest.TestCase): def test_predict_requires_start(self): predictor = mock.MagicMock() - executor = inline_prediction_executor.InlinePredictionExecutor(predictor) + executor = inline_prediction_executor.InlinePredictionExecutor( + predictor, server_model_runner.ServerModelRunner + ) with self.assertRaises(RuntimeError): executor.predict({"placeholder": "input"}) def test_execute_catches_predictor_exception(self): predictor = mock.MagicMock(side_effect=Exception("test error")) - executor = inline_prediction_executor.InlinePredictionExecutor(predictor) + executor = inline_prediction_executor.InlinePredictionExecutor( + predictor, server_model_runner.ServerModelRunner + ) executor.start() with self.assertRaises(RuntimeError): executor.execute({"placeholder": "input"}) def test_execute_calls_predictor(self): predictor = mock.MagicMock(return_value={"placeholder": "output"}) - executor = inline_prediction_executor.InlinePredictionExecutor(predictor) mock_model_runner = mock.create_autospec( server_model_runner.ServerModelRunner, instance=True ) - with mock.patch.object( - server_model_runner, "ServerModelRunner", autospec=True - ) as mock_model_runner_class: - mock_model_runner_class.return_value = mock_model_runner - executor.start() + mock_model_runner_class = mock.create_autospec( + server_model_runner.ServerModelRunner, autospec=True + ) + mock_model_runner_class.return_value = mock_model_runner + executor = inline_prediction_executor.InlinePredictionExecutor( + predictor, mock_model_runner_class + ) + + executor.start() self.assertEqual( executor.execute({"placeholder": "input"}), {"placeholder": "output"}, From 14e0ef668eafef854317d893777bbcea0180a02e Mon Sep 17 00:00:00 2001 From: Google Health Date: Wed, 7 May 2025 10:45:23 -0700 Subject: [PATCH 12/15] Remove config from notebook since it is now specified in the repo directly https://huggingface.co/google/hear-pytorch/blob/main/config.json PiperOrigin-RevId: 755917767 --- ...uick_start_with_hugging_face_pytorch.ipynb | 26 ++----------------- 1 file changed, 2 insertions(+), 24 deletions(-) diff --git a/notebooks/quick_start_with_hugging_face_pytorch.ipynb b/notebooks/quick_start_with_hugging_face_pytorch.ipynb index 1475491..438e6d8 100644 --- a/notebooks/quick_start_with_hugging_face_pytorch.ipynb +++ b/notebooks/quick_start_with_hugging_face_pytorch.ipynb @@ -184,33 +184,11 @@ }, "cell_type": "code", "source": [ - "from transformers import ViTConfig, ViTModel\n", + "from transformers import AutoModel\n", "\n", "\n", "# Load the model directly from Hugging Face Hub\n", - "configuration = ViTConfig(\n", - " image_size=(192, 128),\n", - " hidden_size=1024,\n", - " num_hidden_layers=24,\n", - " num_attention_heads=16,\n", - " intermediate_size=1024 * 4,\n", - " hidden_act=\"gelu_fast\",\n", - " hidden_dropout_prob=0.0,\n", - " attention_probs_dropout_prob=0.0,\n", - " initializer_range=0.02,\n", - " layer_norm_eps=1e-6,\n", - " pooled_dim=512,\n", - " patch_size=16,\n", - " num_channels=1,\n", - " qkv_bias=True,\n", - " encoder_stride=16,\n", - " pooler_act='linear',\n", - " pooler_output_size=512,\n", - ")\n", - "loaded_model = ViTModel.from_pretrained(\n", - " \"google/hear-pytorch\",\n", - " config=configuration\n", - ")" + "loaded_model = AutoModel.from_pretrained(\"google/hear-pytorch\")" ], "outputs": [], "execution_count": null From d0d05f90051df86190c5e07ad29fb24198abd555 Mon Sep 17 00:00:00 2001 From: bramsterling Date: Fri, 9 May 2025 20:13:19 +0000 Subject: [PATCH 13/15] Update serving requirements. PiperOrigin-RevId: 756880540 --- python/serving/serving_framework/requirements.in | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/serving/serving_framework/requirements.in b/python/serving/serving_framework/requirements.in index a79e0d9..e0cbe18 100644 --- a/python/serving/serving_framework/requirements.in +++ b/python/serving/serving_framework/requirements.in @@ -1,9 +1,9 @@ -absl-py~=2.1.0 +absl-py>=2.1.0 flask~=3.0.3 grpcio~=1.68.1 grpcio-status~=1.68.1 gunicorn~=23.0.0 -numpy~=2.0.2 +numpy<=2.0.2 # bypassing faulty version restriction in tritonclient requests~=2.32.3 # TODO: b/375469331 - Enable testing with most current requests-mock release. requests-mock==1.9.3 From 1b9dea36fdb7eedc3a54dc909bababbd7543e143 Mon Sep 17 00:00:00 2001 From: bramsterling Date: Fri, 16 May 2025 21:25:41 +0000 Subject: [PATCH 14/15] Separate tensorflow model runner code from serving framework as a sub-library. This allows extraneous dependencies to be dropped from cases that don't require them. PiperOrigin-RevId: 759755574 --- python/serving/requirements.in | 1 + python/serving/server_gunicorn.py | 2 +- .../inline_prediction_executor_test.py | 28 +++++++++++++++---- .../serving/serving_framework/requirements.in | 3 -- .../serving_framework/tensorflow/__init__.py | 0 .../{ => tensorflow}/inline_model_runner.py | 0 .../inline_model_runner_test.py | 2 +- .../tensorflow/requirements.in | 3 ++ .../{ => tensorflow}/server_model_runner.py | 0 .../server_model_runner_test.py | 2 +- 10 files changed, 30 insertions(+), 11 deletions(-) create mode 100644 python/serving/serving_framework/tensorflow/__init__.py rename python/serving/serving_framework/{ => tensorflow}/inline_model_runner.py (100%) rename python/serving/serving_framework/{ => tensorflow}/inline_model_runner_test.py (99%) create mode 100644 python/serving/serving_framework/tensorflow/requirements.in rename python/serving/serving_framework/{ => tensorflow}/server_model_runner.py (100%) rename python/serving/serving_framework/{ => tensorflow}/server_model_runner_test.py (99%) diff --git a/python/serving/requirements.in b/python/serving/requirements.in index 8c6fc5a..2410d2f 100644 --- a/python/serving/requirements.in +++ b/python/serving/requirements.in @@ -5,3 +5,4 @@ tensorflow~=2.18.0 -r ../data_processing/requirements.in -r serving_framework/requirements.in +-r serving_framework/tensorflow/requirements.in diff --git a/python/serving/server_gunicorn.py b/python/serving/server_gunicorn.py index 1190e2c..a554f3c 100644 --- a/python/serving/server_gunicorn.py +++ b/python/serving/server_gunicorn.py @@ -28,7 +28,7 @@ from serving.serving_framework import inline_prediction_executor from serving.serving_framework import server_gunicorn -from serving.serving_framework import server_model_runner +from serving.serving_framework.tensorflow import server_model_runner from serving import predictor diff --git a/python/serving/serving_framework/inline_prediction_executor_test.py b/python/serving/serving_framework/inline_prediction_executor_test.py index 70b3b09..96ede9a 100644 --- a/python/serving/serving_framework/inline_prediction_executor_test.py +++ b/python/serving/serving_framework/inline_prediction_executor_test.py @@ -12,11 +12,29 @@ # See the License for the specific language governing permissions and # limitations under the License. +from collections.abc import Mapping, Set from unittest import mock +import numpy as np + from absl.testing import absltest from serving.serving_framework import inline_prediction_executor -from serving.serving_framework import server_model_runner +from serving.serving_framework import model_runner + + +class DummyModelRunner(model_runner.ModelRunner): + """Dummy model runner for testing.""" + + def run_model_multiple_output( + self, + model_input: Mapping[str, np.ndarray] | np.ndarray, + *, + model_name: str = "default", + model_version: int | None = None, + model_output_keys: Set[str], + ) -> Mapping[str, np.ndarray]: + del model_name, model_version, model_output_keys + return {"output_0": np.ones((1, 2), dtype=np.float32)} class InlinePredictionExecutorTest(absltest.TestCase): @@ -24,7 +42,7 @@ class InlinePredictionExecutorTest(absltest.TestCase): def test_predict_requires_start(self): predictor = mock.MagicMock() executor = inline_prediction_executor.InlinePredictionExecutor( - predictor, server_model_runner.ServerModelRunner + predictor, DummyModelRunner ) with self.assertRaises(RuntimeError): executor.predict({"placeholder": "input"}) @@ -32,7 +50,7 @@ def test_predict_requires_start(self): def test_execute_catches_predictor_exception(self): predictor = mock.MagicMock(side_effect=Exception("test error")) executor = inline_prediction_executor.InlinePredictionExecutor( - predictor, server_model_runner.ServerModelRunner + predictor, DummyModelRunner ) executor.start() with self.assertRaises(RuntimeError): @@ -41,10 +59,10 @@ def test_execute_catches_predictor_exception(self): def test_execute_calls_predictor(self): predictor = mock.MagicMock(return_value={"placeholder": "output"}) mock_model_runner = mock.create_autospec( - server_model_runner.ServerModelRunner, instance=True + DummyModelRunner, instance=True ) mock_model_runner_class = mock.create_autospec( - server_model_runner.ServerModelRunner, autospec=True + DummyModelRunner, autospec=True ) mock_model_runner_class.return_value = mock_model_runner executor = inline_prediction_executor.InlinePredictionExecutor( diff --git a/python/serving/serving_framework/requirements.in b/python/serving/serving_framework/requirements.in index e0cbe18..33fed85 100644 --- a/python/serving/serving_framework/requirements.in +++ b/python/serving/serving_framework/requirements.in @@ -7,8 +7,5 @@ numpy<=2.0.2 # bypassing faulty version restriction in tritonclient requests~=2.32.3 # TODO: b/375469331 - Enable testing with most current requests-mock release. requests-mock==1.9.3 -tensorflow~=2.18.0 -tensorflow-serving-api~=2.18.0 -tensorflow-io-gcs-filesystem>=0.23.1 setuptools~=75.6.0 typing-extensions~=4.12.2 diff --git a/python/serving/serving_framework/tensorflow/__init__.py b/python/serving/serving_framework/tensorflow/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python/serving/serving_framework/inline_model_runner.py b/python/serving/serving_framework/tensorflow/inline_model_runner.py similarity index 100% rename from python/serving/serving_framework/inline_model_runner.py rename to python/serving/serving_framework/tensorflow/inline_model_runner.py diff --git a/python/serving/serving_framework/inline_model_runner_test.py b/python/serving/serving_framework/tensorflow/inline_model_runner_test.py similarity index 99% rename from python/serving/serving_framework/inline_model_runner_test.py rename to python/serving/serving_framework/tensorflow/inline_model_runner_test.py index 2f6029a..337ac6d 100644 --- a/python/serving/serving_framework/inline_model_runner_test.py +++ b/python/serving/serving_framework/tensorflow/inline_model_runner_test.py @@ -19,7 +19,7 @@ from absl.testing import absltest from absl.testing import parameterized -from serving.serving_framework import inline_model_runner +from serving.serving_framework.tensorflow import inline_model_runner def tensor_equal(first: tf.Tensor, second: tf.Tensor) -> bool: diff --git a/python/serving/serving_framework/tensorflow/requirements.in b/python/serving/serving_framework/tensorflow/requirements.in new file mode 100644 index 0000000..b455760 --- /dev/null +++ b/python/serving/serving_framework/tensorflow/requirements.in @@ -0,0 +1,3 @@ +tensorflow~=2.18.0 +tensorflow-serving-api~=2.18.0 +tensorflow-io-gcs-filesystem>=0.23.1 \ No newline at end of file diff --git a/python/serving/serving_framework/server_model_runner.py b/python/serving/serving_framework/tensorflow/server_model_runner.py similarity index 100% rename from python/serving/serving_framework/server_model_runner.py rename to python/serving/serving_framework/tensorflow/server_model_runner.py diff --git a/python/serving/serving_framework/server_model_runner_test.py b/python/serving/serving_framework/tensorflow/server_model_runner_test.py similarity index 99% rename from python/serving/serving_framework/server_model_runner_test.py rename to python/serving/serving_framework/tensorflow/server_model_runner_test.py index c13b3d8..162ca00 100644 --- a/python/serving/serving_framework/server_model_runner_test.py +++ b/python/serving/serving_framework/tensorflow/server_model_runner_test.py @@ -18,7 +18,7 @@ import tensorflow as tf from absl.testing import absltest -from serving.serving_framework import server_model_runner +from serving.serving_framework.tensorflow import server_model_runner from tensorflow_serving.apis import predict_pb2 from tensorflow_serving.apis import prediction_service_pb2_grpc From 36bbb90879545a2f1106668048b1a12a61ff2a0d Mon Sep 17 00:00:00 2001 From: tiffchen Date: Thu, 3 Jul 2025 00:30:24 +0000 Subject: [PATCH 15/15] Update Github logo link PiperOrigin-RevId: 778693202 --- notebooks/hear_event_detector_demo.ipynb | 4 ++-- notebooks/quick_start_with_hugging_face.ipynb | 4 ++-- notebooks/quick_start_with_hugging_face_pytorch.ipynb | 4 ++-- notebooks/quick_start_with_model_garden.ipynb | 2 +- notebooks/train_data_efficient_classifier.ipynb | 4 ++-- 5 files changed, 9 insertions(+), 9 deletions(-) diff --git a/notebooks/hear_event_detector_demo.ipynb b/notebooks/hear_event_detector_demo.ipynb index 227bc33..03d2140 100644 --- a/notebooks/hear_event_detector_demo.ipynb +++ b/notebooks/hear_event_detector_demo.ipynb @@ -31,12 +31,12 @@ " \u003c/td\u003e \n", " \u003ctd style=\"text-align: center\"\u003e\n", " \u003ca href=\"https://github.com/google-health/hear/blob/master/notebooks/hear_event_detector_demo.ipynb\"\u003e\n", - " \u003cimg alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"\u003e\u003cbr\u003e View on GitHub\n", + " \u003cimg alt=\"GitHub logo\" src=\"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png\" width=\"32px\"\u003e\u003cbr\u003e View on GitHub\n", " \u003c/a\u003e\n", " \u003c/td\u003e\n", " \u003ctd style=\"text-align: center\"\u003e\n", " \u003ca href=\"https://huggingface.co/google/hear\"\u003e\n", - " \u003cimg alt=\"HuggingFace logo\" src=\"https://huggingface.co/front/assets/huggingface_logo-noborder.svg\" width=\"32px\"\u003e\u003cbr\u003e View on HuggingFace\n", + " \u003cimg alt=\"Hugging Face logo\" src=\"https://huggingface.co/front/assets/huggingface_logo-noborder.svg\" width=\"32px\"\u003e\u003cbr\u003e View on Hugging Face\n", " \u003c/a\u003e\n", " \u003c/td\u003e\n", "\u003c/tr\u003e\u003c/tbody\u003e\u003c/table\u003e\n", diff --git a/notebooks/quick_start_with_hugging_face.ipynb b/notebooks/quick_start_with_hugging_face.ipynb index 16728bc..716a340 100644 --- a/notebooks/quick_start_with_hugging_face.ipynb +++ b/notebooks/quick_start_with_hugging_face.ipynb @@ -32,12 +32,12 @@ " \u003c/td\u003e \n", " \u003ctd style=\"text-align: center\"\u003e\n", " \u003ca href=\"https://github.com/google-health/hear/blob/master/notebooks/quick_start_with_hugging_face.ipynb\"\u003e\n", - " \u003cimg alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"\u003e\u003cbr\u003e View on GitHub\n", + " \u003cimg alt=\"GitHub logo\" src=\"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png\" width=\"32px\"\u003e\u003cbr\u003e View on GitHub\n", " \u003c/a\u003e\n", " \u003c/td\u003e\n", " \u003ctd style=\"text-align: center\"\u003e\n", " \u003ca href=\"https://huggingface.co/google/hear\"\u003e\n", - " \u003cimg alt=\"HuggingFace logo\" src=\"https://huggingface.co/front/assets/huggingface_logo-noborder.svg\" width=\"32px\"\u003e\u003cbr\u003e View on HuggingFace\n", + " \u003cimg alt=\"Hugging Face logo\" src=\"https://huggingface.co/front/assets/huggingface_logo-noborder.svg\" width=\"32px\"\u003e\u003cbr\u003e View on Hugging Face\n", " \u003c/a\u003e\n", " \u003c/td\u003e\n", "\u003c/tr\u003e\u003c/tbody\u003e\u003c/table\u003e\n", diff --git a/notebooks/quick_start_with_hugging_face_pytorch.ipynb b/notebooks/quick_start_with_hugging_face_pytorch.ipynb index 438e6d8..838e06c 100644 --- a/notebooks/quick_start_with_hugging_face_pytorch.ipynb +++ b/notebooks/quick_start_with_hugging_face_pytorch.ipynb @@ -33,12 +33,12 @@ " \u003c/td\u003e \n", " \u003ctd style=\"text-align: center\"\u003e\n", " \u003ca href=\"https://github.com/google-health/hear/blob/master/notebooks/quick_start_with_hugging_face_pytorch.ipynb\"\u003e\n", - " \u003cimg alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"\u003e\u003cbr\u003e View on GitHub\n", + " \u003cimg alt=\"GitHub logo\" src=\"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png\" width=\"32px\"\u003e\u003cbr\u003e View on GitHub\n", " \u003c/a\u003e\n", " \u003c/td\u003e\n", " \u003ctd style=\"text-align: center\"\u003e\n", " \u003ca href=\"https://huggingface.co/google/hear-pytorch\"\u003e\n", - " \u003cimg alt=\"HuggingFace logo\" src=\"https://huggingface.co/front/assets/huggingface_logo-noborder.svg\" width=\"32px\"\u003e\u003cbr\u003e View on HuggingFace\n", + " \u003cimg alt=\"Hugging Face logo\" src=\"https://huggingface.co/front/assets/huggingface_logo-noborder.svg\" width=\"32px\"\u003e\u003cbr\u003e View on Hugging Face\n", " \u003c/a\u003e\n", " \u003c/td\u003e\n", "\u003c/tr\u003e\u003c/tbody\u003e\u003c/table\u003e\n", diff --git a/notebooks/quick_start_with_model_garden.ipynb b/notebooks/quick_start_with_model_garden.ipynb index 24f541e..fbd0fea 100644 --- a/notebooks/quick_start_with_model_garden.ipynb +++ b/notebooks/quick_start_with_model_garden.ipynb @@ -40,7 +40,7 @@ " \n", " \n", " \n", - " \"GitHub
View on GitHub\n", + " \"GitHub
View on GitHub\n", " \n", " \n", "" diff --git a/notebooks/train_data_efficient_classifier.ipynb b/notebooks/train_data_efficient_classifier.ipynb index f47200e..4542bd5 100644 --- a/notebooks/train_data_efficient_classifier.ipynb +++ b/notebooks/train_data_efficient_classifier.ipynb @@ -32,12 +32,12 @@ " \u003c/td\u003e \n", " \u003ctd style=\"text-align: center\"\u003e\n", " \u003ca href=\"https://github.com/google-health/hear/blob/master/notebooks/train_data_efficient_classifier.ipynb\"\u003e\n", - " \u003cimg alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"\u003e\u003cbr\u003e View on GitHub\n", + " \u003cimg alt=\"GitHub logo\" src=\"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png\" width=\"32px\"\u003e\u003cbr\u003e View on GitHub\n", " \u003c/a\u003e\n", " \u003c/td\u003e\n", " \u003ctd style=\"text-align: center\"\u003e\n", " \u003ca href=\"https://huggingface.co/google/hear\"\u003e\n", - " \u003cimg alt=\"HuggingFace logo\" src=\"https://huggingface.co/front/assets/huggingface_logo-noborder.svg\" width=\"32px\"\u003e\u003cbr\u003e View on HuggingFace\n", + " \u003cimg alt=\"Hugging Face logo\" src=\"https://huggingface.co/front/assets/huggingface_logo-noborder.svg\" width=\"32px\"\u003e\u003cbr\u003e View on Hugging Face\n", " \u003c/a\u003e\n", " \u003c/td\u003e\n", "\u003c/tr\u003e\u003c/tbody\u003e\u003c/table\u003e\n",