Merge pull request #30 from makaveli10/package_pip

pip package whisper-live.

Dockerfile

@@ -39,6 +39,7 @@ COPY requirements/ /app
 RUN bash setup.sh
 RUN pip install -r server.txt
 
-COPY *py /app
+RUN pip install whisper-live
 
-CMD ["python", "server.py"]
+COPY run_server.py /app
+CMD ["python", "run_server.py"]

README.md

@@ -11,41 +11,43 @@ Unlike traditional speech recognition systems that rely on continuous audio stre
  bash setup.sh
 ```
 
-- To install client requirements
+- Install whisper-live from pip
 ```bash
- pip install -r requirements/client.txt
+ pip install whisper-live
-```
-
-- To install server requirements
-```bash
- pip install -r requirements/server.txt
 ```
 
 ## Getting Started
 - Run the server
-```bash
+```python
- python server.py
+ from whisper_live.server import TranscriptionServer
+ server = TranscriptionServer()
+ server.run("0.0.0.0", 9090)
 ```
 
 - On the client side
     - To transcribe an audio file:
-    ```bash
+    ```python
-      python client.py --audio "audio.wav" --host "localhost" --port "9090" --multilingual --language "hi" --task "transcribe"
+      from whisper_live.client import TranscriptionClient
-                                                                                                                  "translate"
+      client = TranscriptionClient("localhost", 9090, multilingual=True, language="hi", translate=True)
+      client(audio_file_path)
     ```
-    This command transcribes the specified audio file (audio.wav) using the Whisper model. It connects to the server running on localhost at port 9090. It also enables the multilingual feature, allowing transcription in multiple languages. The --language flag specifies the target language for transcription, in this case, Hindi ("hi"). The --task flag is set to "transcribe" to indicate that transcription is the desired task. Also, --task can be set to "translate" to translate source language to English.
+    This command transcribes the specified audio file (audio.wav) using the Whisper model. It connects to the server running on localhost at port 9090. It also enables the multilingual feature, allowing transcription in multiple languages. The language option specifies the target language for transcription, in this case, Hindi ("hi"). The translate option should be set to `True` if we want to translate from the source language to English and `False` if we want to transcribe in the source language.
 
     - To transcribe from microphone:
-    ```bash
+    ```python
-     python client.py --host "localhost" --port "9090" --multilingual --language "en" --task "transcribe"
+      from whisper_live.client import TranscriptionClient
+      client = TranscriptionClient(host, port, multilingual=True, language="hi", translate=True)
+      client()
     ```
     This command captures audio from the microphone and sends it to the server for transcription. It uses the same options as the previous command, enabling the multilingual feature and specifying the target language and task.
 
 
 ## Transcribe audio from browser
 - Run the server
-```bash
+```python
- python server.py
+ from whisper_live.server import TranscriptionServer
+ server = TranscriptionServer()
+ server.run("0.0.0.0", 9090)
 ```
 This would start the websocket server on port ```9090```.
 
@@ -67,9 +69,8 @@ This would start the websocket server on port ```9090```.
 ```
 
 ## Future Work
-- [x] Update Documentation.
-- [x] Keep only a single server implementation i.e. websockets and get rid of the socket implementation in ```server.py```. Also, update ```client.py``` to websockets-client implemenation.
 - [ ] Add translation to other languages on top of transcription.
