Fix healthcheck and streaming with whisper

healthcheck.sh

@@ -11,6 +11,11 @@ else
         # GPU is being utilized, assuming healthy
         exit 0
     else
-        celery --app=celery_app.celeryapp inspect ping -d ${SERVICE_NAME}_worker@$HOSTNAME || exit 1
+        if [ "$SERVICE_MODE" = "websocket" ]
+        then
+            nc -z localhost ${STREAMING_PORT:=80} && exit 0 || exit 1
+        else
+            celery --app=celery_app.celeryapp inspect ping -d ${SERVICE_NAME}_worker@$HOSTNAME || exit 1
+        fi
     fi
 fi

kaldi/Dockerfile

@@ -1,5 +1,5 @@
-FROM python:3.9
-LABEL maintainer="[email protected], [email protected], dgaynullin@linagora.com"
+FROM python:3.8
+LABEL maintainer="[email protected], [email protected], abert@linagora.com"
 
 ARG KALDI_MKL
 
@@ -14,6 +14,8 @@ RUN apt-get update && \
         g++ \
         make \
         cmake \
+        curl \
+        netcat-traditional \
         git \
         zlib1g-dev \
         automake \

kaldi/stt/processing/decoding.py

@@ -1,11 +1,12 @@
 import json
 import re
+from typing import Tuple
 
 from vosk import KaldiRecognizer, Model
 
 from punctuation.recasepunc import apply_recasepunc
 
-def decode(audio: tuple[bytes, int], model: Model, with_metadata: bool, language=None) -> dict:
+def decode(audio: Tuple[bytes, int], model: Model, with_metadata: bool, language=None) -> dict:
     """Transcribe the audio data using the vosk library with the defined model."""
     decoder_result = {"text": "", "confidence-score": 0.0, "words": []}
 

kaldi/stt/processing/streaming.py

@@ -47,7 +47,7 @@ async def wssDecode(ws: WebSocketServerProtocol, model: Model):
         if (isinstance(message, str) and re.match(EOF_REGEX, message)):
             ret = recognizer.FinalResult()
             ret = apply_recasepunc(punctuation_model, ret)
-            await ws.send(json.dumps(ret))
+            await ws.send(ret)
             await ws.close(reason="End of stream")
             break
 
@@ -102,7 +102,7 @@ def ws_streaming(websocket_server: WSServer, model: Model):
         if (isinstance(message, str) and re.match(EOF_REGEX, message)):
             ret = recognizer.FinalResult()
             ret = apply_recasepunc(punctuation_model, ret)
-            websocket_server.send(json.dumps(re.sub("<unk> ", "", ret)))
+            websocket_server.send(re.sub("<unk> ", "", ret))
             websocket_server.close()
             break
         # Audio chunk

test/automated/automated_utils.py

@@ -26,7 +26,14 @@ def parse_env_variables(env_variables):
     # make a dict
     env_variables = env_variables.split()
     env = {}
+    v_opt = False
     for env_variable in env_variables:
+        if env_variable=="-v":
+            v_opt = True
+            continue
+        if v_opt:
+            v_opt = False
+            continue   
         key, value = env_variable.split("=")
         env[key] = value
     return env

test/automated/core.py

@@ -66,7 +66,7 @@ def check_http_server_availability(self, server, pid):
         while elapsed_time < total_wait_time:
             try:
                 response = requests.head(server)
-                if response.status_code == 200:
+                if response.status_code == 200 or response.status_code == 400:
                     self.echo_note(f"Server: {server} is available after {elapsed_time} sec.")
                     return
             except requests.ConnectionError:
