first commit for takway.ai

runtime/python/websocket/README.md

@@ -0,0 +1,127 @@
+# Service with websocket-python
+
+This is a demo using funasr pipeline with websocket python-api. It supports the offline, online, offline/online-2pass unifying speech recognition. 
+
+## For the Server
+
+### Install the modelscope and funasr
+
+```shell
+pip install -U modelscope funasr
+# For the users in China, you could install with the command:
+# pip install -U modelscope funasr -i https://mirror.sjtu.edu.cn/pypi/web/simple
+git clone https://github.com/alibaba/FunASR.git && cd FunASR
+```
+
+### Install the requirements for server
+
+```shell
+cd runtime/python/websocket
+pip install -r requirements_server.txt
+```
+
+### Start server
+
+##### API-reference
+```shell
+python funasr_wss_server.py \
+--port [port id] \
+--asr_model [asr model_name] \
+--asr_model_online [asr model_name] \
+--punc_model [punc model_name] \
+--ngpu [0 or 1] \
+--ncpu [1 or 4] \
+--certfile [path of certfile for ssl] \
+--keyfile [path of keyfile for ssl] 
+```
+##### Usage examples
+```shell
+python funasr_wss_server.py --port 10095
+```
+
+## For the client
+
+Install the requirements for client
+```shell
+git clone https://github.com/alibaba/FunASR.git && cd FunASR
+cd funasr/runtime/python/websocket
+pip install -r requirements_client.txt
+```
+If you want infer from videos, you should install `ffmpeg`
+```shell
+apt-get install -y ffmpeg #ubuntu
+# yum install -y ffmpeg # centos
+# brew install ffmpeg # mac
+# winget install ffmpeg # wins
+pip3 install websockets ffmpeg-python
+```
+
+### Start client
+#### API-reference
+```shell
+python funasr_wss_client.py \
+--host [ip_address] \
+--port [port id] \
+--chunk_size ["5,10,5"=600ms, "8,8,4"=480ms] \
+--chunk_interval [duration of send chunk_size/chunk_interval] \
+--words_max_print [max number of words to print] \
+--audio_in [if set, loadding from wav.scp, else recording from mircrophone] \
+--output_dir [if set, write the results to output_dir] \
+--mode [`online` for streaming asr, `offline` for non-streaming, `2pass` for unifying streaming and non-streaming asr] \
+--thread_num [thread_num for send data]
+```
+
+#### Usage examples
+##### ASR offline client
+Recording from mircrophone
+```shell
+# --chunk_interval, "10": 600/10=60ms, "5"=600/5=120ms, "20": 600/12=30ms
+python funasr_wss_client.py --host "0.0.0.0" --port 10095 --mode offline
+```
+Loadding from wav.scp(kaldi style)
+```shell
+# --chunk_interval, "10": 600/10=60ms, "5"=600/5=120ms, "20": 600/12=30ms
+python funasr_wss_client.py --host "0.0.0.0" --port 10095 --mode offline --audio_in "./data/wav.scp" --output_dir "./results"
+```
+
+##### ASR streaming client
+Recording from mircrophone
+```shell
+# --chunk_size, "5,10,5"=600ms, "8,8,4"=480ms
+python funasr_wss_client.py --host "0.0.0.0" --port 10095 --mode online --chunk_size "5,10,5"
+```
+Loadding from wav.scp(kaldi style)
+```shell
+# --chunk_size, "5,10,5"=600ms, "8,8,4"=480ms
+python funasr_wss_client.py --host "0.0.0.0" --port 10095 --mode online --chunk_size "5,10,5" --audio_in "./data/wav.scp" --output_dir "./results"
+```
+
+##### ASR offline/online 2pass client
+Recording from mircrophone
+```shell
+# --chunk_size, "5,10,5"=600ms, "8,8,4"=480ms
+python funasr_wss_client.py --host "0.0.0.0" --port 10095 --mode 2pass --chunk_size "8,8,4"
+```
+Loadding from wav.scp(kaldi style)
+```shell
+# --chunk_size, "5,10,5"=600ms, "8,8,4"=480ms
+python funasr_wss_client.py --host "0.0.0.0" --port 10095 --mode 2pass --chunk_size "8,8,4" --audio_in "./data/wav.scp" --output_dir "./results"
+```
+
+#### Websocket api
+```shell
+# class Funasr_websocket_recognizer example with 3 step
+# 1.create an recognizer 
+rcg=Funasr_websocket_recognizer(host="127.0.0.1",port="30035",is_ssl=True,mode="2pass")
+# 2.send pcm data to asr engine and get asr result
+text=rcg.feed_chunk(data)
+print("text",text)
+# 3.get last result, set timeout=3
+text=rcg.close(timeout=3)
+print("text",text)
+```
+
+## Acknowledge
+1. This project is maintained by [FunASR community](https://github.com/alibaba-damo-academy/FunASR).
+2. We acknowledge [zhaoming](https://github.com/zhaomingwork/FunASR/tree/fix_bug_for_python_websocket) for contributing the websocket service.
+3. We acknowledge [cgisky1980](https://github.com/cgisky1980/FunASR) for contributing the websocket service of offline model.

runtime/python/websocket/funasr_client_api.py

@@ -0,0 +1,160 @@
+"""
+  Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights
+  Reserved. MIT License  (https://opensource.org/licenses/MIT)
+  
+  2022-2023 by zhaomingwork@qq.com  
+"""
+
+# pip install websocket-client
+import ssl
+from websocket import ABNF
+from websocket import create_connection
+from queue import Queue
+import threading
+import traceback
+import json
+import time
+import numpy as np
+
+
+# class for recognizer in websocket
+class Funasr_websocket_recognizer:
+    """
+    python asr recognizer lib
+
+    """
+
+    def __init__(
+        self,
+        host="127.0.0.1",
+        port="30035",
+        is_ssl=True,
+        chunk_size="0, 10, 5",
+        chunk_interval=10,
+        mode="offline",
+        wav_name="default",
+    ):
+        """
+        host: server host ip
+        port: server port
+        is_ssl: True for wss protocal, False for ws
+        """
+        try:
+            if is_ssl == True:
+                ssl_context = ssl.SSLContext()
+                ssl_context.check_hostname = False
+                ssl_context.verify_mode = ssl.CERT_NONE
+                uri = "wss://{}:{}".format(host, port)
+                ssl_opt = {"cert_reqs": ssl.CERT_NONE}
+            else:
+                uri = "ws://{}:{}".format(host, port)
+                ssl_context = None
+                ssl_opt = None
+            self.host = host
+            self.port = port
+
+            self.msg_queue = Queue()  # used for recognized result text
+
+            print("connect to url", uri)
+            self.websocket = create_connection(uri, ssl=ssl_context, sslopt=ssl_opt)
+
+            self.thread_msg = threading.Thread(
+                target=Funasr_websocket_recognizer.thread_rec_msg, args=(self,)
+            )
+            self.thread_msg.start()
+            chunk_size = [int(x) for x in chunk_size.split(",")]
+            stride = int(60 * chunk_size[1] / chunk_interval / 1000 * 16000 * 2)
+            chunk_num = (len(audio_bytes) - 1) // stride + 1
+
+            message = json.dumps(
+                {
+                    "mode": mode,
+                    "chunk_size": chunk_size,
+                    "encoder_chunk_look_back": 4,
+                    "decoder_chunk_look_back": 1,
+                    "chunk_interval": chunk_interval,
+                    "wav_name": wav_name,
+                    "is_speaking": True,
+                }
+            )
+
+            self.websocket.send(message)
+
+            print("send json", message)
+
+        except Exception as e:
+            print("Exception:", e)
+            traceback.print_exc()
+
+    # threads for rev msg
+    def thread_rec_msg(self):
+        try:
+            while True:
+                msg = self.websocket.recv()
+                if msg is None or len(msg) == 0:
+                    continue
+                msg = json.loads(msg)
+
+                self.msg_queue.put(msg)
+        except Exception as e:
+            print("client closed")
+
+    # feed data to asr engine, wait_time means waiting for result until time out
+    def feed_chunk(self, chunk, wait_time=0.01):
+        try:
+            self.websocket.send(chunk, ABNF.OPCODE_BINARY)
+            # loop to check if there is a message, timeout in 0.01s
+            while True:
+                msg = self.msg_queue.get(timeout=wait_time)
+                if self.msg_queue.empty():
+                    break
+
+            return msg
+        except:
+            return ""
+
+    def close(self, timeout=1):
+        message = json.dumps({"is_speaking": False})
+        self.websocket.send(message)
+        # sleep for timeout seconds to wait for result
+        time.sleep(timeout)
+        msg = ""
+        while not self.msg_queue.empty():
+            msg = self.msg_queue.get()
+
+        self.websocket.close()
+        # only resturn the last msg
+        return msg
+
+
+if __name__ == "__main__":
+
+    print("example for Funasr_websocket_recognizer")
+    import wave
+
+    wav_path = "/Users/zhifu/Downloads/modelscope_models/speech_seaco_paraformer_large_asr_nat-zh-cn-16k-common-vocab8404-pytorch/example/asr_example.wav"
+    with wave.open(wav_path, "rb") as wav_file:
+        params = wav_file.getparams()
+        frames = wav_file.readframes(wav_file.getnframes())
+        audio_bytes = bytes(frames)
+
+    stride = int(60 * 10 / 10 / 1000 * 16000 * 2)
+    chunk_num = (len(audio_bytes) - 1) // stride + 1
+    # create an recognizer
+    rcg = Funasr_websocket_recognizer(
+        host="127.0.0.1", port="10095", is_ssl=True, mode="2pass", chunk_size="0,10,5"
+    )
+    # loop to send chunk
+    for i in range(chunk_num):
+
+        beg = i * stride
+        data = audio_bytes[beg : beg + stride]
+
+        text = rcg.feed_chunk(data, wait_time=0.02)
+        if len(text) > 0:
+            print("text", text)
+        time.sleep(0.05)
+
+    # get last message
+    text = rcg.close(timeout=3)
+    print("text", text)

runtime/python/websocket/funasr_wss_client.py

@@ -0,0 +1,393 @@
+# -*- encoding: utf-8 -*-
+import os
+import time
+import websockets, ssl
+import asyncio
+
+# import threading
+import argparse
+import json
+import traceback
+from multiprocessing import Process
+
+# from funasr.fileio.datadir_writer import DatadirWriter
+
+import logging
+
+logging.basicConfig(level=logging.ERROR)
+
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    "--host", type=str, default="localhost", required=False, help="host ip, localhost, 0.0.0.0"
+)
+parser.add_argument("--port", type=int, default=10095, required=False, help="grpc server port")
+parser.add_argument("--chunk_size", type=str, default="5, 10, 5", help="chunk")
+parser.add_argument("--encoder_chunk_look_back", type=int, default=4, help="chunk")
+parser.add_argument("--decoder_chunk_look_back", type=int, default=0, help="chunk")
+parser.add_argument("--chunk_interval", type=int, default=10, help="chunk")
+parser.add_argument(
+    "--hotword",
+    type=str,
+    default="",
+    help="hotword file path, one hotword perline (e.g.:阿里巴巴 20)",
+)
+parser.add_argument("--audio_in", type=str, default=None, help="audio_in")
+parser.add_argument("--audio_fs", type=int, default=16000, help="audio_fs")
+parser.add_argument(
+    "--send_without_sleep",
+    action="store_true",
+    default=True,
+    help="if audio_in is set, send_without_sleep",
+)
+parser.add_argument("--thread_num", type=int, default=1, help="thread_num")
+parser.add_argument("--words_max_print", type=int, default=10000, help="chunk")
+parser.add_argument("--output_dir", type=str, default=None, help="output_dir")
+parser.add_argument("--ssl", type=int, default=1, help="1 for ssl connect, 0 for no ssl")
+parser.add_argument("--use_itn", type=int, default=1, help="1 for using itn, 0 for not itn")
+parser.add_argument("--mode", type=str, default="2pass", help="offline, online, 2pass")
+
+args = parser.parse_args()
+args.chunk_size = [int(x) for x in args.chunk_size.split(",")]
+print(args)
+# voices = asyncio.Queue()
+from queue import Queue
+
+voices = Queue()
+offline_msg_done = False
+
+if args.output_dir is not None:
+    # if os.path.exists(args.output_dir):
+    #     os.remove(args.output_dir)
+
+    if not os.path.exists(args.output_dir):
+        os.makedirs(args.output_dir)
+
+
+async def record_microphone():
+    is_finished = False
+    import pyaudio
+
+    # print("2")
+    global voices
+    FORMAT = pyaudio.paInt16
+    CHANNELS = 1
+    RATE = 16000
+    chunk_size = 60 * args.chunk_size[1] / args.chunk_interval
+    CHUNK = int(RATE / 1000 * chunk_size)
+
+    p = pyaudio.PyAudio()
+
+    stream = p.open(
+        format=FORMAT, channels=CHANNELS, rate=RATE, input=True, frames_per_buffer=CHUNK
+    )
+    # hotwords
+    fst_dict = {}
+    hotword_msg = ""
+    if args.hotword.strip() != "":
+        if os.path.exists(args.hotword):
+            f_scp = open(args.hotword)
+            hot_lines = f_scp.readlines()
+            for line in hot_lines:
+                words = line.strip().split(" ")
+                if len(words) < 2:
+                    print("Please checkout format of hotwords")
+                    continue
+                try:
+                    fst_dict[" ".join(words[:-1])] = int(words[-1])
+                except ValueError:
+                    print("Please checkout format of hotwords")
+            hotword_msg = json.dumps(fst_dict)
+        else:
+            hotword_msg = args.hotword
+
+    use_itn = True
+    if args.use_itn == 0:
+        use_itn = False
+
+    message = json.dumps(
+        {
+            "mode": args.mode,
+            "chunk_size": args.chunk_size,
+            "chunk_interval": args.chunk_interval,
+            "encoder_chunk_look_back": args.encoder_chunk_look_back,
+            "decoder_chunk_look_back": args.decoder_chunk_look_back,
+            "wav_name": "microphone",
+            "is_speaking": True,
+            "hotwords": hotword_msg,
+            "itn": use_itn,
+        }
+    )
+    # voices.put(message)
+    await websocket.send(message)
+    while True:
+        data = stream.read(CHUNK)
+        message = data
+        # voices.put(message)
+        await websocket.send(message)
+        await asyncio.sleep(0.005)
+
+
+async def record_from_scp(chunk_begin, chunk_size):
+    global voices
+    is_finished = False
+    if args.audio_in.endswith(".scp"):
+        f_scp = open(args.audio_in)
+        wavs = f_scp.readlines()
+    else:
+        wavs = [args.audio_in]
+
+    # hotwords
+    fst_dict = {}
+    hotword_msg = ""
+    if args.hotword.strip() != "":
+        if os.path.exists(args.hotword):
+            f_scp = open(args.hotword)
+            hot_lines = f_scp.readlines()
+            for line in hot_lines:
+                words = line.strip().split(" ")
+                if len(words) < 2:
+                    print("Please checkout format of hotwords")
+                    continue
+                try:
+                    fst_dict[" ".join(words[:-1])] = int(words[-1])
+                except ValueError:
+                    print("Please checkout format of hotwords")
+            hotword_msg = json.dumps(fst_dict)
+        else:
+            hotword_msg = args.hotword
+        print(hotword_msg)
+
+    sample_rate = args.audio_fs
+    wav_format = "pcm"
+    use_itn = True
+    if args.use_itn == 0:
+        use_itn = False
+
+    if chunk_size > 0:
+        wavs = wavs[chunk_begin : chunk_begin + chunk_size]
+    for wav in wavs:
+        wav_splits = wav.strip().split()
+
+        wav_name = wav_splits[0] if len(wav_splits) > 1 else "demo"
+        wav_path = wav_splits[1] if len(wav_splits) > 1 else wav_splits[0]
+        if not len(wav_path.strip()) > 0:
+            continue
+        if wav_path.endswith(".pcm"):
+            with open(wav_path, "rb") as f:
+                audio_bytes = f.read()
+        elif wav_path.endswith(".wav"):
+            import wave
+
+            with wave.open(wav_path, "rb") as wav_file:
+                params = wav_file.getparams()
+                sample_rate = wav_file.getframerate()
+                frames = wav_file.readframes(wav_file.getnframes())
+                audio_bytes = bytes(frames)
+        else:
+            wav_format = "others"
+            with open(wav_path, "rb") as f:
+                audio_bytes = f.read()
+
+        stride = int(60 * args.chunk_size[1] / args.chunk_interval / 1000 * sample_rate * 2)
+        chunk_num = (len(audio_bytes) - 1) // stride + 1
+        # print(stride)
+
+        # send first time
+        message = json.dumps(
+            {
+                "mode": args.mode,
+                "chunk_size": args.chunk_size,
+                "chunk_interval": args.chunk_interval,
+                "encoder_chunk_look_back": args.encoder_chunk_look_back,
+                "decoder_chunk_look_back": args.decoder_chunk_look_back,
+                "audio_fs": sample_rate,
+                "wav_name": wav_name,
+                "wav_format": wav_format,
+                "is_speaking": True,
+                "hotwords": hotword_msg,
+                "itn": use_itn,
+            }
+        )
+
+        # voices.put(message)
+        await websocket.send(message)
+        is_speaking = True
+        for i in range(chunk_num):
+
+            beg = i * stride
+            data = audio_bytes[beg : beg + stride]
+            message = data
+            # voices.put(message)
+            await websocket.send(message)
+            if i == chunk_num - 1:
+                is_speaking = False
+                message = json.dumps({"is_speaking": is_speaking})
+                # voices.put(message)
+                await websocket.send(message)
+
+            sleep_duration = (
+                0.001
+                if args.mode == "offline"
+                else 60 * args.chunk_size[1] / args.chunk_interval / 1000
+            )
+
+            await asyncio.sleep(sleep_duration)
+
+    if not args.mode == "offline":
+        await asyncio.sleep(2)
+    # offline model need to wait for message recved
+
+    if args.mode == "offline":
+        global offline_msg_done
+        while not offline_msg_done:
+            await asyncio.sleep(1)
+
+    await websocket.close()
+
+
+async def message(id):
+    global websocket, voices, offline_msg_done
+    text_print = ""
+    text_print_2pass_online = ""
+    text_print_2pass_offline = ""
+    if args.output_dir is not None:
+        ibest_writer = open(
+            os.path.join(args.output_dir, "text.{}".format(id)), "a", encoding="utf-8"
+        )
+    else:
+        ibest_writer = None
+    try:
+        while True:
+
+            meg = await websocket.recv()
+            meg = json.loads(meg)
+            wav_name = meg.get("wav_name", "demo")
+            text = meg["text"]
+            timestamp = ""
+            offline_msg_done = meg.get("is_final", False)
+            if "timestamp" in meg:
+                timestamp = meg["timestamp"]
+
+            if ibest_writer is not None:
+                if timestamp != "":
+                    text_write_line = "{}\t{}\t{}\n".format(wav_name, text, timestamp)
+                else:
+                    text_write_line = "{}\t{}\n".format(wav_name, text)
+                ibest_writer.write(text_write_line)
+
+            if "mode" not in meg:
+                continue
+            if meg["mode"] == "online":
+                text_print += "{}".format(text)
+                text_print = text_print[-args.words_max_print :]
+                os.system("clear")
+                print("\rpid" + str(id) + ": " + text_print)
+            elif meg["mode"] == "offline":
+                if timestamp != "":
+                    text_print += "{} timestamp: {}".format(text, timestamp)
+                else:
+                    text_print += "{}".format(text)
+
+                # text_print = text_print[-args.words_max_print:]
+                # os.system('clear')
+                print("\rpid" + str(id) + ": " + wav_name + ": " + text_print)
+                offline_msg_done = True
+            else:
+                if meg["mode"] == "2pass-online":
+                    text_print_2pass_online += "{}".format(text)
+                    text_print = text_print_2pass_offline + text_print_2pass_online
+                else:
+                    text_print_2pass_online = ""
+                    text_print = text_print_2pass_offline + "{}".format(text)
+                    text_print_2pass_offline += "{}".format(text)
+                text_print = text_print[-args.words_max_print :]
+                os.system("clear")
+                print("\rpid" + str(id) + ": " + text_print)
+                # offline_msg_done=True
+
+    except Exception as e:
+        print("Exception:", e)
+        # traceback.print_exc()
+        # await websocket.close()
+
+
+async def ws_client(id, chunk_begin, chunk_size):
+    if args.audio_in is None:
+        chunk_begin = 0
+        chunk_size = 1
+    global websocket, voices, offline_msg_done
+
+    for i in range(chunk_begin, chunk_begin + chunk_size):
+        offline_msg_done = False
+        voices = Queue()
+        if args.ssl == 1:
+            ssl_context = ssl.SSLContext()
+            ssl_context.check_hostname = False
+            ssl_context.verify_mode = ssl.CERT_NONE
+            uri = "wss://{}:{}".format(args.host, args.port)
+        else:
+            uri = "ws://{}:{}".format(args.host, args.port)
+            ssl_context = None
+        print("connect to", uri)
+        async with websockets.connect(
+            uri, subprotocols=["binary"], ping_interval=None, ssl=ssl_context
+        ) as websocket:
+            if args.audio_in is not None:
+                task = asyncio.create_task(record_from_scp(i, 1))
+            else:
+                task = asyncio.create_task(record_microphone())
+            task3 = asyncio.create_task(message(str(id) + "_" + str(i)))  # processid+fileid
+            await asyncio.gather(task, task3)
+    exit(0)
+
+
+def one_thread(id, chunk_begin, chunk_size):
+    asyncio.get_event_loop().run_until_complete(ws_client(id, chunk_begin, chunk_size))
+    asyncio.get_event_loop().run_forever()
+
+
+if __name__ == "__main__":
+    # for microphone
+    if args.audio_in is None:
+        p = Process(target=one_thread, args=(0, 0, 0))
+        p.start()
+        p.join()
+        print("end")
+    else:
+        # calculate the number of wavs for each preocess
+        if args.audio_in.endswith(".scp"):
+            f_scp = open(args.audio_in)
+            wavs = f_scp.readlines()
+        else:
+            wavs = [args.audio_in]
+        for wav in wavs:
+            wav_splits = wav.strip().split()
+            wav_name = wav_splits[0] if len(wav_splits) > 1 else "demo"
+            wav_path = wav_splits[1] if len(wav_splits) > 1 else wav_splits[0]
+            audio_type = os.path.splitext(wav_path)[-1].lower()
+
+        total_len = len(wavs)
+        if total_len >= args.thread_num:
+            chunk_size = int(total_len / args.thread_num)
+            remain_wavs = total_len - chunk_size * args.thread_num
+        else:
+            chunk_size = 1
+            remain_wavs = 0
+
+        process_list = []
+        chunk_begin = 0
+        for i in range(args.thread_num):
+            now_chunk_size = chunk_size
+            if remain_wavs > 0:
+                now_chunk_size = chunk_size + 1
+                remain_wavs = remain_wavs - 1
+            # process i handle wavs at chunk_begin and size of now_chunk_size
+            p = Process(target=one_thread, args=(i, chunk_begin, now_chunk_size))
+            chunk_begin = chunk_begin + now_chunk_size
+            p.start()
+            process_list.append(p)
+
+        for i in process_list:
+            p.join()
+
+        print("end")

runtime/python/websocket/funasr_wss_server.py

@@ -0,0 +1,334 @@
+import asyncio
+import json
+import websockets
+import time
+import logging
+import tracemalloc
+import numpy as np
+import argparse
+import ssl
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    "--host", type=str, default="0.0.0.0", required=False, help="host ip, localhost, 0.0.0.0"
+)
+parser.add_argument("--port", type=int, default=10095, required=False, help="grpc server port")
+parser.add_argument(
+    "--asr_model",
+    type=str,
+    default="iic/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch",
+    help="model from modelscope",
+)
+parser.add_argument("--asr_model_revision", type=str, default="v2.0.4", help="")
+parser.add_argument(
+    "--asr_model_online",
+    type=str,
+    default="iic/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-online",
+    help="model from modelscope",
+)
+parser.add_argument("--asr_model_online_revision", type=str, default="v2.0.4", help="")
+parser.add_argument(
+    "--vad_model",
+    type=str,
+    default="iic/speech_fsmn_vad_zh-cn-16k-common-pytorch",
+    help="model from modelscope",
+)
+parser.add_argument("--vad_model_revision", type=str, default="v2.0.4", help="")
+parser.add_argument(
+    "--punc_model",
+    type=str,
+    default="iic/punc_ct-transformer_zh-cn-common-vad_realtime-vocab272727",
+    help="model from modelscope",
+)
+parser.add_argument("--punc_model_revision", type=str, default="v2.0.4", help="")
+parser.add_argument("--ngpu", type=int, default=1, help="0 for cpu, 1 for gpu")
+parser.add_argument("--device", type=str, default="cuda", help="cuda, cpu")
+parser.add_argument("--ncpu", type=int, default=4, help="cpu cores")
+parser.add_argument(
+    "--certfile",
+    type=str,
+    default="../../ssl_key/server.crt",
+    required=False,
+    help="certfile for ssl",
+)
+
+parser.add_argument(
+    "--keyfile",
+    type=str,
+    default="../../ssl_key/server.key",
+    required=False,
+    help="keyfile for ssl",
+)
+args = parser.parse_args()
+
+
+websocket_users = set()
+
+print("model loading")
+from funasr import AutoModel
+
+# asr
+model_asr = AutoModel(
+    model=args.asr_model,
+    model_revision=args.asr_model_revision,
+    ngpu=args.ngpu,
+    ncpu=args.ncpu,
+    device=args.device,
+    disable_pbar=True,
+    disable_log=True,
+)
+# asr
+model_asr_streaming = AutoModel(
+    model=args.asr_model_online,
+    model_revision=args.asr_model_online_revision,
+    ngpu=args.ngpu,
+    ncpu=args.ncpu,
+    device=args.device,
+    disable_pbar=True,
+    disable_log=True,
+)
+# vad
+model_vad = AutoModel(
+    model=args.vad_model,
+    model_revision=args.vad_model_revision,
+    ngpu=args.ngpu,
+    ncpu=args.ncpu,
+    device=args.device,
+    disable_pbar=True,
+    disable_log=True,
+    # chunk_size=60,
+)
+
+if args.punc_model != "":
+    model_punc = AutoModel(
+        model=args.punc_model,
+        model_revision=args.punc_model_revision,
+        ngpu=args.ngpu,
+        ncpu=args.ncpu,
+        device=args.device,
+        disable_pbar=True,
+        disable_log=True,
+    )
+else:
+    model_punc = None
+
+
+print("model loaded! only support one client at the same time now!!!!")
+
+
+async def ws_reset(websocket):
+    print("ws reset now, total num is ", len(websocket_users))
+
+    websocket.status_dict_asr_online["cache"] = {}
+    websocket.status_dict_asr_online["is_final"] = True
+    websocket.status_dict_vad["cache"] = {}
+    websocket.status_dict_vad["is_final"] = True
+    websocket.status_dict_punc["cache"] = {}
+
+    await websocket.close()
+
+
+async def clear_websocket():
+    for websocket in websocket_users:
+        await ws_reset(websocket)
+    websocket_users.clear()
+
+
+async def ws_serve(websocket, path):
+    frames = []
+    frames_asr = []
+    frames_asr_online = []
+    global websocket_users
+    # await clear_websocket()
+    websocket_users.add(websocket)
+    websocket.status_dict_asr = {}
+    websocket.status_dict_asr_online = {"cache": {}, "is_final": False}
+    websocket.status_dict_vad = {"cache": {}, "is_final": False}
+    websocket.status_dict_punc = {"cache": {}}
+    websocket.chunk_interval = 10
+    websocket.vad_pre_idx = 0
+    speech_start = False
+    speech_end_i = -1
+    websocket.wav_name = "microphone"
+    websocket.mode = "2pass"
+    print("new user connected", flush=True)
+
+    try:
+        async for message in websocket:
+            if isinstance(message, str):
+                messagejson = json.loads(message)
+
+                if "is_speaking" in messagejson:
+                    websocket.is_speaking = messagejson["is_speaking"]
+                    websocket.status_dict_asr_online["is_final"] = not websocket.is_speaking
+                if "chunk_interval" in messagejson:
+                    websocket.chunk_interval = messagejson["chunk_interval"]
+                if "wav_name" in messagejson:
+                    websocket.wav_name = messagejson.get("wav_name")
+                if "chunk_size" in messagejson:
+                    chunk_size = messagejson["chunk_size"]
+                    if isinstance(chunk_size, str):
+                        chunk_size = chunk_size.split(",")
+                    websocket.status_dict_asr_online["chunk_size"] = [int(x) for x in chunk_size]
+                if "encoder_chunk_look_back" in messagejson:
+                    websocket.status_dict_asr_online["encoder_chunk_look_back"] = messagejson[
+                        "encoder_chunk_look_back"
+                    ]
+                if "decoder_chunk_look_back" in messagejson:
+                    websocket.status_dict_asr_online["decoder_chunk_look_back"] = messagejson[
+                        "decoder_chunk_look_back"
+                    ]
+                if "hotword" in messagejson:
+                    websocket.status_dict_asr["hotword"] = messagejson["hotword"]
+                if "mode" in messagejson:
+                    websocket.mode = messagejson["mode"]
+
+            websocket.status_dict_vad["chunk_size"] = int(
+                websocket.status_dict_asr_online["chunk_size"][1] * 60 / websocket.chunk_interval
+            )
+            if len(frames_asr_online) > 0 or len(frames_asr) > 0 or not isinstance(message, str):
+                if not isinstance(message, str):
+                    frames.append(message)
+                    duration_ms = len(message) // 32
+                    websocket.vad_pre_idx += duration_ms
+
+                    # asr online
+                    frames_asr_online.append(message)
+                    websocket.status_dict_asr_online["is_final"] = speech_end_i != -1
+                    if (
+                        len(frames_asr_online) % websocket.chunk_interval == 0
+                        or websocket.status_dict_asr_online["is_final"]
+                    ):
+                        if websocket.mode == "2pass" or websocket.mode == "online":
+                            audio_in = b"".join(frames_asr_online)
+                            try:
+                                await async_asr_online(websocket, audio_in)
+                            except:
+                                print(f"error in asr streaming, {websocket.status_dict_asr_online}")
+                        frames_asr_online = []
+                    if speech_start:
+                        frames_asr.append(message)
+                    # vad online
+                    try:
+                        speech_start_i, speech_end_i = await async_vad(websocket, message)
+                    except:
+                        print("error in vad")
+                    if speech_start_i != -1:
+                        speech_start = True
+                        beg_bias = (websocket.vad_pre_idx - speech_start_i) // duration_ms
+                        frames_pre = frames[-beg_bias:]
+                        frames_asr = []
+                        frames_asr.extend(frames_pre)
+                # asr punc offline
+                if speech_end_i != -1 or not websocket.is_speaking:
+                    # print("vad end point")
+                    if websocket.mode == "2pass" or websocket.mode == "offline":
+                        audio_in = b"".join(frames_asr)
+                        try:
+                            await async_asr(websocket, audio_in)
+                        except:
+                            print("error in asr offline")
+                    frames_asr = []
+                    speech_start = False
+                    frames_asr_online = []
+                    websocket.status_dict_asr_online["cache"] = {}
+                    if not websocket.is_speaking:
+                        websocket.vad_pre_idx = 0
+                        frames = []
+                        websocket.status_dict_vad["cache"] = {}
+                    else:
+                        frames = frames[-20:]
+
+    except websockets.ConnectionClosed:
+        print("ConnectionClosed...", websocket_users, flush=True)
+        await ws_reset(websocket)
+        websocket_users.remove(websocket)
+    except websockets.InvalidState:
+        print("InvalidState...")
+    except Exception as e:
+        print("Exception:", e)
+
+
+async def async_vad(websocket, audio_in):
+
+    segments_result = model_vad.generate(input=audio_in, **websocket.status_dict_vad)[0]["value"]
+    # print(segments_result)
+
+    speech_start = -1
+    speech_end = -1
+
+    if len(segments_result) == 0 or len(segments_result) > 1:
+        return speech_start, speech_end
+    if segments_result[0][0] != -1:
+        speech_start = segments_result[0][0]
+    if segments_result[0][1] != -1:
+        speech_end = segments_result[0][1]
+    return speech_start, speech_end
+
+
+async def async_asr(websocket, audio_in):
+    if len(audio_in) > 0:
+        # print(len(audio_in))
+        rec_result = model_asr.generate(input=audio_in, **websocket.status_dict_asr)[0]
+        # print("offline_asr, ", rec_result)
+        if model_punc is not None and len(rec_result["text"]) > 0:
+            # print("offline, before punc", rec_result, "cache", websocket.status_dict_punc)
+            rec_result = model_punc.generate(
+                input=rec_result["text"], **websocket.status_dict_punc
+            )[0]
+            # print("offline, after punc", rec_result)
+        if len(rec_result["text"]) > 0:
+            # print("offline", rec_result)
+            mode = "2pass-offline" if "2pass" in websocket.mode else websocket.mode
+            message = json.dumps(
+                {
+                    "mode": mode,
+                    "text": rec_result["text"],
+                    "wav_name": websocket.wav_name,
+                    "is_final": websocket.is_speaking,
+                }
+            )
+            await websocket.send(message)
+
+
+async def async_asr_online(websocket, audio_in):
+    if len(audio_in) > 0:
+        # print(websocket.status_dict_asr_online.get("is_final", False))
+        rec_result = model_asr_streaming.generate(
+            input=audio_in, **websocket.status_dict_asr_online
+        )[0]
+        # print("online, ", rec_result)
+        if websocket.mode == "2pass" and websocket.status_dict_asr_online.get("is_final", False):
+            return
+            #     websocket.status_dict_asr_online["cache"] = dict()
+        if len(rec_result["text"]):
+            mode = "2pass-online" if "2pass" in websocket.mode else websocket.mode
+            message = json.dumps(
+                {
+                    "mode": mode,
+                    "text": rec_result["text"],
+                    "wav_name": websocket.wav_name,
+                    "is_final": websocket.is_speaking,
+                }
+            )
+            await websocket.send(message)
+
+
+if len(args.certfile) > 0:
+    ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
+
+    # Generate with Lets Encrypt, copied to this location, chown to current user and 400 permissions
+    ssl_cert = args.certfile
+    ssl_key = args.keyfile
+
+    ssl_context.load_cert_chain(ssl_cert, keyfile=ssl_key)
+    start_server = websockets.serve(
+        ws_serve, args.host, args.port, subprotocols=["binary"], ping_interval=None, ssl=ssl_context
+    )
+else:
+    start_server = websockets.serve(
+        ws_serve, args.host, args.port, subprotocols=["binary"], ping_interval=None
+    )
+asyncio.get_event_loop().run_until_complete(start_server)
+asyncio.get_event_loop().run_forever()

runtime/python/websocket/requirements_client.txt

@@ -0,0 +1,2 @@
+websockets
+pyaudio

runtime/python/websocket/requirements_server.txt

@@ -0,0 +1 @@
+websockets