CPU動作もする日本語のストリーミング音声認識モデル「kodama-ja-streaming-small」を Mac で少し試した（ONNX ＋ Python）

はじめに

●ayousanz/kodama-ja-streaming-small · Hugging Face
　https://huggingface.co/ayousanz/kodama-ja-streaming-small

以下は、キャプチャした画像内で書かれている、このモデルに関する説明です。

日本語ストリーミング音声認識（ASR）モデル。moonshine-ai/moonshine-streaming-small（MIT・英語）を ReazonSpeech v2 の全量 35,000 時間でフルファインチューニングしたものです。 CPU オフライン動作・低遅延ストリーミングを想定し、ONNX / ORT の deploy 資産（324MB）を同梱しています。

余談

余談ですが、上で話に出ている VOSK は、2022年ごろに試して以下の記事を書いていたことがあったりします。

●日本語音声のマイク入力をオフラインでリアルタイム音声認識：「VOSK」を JavaScript（Node.js）で扱う - Qiita
　https://qiita.com/youtoy/items/649dcad9ecccf75a9d01

試した内容の動画

先に、今回のお試しの結果（リアルタイムな音声認識を試した時の様子）を掲載します。

しゃべった内容がテキスト化されていることを確認できました。

試した手順

以下、今回の内容を試した手順です。

今回のフォルダ構成

今回、以下のようなフォルダ構成になるようにして進めました。

【ルート】/
├─ models/
│  └─ kodama/
│     ├─ 【必要な JSON ファイル】
│     └─ onnx/
│        └─ 【必要なファイル一式】
└─ python/
   └─ 【必要なファイル一式】

下準備

下準備を進めます。

まずは、「kodama-ja-streaming-small」関連のファイルを用意する環境を準備します。

フォルダ・仮想環境の準備

フォルダと仮想環境を準備します。

mkdir -p models/kodama python

uv venv
source .venv/bin/activate

仮想環境を上記でアクティベートした状態で、以下のインストール・ダウンロードを行います。

インストール・ダウンロード

今回の内容を進めるのに必要なものをインストール・ダウンロードします

uv pip install \
  onnxruntime \
  numpy \
  sounddevice \
  tokenizers \
  huggingface_hub

hf download ayousanz/kodama-ja-streaming-small \
  --local-dir models/kodama \
  --include "onnx/*" \
  --include "*.json"

これで、ひとまずの下準備は完了しました。

「kodama-ja-streaming-small」を使った音声認識

あとは処理用のコードを用意して、それを実行します。

コードの用意

具体的には、以下の内容です。

code python/transcribe_mic.py

import json
import queue
import sys
from pathlib import Path

import numpy as np
import onnxruntime as ort
import sounddevice as sd
from tokenizers import Tokenizer


# ============================================================
# 設定
# ============================================================

SAMPLE_RATE = 16000

# この秒数ごとに独立して書き起こす
CHUNK_SEC = 3.0

# これ未満の音量なら推論しない
# 無音時の「あ」などのハルシネーション対策
SILENCE_RMS = 0.008

TOKENS_PER_SEC = 6.5

NUM_LAYERS = 10
NUM_HEADS = 8
HEAD_DIM = 64

PAD_TO_MULTIPLE_OF = 80


# ============================================================
# パス
# ============================================================

ROOT = Path(__file__).resolve().parent.parent

MODEL_DIR = ROOT / "models" / "kodama"
ONNX_DIR = MODEL_DIR / "onnx"

ENCODER_PATH = (
    ONNX_DIR / "encoder.onnx"
)

CROSS_KV_PATH = (
    ONNX_DIR / "cross_kv_prefill.onnx"
)

DECODER_PATH = (
    ONNX_DIR / "decoder_step_crosskv.int8a.onnx"
)

GENERATION_CONFIG_PATH = (
    MODEL_DIR / "generation_config.json"
)

TOKENIZER_PATH = (
    MODEL_DIR / "tokenizer.json"
)