+- [ ] TensorRT backend for Whisper.
 
 ## Citations
 ```bibtex

requirements/client.txt

@@ -2,4 +2,3 @@ PyAudio
 ffmpeg-python
 scipy
 websocket-client
-onnxruntime

requirements/server.txt

@@ -4,3 +4,4 @@ faster-whisper==0.6.0
 torch==1.10.1
 torchaudio==0.10.1
 websockets
+onnxruntime

run_server.py

@@ -0,0 +1,6 @@
+
+from whisper_live.server import TranscriptionServer
+
+if __name__ == "__main__":
+    server = TranscriptionServer()
+    server.run_server("0.0.0.0", 9090)

setup.py

@@ -0,0 +1,44 @@
+import pathlib
+from setuptools import find_packages, setup
+
+# The directory containing this file
+HERE = pathlib.Path(__file__).parent
+
+# The text of the README file
+README = (HERE / "README.md").read_text()
+
+# This call to setup() does all the work
+setup(name="whisper-live",
+      version="0.0.4",
+      description="A nearly-live implementation of OpenAI's Whisper.",
+      long_description=README,
+      long_description_content_type="text/markdown",
+      include_package_data=True,
+      url="https://github.com/collabora/WhisperLive",
+      author="Collabora Ltd",
+      author_email="vineet.suryan@collabora.com",
+      license="MIT",
+      classifiers=[
+          "License :: OSI Approved :: MIT License",
+          "Programming Language :: Python :: 3.8",
+          "Programming Language :: Python :: 3.9",
+      ],
+      packages=find_packages(
+        exclude=("examples",
+                 "Audio-Transcription-Chrome",
+                 "Audio-Transcription-Firefox",
+                 "requirements",
+                 "whisper-finetuning"
+                )
+      ),
+      install_requires=[
+        "PyAudio",
+        "faster-whisper==0.6.0",
+        "torch",
+        "torchaudio",
+        "websockets",
+        "onnxruntime",
+        "ffmpeg-python",
+        "scipy",
+        "websocket-client",
+      ])

setup.sh

@@ -1,3 +1,3 @@
 #! /bin/bash
 
-apt-get install portaudio19-dev -y
+apt-get install portaudio19-dev ffmpeg -y

whisper_live/client.py

@@ -12,204 +12,6 @@ import json
 import websocket
 
 
-CHUNK = 1024
-FORMAT = pyaudio.paInt16
-CHANNELS = 1
-RATE = 16000
-RECORD_SECONDS = 60000
-START_RECORDING = False
-multilingual = False
-language = None
-
-
-
-def on_message(ws, message):
-    global START_RECORDING, language
-    message = json.loads(message)
-    if message == "SERVER_READY":
-        START_RECORDING = True
-        return
-
-    if isinstance(message, dict):
-        language = message.get("language")
-        lang_prob = message.get("language_prob")
-        print(f"Server detected language {language} with probability {lang_prob}")
-        return
-
-    text = []
-    if len(message):
-        for seg in message:
-            if len(text):
-                if text[-1] != seg["text"]:
-                    text.append(seg["text"])
-            else:
-                text.append(seg["text"])
-    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':
-        os.system('cls')
-    else:
-        os.system('clear')
-    for element in word_list:
-        print(element)
-
-def on_error(ws, error):
-    print(error)
-
-def on_close(ws, close_status_code, close_msg):
-    print("### websocket connection closed ###")
-
-def on_open(ws):
-    global multilingual, language, task
-    print(multilingual, language, task)
-
-    print("Opened connection")
-    ws.send(json.dumps({
-        'multilingual': multilingual[0],
-        'language': language[0],
-        'task': task
-    }))
-
-
-class Client:
-    def __init__(self, host=None, port=None):
-        self.timestamp_offset = 0.0
-        self.audio_bytes = None
-        self.p = pyaudio.PyAudio()
-        self.stream = self.p.open(format=FORMAT,
-                        channels=CHANNELS,
-                        rate=RATE,
-                        input=True,
-                        frames_per_buffer=CHUNK)
-        print(self.p.get_sample_size(FORMAT))
-
-        # 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,
-                                  on_open=on_open,
-                                  on_message=on_message,
-                                  on_error=on_error,
-                                  on_close=on_close)
-        else:
-            print("No host or port specified.")
-            return
-
-        # start websocket client in a thread
-        self.ws_thread = threading.Thread(target=self.client_socket.run_forever)
-        self.ws_thread.setDaemon(True)
-        self.ws_thread.start()
-
-        self.frames = b""
-        print("* recording")
-
-    def send_packet_to_server(self, message):
-        try:
-            self.client_socket.send(message, websocket.ABNF.OPCODE_BINARY)
-        except Exception as e:
-            print(e)
-
-    @staticmethod
-    def bytes_to_float_array(audio_bytes):
-        raw_data = np.frombuffer(
-            buffer=audio_bytes, dtype=np.int16
-        )
-        return raw_data.astype(np.float32) / 32768.0
-    
-    def play_file(self, filename):
-        # read audio and create pyaudio stream
-        self.wf = wave.open(filename, 'rb')
-        self.stream = self.p.open(format=self.p.get_format_from_width(self.wf.getsampwidth()),
-                channels=self.wf.getnchannels(),
-                rate=self.wf.getframerate(),
-                input=True,
-                output=True,
-                frames_per_buffer=CHUNK)
-        try:
-            while True:
-                data = self.wf.readframes(CHUNK)
-                if data==b'': break
-
-                audio_array = Client.bytes_to_float_array(data)
-                self.send_packet_to_server(audio_array.tobytes())
-                self.stream.write(data)
-
-            self.wf.close()
-            self.stream.close()
-
-        except KeyboardInterrupt:
-            print("Keyboard interrupt.")
-
-
-    def get_client_socket(self):
-        return self.client_socket
-    
-    def write_audio_frames_to_file(self, frames, file_name):
-        wf = wave.open(file_name, 'wb')
-        wf.setnchannels(CHANNELS)
-        wf.setsampwidth(2)
-        wf.setframerate(RATE)
-        wf.writeframes(frames)
-        wf.close()
-
-    def record(self, out_file="output_recording.wav"):
-        n_audio_file = 0
-        # create dir for saving audio chunks
-        if not os.path.exists("chunks"):
-            os.makedirs("chunks", exist_ok=True)
-        try:
-            for _ in range(0, int(RATE / CHUNK * RECORD_SECONDS)):
-                data = self.stream.read(CHUNK)
-                self.frames += data
-
-                audio_array = Client.bytes_to_float_array(data)
-                
-                self.send_packet_to_server(audio_array.tobytes())
-
-                # save frames if more than a minute
-                if len(self.frames) > 60*RATE:
-                    t = threading.Thread(
-                        target=self.write_audio_frames_to_file,
-                        args=(self.frames[:], f"chunks/{n_audio_file}.wav", )
-                    )
-                    t.start()
-                    n_audio_file += 1
-                    self.frames = b""
-
-        except KeyboardInterrupt:
-            if len(self.frames):
-                self.write_audio_frames_to_file(
-                    self.frames[:], f"chunks/{n_audio_file}.wav")
-                n_audio_file += 1
-            self.stream.stop_stream()
-            self.stream.close()
-            self.p.terminate()
-
-
-            # combine all the audio files
-            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")]
-        wf = wave.open(out_file, 'wb')
-        wf.setnchannels(CHANNELS)
-        wf.setsampwidth(2)
-        wf.setframerate(RATE)
-        for in_file in input_files:
-            w = wave.open(in_file, 'rb')
-            while True:
-                data = w.readframes(CHUNK)
-                if data==b'': break
-                wf.writeframes(data)
-            w.close()
-            # remove this file
-            os.remove(in_file)
-        wf.close()
-
-
 def resample(file: str, sr: int = 16000):
     """
     # https://github.com/openai/whisper/blob/7858aa9c08d98f75575035ecd6481f462d66ca27/whisper/audio.py#L22
@@ -239,30 +41,235 @@ def resample(file: str, sr: int = 16000):
     return resampled_file
 
 
-if __name__=="__main__":
+class Client:
-    parser = argparse.ArgumentParser()
+    CHUNK = 1024
-    parser.add_argument('--audio', type=str, help='audio file to transcribe')
+    FORMAT = pyaudio.paInt16
-    parser.add_argument('--host', default=None, type=str, help='websocket server address to connect to')
+    CHANNELS = 1
-    parser.add_argument('--port', default=None, type=str, help='websocket server port to connect to')
+    RATE = 16000
-    parser.add_argument('--multilingual', action="store_true", help='use multilingual model')
+    RECORD_SECONDS = 60000
-    parser.add_argument('--language', default=None, type=str, help='languages to use')
+    START_RECORDING = False
-    parser.add_argument(
+    multilingual = False
-        '--task', default="transcribe", type=str, help='task transcribe/translate (translates from any to english)')
+    language = None
-    opt = parser.parse_args()
+    task = "transcribe"
-    print(opt)
-    multilingual=opt.multilingual,
-    language = opt.language if opt.multilingual else "en",
-    task = opt.task
-    c = Client(host=opt.host, port=opt.port)
     
-    # while loop to wait for server to be ready
+    def __init__(self, host=None, port=None, is_multilingual=False, lang=None, translate=False):
-    print("Waiting for server ready ...")
+        Client.multilingual = is_multilingual
-    while not START_RECORDING:
+        Client.language = lang if is_multilingual else "en"
-        pass
+        if translate:
-    print("Server Ready!")
+            Client.task = "translate"
+
+        self.timestamp_offset = 0.0
+        self.audio_bytes = None
+        self.p = pyaudio.PyAudio()
+        self.stream = self.p.open(format=self.FORMAT,
+                        channels=self.CHANNELS,
+                        rate=self.RATE,
+                        input=True,
+                        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,
+                                  on_open=Client.on_open,
+                                  on_message=Client.on_message,
+                                  on_error=Client.on_error,
+                                  on_close=Client.on_close)
+        else:
+            print("[ERROR]: No host or port specified.")
+            return
+
+        # start websocket client in a thread
+        self.ws_thread = threading.Thread(target=self.client_socket.run_forever)
+        self.ws_thread.setDaemon(True)
+        self.ws_thread.start()
+
+        self.frames = b""
+        print("[INFO]: * recording")
+    
+    @staticmethod
+    def on_message(ws, message):
+        message = json.loads(message)
+        if message == "SERVER_READY":
+            Client.START_RECORDING = True
+            return
+
+        if isinstance(message, dict):
+            Client.language = message.get("language")
+            lang_prob = message.get("language_prob")
+            print(f"[INFO]: Server detected language {Client.language} with probability {lang_prob}")
+            return
+
+        text = []
+        if len(message):
+            for seg in message:
+                if len(text):
+                    if text[-1] != seg["text"]:
+                        text.append(seg["text"])
+                else:
+                    text.append(seg["text"])
+        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':
+            os.system('cls')
+        else:
+            os.system('clear')
+        for element in word_list:
+            print(element)
+
+    @staticmethod
+    def on_error(ws, error):
+        print(error)
+
+    @staticmethod
+    def on_close(ws, close_status_code, close_msg):
+        print(f"[INFO]: Websocket connection closed.")
+
+    @staticmethod
+    def on_open(ws):
+        print(Client.multilingual, Client.language, Client.task)
+
+        print("[INFO]: Opened connection")
+        ws.send(json.dumps({
+            'multilingual': Client.multilingual,
+            'language': Client.language,
+            'task': Client.task
+        }))
+
+    @staticmethod
+    def bytes_to_float_array(audio_bytes):
+        raw_data = np.frombuffer(
+            buffer=audio_bytes, dtype=np.int16
+        )
+        return raw_data.astype(np.float32) / 32768.0
+    
+    def send_packet_to_server(self, message):
+        try:
+            self.client_socket.send(message, websocket.ABNF.OPCODE_BINARY)
+        except Exception as e:
+            print(e)
+    
+    def play_file(self, filename):
+        # read audio and create pyaudio stream
+        self.wf = wave.open(filename, 'rb')
+        self.stream = self.p.open(format=self.p.get_format_from_width(self.wf.getsampwidth()),
+                channels=self.wf.getnchannels(),
+                rate=self.wf.getframerate(),
+                input=True,
+                output=True,
+                frames_per_buffer=self.CHUNK)
+        try:
+            while True:
+                data = self.wf.readframes(self.CHUNK)
+                if data==b'': break
+
+                audio_array = Client.bytes_to_float_array(data)
+                self.send_packet_to_server(audio_array.tobytes())
+                self.stream.write(data)
+
+            self.wf.close()
+            self.stream.close()
+
+        except KeyboardInterrupt:
+            self.wf.close()
+            self.stream.stop_stream()
+            self.stream.close()
+            self.p.terminate()
+            self.close_websocket()
+            print("[INFO]: Keyboard interrupt.")
+
+    def close_websocket(self):
+        try:
+            self.client_socket.close()  # Close the WebSocket connection
+        except Exception as e:
+            print("[ERROR]: Error closing WebSocket:", e)
+
+        try:
+            self.ws_thread.join()  # Wait for the WebSocket thread to finish
+        except Exception as e:
+            print("[ERROR:] Error joining WebSocket thread:", e)
+
+    def get_client_socket(self):
+        return self.client_socket
+    
+    def write_audio_frames_to_file(self, frames, file_name):
+        wf = wave.open(file_name, 'wb')
+        wf.setnchannels(self.CHANNELS)
+        wf.setsampwidth(2)
+        wf.setframerate(self.RATE)
+        wf.writeframes(frames)
+        wf.close()
+
+    def record(self, out_file="output_recording.wav"):
+        n_audio_file = 0
+        # create dir for saving audio chunks
+        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)):
+                data = self.stream.read(self.CHUNK)
+                self.frames += data
+
+                audio_array = Client.bytes_to_float_array(data)
+                
+                self.send_packet_to_server(audio_array.tobytes())
+
+                # save frames if more than a minute
+                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", )
+                    )
+                    t.start()
+                    n_audio_file += 1
+                    self.frames = b""
+
+        except KeyboardInterrupt:
+            if len(self.frames):
+                self.write_audio_frames_to_file(
+                    self.frames[:], f"chunks/{n_audio_file}.wav")
+                n_audio_file += 1
+            self.stream.stop_stream()
+            self.stream.close()
+            self.p.terminate()
+            self.close_websocket()
+
+            # combine all the audio files
+            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")]
+        wf = wave.open(out_file, 'wb')
+        wf.setnchannels(self.CHANNELS)
+        wf.setsampwidth(2)
+        wf.setframerate(self.RATE)
+        for in_file in input_files:
+            w = wave.open(in_file, 'rb')
+            while True:
+                data = w.readframes(self.CHUNK)
+                if data==b'': break
+                wf.writeframes(data)
+            w.close()
+            # remove this file
+            os.remove(in_file)
+        wf.close()
+
+
+class TranscriptionClient:
+    def __init__(self, host, port, is_multilingual=False, lang=None, translate=False):
+        self.client = Client(host, port, is_multilingual, lang, translate)
+        
+    def __call__(self, audio=None):
+        print("[INFO]: Waiting for server ready ...")
+        while not Client.START_RECORDING:
+            pass
+        print("[INFO]: Server Ready!")
+        if audio is not None:
+            resampled_file = resample(audio)
+            self.client.play_file(resampled_file)
+        else:
+            self.client.record()
 
-    if opt.audio is not None:
-        resampled_file = resample(opt.audio)
-        c.play_file(resampled_file)
-    else:
-        c.record()

whisper_live/server.py

@@ -10,48 +10,70 @@ logging.basicConfig(level = logging.INFO)
 
 from collections import deque
 from dataclasses import dataclass
+from websockets.sync.server import serve
 
 import torch
 import numpy as np
-from websockets.sync import server
+from whisper_live.transcriber import WhisperModel
-from websockets.sync.server import serve
-from transcriber import WhisperModel
 
 
-clients = {}
+class TranscriptionServer:
-SERVER_READY = "SERVER_READY"
-
-def recv_audio(websocket):
     """
-    Receive audio chunks from client in an infinite loop.
+    Represents a transcription server that handles incoming audio from clients.
+
+    Attributes:
+        clients (dict): A dictionary to store connected clients.
     """
-    global clients
-    options = websocket.recv()
-    options = json.loads(options)
-    client = ServeClient(
-        websocket, 
-        multilingual=options["multilingual"],
-        language=options["language"],
-        task=options["task"]
-    )
 
-    clients[websocket] = client
+    def __init__(self):
+        self.clients = {}
 
-    while True:
+    def recv_audio(self, websocket):
-        try:
+        """
-            frame_data = websocket.recv()
+        Receive audio chunks from a client in an infinite loop.
-            frame_np = np.frombuffer(frame_data, np.float32)
-            clients[websocket].add_frames(frame_np)
 
-        except Exception as e:
+        Args:
-            clients[websocket].cleanup()
+            websocket (WebSocket): The WebSocket connection for the client.
-            clients.pop(websocket)
+        """
-            logging.info("Connection Closed.")
+        options = websocket.recv()
-            break
+        options = json.loads(options)
+        client = ServeClient(
+            websocket,
+            multilingual=options["multilingual"],
+            language=options["language"],
+            task=options["task"],
+        )
+
+        self.clients[websocket] = client
+
+        while True:
+            try:
+                frame_data = websocket.recv()
+                frame_np = np.frombuffer(frame_data, np.float32)
+                self.clients[websocket].add_frames(frame_np)
+
+            except Exception as e:
+                self.clients[websocket].cleanup()
+                self.clients.pop(websocket)
+                logging.info("Connection Closed.")
+                break
+
+    def run(self, host, port):
+        """
+        Run the transcription server.
+
+        Args:
+            host (str): The host address to bind the server.
+            port (int): The port number to bind the server.
+        """
+        with serve(self.recv_audio, host, port) as server:
+            server.serve_forever()
 
 
 class ServeClient:
     RATE = 16000
+    SERVER_READY = "SERVER_READY"
+
     def __init__(self, websocket, task="transcribe", device=None, multilingual=False, language=None):
         self.data = b""
         self.frames = b""
@@ -94,7 +116,7 @@ class ServeClient:
         self.websocket = websocket
         self.trans_thread = threading.Thread(target=self.speech_to_text)
         self.trans_thread.start()
-        self.websocket.send(json.dumps(SERVER_READY))
+        self.websocket.send(json.dumps(self.SERVER_READY))
     
     def fill_output(self, output):
         """
@@ -303,7 +325,3 @@ class ServeClient:
         self.exit = True
         self.transcriber.destroy()
 
-
-if __name__ == "__main__":
-    with serve(recv_audio, "0.0.0.0", 9090) as server:
-        server.serve_forever()