bidi streaming impl

Author: pskrunner14

protos/kaldi_serve.proto

@@ -2,12 +2,16 @@ syntax = "proto3";
 package kaldi_serve;
 
 service KaldiServe {
-  // Performs synchronous non-streaming speech recognition;
+  // Performs synchronous non-streaming speech recognition.
   rpc Recognize(RecognizeRequest) returns (RecognizeResponse) {}
 
   // Performs synchronous client-to-server streaming speech recognition: 
   //    receive results after all audio has been streamed and processed.
   rpc StreamingRecognize(stream RecognizeRequest) returns (RecognizeResponse) {}
+
+  // Performs synchronous bidirectional streaming speech recognition: 
+  //    receive results as the audio is being streamed and processed.
+  rpc BidiStreamingRecognize(stream RecognizeRequest) returns (stream RecognizeResponse) {}
 }
 
 message RecognizeRequest {

python/kaldi_serve/core.py

@@ -19,10 +19,18 @@ def recognize(self, config: RecognitionConfig, audio, uuid: str, timeout=None):
         request = RecognizeRequest(config=config, audio=audio, uuid=uuid)
         return self._client.Recognize(request, timeout=timeout)
 
-    def streaming_recognize(self, config: RecognitionConfig, audio_chunks, uuid: str, timeout=None):
-        request_gen = (RecognizeRequest(config=config, audio=chunk, uuid=uuid) for chunk in audio_chunks)
+    def streaming_recognize(self, config: RecognitionConfig, audio_chunks_gen, uuid: str, timeout=None):
+        request_gen = (RecognizeRequest(config=config, audio=chunk, uuid=uuid) for chunk in audio_chunks_gen)
         return self._client.StreamingRecognize(request_gen, timeout=timeout)
 
-    def streaming_recognize_raw(self, audio_params, uuid: str, timeout=None):
-        request_gen = (RecognizeRequest(config=config, audio=chunk, uuid=uuid) for config, chunk in audio_params)
-        return self._client.StreamingRecognize(request_gen, timeout=timeout)
+    def streaming_recognize_raw(self, audio_params_gen, uuid: str, timeout=None):
+        request_gen = (RecognizeRequest(config=config, audio=chunk, uuid=uuid) for config, chunk in audio_params_gen)
+        return self._client.StreamingRecognize(request_gen, timeout=timeout)
+
+    def bidi_streaming_recognize(self, config: RecognitionConfig, audio_chunks_gen, uuid: str, timeout=None):
+        request_gen = (RecognizeRequest(config=config, audio=chunk, uuid=uuid) for chunk in audio_chunks_gen)
+        return self._client.BidiStreamingRecognize(request_gen, timeout=timeout)
+
+    def bidi_streaming_recognize_raw(self, audio_params_gen, uuid: str, timeout=None):
+        request_gen = (RecognizeRequest(config=config, audio=chunk, uuid=uuid) for config, chunk in audio_params_gen)
+        return self._client.BidiStreamingRecognize(request_gen, timeout=timeout)

python/kaldi_serve/kaldi_serve_pb2.py

@@ -24,14 +24,19 @@ def __init__(self, channel):
         request_serializer=kaldi__serve__pb2.RecognizeRequest.SerializeToString,
         response_deserializer=kaldi__serve__pb2.RecognizeResponse.FromString,
         )
+    self.BidiStreamingRecognize = channel.stream_stream(
+        '/kaldi_serve.KaldiServe/BidiStreamingRecognize',
+        request_serializer=kaldi__serve__pb2.RecognizeRequest.SerializeToString,
+        response_deserializer=kaldi__serve__pb2.RecognizeResponse.FromString,
+        )
 
 
 class KaldiServeServicer(object):
   # missing associated documentation comment in .proto file
   pass
 
   def Recognize(self, request, context):
