KoeAct-JP v0.1.1: koeact.runtime applies the evaluation preprocessing by default

55ca644 verified 4 days ago

9.56 kB

# SPDX-License-Identifier: Apache-2.0
# Copyright 2026 ayutaz
"""KoeAct-JP inference from the released files: 16 kHz speech -> gated Action JSON.
The release directory holds ``model.safetensors`` (Ralomi's input layer and blocks 1..4 as
released, the fine-tuned blocks 5..8, the CTC classifier and the decoder) and ``config.json``
(the ``p_act`` threshold). Nothing is downloaded. The outputs equal ``koeact.infer``'s torch
engine (``tests/test_runtime.py``); this module imports no training or evaluation code, so it
runs from the bundled ``code/`` directory alone (``koeact.release``).
    python -m koeact.runtime --model <release dir> a.wav b.wav
    python -m koeact.runtime --model <release dir> --mic
``from_device_line`` turns a CoreS3 ``ACTION ... p_act=... ids=...`` line (the firmware stops
at the decoder's piece ids) into the same gated JSON; ``follow`` does it for every recognition in
a stream of serial lines (``--serial <port>``, needs pyserial; touch the screen to record 4 s).
"""

import argparse
import json
import math
import re
from pathlib import Path

import numpy as np

from koeact.action import parse_output, to_json
from koeact.features.logmel import HOP, MAX_SECONDS, SAMPLE_RATE, WINDOW, log_mel
from koeact.model import pieces

FORMAT = "koeact-runtime-1"
FIRST_BLOCK = 4  # the decoder reads layer 4; blocks 5..8 belong to the fine-tuned head
MIN_SAMPLES = WINDOW + 3 * HOP  # Ralomi needs four log-Mel frames
CLEAN_RMS_DBFS = -23.0
CLEAN_LEAD, CLEAN_TAIL = 0.5, 0.3  # seconds of silence before and after
PEAK = 0.99
PAIR = re.compile(r"(\w+)=(-?\d+(?:\.\d+)?)")


def build_net(ft_blocks: int):
    """Ralomi's lower half (input layer, blocks 1..4) and the plan B head over it."""
    from torch import nn

    from koeact.model.plan_b import DecoderB
    from koeact.model.ralomi_torch import RalomiTorch

    class KoeActNet(nn.Module):
        def __init__(self) -> None:
            super().__init__()
            ralomi = RalomiTorch()
            self.subsampling_projection = ralomi.subsampling_projection
            self.input_normalization = ralomi.input_normalization
            self.lower = nn.ModuleList(ralomi.blocks[:FIRST_BLOCK])
            self.head = DecoderB(ralomi, ft_blocks)

    return KoeActNet()


def _set_level(x: np.ndarray, dbfs: float) -> np.ndarray:
    rms = float(np.sqrt(np.mean(np.square(x, dtype=np.float64)))) if len(x) else 0.0
    if not 0 < rms < math.inf:  # silence stays silence; NaN or inf is left as it is
        return x
    y = x * (10 ** ((dbfs - 20 * math.log10(rms)) / 20))
    peak = float(np.abs(y).max())
    return y * (PEAK / peak) if peak > PEAK else y


def clean(audio: np.ndarray) -> np.ndarray:
    """``koeact.features.augment.clean``: -23 dBFS RMS, 0.5 s of silence before, 0.3 s after."""
    lead = np.zeros(round(CLEAN_LEAD * SAMPLE_RATE), np.float32)
    tail = np.zeros(round(CLEAN_TAIL * SAMPLE_RATE), np.float32)
    return np.concatenate([lead, _set_level(audio, CLEAN_RMS_DBFS), tail]).astype(np.float32)


_clean = clean  # ``KoeAct.__call__`` has a ``clean`` parameter


def prepare(audio: np.ndarray) -> np.ndarray:
    """Cut to 15 s and pad short audio to four log-Mel frames (as the evaluations did)."""
    x = audio[: MAX_SECONDS * SAMPLE_RATE]
    if len(x) < MIN_SAMPLES:
        x = np.pad(x, (0, MIN_SAMPLES - len(x)))
    return x


