From 58c59b1b3b45181a30eeb135d8a56be2942fdfd9 Mon Sep 17 00:00:00 2001
From: hnluo <haoneng.lhn@alibaba-inc.com>
Date: 星期四, 27 四月 2023 17:25:27 +0800
Subject: [PATCH] Merge pull request #432 from alibaba-damo-academy/dev_websocket
---
funasr/modules/embedding.py | 13
/dev/null | 185 ---------
funasr/models/e2e_asr_paraformer.py | 4
funasr/models/encoder/sanm_encoder.py | 24 +
funasr/bin/asr_inference_paraformer_streaming.py | 410 +++++---------------
funasr/models/predictor/cif.py | 128 ++++--
funasr/runtime/python/websocket/README.md | 30 +
funasr/runtime/python/websocket/parse_args.py | 35 +
funasr/runtime/python/websocket/ws_client.py | 182 +++++++++
funasr/runtime/python/websocket/ws_server_online.py | 108 +++++
10 files changed, 564 insertions(+), 555 deletions(-)
diff --git a/funasr/bin/asr_inference_paraformer_streaming.py b/funasr/bin/asr_inference_paraformer_streaming.py
index 821f694..ff8bb8c 100644
--- a/funasr/bin/asr_inference_paraformer_streaming.py
+++ b/funasr/bin/asr_inference_paraformer_streaming.py
@@ -8,6 +8,7 @@
import codecs
import tempfile
import requests
+import yaml
from pathlib import Path
from typing import Optional
from typing import Sequence
@@ -19,7 +20,6 @@
import numpy as np
import torch
-import torchaudio
from typeguard import check_argument_types
from funasr.fileio.datadir_writer import DatadirWriter
@@ -40,10 +40,11 @@
from funasr.utils.types import str2triple_str
from funasr.utils.types import str_or_none
from funasr.utils import asr_utils, wav_utils, postprocess_utils
-from funasr.models.frontend.wav_frontend import WavFrontend
-from funasr.models.e2e_asr_paraformer import BiCifParaformer, ContextualParaformer
+from funasr.models.frontend.wav_frontend import WavFrontend, WavFrontendOnline
from funasr.export.models.e2e_asr_paraformer import Paraformer as Paraformer_export
+
np.set_printoptions(threshold=np.inf)
+
class Speech2Text:
"""Speech2Text class
@@ -89,7 +90,7 @@
)
frontend = None
if asr_train_args.frontend is not None and asr_train_args.frontend_conf is not None:
- frontend = WavFrontend(cmvn_file=cmvn_file, **asr_train_args.frontend_conf)
+ frontend = WavFrontendOnline(cmvn_file=cmvn_file, **asr_train_args.frontend_conf)
logging.info("asr_model: {}".format(asr_model))
logging.info("asr_train_args: {}".format(asr_train_args))
@@ -189,8 +190,7 @@
@torch.no_grad()
def __call__(
- self, cache: dict, speech: Union[torch.Tensor, np.ndarray], speech_lengths: Union[torch.Tensor, np.ndarray] = None,
- begin_time: int = 0, end_time: int = None,
+ self, cache: dict, speech: Union[torch.Tensor], speech_lengths: Union[torch.Tensor] = None
):
"""Inference
@@ -201,38 +201,59 @@
"""
assert check_argument_types()
-
- # Input as audio signal
- if isinstance(speech, np.ndarray):
- speech = torch.tensor(speech)
- if self.frontend is not None:
- feats, feats_len = self.frontend.forward(speech, speech_lengths)
- feats = to_device(feats, device=self.device)
- feats_len = feats_len.int()
- self.asr_model.frontend = None
+ results = []
+ cache_en = cache["encoder"]
+ if speech.shape[1] < 16 * 60 and cache_en["is_final"]:
+ cache_en["tail_chunk"] = True
+ feats = cache_en["feats"]
+ feats_len = torch.tensor([feats.shape[1]])
+ results = self.infer(feats, feats_len, cache)
+ return results
else:
- feats = speech
- feats_len = speech_lengths
- lfr_factor = max(1, (feats.size()[-1] // 80) - 1)
- feats_len = cache["encoder"]["stride"] + cache["encoder"]["pad_left"] + cache["encoder"]["pad_right"]
- feats = feats[:,cache["encoder"]["start_idx"]:cache["encoder"]["start_idx"]+feats_len,:]
- feats_len = torch.tensor([feats_len])
- batch = {"speech": feats, "speech_lengths": feats_len, "cache": cache}
+ if self.frontend is not None:
+ feats, feats_len = self.frontend.forward(speech, speech_lengths, cache_en["is_final"])
+ feats = to_device(feats, device=self.device)
+ feats_len = feats_len.int()
+ self.asr_model.frontend = None
+ else:
+ feats = speech
+ feats_len = speech_lengths
- # a. To device
+ if feats.shape[1] != 0:
+ if cache_en["is_final"]:
+ if feats.shape[1] + cache_en["chunk_size"][2] < cache_en["chunk_size"][1]:
+ cache_en["last_chunk"] = True
+ else:
+ # first chunk
+ feats_chunk1 = feats[:, :cache_en["chunk_size"][1], :]
+ feats_len = torch.tensor([feats_chunk1.shape[1]])
+ results_chunk1 = self.infer(feats_chunk1, feats_len, cache)
+
+ # last chunk
+ cache_en["last_chunk"] = True
+ feats_chunk2 = feats[:, -(feats.shape[1] + cache_en["chunk_size"][2] - cache_en["chunk_size"][1]):, :]
+ feats_len = torch.tensor([feats_chunk2.shape[1]])
+ results_chunk2 = self.infer(feats_chunk2, feats_len, cache)
+
+ return ["".join(results_chunk1 + results_chunk2)]
+
+ results = self.infer(feats, feats_len, cache)
+
+ return results
+
+ @torch.no_grad()
+ def infer(self, feats: Union[torch.Tensor], feats_len: Union[torch.Tensor], cache: List = None):
+ batch = {"speech": feats, "speech_lengths": feats_len}
batch = to_device(batch, device=self.device)
-
# b. Forward Encoder
- enc, enc_len = self.asr_model.encode_chunk(feats, feats_len, cache)
+ enc, enc_len = self.asr_model.encode_chunk(feats, feats_len, cache=cache)
if isinstance(enc, tuple):
enc = enc[0]
# assert len(enc) == 1, len(enc)
enc_len_batch_total = torch.sum(enc_len).item() * self.encoder_downsampling_factor
predictor_outs = self.asr_model.calc_predictor_chunk(enc, cache)
- pre_acoustic_embeds, pre_token_length, alphas, pre_peak_index = predictor_outs[0], predictor_outs[1], \
- predictor_outs[2], predictor_outs[3]
- pre_token_length = pre_token_length.floor().long()
+ pre_acoustic_embeds, pre_token_length= predictor_outs[0], predictor_outs[1]
if torch.max(pre_token_length) < 1:
return []
decoder_outs = self.asr_model.cal_decoder_with_predictor_chunk(enc, pre_acoustic_embeds, cache)
@@ -279,163 +300,9 @@
text = self.tokenizer.tokens2text(token)
else:
text = None
-
- results.append((text, token, token_int, hyp, enc_len_batch_total, lfr_factor))
+ results.append(text)
# assert check_return_type(results)
- return results
-
-
-class Speech2TextExport:
- """Speech2TextExport class
-
- """
-
- def __init__(
- self,
- asr_train_config: Union[Path, str] = None,
- asr_model_file: Union[Path, str] = None,
- cmvn_file: Union[Path, str] = None,
- lm_train_config: Union[Path, str] = None,
- lm_file: Union[Path, str] = None,
- token_type: str = None,
- bpemodel: str = None,
- device: str = "cpu",
- maxlenratio: float = 0.0,
- minlenratio: float = 0.0,
- dtype: str = "float32",
- beam_size: int = 20,
- ctc_weight: float = 0.5,
- lm_weight: float = 1.0,
- ngram_weight: float = 0.9,
- penalty: float = 0.0,
- nbest: int = 1,
- frontend_conf: dict = None,
- hotword_list_or_file: str = None,
- **kwargs,
- ):
-
- # 1. Build ASR model
- asr_model, asr_train_args = ASRTask.build_model_from_file(
- asr_train_config, asr_model_file, cmvn_file, device
- )
- frontend = None
- if asr_train_args.frontend is not None and asr_train_args.frontend_conf is not None:
- frontend = WavFrontend(cmvn_file=cmvn_file, **asr_train_args.frontend_conf)
-
- logging.info("asr_model: {}".format(asr_model))
- logging.info("asr_train_args: {}".format(asr_train_args))
- asr_model.to(dtype=getattr(torch, dtype)).eval()
-
- token_list = asr_model.token_list
-
- logging.info(f"Decoding device={device}, dtype={dtype}")
-
- # 5. [Optional] Build Text converter: e.g. bpe-sym -> Text
- if token_type is None:
- token_type = asr_train_args.token_type
- if bpemodel is None:
- bpemodel = asr_train_args.bpemodel
-
- if token_type is None:
- tokenizer = None
- elif token_type == "bpe":
- if bpemodel is not None:
- tokenizer = build_tokenizer(token_type=token_type, bpemodel=bpemodel)
- else:
- tokenizer = None
- else:
- tokenizer = build_tokenizer(token_type=token_type)
- converter = TokenIDConverter(token_list=token_list)
- logging.info(f"Text tokenizer: {tokenizer}")
-
- # self.asr_model = asr_model
- self.asr_train_args = asr_train_args
- self.converter = converter
- self.tokenizer = tokenizer
-
- self.device = device
- self.dtype = dtype
- self.nbest = nbest
- self.frontend = frontend
-
- model = Paraformer_export(asr_model, onnx=False)
- self.asr_model = model
-
- @torch.no_grad()
- def __call__(
- self, speech: Union[torch.Tensor, np.ndarray], speech_lengths: Union[torch.Tensor, np.ndarray] = None
- ):
- """Inference
-
- Args:
- speech: Input speech data
- Returns:
- text, token, token_int, hyp
-
- """
- assert check_argument_types()
-
- # Input as audio signal
- if isinstance(speech, np.ndarray):
- speech = torch.tensor(speech)
-
- if self.frontend is not None:
- feats, feats_len = self.frontend.forward(speech, speech_lengths)
- feats = to_device(feats, device=self.device)
- feats_len = feats_len.int()
- self.asr_model.frontend = None
- else:
- feats = speech
- feats_len = speech_lengths
-
- enc_len_batch_total = feats_len.sum()
- lfr_factor = max(1, (feats.size()[-1] // 80) - 1)
- batch = {"speech": feats, "speech_lengths": feats_len}
-
- # a. To device
- batch = to_device(batch, device=self.device)
-
- decoder_outs = self.asr_model(**batch)
- decoder_out, ys_pad_lens = decoder_outs[0], decoder_outs[1]
-
- results = []
- b, n, d = decoder_out.size()
- for i in range(b):
- am_scores = decoder_out[i, :ys_pad_lens[i], :]
-
- yseq = am_scores.argmax(dim=-1)
- score = am_scores.max(dim=-1)[0]
- score = torch.sum(score, dim=-1)
- # pad with mask tokens to ensure compatibility with sos/eos tokens
- yseq = torch.tensor(
- yseq.tolist(), device=yseq.device
- )
- nbest_hyps = [Hypothesis(yseq=yseq, score=score)]
-
- for hyp in nbest_hyps:
- assert isinstance(hyp, (Hypothesis)), type(hyp)
-
- # remove sos/eos and get results
- last_pos = -1
- if isinstance(hyp.yseq, list):
- token_int = hyp.yseq[1:last_pos]
- else:
- token_int = hyp.yseq[1:last_pos].tolist()
-
- # remove blank symbol id, which is assumed to be 0
- token_int = list(filter(lambda x: x != 0 and x != 2, token_int))
-
- # Change integer-ids to tokens
- token = self.converter.ids2tokens(token_int)
-
- if self.tokenizer is not None:
- text = self.tokenizer.tokens2text(token)
- else:
- text = None
-
- results.append((text, token, token_int, hyp, enc_len_batch_total, lfr_factor))
-
return results
@@ -536,8 +403,6 @@
**kwargs,
):
assert check_argument_types()
- ncpu = kwargs.get("ncpu", 1)
- torch.set_num_threads(ncpu)
if word_lm_train_config is not None:
raise NotImplementedError("Word LM is not implemented")
@@ -580,11 +445,9 @@
penalty=penalty,
nbest=nbest,
)
- if export_mode:
- speech2text = Speech2TextExport(**speech2text_kwargs)
- else:
- speech2text = Speech2Text(**speech2text_kwargs)
-
+
+ speech2text = Speech2Text(**speech2text_kwargs)
+
def _load_bytes(input):
middle_data = np.frombuffer(input, dtype=np.int16)
middle_data = np.asarray(middle_data)
@@ -599,7 +462,46 @@
offset = i.min + abs_max
array = np.frombuffer((middle_data.astype(dtype) - offset) / abs_max, dtype=np.float32)
return array
-
+
+ def _read_yaml(yaml_path: Union[str, Path]) -> Dict:
+ if not Path(yaml_path).exists():
+ raise FileExistsError(f'The {yaml_path} does not exist.')
+
+ with open(str(yaml_path), 'rb') as f:
+ data = yaml.load(f, Loader=yaml.Loader)
+ return data
+
+ def _prepare_cache(cache: dict = {}, chunk_size=[5,10,5], batch_size=1):
+ if len(cache) > 0:
+ return cache
+ config = _read_yaml(asr_train_config)
+ enc_output_size = config["encoder_conf"]["output_size"]
+ feats_dims = config["frontend_conf"]["n_mels"] * config["frontend_conf"]["lfr_m"]
+ cache_en = {"start_idx": 0, "cif_hidden": torch.zeros((batch_size, 1, enc_output_size)),
+ "cif_alphas": torch.zeros((batch_size, 1)), "chunk_size": chunk_size, "last_chunk": False,
+ "feats": torch.zeros((batch_size, chunk_size[0] + chunk_size[2], feats_dims)), "tail_chunk": False}
+ cache["encoder"] = cache_en
+
+ cache_de = {"decode_fsmn": None}
+ cache["decoder"] = cache_de
+
+ return cache
+
+ def _cache_reset(cache: dict = {}, chunk_size=[5,10,5], batch_size=1):
+ if len(cache) > 0:
+ config = _read_yaml(asr_train_config)
+ enc_output_size = config["encoder_conf"]["output_size"]
+ feats_dims = config["frontend_conf"]["n_mels"] * config["frontend_conf"]["lfr_m"]
+ cache_en = {"start_idx": 0, "cif_hidden": torch.zeros((batch_size, 1, enc_output_size)),
+ "cif_alphas": torch.zeros((batch_size, 1)), "chunk_size": chunk_size, "last_chunk": False,
+ "feats": torch.zeros((batch_size, chunk_size[0] + chunk_size[2], feats_dims)), "tail_chunk": False}
+ cache["encoder"] = cache_en
+
+ cache_de = {"decode_fsmn": None}
+ cache["decoder"] = cache_de
+
+ return cache
+
def _forward(
data_path_and_name_and_type,
raw_inputs: Union[np.ndarray, torch.Tensor] = None,
@@ -610,123 +512,35 @@
):
# 3. Build data-iterator
+ if data_path_and_name_and_type is not None and data_path_and_name_and_type[2] == "bytes":
+ raw_inputs = _load_bytes(data_path_and_name_and_type[0])
+ raw_inputs = torch.tensor(raw_inputs)
+ if data_path_and_name_and_type is None and raw_inputs is not None:
+ if isinstance(raw_inputs, np.ndarray):
+ raw_inputs = torch.tensor(raw_inputs)
is_final = False
cache = {}
+ chunk_size = [5, 10, 5]
if param_dict is not None and "cache" in param_dict:
cache = param_dict["cache"]
if param_dict is not None and "is_final" in param_dict:
is_final = param_dict["is_final"]
+ if param_dict is not None and "chunk_size" in param_dict:
+ chunk_size = param_dict["chunk_size"]
- if data_path_and_name_and_type is not None and data_path_and_name_and_type[2] == "bytes":
- raw_inputs = _load_bytes(data_path_and_name_and_type[0])
- raw_inputs = torch.tensor(raw_inputs)
- if data_path_and_name_and_type is not None and data_path_and_name_and_type[2] == "sound":
- raw_inputs = torchaudio.load(data_path_and_name_and_type[0])[0][0]
- is_final = True
- if data_path_and_name_and_type is None and raw_inputs is not None:
- if isinstance(raw_inputs, np.ndarray):
- raw_inputs = torch.tensor(raw_inputs)
# 7 .Start for-loop
# FIXME(kamo): The output format should be discussed about
+ raw_inputs = torch.unsqueeze(raw_inputs, axis=0)
+ input_lens = torch.tensor([raw_inputs.shape[1]])
asr_result_list = []
- results = []
- asr_result = ""
- wait = True
- if len(cache) == 0:
- cache["encoder"] = {"start_idx": 0, "pad_left": 0, "stride": 10, "pad_right": 5, "cif_hidden": None, "cif_alphas": None, "is_final": is_final, "left": 0, "right": 0}
- cache_de = {"decode_fsmn": None}
- cache["decoder"] = cache_de
- cache["first_chunk"] = True
- cache["speech"] = []
- cache["accum_speech"] = 0
- if raw_inputs is not None:
- if len(cache["speech"]) == 0:
- cache["speech"] = raw_inputs
- else:
- cache["speech"] = torch.cat([cache["speech"], raw_inputs], dim=0)
- cache["accum_speech"] += len(raw_inputs)
- while cache["accum_speech"] >= 960:
- if cache["first_chunk"]:
- if cache["accum_speech"] >= 14400:
- speech = torch.unsqueeze(cache["speech"], axis=0)
- speech_length = torch.tensor([len(cache["speech"])])
- cache["encoder"]["pad_left"] = 5
- cache["encoder"]["pad_right"] = 5
- cache["encoder"]["stride"] = 10
- cache["encoder"]["left"] = 5
- cache["encoder"]["right"] = 0
- results = speech2text(cache, speech, speech_length)
- cache["accum_speech"] -= 4800
- cache["first_chunk"] = False
- cache["encoder"]["start_idx"] = -5
- cache["encoder"]["is_final"] = False
- wait = False
- else:
- if is_final:
- cache["encoder"]["stride"] = len(cache["speech"]) // 960
- cache["encoder"]["pad_left"] = 0
- cache["encoder"]["pad_right"] = 0
- speech = torch.unsqueeze(cache["speech"], axis=0)
- speech_length = torch.tensor([len(cache["speech"])])
- results = speech2text(cache, speech, speech_length)
- cache["accum_speech"] = 0
- wait = False
- else:
- break
- else:
- if cache["accum_speech"] >= 19200:
- cache["encoder"]["start_idx"] += 10
- cache["encoder"]["stride"] = 10
- cache["encoder"]["pad_left"] = 5
- cache["encoder"]["pad_right"] = 5
- cache["encoder"]["left"] = 0
- cache["encoder"]["right"] = 0
- speech = torch.unsqueeze(cache["speech"], axis=0)
- speech_length = torch.tensor([len(cache["speech"])])
- results = speech2text(cache, speech, speech_length)
- cache["accum_speech"] -= 9600
- wait = False
- else:
- if is_final:
- cache["encoder"]["is_final"] = True
- if cache["accum_speech"] >= 14400:
- cache["encoder"]["start_idx"] += 10
- cache["encoder"]["stride"] = 10
- cache["encoder"]["pad_left"] = 5
- cache["encoder"]["pad_right"] = 5
- cache["encoder"]["left"] = 0
- cache["encoder"]["right"] = cache["accum_speech"] // 960 - 15
- speech = torch.unsqueeze(cache["speech"], axis=0)
- speech_length = torch.tensor([len(cache["speech"])])
- results = speech2text(cache, speech, speech_length)
- cache["accum_speech"] -= 9600
- wait = False
- else:
- cache["encoder"]["start_idx"] += 10
- cache["encoder"]["stride"] = cache["accum_speech"] // 960 - 5
- cache["encoder"]["pad_left"] = 5
- cache["encoder"]["pad_right"] = 0
- cache["encoder"]["left"] = 0
- cache["encoder"]["right"] = 0
- speech = torch.unsqueeze(cache["speech"], axis=0)
- speech_length = torch.tensor([len(cache["speech"])])
- results = speech2text(cache, speech, speech_length)
- cache["accum_speech"] = 0
- wait = False
- else:
- break
-
- if len(results) >= 1:
- asr_result += results[0][0]
- if asr_result == "":
- asr_result = "sil"
- if wait:
- asr_result = "waiting_for_more_voice"
- item = {'key': "utt", 'value': asr_result}
- asr_result_list.append(item)
- else:
- return []
+ cache = _prepare_cache(cache, chunk_size=chunk_size, batch_size=1)
+ cache["encoder"]["is_final"] = is_final
+ asr_result = speech2text(cache, raw_inputs, input_lens)
+ item = {'key': "utt", 'value': asr_result}
+ asr_result_list.append(item)
+ if is_final:
+ cache = _cache_reset(cache, chunk_size=chunk_size, batch_size=1)
return asr_result_list
return _forward
@@ -920,5 +734,3 @@
#
# rec_result = inference_16k_pipline(audio_in='https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav')
# print(rec_result)
-
-
diff --git a/funasr/models/e2e_asr_paraformer.py b/funasr/models/e2e_asr_paraformer.py
index 699d85f..d02783f 100644
--- a/funasr/models/e2e_asr_paraformer.py
+++ b/funasr/models/e2e_asr_paraformer.py
@@ -712,9 +712,9 @@
def calc_predictor_chunk(self, encoder_out, cache=None):
- pre_acoustic_embeds, pre_token_length, alphas, pre_peak_index = \
+ pre_acoustic_embeds, pre_token_length = \
self.predictor.forward_chunk(encoder_out, cache["encoder"])
- return pre_acoustic_embeds, pre_token_length, alphas, pre_peak_index
+ return pre_acoustic_embeds, pre_token_length
def cal_decoder_with_predictor_chunk(self, encoder_out, sematic_embeds, cache=None):
decoder_outs = self.decoder.forward_chunk(
diff --git a/funasr/models/encoder/sanm_encoder.py b/funasr/models/encoder/sanm_encoder.py
index f2502bb..969ddad 100644
--- a/funasr/models/encoder/sanm_encoder.py
+++ b/funasr/models/encoder/sanm_encoder.py
@@ -6,9 +6,11 @@
import logging
import torch
import torch.nn as nn
+import torch.nn.functional as F
from funasr.modules.streaming_utils.chunk_utilis import overlap_chunk
from typeguard import check_argument_types
import numpy as np
+from funasr.torch_utils.device_funcs import to_device
from funasr.modules.nets_utils import make_pad_mask
from funasr.modules.attention import MultiHeadedAttention, MultiHeadedAttentionSANM, MultiHeadedAttentionSANMwithMask
from funasr.modules.embedding import SinusoidalPositionEncoder, StreamSinusoidalPositionEncoder
@@ -349,6 +351,23 @@
return (xs_pad, intermediate_outs), olens, None
return xs_pad, olens, None
+ def _add_overlap_chunk(self, feats: np.ndarray, cache: dict = {}):
+ if len(cache) == 0:
+ return feats
+ # process last chunk
+ cache["feats"] = to_device(cache["feats"], device=feats.device)
+ overlap_feats = torch.cat((cache["feats"], feats), dim=1)
+ if cache["is_final"]:
+ cache["feats"] = overlap_feats[:, -cache["chunk_size"][0]:, :]
+ if not cache["last_chunk"]:
+ padding_length = sum(cache["chunk_size"]) - overlap_feats.shape[1]
+ overlap_feats = overlap_feats.transpose(1, 2)
+ overlap_feats = F.pad(overlap_feats, (0, padding_length))
+ overlap_feats = overlap_feats.transpose(1, 2)
+ else:
+ cache["feats"] = overlap_feats[:, -(cache["chunk_size"][0] + cache["chunk_size"][2]):, :]
+ return overlap_feats
+
def forward_chunk(self,
xs_pad: torch.Tensor,
ilens: torch.Tensor,
@@ -360,7 +379,10 @@
xs_pad = xs_pad
else:
xs_pad = self.embed(xs_pad, cache)
-
+ if cache["tail_chunk"]:
+ xs_pad = cache["feats"]
+ else:
+ xs_pad = self._add_overlap_chunk(xs_pad, cache)
encoder_outs = self.encoders0(xs_pad, None, None, None, None)
xs_pad, masks = encoder_outs[0], encoder_outs[1]
intermediate_outs = []
diff --git a/funasr/models/predictor/cif.py b/funasr/models/predictor/cif.py
index a5273f8..c59e245 100644
--- a/funasr/models/predictor/cif.py
+++ b/funasr/models/predictor/cif.py
@@ -2,6 +2,7 @@
from torch import nn
import logging
import numpy as np
+from funasr.torch_utils.device_funcs import to_device
from funasr.modules.nets_utils import make_pad_mask
from funasr.modules.streaming_utils.utils import sequence_mask
@@ -200,7 +201,7 @@
return acoustic_embeds, token_num, alphas, cif_peak
def forward_chunk(self, hidden, cache=None):
- b, t, d = hidden.size()
+ batch_size, len_time, hidden_size = hidden.shape
h = hidden
context = h.transpose(1, 2)
queries = self.pad(context)
@@ -211,58 +212,81 @@
alphas = torch.nn.functional.relu(alphas * self.smooth_factor - self.noise_threshold)
alphas = alphas.squeeze(-1)
- mask_chunk_predictor = None
- if cache is not None:
- mask_chunk_predictor = None
- mask_chunk_predictor = torch.zeros_like(alphas)
- mask_chunk_predictor[:, cache["pad_left"]:cache["stride"] + cache["pad_left"]] = 1.0
-
- if mask_chunk_predictor is not None:
- alphas = alphas * mask_chunk_predictor
-
- if cache is not None:
- if cache["is_final"]:
- alphas[:, cache["stride"] + cache["pad_left"] - 1] += 0.45
- if cache["cif_hidden"] is not None:
- hidden = torch.cat((cache["cif_hidden"], hidden), 1)
- if cache["cif_alphas"] is not None:
- alphas = torch.cat((cache["cif_alphas"], alphas), -1)
- token_num = alphas.sum(-1)
- acoustic_embeds, cif_peak = cif(hidden, alphas, self.threshold)
- len_time = alphas.size(-1)
- last_fire_place = len_time - 1
- last_fire_remainds = 0.0
- pre_alphas_length = 0
- last_fire = False
-
- mask_chunk_peak_predictor = None
- if cache is not None:
- mask_chunk_peak_predictor = None
- mask_chunk_peak_predictor = torch.zeros_like(cif_peak)
- if cache["cif_alphas"] is not None:
- pre_alphas_length = cache["cif_alphas"].size(-1)
- mask_chunk_peak_predictor[:, :pre_alphas_length] = 1.0
- mask_chunk_peak_predictor[:, pre_alphas_length + cache["pad_left"]:pre_alphas_length + cache["stride"] + cache["pad_left"]] = 1.0
-
- if mask_chunk_peak_predictor is not None:
- cif_peak = cif_peak * mask_chunk_peak_predictor.squeeze(-1)
-
- for i in range(len_time):
- if cif_peak[0][len_time - 1 - i] > self.threshold or cif_peak[0][len_time - 1 - i] == self.threshold:
- last_fire_place = len_time - 1 - i
- last_fire_remainds = cif_peak[0][len_time - 1 - i] - self.threshold
- last_fire = True
- break
- if last_fire:
- last_fire_remainds = torch.tensor([last_fire_remainds], dtype=alphas.dtype).to(alphas.device)
- cache["cif_hidden"] = hidden[:, last_fire_place:, :]
- cache["cif_alphas"] = torch.cat((last_fire_remainds.unsqueeze(0), alphas[:, last_fire_place+1:]), -1)
- else:
- cache["cif_hidden"] = hidden
- cache["cif_alphas"] = alphas
- token_num_int = token_num.floor().type(torch.int32).item()
- return acoustic_embeds[:, 0:token_num_int, :], token_num, alphas, cif_peak
+ token_length = []
+ list_fires = []
+ list_frames = []
+ cache_alphas = []
+ cache_hiddens = []
+
+ if cache is not None and "chunk_size" in cache:
+ alphas[:, :cache["chunk_size"][0]] = 0.0
+ alphas[:, sum(cache["chunk_size"][:2]):] = 0.0
+ if cache is not None and "cif_alphas" in cache and "cif_hidden" in cache:
+ cache["cif_hidden"] = to_device(cache["cif_hidden"], device=hidden.device)
+ cache["cif_alphas"] = to_device(cache["cif_alphas"], device=alphas.device)
+ hidden = torch.cat((cache["cif_hidden"], hidden), dim=1)
+ alphas = torch.cat((cache["cif_alphas"], alphas), dim=1)
+ if cache is not None and "last_chunk" in cache and cache["last_chunk"]:
+ tail_hidden = torch.zeros((batch_size, 1, hidden_size), device=hidden.device)
+ tail_alphas = torch.tensor([[self.tail_threshold]], device=alphas.device)
+ tail_alphas = torch.tile(tail_alphas, (batch_size, 1))
+ hidden = torch.cat((hidden, tail_hidden), dim=1)
+ alphas = torch.cat((alphas, tail_alphas), dim=1)
+
+ len_time = alphas.shape[1]
+ for b in range(batch_size):
+ integrate = 0.0
+ frames = torch.zeros((hidden_size), device=hidden.device)
+ list_frame = []
+ list_fire = []
+ for t in range(len_time):
+ alpha = alphas[b][t]
+ if alpha + integrate < self.threshold:
+ integrate += alpha
+ list_fire.append(integrate)
+ frames += alpha * hidden[b][t]
+ else:
+ frames += (self.threshold - integrate) * hidden[b][t]
+ list_frame.append(frames)
+ integrate += alpha
+ list_fire.append(integrate)
+ integrate -= self.threshold
+ frames = integrate * hidden[b][t]
+
+ cache_alphas.append(integrate)
+ if integrate > 0.0:
+ cache_hiddens.append(frames / integrate)
+ else:
+ cache_hiddens.append(frames)
+
+ token_length.append(torch.tensor(len(list_frame), device=alphas.device))
+ list_fires.append(list_fire)
+ list_frames.append(list_frame)
+
+ cache["cif_alphas"] = torch.stack(cache_alphas, axis=0)
+ cache["cif_alphas"] = torch.unsqueeze(cache["cif_alphas"], axis=0)
+ cache["cif_hidden"] = torch.stack(cache_hiddens, axis=0)
+ cache["cif_hidden"] = torch.unsqueeze(cache["cif_hidden"], axis=0)
+
+ max_token_len = max(token_length)
+ if max_token_len == 0:
+ return hidden, torch.stack(token_length, 0)
+ list_ls = []
+ for b in range(batch_size):
+ pad_frames = torch.zeros((max_token_len - token_length[b], hidden_size), device=alphas.device)
+ if token_length[b] == 0:
+ list_ls.append(pad_frames)
+ else:
+ list_frames[b] = torch.stack(list_frames[b])
+ list_ls.append(torch.cat((list_frames[b], pad_frames), dim=0))
+
+ cache["cif_alphas"] = torch.stack(cache_alphas, axis=0)
+ cache["cif_alphas"] = torch.unsqueeze(cache["cif_alphas"], axis=0)
+ cache["cif_hidden"] = torch.stack(cache_hiddens, axis=0)
+ cache["cif_hidden"] = torch.unsqueeze(cache["cif_hidden"], axis=0)
+ return torch.stack(list_ls, 0), torch.stack(token_length, 0)
+
def tail_process_fn(self, hidden, alphas, token_num=None, mask=None):
b, t, d = hidden.size()
diff --git a/funasr/modules/embedding.py b/funasr/modules/embedding.py
index c347e24..aaac80a 100644
--- a/funasr/modules/embedding.py
+++ b/funasr/modules/embedding.py
@@ -425,21 +425,14 @@
return encoding.type(dtype)
def forward(self, x, cache=None):
- start_idx = 0
- pad_left = 0
- pad_right = 0
batch_size, timesteps, input_dim = x.size()
+ start_idx = 0
if cache is not None:
start_idx = cache["start_idx"]
- pad_left = cache["left"]
- pad_right = cache["right"]
+ cache["start_idx"] += timesteps
positions = torch.arange(1, timesteps+start_idx+1)[None, :]
position_encoding = self.encode(positions, input_dim, x.dtype).to(x.device)
- outputs = x + position_encoding[:, start_idx: start_idx + timesteps]
- outputs = outputs.transpose(1, 2)
- outputs = F.pad(outputs, (pad_left, pad_right))
- outputs = outputs.transpose(1, 2)
- return outputs
+ return x + position_encoding[:, start_idx: start_idx + timesteps]
class StreamingRelPositionalEncoding(torch.nn.Module):
"""Relative positional encoding.
diff --git a/funasr/runtime/python/websocket/ASR_client.py b/funasr/runtime/python/websocket/ASR_client.py
deleted file mode 100644
index fe67981..0000000
--- a/funasr/runtime/python/websocket/ASR_client.py
+++ /dev/null
@@ -1,100 +0,0 @@
-import pyaudio
-# import websocket #鍖哄埆鏈嶅姟绔繖閲屾槸 websocket-client搴�
-import time
-import websockets
-import asyncio
-from queue import Queue
-# import threading
-import argparse
-
-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=int,
- default=300,
- help="ms")
-
-args = parser.parse_args()
-
-voices = Queue()
-
-
-
-# 鍏朵粬鍑芥暟鍙互閫氳繃璋冪敤send(data)鏉ュ彂閫佹暟鎹紝渚嬪锛�
-async def record():
- #print("2")
- global voices
- FORMAT = pyaudio.paInt16
- CHANNELS = 1
- RATE = 16000
- CHUNK = int(RATE / 1000 * args.chunk_size)
-
- p = pyaudio.PyAudio()
-
- stream = p.open(format=FORMAT,
- channels=CHANNELS,
- rate=RATE,
- input=True,
- frames_per_buffer=CHUNK)
-
- while True:
-
- data = stream.read(CHUNK)
-
- voices.put(data)
- #print(voices.qsize())
-
- await asyncio.sleep(0.01)
-
-
-
-async def ws_send():
- global voices
- global websocket
- print("started to sending data!")
- while True:
- while not voices.empty():
- data = voices.get()
- voices.task_done()
- try:
- await websocket.send(data) # 閫氳繃ws瀵硅薄鍙戦�佹暟鎹�
- except Exception as e:
- print('Exception occurred:', e)
- await asyncio.sleep(0.01)
- await asyncio.sleep(0.01)
-
-
-
-async def message():
- global websocket
- while True:
- try:
- print(await websocket.recv())
- except Exception as e:
- print("Exception:", e)
-
-
-
-async def ws_client():
- global websocket # 瀹氫箟涓�涓叏灞�鍙橀噺ws锛岀敤浜庝繚瀛榳ebsocket杩炴帴瀵硅薄
- # uri = "ws://11.167.134.197:8899"
- uri = "ws://{}:{}".format(args.host, args.port)
- #ws = await websockets.connect(uri, subprotocols=["binary"]) # 鍒涘缓涓�涓暱杩炴帴
- async for websocket in websockets.connect(uri, subprotocols=["binary"], ping_interval=None):
- task = asyncio.create_task(record()) # 鍒涘缓涓�涓悗鍙颁换鍔″綍闊�
- task2 = asyncio.create_task(ws_send()) # 鍒涘缓涓�涓悗鍙颁换鍔″彂閫�
- task3 = asyncio.create_task(message()) # 鍒涘缓涓�涓悗鍙版帴鏀舵秷鎭殑浠诲姟
- await asyncio.gather(task, task2, task3)
-
-
-asyncio.get_event_loop().run_until_complete(ws_client()) # 鍚姩鍗忕▼
-asyncio.get_event_loop().run_forever()
diff --git a/funasr/runtime/python/websocket/ASR_server.py b/funasr/runtime/python/websocket/ASR_server.py
deleted file mode 100644
index 827df7b..0000000
--- a/funasr/runtime/python/websocket/ASR_server.py
+++ /dev/null
@@ -1,185 +0,0 @@
-import asyncio
-import websockets
-import time
-from queue import Queue
-import threading
-import argparse
-
-from modelscope.pipelines import pipeline
-from modelscope.utils.constant import Tasks
-from modelscope.utils.logger import get_logger
-import logging
-import tracemalloc
-tracemalloc.start()
-
-logger = get_logger(log_level=logging.CRITICAL)
-logger.setLevel(logging.CRITICAL)
-
-
-websocket_users = set() #缁存姢瀹㈡埛绔垪琛�
-
-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="damo/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch",
- help="model from modelscope")
-parser.add_argument("--vad_model",
- type=str,
- default="damo/speech_fsmn_vad_zh-cn-16k-common-pytorch",
- help="model from modelscope")
-
-parser.add_argument("--punc_model",
- type=str,
- default="",
- help="model from modelscope")
-parser.add_argument("--ngpu",
- type=int,
- default=1,
- help="0 for cpu, 1 for gpu")
-
-args = parser.parse_args()
-
-print("model loading")
-
-
-# vad
-inference_pipeline_vad = pipeline(
- task=Tasks.voice_activity_detection,
- model=args.vad_model,
- model_revision=None,
- output_dir=None,
- batch_size=1,
- mode='online',
- ngpu=args.ngpu,
-)
-# param_dict_vad = {'in_cache': dict(), "is_final": False}
-
-# asr
-param_dict_asr = {}
-# param_dict["hotword"] = "灏忎簲 灏忎簲鏈�" # 璁剧疆鐑瘝锛岀敤绌烘牸闅斿紑
-inference_pipeline_asr = pipeline(
- task=Tasks.auto_speech_recognition,
- model=args.asr_model,
- param_dict=param_dict_asr,
- ngpu=args.ngpu,
-)
-if args.punc_model != "":
- # param_dict_punc = {'cache': list()}
- inference_pipeline_punc = pipeline(
- task=Tasks.punctuation,
- model=args.punc_model,
- model_revision=None,
- ngpu=args.ngpu,
- )
-else:
- inference_pipeline_punc = None
-
-print("model loaded")
-
-
-
-async def ws_serve(websocket, path):
- #speek = Queue()
- frames = [] # 瀛樺偍鎵�鏈夌殑甯ф暟鎹�
- buffer = [] # 瀛樺偍缂撳瓨涓殑甯ф暟鎹紙鏈�澶氫袱涓墖娈碉級
- RECORD_NUM = 0
- global websocket_users
- speech_start, speech_end = False, False
- # 璋冪敤asr鍑芥暟
- websocket.param_dict_vad = {'in_cache': dict(), "is_final": False}
- websocket.param_dict_punc = {'cache': list()}
- websocket.speek = Queue() #websocket 娣诲姞杩涢槦鍒楀璞� 璁゛sr璇诲彇璇煶鏁版嵁鍖�
- websocket.send_msg = Queue() #websocket 娣诲姞涓槦鍒楀璞� 璁﹚s鍙戦�佹秷鎭埌瀹㈡埛绔�
- websocket_users.add(websocket)
- ss = threading.Thread(target=asr, args=(websocket,))
- ss.start()
-
- try:
- async for message in websocket:
- #voices.put(message)
- #print("put")
- #await websocket.send("123")
- buffer.append(message)
- if len(buffer) > 2:
- buffer.pop(0) # 濡傛灉缂撳瓨瓒呰繃涓や釜鐗囨锛屽垯鍒犻櫎鏈�鏃╃殑涓�涓�
-
- if speech_start:
- frames.append(message)
- RECORD_NUM += 1
- speech_start_i, speech_end_i = vad(message, websocket)
- #print(speech_start_i, speech_end_i)
- if speech_start_i:
- speech_start = speech_start_i
- frames = []
- frames.extend(buffer) # 鎶婁箣鍓�2涓闊虫暟鎹揩鍔犲叆
- if speech_end_i or RECORD_NUM > 300:
- speech_start = False
- audio_in = b"".join(frames)
- websocket.speek.put(audio_in)
- frames = [] # 娓呯┖鎵�鏈夌殑甯ф暟鎹�
- buffer = [] # 娓呯┖缂撳瓨涓殑甯ф暟鎹紙鏈�澶氫袱涓墖娈碉級
- RECORD_NUM = 0
- if not websocket.send_msg.empty():
- await websocket.send(websocket.send_msg.get())
- websocket.send_msg.task_done()
-
-
- except websockets.ConnectionClosed:
- print("ConnectionClosed...", websocket_users) # 閾炬帴鏂紑
- websocket_users.remove(websocket)
- except websockets.InvalidState:
- print("InvalidState...") # 鏃犳晥鐘舵��
- except Exception as e:
- print("Exception:", e)
-
-
-def asr(websocket): # ASR鎺ㄧ悊
- global inference_pipeline_asr, inference_pipeline_punc
- # global param_dict_punc
- global websocket_users
- while websocket in websocket_users:
- if not websocket.speek.empty():
- audio_in = websocket.speek.get()
- websocket.speek.task_done()
- if len(audio_in) > 0:
- rec_result = inference_pipeline_asr(audio_in=audio_in)
- if inference_pipeline_punc is not None and 'text' in rec_result:
- rec_result = inference_pipeline_punc(text_in=rec_result['text'], param_dict=websocket.param_dict_punc)
- # print(rec_result)
- if "text" in rec_result:
- websocket.send_msg.put(rec_result["text"]) # 瀛樺叆鍙戦�侀槦鍒� 鐩存帴璋冪敤send鍙戦�佷笉浜�
-
- time.sleep(0.1)
-
-def vad(data, websocket): # VAD鎺ㄧ悊
- global inference_pipeline_vad
- #print(type(data))
- # print(param_dict_vad)
- segments_result = inference_pipeline_vad(audio_in=data, param_dict=websocket.param_dict_vad)
- # print(segments_result)
- # print(param_dict_vad)
- speech_start = False
- speech_end = False
-
- if len(segments_result) == 0 or len(segments_result["text"]) > 1:
- return speech_start, speech_end
- if segments_result["text"][0][0] != -1:
- speech_start = True
- if segments_result["text"][0][1] != -1:
- speech_end = True
- return speech_start, speech_end
-
-
-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()
\ No newline at end of file
diff --git a/funasr/runtime/python/websocket/README.md b/funasr/runtime/python/websocket/README.md
index 73f8aeb..723782f 100644
--- a/funasr/runtime/python/websocket/README.md
+++ b/funasr/runtime/python/websocket/README.md
@@ -5,7 +5,7 @@
## For the Server
-Install the modelscope and funasr
+### Install the modelscope and funasr
```shell
pip install -U modelscope funasr
@@ -14,18 +14,34 @@
git clone https://github.com/alibaba/FunASR.git && cd FunASR
```
-Install the requirements for server
+### Install the requirements for server
```shell
cd funasr/runtime/python/websocket
pip install -r requirements_server.txt
```
-Start server
+### Start server
+#### ASR offline server
+[//]: # (```shell)
+
+[//]: # (python ws_server_online.py --host "0.0.0.0" --port 10095 --asr_model "damo/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch")
+
+[//]: # (```)
+#### ASR streaming server
```shell
-python ASR_server.py --host "0.0.0.0" --port 10095 --asr_model "damo/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch"
+python ws_server_online.py --host "0.0.0.0" --port 10095
```
+####
+
+#### ASR offline/online 2pass server
+
+[//]: # (```shell)
+
+[//]: # (python ws_server_online.py --host "0.0.0.0" --port 10095 --asr_model "damo/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch")
+
+[//]: # (```)
## For the client
@@ -39,8 +55,10 @@
Start client
```shell
-python ASR_client.py --host "127.0.0.1" --port 10095 --chunk_size 300
+# --chunk_size, "5,10,5"=600ms, "8,8,4"=480ms
+python ws_client.py --host "127.0.0.1" --port 10096 --chunk_size "5,10,5"
```
## Acknowledge
-1. We acknowledge [cgisky1980](https://github.com/cgisky1980/FunASR) for contributing the websocket service.
+1. This project is maintained by [FunASR community](https://github.com/alibaba-damo-academy/FunASR).
+2. We acknowledge [cgisky1980](https://github.com/cgisky1980/FunASR) for contributing the websocket service.
diff --git a/funasr/runtime/python/websocket/parse_args.py b/funasr/runtime/python/websocket/parse_args.py
new file mode 100644
index 0000000..2528a76
--- /dev/null
+++ b/funasr/runtime/python/websocket/parse_args.py
@@ -0,0 +1,35 @@
+# -*- encoding: utf-8 -*-
+import argparse
+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="damo/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch",
+ help="model from modelscope")
+parser.add_argument("--asr_model_online",
+ type=str,
+ default="damo/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-online",
+ help="model from modelscope")
+parser.add_argument("--vad_model",
+ type=str,
+ default="damo/speech_fsmn_vad_zh-cn-16k-common-pytorch",
+ help="model from modelscope")
+parser.add_argument("--punc_model",
+ type=str,
+ default="damo/punc_ct-transformer_zh-cn-common-vad_realtime-vocab272727",
+ help="model from modelscope")
+parser.add_argument("--ngpu",
+ type=int,
+ default=1,
+ help="0 for cpu, 1 for gpu")
+
+args = parser.parse_args()
\ No newline at end of file
diff --git a/funasr/runtime/python/websocket/ws_client.py b/funasr/runtime/python/websocket/ws_client.py
new file mode 100644
index 0000000..8bbf103
--- /dev/null
+++ b/funasr/runtime/python/websocket/ws_client.py
@@ -0,0 +1,182 @@
+# -*- encoding: utf-8 -*-
+import os
+import time
+import websockets
+import asyncio
+# import threading
+import argparse
+import json
+
+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("--chunk_interval",
+ type=int,
+ default=10,
+ help="chunk")
+parser.add_argument("--audio_in",
+ type=str,
+ default=None,
+ help="audio_in")
+
+args = parser.parse_args()
+args.chunk_size = [int(x) for x in args.chunk_size.split(",")]
+
+# voices = asyncio.Queue()
+from queue import Queue
+voices = Queue()
+
+# 鍏朵粬鍑芥暟鍙互閫氳繃璋冪敤send(data)鏉ュ彂閫佹暟鎹紝渚嬪锛�
+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)
+ is_speaking = True
+ while True:
+
+ data = stream.read(CHUNK)
+ data = data.decode('ISO-8859-1')
+ message = json.dumps({"chunk_size": args.chunk_size, "chunk_interval": args.chunk_interval, "audio": data, "is_speaking": is_speaking, "is_finished": is_finished})
+
+ voices.put(message)
+ #print(voices.qsize())
+
+ await asyncio.sleep(0.005)
+
+# 鍏朵粬鍑芥暟鍙互閫氳繃璋冪敤send(data)鏉ュ彂閫佹暟鎹紝渚嬪锛�
+async def record_from_scp():
+ import wave
+ 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]
+ for wav in wavs:
+ wav_splits = wav.strip().split()
+ wav_path = wav_splits[1] if len(wav_splits) > 1 else wav_splits[0]
+ # bytes_f = open(wav_path, "rb")
+ # bytes_data = bytes_f.read()
+ with wave.open(wav_path, "rb") as wav_file:
+ # 鑾峰彇闊抽鍙傛暟
+ params = wav_file.getparams()
+ # 鑾峰彇澶翠俊鎭殑闀垮害
+ # header_length = wav_file.getheaders()[0][1]
+ # 璇诲彇闊抽甯ф暟鎹紝璺宠繃澶翠俊鎭�
+ # wav_file.setpos(header_length)
+ frames = wav_file.readframes(wav_file.getnframes())
+
+ # 灏嗛煶棰戝抚鏁版嵁杞崲涓哄瓧鑺傜被鍨嬬殑鏁版嵁
+ audio_bytes = bytes(frames)
+ # stride = int(args.chunk_size/1000*16000*2)
+ stride = int(60*args.chunk_size[1]/args.chunk_interval/1000*16000*2)
+ chunk_num = (len(audio_bytes)-1)//stride + 1
+ # print(stride)
+ is_speaking = True
+ for i in range(chunk_num):
+ if i == chunk_num-1:
+ is_speaking = False
+ beg = i*stride
+ data = audio_bytes[beg:beg+stride]
+ data = data.decode('ISO-8859-1')
+ message = json.dumps({"chunk_size": args.chunk_size, "chunk_interval": args.chunk_interval, "is_speaking": is_speaking, "audio": data, "is_finished": is_finished})
+ voices.put(message)
+ # print("data_chunk: ", len(data_chunk))
+ # print(voices.qsize())
+
+ await asyncio.sleep(60*args.chunk_size[1]/args.chunk_interval/1000)
+
+ is_finished = True
+ message = json.dumps({"is_finished": is_finished})
+ voices.put(message)
+
+async def ws_send():
+ global voices
+ global websocket
+ print("started to sending data!")
+ while True:
+ while not voices.empty():
+ data = voices.get()
+ voices.task_done()
+ try:
+ await websocket.send(data) # 閫氳繃ws瀵硅薄鍙戦�佹暟鎹�
+ except Exception as e:
+ print('Exception occurred:', e)
+ await asyncio.sleep(0.005)
+ await asyncio.sleep(0.005)
+
+
+
+async def message():
+ global websocket
+ text_print = ""
+ while True:
+ try:
+ meg = await websocket.recv()
+ meg = json.loads(meg)
+ # print(meg, end = '')
+ # print("\r")
+ text = meg["text"][0]
+ text_print += text
+ text_print = text_print[-55:]
+ os.system('clear')
+ print("\r"+text_print)
+ except Exception as e:
+ print("Exception:", e)
+
+
+async def print_messge():
+ global websocket
+ while True:
+ try:
+ meg = await websocket.recv()
+ meg = json.loads(meg)
+ print(meg)
+ except Exception as e:
+ print("Exception:", e)
+
+
+async def ws_client():
+ global websocket # 瀹氫箟涓�涓叏灞�鍙橀噺ws锛岀敤浜庝繚瀛榳ebsocket杩炴帴瀵硅薄
+ # uri = "ws://11.167.134.197:8899"
+ uri = "ws://{}:{}".format(args.host, args.port)
+ #ws = await websockets.connect(uri, subprotocols=["binary"]) # 鍒涘缓涓�涓暱杩炴帴
+ async for websocket in websockets.connect(uri, subprotocols=["binary"], ping_interval=None):
+ if args.audio_in is not None:
+ task = asyncio.create_task(record_from_scp()) # 鍒涘缓涓�涓悗鍙颁换鍔″綍闊�
+ else:
+ task = asyncio.create_task(record_microphone()) # 鍒涘缓涓�涓悗鍙颁换鍔″綍闊�
+ task2 = asyncio.create_task(ws_send()) # 鍒涘缓涓�涓悗鍙颁换鍔″彂閫�
+ task3 = asyncio.create_task(message()) # 鍒涘缓涓�涓悗鍙版帴鏀舵秷鎭殑浠诲姟
+ await asyncio.gather(task, task2, task3)
+
+
+asyncio.get_event_loop().run_until_complete(ws_client()) # 鍚姩鍗忕▼
+asyncio.get_event_loop().run_forever()
diff --git a/funasr/runtime/python/websocket/ws_server_online.py b/funasr/runtime/python/websocket/ws_server_online.py
new file mode 100644
index 0000000..7ef0e21
--- /dev/null
+++ b/funasr/runtime/python/websocket/ws_server_online.py
@@ -0,0 +1,108 @@
+import asyncio
+import json
+import websockets
+import time
+from queue import Queue
+import threading
+import logging
+import tracemalloc
+import numpy as np
+
+from parse_args import args
+from modelscope.pipelines import pipeline
+from modelscope.utils.constant import Tasks
+from modelscope.utils.logger import get_logger
+from funasr_onnx.utils.frontend import load_bytes
+
+tracemalloc.start()
+
+logger = get_logger(log_level=logging.CRITICAL)
+logger.setLevel(logging.CRITICAL)
+
+
+websocket_users = set()
+
+
+print("model loading")
+
+inference_pipeline_asr_online = pipeline(
+ task=Tasks.auto_speech_recognition,
+ model=args.asr_model_online,
+ model_revision='v1.0.4')
+
+print("model loaded")
+
+
+
+async def ws_serve(websocket, path):
+ frames_online = []
+ global websocket_users
+ websocket.send_msg = Queue()
+ websocket_users.add(websocket)
+ websocket.param_dict_asr_online = {"cache": dict()}
+ websocket.speek_online = Queue()
+ ss_online = threading.Thread(target=asr_online, args=(websocket,))
+ ss_online.start()
+
+ try:
+ async for message in websocket:
+ message = json.loads(message)
+ is_finished = message["is_finished"]
+ if not is_finished:
+ audio = bytes(message['audio'], 'ISO-8859-1')
+
+ is_speaking = message["is_speaking"]
+ websocket.param_dict_asr_online["is_final"] = not is_speaking
+
+ websocket.param_dict_asr_online["chunk_size"] = message["chunk_size"]
+
+
+ frames_online.append(audio)
+
+ if len(frames_online) % message["chunk_interval"] == 0 or not is_speaking:
+
+ audio_in = b"".join(frames_online)
+ websocket.speek_online.put(audio_in)
+ frames_online = []
+
+ if not websocket.send_msg.empty():
+ await websocket.send(websocket.send_msg.get())
+ websocket.send_msg.task_done()
+
+
+ except websockets.ConnectionClosed:
+ print("ConnectionClosed...", websocket_users) # 閾炬帴鏂紑
+ websocket_users.remove(websocket)
+ except websockets.InvalidState:
+ print("InvalidState...") # 鏃犳晥鐘舵��
+ except Exception as e:
+ print("Exception:", e)
+
+
+
+def asr_online(websocket): # ASR鎺ㄧ悊
+ global websocket_users
+ while websocket in websocket_users:
+ if not websocket.speek_online.empty():
+ audio_in = websocket.speek_online.get()
+ websocket.speek_online.task_done()
+ if len(audio_in) > 0:
+ # print(len(audio_in))
+ audio_in = load_bytes(audio_in)
+ rec_result = inference_pipeline_asr_online(audio_in=audio_in,
+ param_dict=websocket.param_dict_asr_online)
+ if websocket.param_dict_asr_online["is_final"]:
+ websocket.param_dict_asr_online["cache"] = dict()
+
+ if "text" in rec_result:
+ if rec_result["text"] != "sil" and rec_result["text"] != "waiting_for_more_voice":
+ print(rec_result["text"])
+ message = json.dumps({"mode": "online", "text": rec_result["text"]})
+ websocket.send_msg.put(message)
+
+ time.sleep(0.005)
+
+
+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()
\ No newline at end of file
--
Gitblit v1.9.1