From 8e9baaea8fe207dcd3f9b8dcf66e8ab0357e042c Mon Sep 17 00:00:00 2001
From: makaveli10 <suryanvineet47@gmail.com>
Date: Mon, 17 Jul 2023 22:06:59 +0530
Subject: [PATCH] add multilingual, task option to python client & server

---
 client.py | 21 +++++++++++++++++++++
 server.py | 36 +++++++++++++++++++++++-------------
 2 files changed, 44 insertions(+), 13 deletions(-)

diff --git a/client.py b/client.py
index 62dfd83..b38f7bd 100644
--- a/client.py
+++ b/client.py
@@ -18,6 +18,9 @@ CHANNELS = 1
 RATE = 16000
 RECORD_SECONDS = 60000
 START_RECORDING = False
+multilingual = False
+language = None
+
 
 
 def on_message(ws, message):
@@ -53,7 +56,16 @@ 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:
@@ -226,13 +238,22 @@ if __name__=="__main__":
     parser.add_argument('--audio', type=str, help='audio file to transcribe')
     parser.add_argument('--host', default=None, type=str, help='websocket server address to connect to')
     parser.add_argument('--port', default=None, type=str, help='websocket server port to connect to')
+    parser.add_argument('--multilingual', action="store_true", help='use multilingual model')
+    parser.add_argument('--language', default=None, type=str, help='languages to use')
+    parser.add_argument(
+        '--task', default="transcribe", type=str, help='task transcribe/translate (translates from any to english)')
     opt = parser.parse_args()
+    print(opt)
+    multilingual=opt.multilingual,
+    language = opt.language,
+    task = opt.task
     c = Client(host=opt.host, port=opt.port)
 
     # while loop to wait for server to be ready
     print("Waiting for server ready ...")
     while not START_RECORDING:
         pass
+    print("Server Ready!")
     if os.name=='nt':
         os.system('cls')
     else:
diff --git a/server.py b/server.py
index f189214..67543bb 100644
--- a/server.py
+++ b/server.py
@@ -26,14 +26,20 @@ def recv_audio(websocket):
     Receive audio chunks from client in an infinite loop.
     """
     global clients
-    client = ServeClient(websocket)
+    options = websocket.recv()
+    options = json.loads(options)
+    client = ServeClient(
+        websocket, 
+        multilingual=options["multilingual"],
+        language=options["language"],
+        task=options["task"]
+    )
+
     clients[websocket] = client
+    
     while True:
         try:
             frame_data = websocket.recv()
-            if isinstance(frame_data, str):
-                logging.info(frame_data)
-                continue
             frame_np = np.frombuffer(frame_data, np.float32)
             clients[websocket].add_frames(frame_np)
             
@@ -46,15 +52,16 @@ def recv_audio(websocket):
 
 class ServeClient:
     RATE = 16000
-    def __init__(self, websocket, topic=None, device=None):
-        self.payload_size = struct.calcsize("Q")
+    def __init__(self, websocket, task="transcribe", device=None, multilingual=False, language=None):
         self.data = b""
         self.frames = b""
+        self.language = language
+        self.task = task
         self.transcriber = WhisperModel(
-            "small.en", 
-            device="cuda",
+            "small" if multilingual else "small.en", 
+            device=device if device else "cuda",
             compute_type="float16", 
-            local_files_only=False
+            local_files_only=False,
         )
         
         # voice activity detection model
@@ -83,9 +90,6 @@ class ServeClient:
         self.wrapper = textwrap.TextWrapper(width=50)
         self.pick_previous_segments = 2
 
-        # setup mqtt
-        self.topic = topic
-
         # threading
         self.websocket = websocket
         self.trans_thread = threading.Thread(target=self.speech_to_text)
@@ -164,7 +168,13 @@ class ServeClient:
                     initial_prompt = None
                 
                 # whisper transcribe with prompt
-                result = self.transcriber.transcribe(input_sample, initial_prompt=initial_prompt)
+                result = self.transcriber.transcribe(
+                    input_sample, 
+                    initial_prompt=initial_prompt,
+                    language=self.language,
+                    task=self.task
+                )
+
                 if len(result):
                     self.t_start = None
                     last_segment = self.update_segments(result, duration)