model_demo/inference/audio.py

import numpy as np


class AudioStream:
    """
    Class to mimic streaming audio input.
    """

    def __init__(
        self, audio: np.ndarray, min_samples_per_step: int, max_samples_per_step: int
    ):
        self.audio = audio
        self.min_samples_per_step = min_samples_per_step
        self.max_samples_per_step = max_samples_per_step
        self.current_idx = 0
        self.can_step = True

    def step(self) -> np.ndarray:
        if not self.can_step:
            raise StopIteration("End of audio stream")
        start_idx = self.current_idx
        if self.min_samples_per_step == self.max_samples_per_step:
            samples_per_step = self.min_samples_per_step
        else:
            samples_per_step = np.random.randint(
                self.min_samples_per_step, self.max_samples_per_step, (1,)
            ).item()
        end_idx = min(start_idx + samples_per_step, len(self.audio))
        audio_chunk = self.audio[start_idx:end_idx]
        self.current_idx = end_idx
        if end_idx >= len(self.audio):
            self.can_step = False
        return audio_chunk