Local Whisper transcription (#10)

* First pass at Whisper transcription

* deletions

* Revise based on feedback, add autopep8

requirements.txt

@@ -1,3 +1,4 @@
+autopep8==2.0.4
 build==1.0.3
 packaging==23.2
-pyproject_hooks==1.0.0
+pyproject_hooks==1.0.0

src/dailyai/requirements.txt

@@ -1,2 +1,3 @@
 Pillow==10.1.0
-typing_extensions==4.9.0
+typing_extensions==4.9.0
+faster-whisper==0.10.0

src/dailyai/services/ai_services.py

@@ -1,5 +1,9 @@
+import array
 import asyncio
+import io
 import logging
+import math
+import wave
 
 from dailyai.queue_frame import (
     AudioQueueFrame,
@@ -11,7 +15,7 @@ from dailyai.queue_frame import (
 )
 
 from abc import abstractmethod
-from typing import AsyncGenerator, AsyncIterable, Iterable
+from typing import AsyncGenerator, AsyncIterable, BinaryIO, Iterable
 from dataclasses import dataclass
 
 
@@ -150,9 +154,40 @@ class ImageGenService(AIService):
         (url, image_data) = await self.run_image_gen(frame.text)
         yield ImageQueueFrame(url, image_data)
 
+class STTService(AIService):
+    """STTService is a base class for speech-to-text services."""
+
+    _frame_rate: int
+    def __init__(self, frame_rate: int = 16000, **kwargs):
+        super().__init__(**kwargs)
+        self._frame_rate = frame_rate
+
+
+    @abstractmethod
+    async def run_stt(self, audio: BinaryIO) -> str:
+        """Returns transcript as a string"""
+        pass
+
+    async def process_frame(self, frame: QueueFrame) -> AsyncGenerator[QueueFrame, None]:
+        """Processes a frame of audio data, either buffering or transcribing it."""
+        if not isinstance(frame, AudioQueueFrame):
+            return
+        
+        data = frame.data
+        content = io.BufferedRandom(io.BytesIO())
+        ww = wave.open(self._content, "wb")
+        ww.setnchannels(1)
+        ww.setsampwidth(2)
+        ww.setframerate(self._frame_rate)
+        ww.writeframesraw(data)
+        ww.close()
+        content.seek(0)
+        text = await self.run_stt(content)
+        yield TextQueueFrame(text)
 
 @dataclass
 class AIServiceConfig:
     tts: TTSService
     image: ImageGenService
     llm: LLMService
+    stt: STTService

src/dailyai/services/daily_transport_service.py

@@ -31,6 +31,10 @@ from daily import (
 class DailyTransportService(EventHandler):
     _daily_initialized = False
     _lock = threading.Lock()
+
+    speaker_enabled: bool
+    speaker_sample_rate: int
+
     def __init__(
         self,
         room_url: str,
@@ -38,6 +42,9 @@ class DailyTransportService(EventHandler):
         bot_name: str,
         duration: float = 10,
         min_others_count: int = 1,
+        start_transcription: bool = True,
+        speaker_enabled: bool = False,
+        speaker_sample_rate: int = 16000,
     ):
         super().__init__()
         self.bot_name: str = bot_name
@@ -46,6 +53,7 @@ class DailyTransportService(EventHandler):
         self.duration: float = duration
         self.expiration = time.time() + duration * 60
         self.min_others_count = min_others_count
+        self.start_transcription = start_transcription
 
         # This queue is used to marshal frames from the async send queue to the thread that emits audio & video.
         # We need this to maintain the asynchronous behavior of asyncio queues -- to give async functions
@@ -61,6 +69,8 @@ class DailyTransportService(EventHandler):
         self.camera_width = 1024
         self.camera_height = 768
         self.camera_enabled = False
+        self.speaker_enabled = speaker_enabled
+        self.speaker_sample_rate = speaker_sample_rate
 
         self.send_queue = asyncio.Queue()
         self.receive_queue = asyncio.Queue()
@@ -144,9 +154,11 @@ class DailyTransportService(EventHandler):
                 "camera", width=self.camera_width, height=self.camera_height, color_format="RGB"
             )
 
-        self.speaker: VirtualSpeakerDevice = Daily.create_speaker_device(
-            "speaker", sample_rate=16000, channels=1
-        )
+        if self.speaker_enabled:
+            self.speaker: VirtualSpeakerDevice = Daily.create_speaker_device(
+                "speaker", sample_rate=self.speaker_sample_rate, channels=1
+            )
+            Daily.select_speaker_device("speaker")
 
         self.image: bytes | None = None
         self.camera_thread = Thread(target=self.run_camera, daemon=True)
@@ -156,8 +168,6 @@ class DailyTransportService(EventHandler):
         self.frame_consumer_thread = Thread(target=self.frame_consumer, daemon=True)
         self.frame_consumer_thread.start()
 
-        Daily.select_speaker_device("speaker")
-
         self.client.set_user_name(self.bot_name)
         self.client.join(self.room_url, self.token, completion=self.call_joined)
         self.my_participant_id = self.client.participants()["local"]["id"]
@@ -201,9 +211,20 @@ class DailyTransportService(EventHandler):
             }
         )
 
