Tidy up the client python file

whisper_live/client.py

@@ -1,5 +1,4 @@
 import os
-import argparse
 import wave
 
 import numpy as np
@@ -44,47 +43,58 @@ def resample(file: str, sr: int = 16000):
 
 
 class Client:
-    CHUNK = 1024
+    INSTANCES = {}
-    FORMAT = pyaudio.paInt16
-    CHANNELS = 1
-    RATE = 16000
-    RECORD_SECONDS = 60000
-    RECORDING = False
-    multilingual = False
-    language = None
-    task = "transcribe"
-    uid = str(uuid.uuid4())
-    WAITING = False
-    LAST_RESPONSE_RECIEVED = None
-    DISCONNECT_IF_NO_RESPONSE_FOR = 15
 
-    def __init__(self, host=None, port=None, is_multilingual=False, lang=None, translate=False):
+    def __init__(
-        Client.multilingual = is_multilingual
+        self, host=None, port=None, is_multilingual=False, lang=None, translate=False
-        Client.language = lang if is_multilingual else "en"
+    ):
+        self.chunk = 1024
+        self.format = pyaudio.paInt16
+        self.channels = 1
+        self.rate = 16000
+        self.record_seconds = 60000
+        self.recording = False
+        self.multilingual = False
+        self.language = None
+        self.task = "transcribe"
+        self.uid = str(uuid.uuid4())
+        self.waiting = False
+        self.last_response_recieved = None
+        self.disconnect_if_no_response_for = 15
+        self.multilingual = is_multilingual
+        self.language = lang if is_multilingual else "en"
         if translate:
-            Client.task = "translate"
+            self.task = "translate"
 
         self.timestamp_offset = 0.0
         self.audio_bytes = None
         self.p = pyaudio.PyAudio()
-        self.stream = self.p.open(format=self.FORMAT,
+        self.stream = self.p.open(
-                        channels=self.CHANNELS,
+            format=self.format,
-                        rate=self.RATE,
+            channels=self.channels,
+            rate=self.rate,
             input=True,
-                        frames_per_buffer=self.CHUNK)
+            frames_per_buffer=self.chunk,
+        )
 
         # create websocket connection
         if host is not None and port is not None:
             socket_url = f"ws://{host}:{port}"
-            self.client_socket = websocket.WebSocketApp(socket_url,
+            self.client_socket = websocket.WebSocketApp(
-                                  on_open=Client.on_open,
+                socket_url,
-                                  on_message=Client.on_message,
+                on_open=lambda ws: self.on_open(ws),
-                                  on_error=Client.on_error,
+                on_message=lambda ws, message: self.on_message(ws, message),
-                                  on_close=Client.on_close)
+                on_error=lambda ws, error: self.on_error(ws, error),
+                on_close=lambda ws, close_status_code, close_msg: self.on_close(
+                    ws, close_status_code, close_msg
+                ),
+            )
         else:
             print("[ERROR]: No host or port specified.")
             return
 
+        Client.INSTANCES[self.uid] = self
+
         # start websocket client in a thread
         self.ws_thread = threading.Thread(target=self.client_socket.run_forever)
         self.ws_thread.setDaemon(True)
@@ -93,30 +103,34 @@ class Client:
         self.frames = b""
         print("[INFO]: * recording")
 
-    @staticmethod
+    def on_message(self, ws, message):
-    def on_message(ws, message):
+        self.last_response_recieved = time.time()
-        Client.LAST_RESPONSE_RECIEVED = time.time()
         message = json.loads(message)
-        if message.get('uid')!=Client.uid:
+
+        if self.uid != message.get("uid"):
             print("[ERROR]: invalid client uid")
             return
 
         if "status" in message.keys() and message["status"] == "WAIT":
-            Client.WAITING = True
+            self.waiting = True
-            print(f"[INFO]:Server is full. Estimated wait time {round(message['message'])} minutes.")
+            print(
+                f"[INFO]:Server is full. Estimated wait time {round(message['message'])} minutes."
+            )
 
         if "message" in message.keys() and message["message"] == "DISCONNECT":
             print("[INFO]: Server overtime disconnected.")
-            Client.RECORDING = False
+            self.recording = False
 
         if "message" in message.keys() and message["message"] == "SERVER_READY":
-            Client.RECORDING = True
+            self.recording = True
             return
 
         if "language" in message.keys():