def resample(x: np.ndarray, sr: int) -> np.ndarray:
    if sr == SAMPLE_RATE:
        return x
    from math import gcd

    from scipy.signal import resample_poly

    g = gcd(SAMPLE_RATE, sr)
    return resample_poly(x, SAMPLE_RATE // g, sr // g).astype(np.float32)


def load_audio(path: Path) -> np.ndarray:
    """Mono float32 at 16 kHz (other rates are resampled; channels are averaged)."""
    import soundfile as sf

    x, sr = sf.read(path, dtype="float32", always_2d=True)
    return resample(x.mean(axis=1), sr)


class KoeAct:
    def __init__(self, net, threshold: float) -> None:
        self.net, self.threshold = net.eval(), threshold
        self.grammar = pieces.ActionGrammar()

    @classmethod
    def from_dir(cls, path: Path | str, threads: int = 0) -> "KoeAct":
        import torch
        from safetensors.torch import load_file

        if threads:
            torch.set_num_threads(threads)
        path = Path(path)
        config = json.loads((path / "config.json").read_text("utf-8"))
        if config.get("format") != FORMAT:
            raise ValueError(f"{path}: not a {FORMAT} release ({config.get('format')!r})")
        net = build_net(config["ft_blocks"])
        net.load_state_dict(load_file(path / "model.safetensors"))
        return cls(net, config["threshold"])

    def _frames(self, features: np.ndarray):
        """``[1, steps, 2, 320]``: layer 4 (rounded to fp16, as the training cache) and zeros
        for layer 8, which the head recomputes from layer 4."""
        import torch

        from koeact.model.ralomi_torch import SUBSAMPLING

        x = torch.from_numpy(features)[None]
        steps = x.shape[1] // SUBSAMPLING
        h = x[:, : steps * SUBSAMPLING].reshape(1, steps, -1)
        h = torch.relu(self.net.input_normalization(self.net.subsampling_projection(h)))
        for block in self.net.lower:
            h = block(h)
        h4 = h[0].numpy()
        frames = np.stack([h4, np.zeros_like(h4)], axis=1).astype(np.float16)
        return torch.from_numpy(frames.astype(np.float32))[None]

    def __call__(self, audio: np.ndarray, clean: bool = True) -> dict:
        """One utterance (float32, 16 kHz) -> ``output`` (gated JSON), ``raw``, ``p_act``.
        ``clean`` applies the training preprocessing first (v0.1.0 did not)."""
        import torch

        from koeact.model import decoder_a2

        if clean:
            audio = _clean(audio)
        with torch.no_grad():
            x = self._frames(log_mel(prepare(audio)))
            mask = torch.ones(x.shape[:2], dtype=torch.bool)
            d = decoder_a2.greedy(self.net.head, x, mask, self.grammar)[0]
        return gate(to_json(d.calls), round(d.p_act, 6), self.threshold)


def gate(raw: str, p_act: float, threshold: float) -> dict:
    accepted = p_act >= threshold  # below the threshold the answer is "[]"
    return {"output": raw if accepted else "[]", "raw": raw, "p_act": p_act,
            "accepted": accepted}  # fmt: skip


def from_device_line(line: str, threshold: float) -> dict:
    """A CoreS3 ``ACTION``/``A2`` line (``p_act=`` and the piece ids after ``ids=``) -> JSON."""
    head, _, ids = line.partition(" ids=")
    fields = dict(PAIR.findall(head))
    calls = parse_output(pieces.decode_text([int(i) for i in ids.split()])).calls or []
    return gate(to_json(calls), float(fields["p_act"]), threshold)


def follow(lines, threshold: float):
    """Gated JSON for each microphone recognition (``A2 ... real=1 ... ids=``) or ``ACTION`` line;
    ``real=0`` lines are timings over stand-in frames and are skipped."""
    for line in lines:
        if " ids=" not in line:
            continue
        if line.startswith("ACTION ") or (line.startswith("A2 ") and " real=1 " in line):
            yield from_device_line(line, threshold)


def serial_lines(port: str):
    import serial

    with serial.Serial(port, 115200, timeout=1) as s:
        while True:
            raw = s.readline()
            if raw:
                yield raw.decode("utf-8", "replace").rstrip()


def record_utterance() -> np.ndarray:
    """One utterance from the default microphone: Enter starts, Enter stops."""
    import sounddevice as sd

    sr = int(sd.query_devices(kind="input")["default_samplerate"])
    chunks: list[np.ndarray] = []
    with sd.InputStream(samplerate=sr, channels=1, dtype="float32",
                        callback=lambda data, *_: chunks.append(data[:, 0].copy())):  # fmt: skip
        input("  ● 話してください（終わったら Enter） / speak, then press Enter")
    return resample(np.concatenate(chunks) if chunks else np.zeros(0, np.float32), sr)


def main() -> None:
    p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    p.add_argument("wav", nargs="*", type=Path)
    p.add_argument("--model", type=Path, default=Path("."), help="the release directory")
    p.add_argument("--threads", type=int, default=0)
    p.add_argument("--mic", action="store_true", help="Enter starts and stops; q quits")
    p.add_argument(
        "--no-clean",
        action="store_true",
        help="feed the audio as is (v0.1.0); the default levels it and pads silence",
    )
    p.add_argument("--serial", help="a CoreS3 with the KoeAct firmware (reads config.json only)")
    args = p.parse_args()
    if args.serial:
        threshold = json.loads((args.model / "config.json").read_text("utf-8"))["threshold"]
        for r in follow(serial_lines(args.serial), threshold):
            print(json.dumps(r, ensure_ascii=False), flush=True)
        return
    model = KoeAct.from_dir(args.model, args.threads)
    if args.mic:
        while input("Enter で録音（q で終了） / Enter to record, q to quit > ").strip() != "q":
            r = model(record_utterance(), clean=not args.no_clean)
            print(f"  {r['output']}  (p_act {r['p_act']:.3f})")
        return
    for w in args.wav:
        x = load_audio(w)
        r = model(x, clean=not args.no_clean)
        print(json.dumps({"file": str(w), "output": r["output"], "p_act": r["p_act"],
                          "accepted": r["accepted"]}, ensure_ascii=False), flush=True)  # fmt: skip


if __name__ == "__main__":
    main()

	# SPDX-License-Identifier: Apache-2.0
	# Copyright 2026 ayutaz
	"""KoeAct-JP inference from the released files: 16 kHz speech -> gated Action JSON.

	The release directory holds ``model.safetensors`` (Ralomi's input layer and blocks 1..4 as
	released, the fine-tuned blocks 5..8, the CTC classifier and the decoder) and ``config.json``
	(the ``p_act`` threshold). Nothing is downloaded. The outputs equal ``koeact.infer``'s torch
	engine (``tests/test_runtime.py``); this module imports no training or evaluation code, so it
	runs from the bundled ``code/`` directory alone (``koeact.release``).

	python -m koeact.runtime --model <release dir> a.wav b.wav
	python -m koeact.runtime --model <release dir> --mic

	``from_device_line`` turns a CoreS3 ``ACTION ... p_act=... ids=...`` line (the firmware stops
	at the decoder's piece ids) into the same gated JSON; ``follow`` does it for every recognition in
	a stream of serial lines (``--serial <port>``, needs pyserial; touch the screen to record 4 s).
	"""

	import argparse
	import json
	import math
	import re
	from pathlib import Path

	import numpy as np

	from koeact.action import parse_output, to_json
	from koeact.features.logmel import HOP, MAX_SECONDS, SAMPLE_RATE, WINDOW, log_mel
	from koeact.model import pieces

	FORMAT = "koeact-runtime-1"
	FIRST_BLOCK = 4 # the decoder reads layer 4; blocks 5..8 belong to the fine-tuned head
	MIN_SAMPLES = WINDOW + 3 * HOP # Ralomi needs four log-Mel frames
	CLEAN_RMS_DBFS = -23.0
	CLEAN_LEAD, CLEAN_TAIL = 0.5, 0.3 # seconds of silence before and after
	PEAK = 0.99
	PAIR = re.compile(r"(\w+)=(-?\d+(?:\.\d+)?)")


	def build_net(ft_blocks: int):
	"""Ralomi's lower half (input layer, blocks 1..4) and the plan B head over it."""
	from torch import nn

	from koeact.model.plan_b import DecoderB
	from koeact.model.ralomi_torch import RalomiTorch

	class KoeActNet(nn.Module):
	def __init__(self) -> None:
	super().__init__()
	ralomi = RalomiTorch()
	self.subsampling_projection = ralomi.subsampling_projection
	self.input_normalization = ralomi.input_normalization
	self.lower = nn.ModuleList(ralomi.blocks[:FIRST_BLOCK])
	self.head = DecoderB(ralomi, ft_blocks)

	return KoeActNet()


	def _set_level(x: np.ndarray, dbfs: float) -> np.ndarray:
	rms = float(np.sqrt(np.mean(np.square(x, dtype=np.float64)))) if len(x) else 0.0
	if not 0 < rms < math.inf: # silence stays silence; NaN or inf is left as it is
	return x
	y = x * (10 ** ((dbfs - 20 * math.log10(rms)) / 20))
	peak = float(np.abs(y).max())
	return y * (PEAK / peak) if peak > PEAK else y


	def clean(audio: np.ndarray) -> np.ndarray:
	"""``koeact.features.augment.clean``: -23 dBFS RMS, 0.5 s of silence before, 0.3 s after."""
	lead = np.zeros(round(CLEAN_LEAD * SAMPLE_RATE), np.float32)
	tail = np.zeros(round(CLEAN_TAIL * SAMPLE_RATE), np.float32)
	return np.concatenate([lead, _set_level(audio, CLEAN_RMS_DBFS), tail]).astype(np.float32)


	_clean = clean # ``KoeAct.__call__`` has a ``clean`` parameter


	def prepare(audio: np.ndarray) -> np.ndarray:
	"""Cut to 15 s and pad short audio to four log-Mel frames (as the evaluations did)."""
	x = audio[: MAX_SECONDS * SAMPLE_RATE]
	if len(x) < MIN_SAMPLES:
	x = np.pad(x, (0, MIN_SAMPLES - len(x)))
	return x


	def resample(x: np.ndarray, sr: int) -> np.ndarray:
	if sr == SAMPLE_RATE:
	return x
	from math import gcd

	from scipy.signal import resample_poly

	g = gcd(SAMPLE_RATE, sr)
	return resample_poly(x, SAMPLE_RATE // g, sr // g).astype(np.float32)


	def load_audio(path: Path) -> np.ndarray:
	"""Mono float32 at 16 kHz (other rates are resampled; channels are averaged)."""
	import soundfile as sf

	x, sr = sf.read(path, dtype="float32", always_2d=True)
	return resample(x.mean(axis=1), sr)


	class KoeAct:
	def __init__(self, net, threshold: float) -> None:
	self.net, self.threshold = net.eval(), threshold
	self.grammar = pieces.ActionGrammar()

	@classmethod
	def from_dir(cls, path: Path \| str, threads: int = 0) -> "KoeAct":
	import torch
	from safetensors.torch import load_file

	if threads:
	torch.set_num_threads(threads)
	path = Path(path)
	config = json.loads((path / "config.json").read_text("utf-8"))
	if config.get("format") != FORMAT:
	raise ValueError(f"{path}: not a {FORMAT} release ({config.get('format')!r})")
	net = build_net(config["ft_blocks"])
	net.load_state_dict(load_file(path / "model.safetensors"))
	return cls(net, config["threshold"])

	def _frames(self, features: np.ndarray):
	"""``[1, steps, 2, 320]``: layer 4 (rounded to fp16, as the training cache) and zeros
	for layer 8, which the head recomputes from layer 4."""
	import torch

	from koeact.model.ralomi_torch import SUBSAMPLING

	x = torch.from_numpy(features)[None]
	steps = x.shape[1] // SUBSAMPLING
	h = x[:, : steps * SUBSAMPLING].reshape(1, steps, -1)
	h = torch.relu(self.net.input_normalization(self.net.subsampling_projection(h)))
	for block in self.net.lower:
	h = block(h)
	h4 = h[0].numpy()
	frames = np.stack([h4, np.zeros_like(h4)], axis=1).astype(np.float16)
	return torch.from_numpy(frames.astype(np.float32))[None]

	def __call__(self, audio: np.ndarray, clean: bool = True) -> dict:
	"""One utterance (float32, 16 kHz) -> ``output`` (gated JSON), ``raw``, ``p_act``.

	``clean`` applies the training preprocessing first (v0.1.0 did not)."""
	import torch

	from koeact.model import decoder_a2

	if clean:
	audio = _clean(audio)
	with torch.no_grad():
	x = self._frames(log_mel(prepare(audio)))
	mask = torch.ones(x.shape[:2], dtype=torch.bool)
	d = decoder_a2.greedy(self.net.head, x, mask, self.grammar)[0]
	return gate(to_json(d.calls), round(d.p_act, 6), self.threshold)


	def gate(raw: str, p_act: float, threshold: float) -> dict:
	accepted = p_act >= threshold # below the threshold the answer is "[]"
	return {"output": raw if accepted else "[]", "raw": raw, "p_act": p_act,
	"accepted": accepted} # fmt: skip


	def from_device_line(line: str, threshold: float) -> dict:
	"""A CoreS3 ``ACTION``/``A2`` line (``p_act=`` and the piece ids after ``ids=``) -> JSON."""
	head, _, ids = line.partition(" ids=")
	fields = dict(PAIR.findall(head))
	calls = parse_output(pieces.decode_text([int(i) for i in ids.split()])).calls or []
	return gate(to_json(calls), float(fields["p_act"]), threshold)


	def follow(lines, threshold: float):
	"""Gated JSON for each microphone recognition (``A2 ... real=1 ... ids=``) or ``ACTION`` line;
	``real=0`` lines are timings over stand-in frames and are skipped."""
	for line in lines:
	if " ids=" not in line:
	continue
	if line.startswith("ACTION ") or (line.startswith("A2 ") and " real=1 " in line):
	yield from_device_line(line, threshold)


	def serial_lines(port: str):
	import serial

	with serial.Serial(port, 115200, timeout=1) as s:
	while True:
	raw = s.readline()
	if raw:
	yield raw.decode("utf-8", "replace").rstrip()


	def record_utterance() -> np.ndarray:
	"""One utterance from the default microphone: Enter starts, Enter stops."""
	import sounddevice as sd

	sr = int(sd.query_devices(kind="input")["default_samplerate"])
	chunks: list[np.ndarray] = []
	with sd.InputStream(samplerate=sr, channels=1, dtype="float32",
	callback=lambda data, *_: chunks.append(data[:, 0].copy())): # fmt: skip
	input(" ● 話してください（終わったら Enter） / speak, then press Enter")
	return resample(np.concatenate(chunks) if chunks else np.zeros(0, np.float32), sr)


	def main() -> None:
	p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
	p.add_argument("wav", nargs="*", type=Path)
	p.add_argument("--model", type=Path, default=Path("."), help="the release directory")
	p.add_argument("--threads", type=int, default=0)
	p.add_argument("--mic", action="store_true", help="Enter starts and stops; q quits")
	p.add_argument(
	"--no-clean",
	action="store_true",
	help="feed the audio as is (v0.1.0); the default levels it and pads silence",
	)
	p.add_argument("--serial", help="a CoreS3 with the KoeAct firmware (reads config.json only)")
	args = p.parse_args()
	if args.serial:
	threshold = json.loads((args.model / "config.json").read_text("utf-8"))["threshold"]
	for r in follow(serial_lines(args.serial), threshold):
	print(json.dumps(r, ensure_ascii=False), flush=True)
	return
	model = KoeAct.from_dir(args.model, args.threads)
	if args.mic:
	while input("Enter で録音（q で終了） / Enter to record, q to quit > ").strip() != "q":
	r = model(record_utterance(), clean=not args.no_clean)
	print(f" {r['output']} (p_act {r['p_act']:.3f})")
	return
	for w in args.wav:
	x = load_audio(w)
	r = model(x, clean=not args.no_clean)
	print(json.dumps({"file": str(w), "output": r["output"], "p_act": r["p_act"],
	"accepted": r["accepted"]}, ensure_ascii=False), flush=True) # fmt: skip


	if __name__ == "__main__":
	main()