funasr_wss_server.py: hotwords in the websocket protocol's JSON format reach the model as <unk> #3879

if "hotwords" in messagejson:
    hotword_data = messagejson["hotwords"]
    websocket.status_dict_asr["hotword"] = hotword_data
    websocket.status_dict_asr_online["hotword"] = hotword_data

import asyncio
import contextlib
import importlib.util
import io
import json
import os
import sys
import tempfile
import types

import funasr
import torch
from funasr.models.seaco_paraformer.model import SeacoParaformer

# Load the real server; only AutoModel (called when the module is imported) is replaced.
funasr.AutoModel = lambda **kwargs: types.SimpleNamespace(**kwargs)
sys.argv = ["funasr_wss_server.py"]
spec = importlib.util.spec_from_file_location(
    "funasr_wss_server", "runtime/python/websocket/funasr_wss_server.py"
)
server = importlib.util.module_from_spec(spec)
spec.loader.exec_module(server)

received = []
server.model_asr = types.SimpleNamespace(
    generate=lambda input, **kwargs: received.append(kwargs["hotword"]) or [{"text": ""}]
)
server.model_vad = types.SimpleNamespace(generate=lambda input, **kwargs: [{"value": []}])
server.model_sv = types.SimpleNamespace(
    generate=lambda input, **kwargs: [{"spk_embedding": torch.zeros(1, 192)}]
)


class Client:
    """One offline request: the config message, 0.1 s of audio, end of input."""

    def __init__(self, hotwords):
        config = {"mode": "offline", "chunk_size": [5, 10, 5], "chunk_interval": 10,
                  "wav_name": "demo", "is_speaking": True, "hotwords": hotwords}
        self.messages = iter([json.dumps(config, ensure_ascii=False), b"\x01\x00" * 1600,
                              json.dumps({"is_speaking": False, "is_end": True})])
        self.path = None

    def __aiter__(self):
        return self

    async def __anext__(self):
        try:
            return next(self.messages)
        except StopIteration:
            raise StopAsyncIteration

    async def send(self, message):
        pass

    async def close(self):
        pass


# SeACo-Paraformer's own hotword parser, with a toy vocabulary and a character seg_dict.
tokens = ["<blank>", "<s>", "</s>", "<unk>"] + list("阿里巴通义实验室")
model_dir = tempfile.mkdtemp()
with open(os.path.join(model_dir, "seg_dict"), "w", encoding="utf-8") as f:
    f.write("".join(f"{t} {t}\n" for t in tokens[4:]))
tokenizer = types.SimpleNamespace(tokens2ids=lambda ws: [tokens.index(w) if w in tokens else 3 for w in ws])
frontend = types.SimpleNamespace(cmvn_file=os.path.join(model_dir, "am.mvn"))

for hotwords in [
    '{"阿里巴巴":20,"通义实验室":30}',  # HTML5, Java, Android and C++ clients
    json.dumps({"阿里巴巴": 20, "通义实验室": 30}, ensure_ascii=False),  # funasr_wss_client.py
    "阿里巴巴 通义实验室",  # what AutoModel expects
]:
    received.clear()
    with contextlib.redirect_stdout(io.StringIO()):  # the server prints every message
        asyncio.run(server.ws_serve(Client(hotwords)))
        ids = SeacoParaformer.generate_hotwords_list(
            types.SimpleNamespace(sos=1), received[0], tokenizer=tokenizer, frontend=frontend
        )
    print(f"client sends {hotwords!r}")
    print(f"  model.generate(hotword={received[0]!r})")
    print(f"  hotword ids: {ids}")
server.EXECUTOR.shutdown()