-            Client.language = message.get("language")
+            self.language = message.get("language")
             lang_prob = message.get("language_prob")
-            print(f"[INFO]: Server detected language {Client.language} with probability {lang_prob}")
+            print(
+                f"[INFO]: Server detected language {self.language} with probability {lang_prob}"
+            )
             return
 
         if "segments" not in message.keys():
@@ -126,48 +140,47 @@ class Client:
         text = []
         if len(message):
             for seg in message:
-                if len(text):
+                if text and text[-1] == seg["text"]:
-                    if text[-1] != seg["text"]:
+                    # already got it
-                        text.append(seg["text"])
+                    continue
-                else:
                 text.append(seg["text"])
+        # keep only last 3
         if len(text) > 3:
             text = text[-3:]
         wrapper = textwrap.TextWrapper(width=60)
         word_list = wrapper.wrap(text="".join(text))
         # Print each line.
-        if os.name=='nt':
+        if os.name == "nt":
-            os.system('cls')
+            os.system("cls")
         else:
-            os.system('clear')
+            os.system("clear")
         for element in word_list:
             print(element)
 
-    @staticmethod
+    def on_error(self, ws, error):
-    def on_error(ws, error):
         print(error)
 
-    @staticmethod
+    def on_close(self, ws, close_status_code, close_msg):
-    def on_close(ws, close_status_code, close_msg):
+        print(f"[INFO]: Websocket connection closed: {close_status_code}: {close_msg}")
-        print(f"[INFO]: Websocket connection closed.")
 
-    @staticmethod
+    def on_open(self, ws):
-    def on_open(ws):
+        print(self.multilingual, self.language, self.task)
-        print(Client.multilingual, Client.language, Client.task)
 
         print("[INFO]: Opened connection")
-        ws.send(json.dumps({
+        ws.send(
-            'uid': Client.uid,
+            json.dumps(
-            'multilingual': Client.multilingual,
+                {
-            'language': Client.language,
+                    "uid": self.uid,
-            'task': Client.task
+                    "multilingual": self.multilingual,
-        }))
+                    "language": self.language,
+                    "task": self.task,
+                }
+            )
+        )
 
     @staticmethod
     def bytes_to_float_array(audio_bytes):
-        raw_data = np.frombuffer(
+        raw_data = np.frombuffer(buffer=audio_bytes, dtype=np.int16)
-            buffer=audio_bytes, dtype=np.int16
-        )
         return raw_data.astype(np.float32) / 32768.0
 
     def send_packet_to_server(self, message):
@@ -178,31 +191,36 @@ class Client:
 
     def play_file(self, filename):
         # read audio and create pyaudio stream
-        self.wf = wave.open(filename, 'rb')
+        with wave.open(filename, "rb") as wavfile:
-        self.stream = self.p.open(format=self.p.get_format_from_width(self.wf.getsampwidth()),
+            self.stream = self.p.open(
-                channels=self.wf.getnchannels(),
+                format=self.p.get_format_from_width(wavfile.getsampwidth()),
-                rate=self.wf.getframerate(),
+                channels=wavfile.getnchannels(),
+                rate=wavfile.getframerate(),
                 input=True,
                 output=True,
-                frames_per_buffer=self.CHUNK)
+                frames_per_buffer=self.chunk,
+            )
             try:
-            while Client.RECORDING:
+                while self.recording:
-                data = self.wf.readframes(self.CHUNK)
+                    data = wavfile.readframes(self.chunk)
-                if data==b'': break
+                    if data == b"":
+                        break
 
-                audio_array = Client.bytes_to_float_array(data)
+                    audio_array = self.bytes_to_float_array(data)
                     self.send_packet_to_server(audio_array.tobytes())
                     self.stream.write(data)
 
-            self.wf.close()
+                wavfile.close()
-            elapsed_time = time.time() - self.LAST_RESPONSE_RECIEVED
+
-            while elapsed_time < self.DISCONNECT_IF_NO_RESPONSE_FOR:
+                assert self.last_response_recieved
+                elapsed_time = time.time() - self.last_response_recieved
+                while elapsed_time < self.disconnect_if_no_response_for:
                     continue
                 self.stream.close()
                 self.close_websocket()
 
             except KeyboardInterrupt:
-            self.wf.close()
+                wavfile.close()
                 self.stream.stop_stream()
                 self.stream.close()
                 self.p.terminate()
@@ -224,12 +242,12 @@ class Client:
         return self.client_socket
 
     def write_audio_frames_to_file(self, frames, file_name):
-        wf = wave.open(file_name, 'wb')
+        with wave.open(file_name, "wb") as wavfile:
-        wf.setnchannels(self.CHANNELS)
+            wavfile: wave.Wave_write
-        wf.setsampwidth(2)
+            wavfile.setnchannels(self.channels)
-        wf.setframerate(self.RATE)
+            wavfile.setsampwidth(2)
-        wf.writeframes(frames)
+            wavfile.setframerate(self.rate)
-        wf.close()
+            wavfile.writeframes(frames)
 
     def record(self, out_file="output_recording.wav"):
         n_audio_file = 0
@@ -237,9 +255,10 @@ class Client:
         if not os.path.exists("chunks"):
             os.makedirs("chunks", exist_ok=True)
         try:
-            for _ in range(0, int(self.RATE / self.CHUNK * self.RECORD_SECONDS)):
+            for _ in range(0, int(self.rate / self.chunk * self.record_seconds)):
-                if not Client.RECORDING: break
+                if not self.recording:
-                data = self.stream.read(self.CHUNK)
+                    break
+                data = self.stream.read(self.chunk)
                 self.frames += data
 
                 audio_array = Client.bytes_to_float_array(data)
@@ -247,10 +266,13 @@ class Client:
                 self.send_packet_to_server(audio_array.tobytes())
 
                 # save frames if more than a minute
-                if len(self.frames) > 60*self.RATE:
+                if len(self.frames) > 60 * self.rate:
                     t = threading.Thread(
                         target=self.write_audio_frames_to_file,
-                        args=(self.frames[:], f"chunks/{n_audio_file}.wav", )
+                        args=(
+                            self.frames[:],
+                            f"chunks/{n_audio_file}.wav",
+                        ),
                     )
                     t.start()
                     n_audio_file += 1
@@ -259,7 +281,8 @@ class Client:
         except KeyboardInterrupt:
             if len(self.frames):
                 self.write_audio_frames_to_file(
-                    self.frames[:], f"chunks/{n_audio_file}.wav")
+                    self.frames[:], f"chunks/{n_audio_file}.wav"
+                )
                 n_audio_file += 1
             self.stream.stop_stream()
             self.stream.close()
@@ -270,21 +293,26 @@ class Client:
             self.write_output_recording(n_audio_file, out_file)
 
     def write_output_recording(self, n_audio_file, out_file):
-        input_files = [f"chunks/{i}.wav" for i in range(n_audio_file) if os.path.exists(f"chunks/{i}.wav")]
+        input_files = [
-        wf = wave.open(out_file, 'wb')
+            f"chunks/{i}.wav"
-        wf.setnchannels(self.CHANNELS)
+            for i in range(n_audio_file)
-        wf.setsampwidth(2)
+            if os.path.exists(f"chunks/{i}.wav")
-        wf.setframerate(self.RATE)
+        ]
+        with wave.open(out_file, "wb") as wavfile:
+            wavfile: wave.Wave_write
+            wavfile.setnchannels(self.channels)
+            wavfile.setsampwidth(2)
+            wavfile.setframerate(self.rate)
             for in_file in input_files:
-            w = wave.open(in_file, 'rb')
+                with wave.open(in_file, "rb") as wav_in:
                     while True:
-                data = w.readframes(self.CHUNK)
+                        data = wav_in.readframes(self.chunk)
-                if data==b'': break
+                        if data == b"":
-                wf.writeframes(data)
+                            break
-            w.close()
+                        wavfile.writeframes(data)
                 # remove this file
                 os.remove(in_file)
-        wf.close()
+        wavfile.close()
 
 
 class TranscriptionClient:
@@ -293,8 +321,8 @@ class TranscriptionClient:
 
     def __call__(self, audio=None):
         print("[INFO]: Waiting for server ready ...")
-        while not Client.RECORDING:
+        while not self.client.recording:
-            if Client.WAITING:
+            if self.client.waiting:
                 self.client.close_websocket()
                 return
             pass
@@ -304,4 +332,3 @@ class TranscriptionClient:
             self.client.play_file(resampled_file)
         else:
             self.client.record()
-