-        if self.token:
+        if self.token and self.start_transcription:
             self.client.start_transcription(self.transcription_settings)
 
+
+    def _receive_audio(self):
+        """Receive audio from the Daily call and put it on the receive queue"""
+        seconds = 1
+        desired_frame_count = self.speaker_sample_rate * seconds
+        while True:
+            buffer = self.speaker.read_frames(desired_frame_count)
+            if len(buffer) > 0:
+                frame = AudioQueueFrame(buffer)
+                asyncio.run_coroutine_threadsafe(self.receive_queue.put(frame), self.loop)
+
     async def get_receive_frames(self):
         while True:
             frame = await self.receive_queue.get()
@@ -266,6 +287,9 @@ class DailyTransportService(EventHandler):
 
     def call_joined(self, join_data, client_error):
         self.logger.info(f"Call_joined: {join_data}, {client_error}")
+        if self.speaker_enabled:
+            t = Thread(target=self._receive_audio, daemon=True)
+            t.start()
 
     def on_error(self, error):
         self.logger.error(f"on_error: {error}")

src/dailyai/services/local_stt_service.py

@@ -0,0 +1,72 @@
+import array
+import io
+import math
+from typing import AsyncGenerator
+import wave
+from dailyai.queue_frame import AudioQueueFrame, QueueFrame, TextQueueFrame
+from dailyai.services.ai_services import STTService
+
+
+class LocalSTTService(STTService):
+    _content: io.BufferedRandom
+    _wave: wave.Wave_write
+    _current_silence_frames: int
+
+    # Configuration
+    _min_rms: int
+    _max_silence_frames: int
+    _frame_rate: int
+
+    def __init__(self,
+                 min_rms: int = 400,
+                 max_silence_frames: int = 3,
+                 frame_rate: int = 16000,
+                 **kwargs):
+        super().__init__(frame_rate, **kwargs)
+        self._current_silence_frames = 0
+        self._min_rms = min_rms
+        self._max_silence_frames = max_silence_frames
+        self._frame_rate = frame_rate
+        self._new_wave()
+
+    def _new_wave(self):
+        """Creates a new wave object and content buffer."""
+        self._content = io.BufferedRandom(io.BytesIO())
+        ww = wave.open(self._content, "wb")
+        ww.setnchannels(1)
+        ww.setsampwidth(2)
+        ww.setframerate(self._frame_rate)
+        self._wave = ww
+
+    async def process_frame(self, frame: QueueFrame) -> AsyncGenerator[QueueFrame, None]:
+        """Processes a frame of audio data, either buffering or transcribing it."""
+        if not isinstance(frame, AudioQueueFrame):
+            return
+
+        data = frame.data
+        # Try to filter out empty background noise
+        # (Very rudimentary approach, can be improved)
+        rms = self._get_volume(data)
+        if rms >= self._min_rms:
+            # If volume is high enough, write new data to wave file
+            self._wave.writeframesraw(data)
+
+        # If buffer is not empty and we detect a 3-frame pause in speech,
+        # transcribe the audio gathered so far.
+        if self._content.tell() > 0 and self._current_silence_frames > self._max_silence_frames:
+            self._current_silence_frames = 0
+            self._wave.close()
+            self._content.seek(0)
+            text = await self.run_stt(self._content)
+            self._new_wave()
+            yield TextQueueFrame(text)
+        # If we get this far, this is a frame of silence
+        self._current_silence_frames += 1
+
+    def _get_volume(self, audio: bytes) -> float:
+        # https://docs.python.org/3/library/array.html
+        audio_array = array.array('h', audio)
+        squares = [sample**2 for sample in audio_array]
+        mean = sum(squares) / len(audio_array)
+        rms = math.sqrt(mean)
+        return rms