# ============================================================
# generation_config.json
# ============================================================

generation_config = json.loads(
    GENERATION_CONFIG_PATH.read_text(
        encoding="utf-8"
    )
)

BOS = generation_config.get(
    "decoder_start_token_id",
    generation_config.get(
        "bos_token_id",
        1,
    ),
)

EOS = generation_config.get(
    "eos_token_id",
    2,
)

if isinstance(BOS, list):
    BOS = BOS[0]

if isinstance(EOS, list):
    EOS_IDS = {
        int(x)
        for x in EOS
    }
else:
    EOS_IDS = {
        int(EOS)
    }

BOS = int(BOS)


# ============================================================
# Tokenizer
# ============================================================

tokenizer = Tokenizer.from_file(
    str(TOKENIZER_PATH)
)


# ============================================================
# ONNX Runtime
# ============================================================

options = ort.SessionOptions()

options.graph_optimization_level = (
    ort.GraphOptimizationLevel.ORT_ENABLE_ALL
)

# Apple Silicon CPU
options.intra_op_num_threads = 4
options.inter_op_num_threads = 1


def load_model(path):

    return ort.InferenceSession(
        str(path),
        sess_options=options,
        providers=[
            "CPUExecutionProvider"
        ],
    )


print("Kodamaを読み込み中...")

encoder = load_model(
    ENCODER_PATH
)

cross_kv_model = load_model(
    CROSS_KV_PATH
)

decoder = load_model(
    DECODER_PATH
)

print("モデル読み込み完了")
print()


# ============================================================
# 音声前処理
# ============================================================

def prepare_audio(audio):

    audio = np.asarray(
        audio,
        dtype=np.float32,
    ).reshape(-1)

    real_length = len(audio)

    # 80サンプル単位にpadding
    pad_length = (
        -real_length
    ) % PAD_TO_MULTIPLE_OF

    if pad_length:

        audio = np.pad(
            audio,
            (
                0,
                pad_length,
            ),
            mode="constant",
        )

    input_values = audio[
        None,
        :
    ].astype(
        np.float32,
        copy=False,
    )

    attention_mask = np.zeros(
        (
            1,
            len(audio),
        ),
        dtype=np.int64,
    )

    attention_mask[
        0,
        :real_length
    ] = 1

    return (
        input_values,
        attention_mask,
    )


# ============================================================
# Encoder
# ============================================================

def run_encoder(audio):

    (
        input_values,
        attention_mask,
    ) = prepare_audio(
        audio
    )

    outputs = encoder.run(
        None,
        {
            "input_values":
                input_values,

            "attention_mask":
                attention_mask,
        },
    )

    return (
        outputs[0],
        outputs[1],
    )


# ============================================================
# Cross KV
# ============================================================

def run_cross_kv(
    encoder_hidden_states,
):

    names = [
        output.name
        for output
        in cross_kv_model.get_outputs()
    ]

    outputs = cross_kv_model.run(
        None,
        {
            "encoder_hidden_states":
                encoder_hidden_states
        },
    )

    return dict(
        zip(
            names,
            outputs,
        )
    )


# ============================================================
# Decoder
# ============================================================