@@ -161,9 +161,11 @@ def run_test(self, docker_image="whisper/Dockerfile.ctranslate2", serving="http"
                 cmd += f' -F "language={language}"'
             self.echo_command(cmd)
             r = self.transcribe(cmd, regex, test_file, "Error transcription", "HTTP route 'transcribe'")
+        elif serving == "websocket":
+            r=self.check_http_server_availability("http://localhost:8080", pid)
             if r:
                 return self.report_failure(r, expect_failure=expect_failure)
-            cmd = f"python3 {TESTDIR}/test_streaming.py --audio_file {test_file}"
+            cmd = f"python3 {TESTDIR}/test_streaming.py --audio_file {test_file} -v --stream_duration 1 --stream_wait 0.0"
             if language:
                 cmd += f" --language {language}"
             self.echo_command(cmd)

test/automated/whisper.py

@@ -21,8 +21,6 @@ def generate_whisper_test_setups(
     else:
         dockerfiles = [
             "whisper/Dockerfile.ctranslate2",
-            "whisper/Dockerfile.ctranslate2.cpu",
-            "whisper/Dockerfile.torch",
             "whisper/Dockerfile.torch.cpu",
         ]
 
@@ -53,7 +51,7 @@ def test_04_integration_cpu(self, setup):
         copy_env_file("whisper/.envdefault", env_variables)
         self.run_test(dockerfile, serving=serving, env_variables=env_variables)
 
-    @idata(generate_whisper_test_setups(device="cuda", vads=[None, "silero"]))
+    @idata(generate_whisper_test_setups(device="cuda", vads=[None]))
     def test_05_integration_cuda(self, setup):
         dockerfile, serving, env_variables = setup
         copy_env_file("whisper/.envdefault", env_variables)
@@ -65,12 +63,23 @@ def test_06_integration_nodevice(self, setup):
         copy_env_file("whisper/.envdefault", env_variables)
         self.run_test(dockerfile, serving=serving, env_variables=env_variables)
 
+    def test_03_model(self):
+        env_variables = "MODEL=small LANGUAGE=fr"
+        copy_env_file("whisper/.envdefault", env_variables)
+        self.run_test(env_variables=env_variables)
+
+    def test_03_websocket(self):
+        env_variables = "MODEL=small LANGUAGE=fr"
+        copy_env_file("whisper/.envdefault", env_variables)
+        self.run_test(serving="websocket", env_variables=env_variables)
+
     def test_02_failures_cuda_on_cpu_dockerfile(self):
         env_variables = "MODEL=tiny  DEVICE=cuda"
         dockerfile = "whisper/Dockerfile.ctranslate2.cpu"
         copy_env_file("whisper/.envdefault", env_variables)
         self.assertIn(
-            "CUDA failed with error named symbol not found",
+            # "CUDA failed with error named symbol not found",
+            "Cannot load symbol",
             self.run_test(dockerfile, env_variables=env_variables, expect_failure=True),
         )
 
@@ -85,16 +94,11 @@ def test_02_failure_not_existing_file(self):
             )
         self.cleanup()
 
-    def test_03_model(self):
-        env_variables = "MODEL=small LANGUAGE=fr"
-        copy_env_file("whisper/.envdefault", env_variables)
-        self.run_test(env_variables=env_variables)
-
-    def test_01_failure_wrong_language(self):
-        env_variables = "MODEL=tiny LANGUAGE=whatever"
+    def test_02_failure_wrong_vad(self):  # doesnt work anymore because server dont stop immediately (after worker crash) and will hang between 0 and more than 10mins
+        env_variables = "VAD=whatever MODEL=tiny LANGUAGE=fr"
         copy_env_file("whisper/.envdefault", env_variables)
         self.assertIn(
-            "ValueError: Language \'whatever\' is not available",
+            "Got unexpected VAD method whatever",
             self.run_test(env_variables=env_variables, expect_failure=True),
         )
 
@@ -133,11 +137,11 @@ def test_01_language_over_config_celery(self):
         copy_env_file("whisper/.envdefault", env_variables)
         self.run_test(serving="task", env_variables=env_variables, language="fr")
 
-    def test_02_failure_wrong_vad(self):  # doesnt work anymore because server dont stop immediately (after worker crash) and will hang between 0 and more than 10mins
-        env_variables = "VAD=whatever MODEL=tiny LANGUAGE=fr"
+    def test_01_failure_wrong_language(self):
+        env_variables = "MODEL=tiny LANGUAGE=whatever"
         copy_env_file("whisper/.envdefault", env_variables)
         self.assertIn(
-            "Got unexpected VAD method whatever",
+            "ValueError: Language \'whatever\' is not available",
             self.run_test(env_variables=env_variables, expect_failure=True),
         )