src/dailyai/services/whisper_ai_services.py

@@ -0,0 +1,55 @@
+"""This module implements Whisper transcription with a locally-downloaded model."""
+import asyncio
+from enum import Enum
+import logging
+from typing import BinaryIO
+from faster_whisper import WhisperModel
+from dailyai.services.local_stt_service import LocalSTTService
+
+
+class Model(Enum):
+    """Class of basic Whisper model selection options"""
+    TINY = "tiny"
+    BASE = "base"
+    MEDIUM = "medium"
+    LARGE = "large-v3"
+    DISTIL_LARGE_V2 = "Systran/faster-distil-whisper-large-v2"
+    DISTIL_MEDIUM_EN = "Systran/faster-distil-whisper-medium.en"
+
+
+class WhisperSTTService(LocalSTTService):
+    """Class to transcribe audio with a locally-downloaded Whisper model"""
+    _model: WhisperModel
+
+    # Model configuration
+    _model_name: Model
+    _device: str
+    _compute_type: str
+
+    def __init__(self, model_name: Model = Model.DISTIL_MEDIUM_EN,
+                 device: str = "auto",
+                 compute_type: str = "default"):
+
+        super().__init__()
+        self.logger: logging.Logger = logging.getLogger("dailyai")
+        self._model_name = model_name
+        self._device = device
+        self._compute_type = compute_type
+        self._load()
+
+    def _load(self):
+        """Loads the Whisper model. Note that if this is the first time
+        this model is being run, it will take time to download."""
+        model = WhisperModel(
+            self._model_name.value,
+            device=self._device,
+            compute_type=self._compute_type)
+        self._model = model
+
+    async def run_stt(self, audio: BinaryIO = None) -> str:
+        """Transcribes given audio using Whisper"""
+        segments, _ = await asyncio.to_thread(self._model.transcribe, audio)
+        res: str = ""
+        for segment in segments:
+            res += f"{segment.text} "
+        return res

src/samples/foundational/07-whisper-transcription.py

@@ -0,0 +1,45 @@
+import argparse
+import asyncio
+from threading import Thread
+
+from dailyai.services.daily_transport_service import DailyTransportService
+from dailyai.services.whisper_ai_services import WhisperSTTService
+
+
+async def main(room_url: str):
+    global transport
+    global stt
+
+    transport = DailyTransportService(
+        room_url,
+        None,
+        "Transcription bot",
+    )
+    transport.mic_enabled = False
+    transport.camera_enabled = False
+    transport.speaker_enabled = True
+    stt = WhisperSTTService()
+    transcription_output_queue = asyncio.Queue()
+
+    async def handle_transcription():
+        print("`````````TRANSCRIPTION`````````")
+        while True:
+            item = await transcription_output_queue.get()
+            print(item.text)
+
+    async def handle_speaker():
+        await stt.run_to_queue(
+            transcription_output_queue,
+            transport.get_receive_frames()
+        )
+    await asyncio.gather(transport.run(), handle_speaker(), handle_transcription())
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Simple Daily Bot Sample")
+    parser.add_argument(
+        "-u", "--url", type=str, required=True, help="URL of the Daily room to join"
+    )
+
+    args, unknown = parser.parse_known_args()
+    asyncio.run(main(args.url))