def decode(
    encoder_attention_mask,
    cross_kv,
    max_tokens,
):

    current_token = BOS

    tokens = []

    # self attention cache
    cache = {}

    for layer in range(
        NUM_LAYERS
    ):

        shape = (
            1,
            NUM_HEADS,
            0,
            HEAD_DIM,
        )

        cache[
            f"past_self_k_{layer}"
        ] = np.zeros(
            shape,
            dtype=np.float32,
        )

        cache[
            f"past_self_v_{layer}"
        ] = np.zeros(
            shape,
            dtype=np.float32,
        )

    output_names = [
        output.name
        for output
        in decoder.get_outputs()
    ]

    # token生成
    for _ in range(
        max_tokens
    ):

        feeds = {
            "decoder_input_ids":
                np.array(
                    [[current_token]],
                    dtype=np.int64,
                ),

            "encoder_attention_mask":
                encoder_attention_mask,

            **cache,
        }

        # Cross Attention KV
        for layer in range(
            NUM_LAYERS
        ):

            feeds[
                f"cross_k_{layer}"
            ] = cross_kv[
                f"cross_k_{layer}"
            ]

            feeds[
                f"cross_v_{layer}"
            ] = cross_kv[
                f"cross_v_{layer}"
            ]

        outputs = decoder.run(
            None,
            feeds,
        )

        result = dict(
            zip(
                output_names,
                outputs,
            )
        )

        logits = result[
            "logits"
        ]

        next_token = int(
            np.argmax(
                logits[
                    0,
                    -1,
                    :
                ]
            )
        )

        if next_token in EOS_IDS:
            break

        tokens.append(
            next_token
        )

        current_token = (
            next_token
        )

        # self attention cache更新
        new_cache = {}

        for layer in range(
            NUM_LAYERS
        ):

            new_cache[
                f"past_self_k_{layer}"
            ] = result[
                f"new_self_k_{layer}"
            ]

            new_cache[
                f"past_self_v_{layer}"
            ] = result[
                f"new_self_v_{layer}"
            ]

        cache = new_cache

    return tokenizer.decode(
        tokens,
        skip_special_tokens=True,
    )


# ============================================================
# ASR
# ============================================================

def transcribe(audio):

    (
        hidden,
        encoder_mask,
    ) = run_encoder(
        audio
    )

    cross_kv = run_cross_kv(
        hidden
    )

    seconds = (
        len(audio)
        / SAMPLE_RATE
    )

    max_tokens = max(
        1,
        int(
            seconds
            * TOKENS_PER_SEC
        )
        + 1,
    )

    return decode(
        encoder_mask,
        cross_kv,
        max_tokens,
    )


# ============================================================
# マイク
# ============================================================

audio_queue = queue.Queue()


def audio_callback(
    indata,
    frames,
    time_info,
    status,
):

    if status:

        print(
            status,
            file=sys.stderr,
        )

    audio_queue.put(
        indata[:, 0].copy()
    )


# 0.1秒ごとにマイクから受け取る
BLOCK_SIZE = int(
    SAMPLE_RATE
    * 0.1
)

# 3秒分
CHUNK_SAMPLES = int(
    SAMPLE_RATE
    * CHUNK_SEC
)


# ============================================================
# Main
# ============================================================

print(
    f"{CHUNK_SEC:.1f}秒ごとに文字起こしします"
)

print(
    "Ctrl+C で終了"
)

print()


buffer = np.empty(
    0,
    dtype=np.float32,
)


try:

    with sd.InputStream(
        samplerate=SAMPLE_RATE,
        channels=1,
        dtype="float32",
        blocksize=BLOCK_SIZE,
        callback=audio_callback,
    ):

        while True:

            block = (
                audio_queue.get()
            )

            buffer = np.concatenate(
                (
                    buffer,
                    block,
                )
            )

            # 3秒分たまるまで待つ
            if (
                len(buffer)
                < CHUNK_SAMPLES
            ):
                continue

            # ちょうど3秒を取り出す
            audio = buffer[
                :CHUNK_SAMPLES
            ]

            # 余った部分は次回へ
            buffer = buffer[
                CHUNK_SAMPLES:
            ]

            # 音量確認
            rms = float(
                np.sqrt(
                    np.mean(
                        np.square(
                            audio
                        )
                    )
                )
            )

            # 無音ならKodamaを呼ばない
            if rms < SILENCE_RMS:
                continue

            # 3秒分を独立して認識
            text = transcribe(
                audio
            )

            if text.strip():

                print(
                    text,
                    flush=True,
                )


except KeyboardInterrupt:

    print(
        "\n終了"
    )

音声認識を試す

あとは下記のコマンドを実行して、音声認識を試します。

python python/transcribe_mic.py

試した結果は、以下の通りです。