-    """Performs synchronous non-streaming speech recognition;
+    """Performs synchronous non-streaming speech recognition.
     """
     context.set_code(grpc.StatusCode.UNIMPLEMENTED)
     context.set_details('Method not implemented!')
@@ -45,6 +50,14 @@ def StreamingRecognize(self, request_iterator, context):
     context.set_details('Method not implemented!')
     raise NotImplementedError('Method not implemented!')
 
+  def BidiStreamingRecognize(self, request_iterator, context):
+    """Performs synchronous bidirectional streaming speech recognition: 
+    receive results as the audio is being streamed and processed.
+    """
+    context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+    context.set_details('Method not implemented!')
+    raise NotImplementedError('Method not implemented!')
+
 
 def add_KaldiServeServicer_to_server(servicer, server):
   rpc_method_handlers = {
@@ -58,6 +71,11 @@ def add_KaldiServeServicer_to_server(servicer, server):
           request_deserializer=kaldi__serve__pb2.RecognizeRequest.FromString,
           response_serializer=kaldi__serve__pb2.RecognizeResponse.SerializeToString,
       ),
+      'BidiStreamingRecognize': grpc.stream_stream_rpc_method_handler(
+          servicer.BidiStreamingRecognize,
+          request_deserializer=kaldi__serve__pb2.RecognizeRequest.FromString,
+          response_serializer=kaldi__serve__pb2.RecognizeResponse.SerializeToString,
+      ),
   }
   generic_handler = grpc.method_handlers_generic_handler(
       'kaldi_serve.KaldiServe', rpc_method_handlers)

python/kaldi_serve/kaldi_serve_pb2_grpc.py

@@ -34,8 +34,8 @@ def chunks_from_mic(secs: int, frame_rate: int, channels: int):
     p = pyaudio.PyAudio()
     sample_format = pyaudio.paInt16
 
-    # This is in samples not seconds
-    chunk_size = 4000
+    # 8k samples ~ 1sec of audio
+    chunk_size = 8000
 
     stream = p.open(format=sample_format,
                     channels=channels,
@@ -45,6 +45,7 @@ def chunks_from_mic(secs: int, frame_rate: int, channels: int):
 
     sample_width = p.get_sample_size(sample_format)
 
+    print('recording...')
     for _ in range(0, int(frame_rate / chunk_size * secs)):
         # The right way probably is to not send headers at all and let the
         # server side's chunk handler maintain state, taking data from

python/kaldi_serve/utils.py

@@ -56,7 +56,7 @@ def parse_response(response):
     return output
 
 
-def transcribe_chunks(client, audio_chunks, model: str, language_code: str, raw: bool=False, word_level: bool=False):
+def transcribe_chunks_streaming(client, audio_chunks, model: str, language_code: str, raw: bool=False, word_level: bool=False):
     """
     Transcribe the given audio chunks
     """
@@ -66,7 +66,6 @@ def transcribe_chunks(client, audio_chunks, model: str, language_code: str, raw:
 
     try:
         if raw:
-            print('streaming raw')
             config = lambda chunk_len: RecognitionConfig(
                 sample_rate_hertz=SR,
                 encoding=encoding,
@@ -80,7 +79,6 @@ def transcribe_chunks(client, audio_chunks, model: str, language_code: str, raw:
             audio_params = [(config(len(chunk)), RecognitionAudio(content=chunk)) for chunk in audio_chunks]
             response = client.streaming_recognize_raw(audio_params, uuid="")
         else:
-            print('streaming with headers')
             audio = (RecognitionAudio(content=chunk) for chunk in audio_chunks)
             config = RecognitionConfig(
                 sample_rate_hertz=SR,
@@ -97,6 +95,38 @@ def transcribe_chunks(client, audio_chunks, model: str, language_code: str, raw:
 
     pprint(parse_response(response))
 
+def transcribe_chunks_bidi_streaming(client, audio_chunks, model: str, language_code: str, word_level: bool=False):
+    """
+    Transcribe the given audio chunks
+    """
+
+    response = {}
+    encoding = RecognitionConfig.AudioEncoding.LINEAR16
+
+    try:
+        config = lambda chunk_len: RecognitionConfig(
+            sample_rate_hertz=SR,
+            encoding=encoding,
+            language_code=language_code,
+            max_alternatives=10,
+            model=model,
+            raw=True,
+            word_level=word_level,
+            data_bytes=chunk_len
+        )
+
+        def audio_params_gen(audio_chunks_gen):
+            for chunk in audio_chunks_gen:
+                yield config(len(chunk)), RecognitionAudio(content=chunk)
+
+        response_gen = client.bidi_streaming_recognize_raw(audio_params_gen(audio_chunks), uuid="")
+    except Exception as e:
+        traceback.print_exc()
+        print(f'error: {str(e)}')
+
+    for response in response_gen:
+        pprint(parse_response(response))
+
 
 def decode_files(client, audio_paths: List[str], model: str, language_code: str, raw: bool=False, pcm: bool=False, word_level: bool=False):
     """
@@ -105,7 +135,7 @@ def decode_files(client, audio_paths: List[str], model: str, language_code: str,
     chunked_audios = [chunks_from_file(x, chunk_size=random.randint(1, 3), raw=raw, pcm=pcm) for x in audio_paths]
 
     threads = [
-        threading.Thread(target=transcribe_chunks, args=(client, chunks, model, language_code, raw, word_level))
+        threading.Thread(target=transcribe_chunks_streaming, args=(client, chunks, model, language_code, raw, word_level))
         for chunks in chunked_audios
     ]
 
@@ -126,6 +156,6 @@ def decode_files(client, audio_paths: List[str], model: str, language_code: str,
     word_level = args["--word-level"]
 
     if args["mic"]:
-        transcribe_chunks(client, chunks_from_mic(int(args["--n-secs"]), SR, 1), model, language_code, raw, word_level)
+        transcribe_chunks_bidi_streaming(client, chunks_from_mic(int(args["--n-secs"]), SR, 1), model, language_code, word_level)
     else:
         decode_files(client, args["<file>"], model, language_code, raw, pcm, word_level)

python/scripts/example_client.py

@@ -177,7 +177,8 @@ class Decoder final {
                              kaldi::SingleUtteranceNnet3Decoder &,
                              const std::size_t &,
                              utterance_results_t &,
-                             const bool &) const;
+                             const bool &,
+                             const bool & = false) const;
 };
 
 Decoder::Decoder(const kaldi::BaseFloat &beam,
@@ -532,9 +533,13 @@ void Decoder::decode_stream_final(kaldi::OnlineNnet2FeaturePipeline &feature_pip
                                   kaldi::SingleUtteranceNnet3Decoder &decoder,
                                   const std::size_t &n_best,
                                   utterance_results_t &results,
-                                  const bool &word_level) const {
-    feature_pipeline.InputFinished();
-    decoder.FinalizeDecoding();
+                                  const bool &word_level,
+                                  const bool &bidi_streaming) const {
+
+    if (!bidi_streaming) {
+        feature_pipeline.InputFinished();
+        decoder.FinalizeDecoding();
+    }
 
     if (decoder.NumFramesDecoded() == 0) {
         KALDI_WARN << "audio may be empty :: decoded no frames";