145 lines
4.7 KiB
Python
145 lines
4.7 KiB
Python
import asyncio
|
|
import time
|
|
import queue
|
|
from google.cloud import speech
|
|
import threading
|
|
|
|
from vocode.streaming.models.audio_encoding import AudioEncoding
|
|
from vocode.streaming.transcriber.base_transcriber import (
|
|
BaseTranscriber,
|
|
Transcription,
|
|
)
|
|
from vocode.streaming.models.transcriber import GoogleTranscriberConfig
|
|
from vocode.streaming.utils import create_loop_in_thread
|
|
|
|
|
|
class GoogleTranscriber(BaseTranscriber):
|
|
def __init__(self, transcriber_config: GoogleTranscriberConfig):
|
|
super().__init__(transcriber_config)
|
|
self._queue = queue.Queue()
|
|
self._ended = False
|
|
self.google_streaming_config = self.create_google_streaming_config()
|
|
self.client = speech.SpeechClient()
|
|
self.warmed_up = False
|
|
self.is_ready = False
|
|
if self.transcriber_config.endpointing_config:
|
|
raise Exception("Google endpointing config not supported yet")
|
|
self.event_loop = asyncio.new_event_loop()
|
|
self.thread = threading.Thread(
|
|
name="google_transcriber",
|
|
target=create_loop_in_thread,
|
|
args=(self.event_loop, self.process()),
|
|
)
|
|
|
|
def create_google_streaming_config(self):
|
|
extra_params = {}
|
|
if self.transcriber_config.model:
|
|
extra_params["model"] = self.transcriber_config.model
|
|
extra_params["use_enhanced"] = True
|
|
|
|
if self.transcriber_config.audio_encoding == AudioEncoding.LINEAR16:
|
|
google_audio_encoding = speech.RecognitionConfig.AudioEncoding.LINEAR16
|
|
elif self.transcriber_config.audio_encoding == AudioEncoding.MULAW:
|
|
google_audio_encoding = speech.RecognitionConfig.AudioEncoding.MULAW
|
|
|
|
return speech.StreamingRecognitionConfig(
|
|
config=speech.RecognitionConfig(
|
|
encoding=google_audio_encoding,
|
|
sample_rate_hertz=self.transcriber_config.sampling_rate,
|
|
language_code="en-US",
|
|
**extra_params
|
|
),
|
|
interim_results=True,
|
|
)
|
|
|
|
async def ready(self):
|
|
if not self.transcriber_config.should_warmup_model:
|
|
return True
|
|
while not self.warmed_up:
|
|
await asyncio.sleep(0.1)
|
|
return self.is_ready
|
|
|
|
def warmup(self):
|
|
warmup_bytes = self.get_warmup_bytes()
|
|
|
|
def stream():
|
|
chunk_size = self.transcriber_config.sampling_rate * 2
|
|
for i in range(len(warmup_bytes) // chunk_size):
|
|
yield speech.StreamingRecognizeRequest(
|
|
audio_content=warmup_bytes[i * chunk_size : (i + 1) * chunk_size]
|
|
)
|
|
time.sleep(0.01)
|
|
|
|
for _ in self.client.streaming_recognize(
|
|
self.google_streaming_config, stream()
|
|
):
|
|
pass
|
|
self.warmed_up = True
|
|
self.is_ready = True
|
|
|
|
async def run(self):
|
|
self.thread.start()
|
|
|
|
async def process(self):
|
|
if self.transcriber_config.should_warmup_model:
|
|
self.warmup()
|
|
stream = self.generator()
|
|
requests = (
|
|
speech.StreamingRecognizeRequest(audio_content=content)
|
|
for content in stream
|
|
)
|
|
responses = self.client.streaming_recognize(
|
|
self.google_streaming_config, requests
|
|
)
|
|
await self.process_responses_loop(responses)
|
|
|
|
def terminate(self):
|
|
self._ended = True
|
|
|
|
def send_audio(self, chunk: bytes):
|
|
self._queue.put(chunk, block=False)
|
|
|
|
async def process_responses_loop(self, responses):
|
|
for response in responses:
|
|
await self._on_response(response)
|
|
|
|
if self._ended:
|
|
break
|
|
|
|
async def _on_response(self, response):
|
|
if not response.results:
|
|
return
|
|
|
|
result = response.results[0]
|
|
if not result.alternatives:
|
|
return
|
|
|
|
top_choice = result.alternatives[0]
|
|
message = top_choice.transcript
|
|
confidence = top_choice.confidence
|
|
|
|
return await self.on_response(
|
|
Transcription(message, confidence, result.is_final)
|
|
)
|
|
|
|
def generator(self):
|
|
while not self._ended:
|
|
# Use a blocking get() to ensure there's at least one chunk of
|
|
# data, and stop iteration if the chunk is None, indicating the
|
|
# end of the audio stream.
|
|
chunk = self._queue.get()
|
|
if chunk is None:
|
|
return
|
|
data = [chunk]
|
|
|
|
# Now consume whatever other data's still buffered.
|
|
while True:
|
|
try:
|
|
chunk = self._queue.get(block=False)
|
|
if chunk is None:
|
|
return
|
|
data.append(chunk)
|
|
except queue.Empty:
|
|
break
|
|
|
|
yield b"".join(data)
|