Skip to content

Add CI to export Paraformer to RKNN - #2692

Merged
csukuangfj merged 2 commits into
k2-fsa:masterfrom
csukuangfj:ci-paraformer-rknn
Oct 15, 2025
Merged

csukuangfj merged 2 commits into
k2-fsa:masterfrom
csukuangfj:ci-paraformer-rknn

Conversation

@csukuangfj

@csukuangfj csukuangfj commented Oct 15, 2025 •

Copy link
Copy Markdown
Collaborator

you can find them at
https://github.com/k2-fsa/sherpa-onnx/releases/tag/asr-models

you can use the following python code to test the exported rknn models on your board
test_on_rk3588_board_paraformer.py

(for the model with 15 seconds fixed input)

Screenshot 2025-10-15 at 16 33 06
#!/usr/bin/env python3

try:
    from rknnlite.api import RKNNLite
except:
    print("Please run this file on your board (linux + aarch64 + npu)")
    print("You need to install rknn_toolkit_lite2")
    print(
        " from https://github.com/airockchip/rknn-toolkit2/tree/master/rknn-toolkit-lite2/packages"
    )
    print(
        "https://github.com/airockchip/rknn-toolkit2/blob/v2.1.0/rknn-toolkit-lite2/packages/rknn_toolkit_lite2-2.1.0-cp310-cp310-linux_aarch64.whl"
    )
    print("is known to work")
    raise
import json
import time
from pathlib import Path
from typing import Tuple

import kaldi_native_fbank as knf
import numpy as np
import soundfile as sf


def read_tokens():
    with open("./tokens.json", encoding="utf-8") as f:
        tokens = json.load(f)
    id2token = dict()
    for i, t in enumerate(tokens):
        id2token[i] = t
    return id2token


def load_cmvn():
    neg_mean = None
    inv_stddev = None

    with open("am.mvn") as f:
        for line in f:
            if not line.startswith("<LearnRateCoef>"):
                continue
            t = line.split()[3:-1]

            if neg_mean is None:
                neg_mean = ",".join(t)
            else:
                inv_stddev = ",".join(t)
    neg_mean = map(float, neg_mean.split(","))
    inv_stddev = map(float, inv_stddev.split(","))

    return np.array(list(neg_mean), dtype=np.float32), np.array(
        list(inv_stddev), dtype=np.float32
    )


def compute_feat(filename):
    data, sample_rate = sf.read(
        filename,
        always_2d=True,
        dtype="float32",
    )
    data = data[:, 0]  # use only the first channel
    samples = np.ascontiguousarray(data)

    assert sample_rate == 16000, sample_rate

    opts = knf.FbankOptions()
    opts.frame_opts.dither = 0
    opts.frame_opts.snip_edges = False
    opts.frame_opts.samp_freq = sample_rate
    opts.mel_opts.num_bins = 80

    online_fbank = knf.OnlineFbank(opts)
    online_fbank.accept_waveform(sample_rate, (samples * 32768).tolist())
    online_fbank.input_finished()

    features = np.stack(
        [online_fbank.get_frame(i) for i in range(online_fbank.num_frames_ready)]
    )
    assert features.data.contiguous is True
    assert features.dtype == np.float32, features.dtype
    print("features sum", features.sum(), features.size, features.shape)

    window_size = 7  # lfr_m
    window_shift = 6  # lfr_n

    T = (features.shape[0] - window_size) // window_shift + 1
    features = np.lib.stride_tricks.as_strided(
        features,
        shape=(T, features.shape[1] * window_size),
        strides=((window_shift * features.shape[1]) * 4, 4),
    )
    neg_mean, inv_std = load_cmvn()
    features = (features + neg_mean) * inv_std
    return features


def init_model(filename, target_platform="rk3588"):
    if not Path(filename).is_file():
        exit(f"{filename} does not exist")

    rknn_lite = RKNNLite(verbose=False)
    ret = rknn_lite.load_rknn(path=filename)
    if ret != 0:
        exit(f"Load model {filename} failed!")

    ret = rknn_lite.init_runtime(core_mask=RKNNLite.NPU_CORE_0)
    if ret != 0:
        exit(f"Failed to init rknn runtime for {filename}")
    return rknn_lite


class RKNNModel:
    def __init__(self, model: str, target_platform="rk3588"):
        self.model = init_model(model)

    def release(self):
        self.model.release()

    def __call__(self, x: np.ndarray):
        """
        Args:
          x: (1, 560, 560), np.float32
        Returns:
          prob:
          next_h:
          next_c
        """
        x_length = np.array([x.shape[1]], dtype=np.int32)
        out = self.model.inference(inputs=[x, x_length])
        return out


def main():
    model = RKNNModel(model="./model.rknn")
    for i in range(1):
        test(model)


def test(model):

    id2token = read_tokens()
    print("len id2token", len(id2token))

    start = time.time()
    features = compute_feat("./1.wav")

    print(features.shape)  # (85, 560)
    padding = np.ones((560 - features.shape[0], 560), dtype=np.float32) * features[0][0]

    features = np.concatenate([features, padding], axis=0)
    features = features[None]
    features_length = np.array([features.shape[1]], dtype=np.int32)

    print(features.shape)
    outputs = model(features)
    print("outputs", len(outputs))

    print(type(outputs[0]))
    print(type(outputs[1]))
    print(outputs[0].dtype, outputs[1].dtype)
    print(outputs[0].shape, outputs[1].shape)

    log_probs = outputs[0][0]
    token_num = int(outputs[1][0].item())
    print(token_num, type(token_num))
    print(log_probs.shape)
    print(log_probs)

    y = log_probs.argmax(axis=-1)[:token_num]
    print(y)

    return


if __name__ == "__main__":
    main()

Summary by CodeRabbit

  • New Features
    • Automated creation and release of Paraformer RKNN packages for multiple RK356x targets, including model files, tokens, README, and test audio.
    • Example script now downloads required resources and converts encoder/predictor/decoder ONNX models to RKNN in one flow.
  • Improvements
    • More robust model loading that supports checkpoints saved in different formats.
  • Chores
    • Introduced CI workflow to build, package, and publish RKNN artifacts across platforms and frameworks.

@coderabbitai

coderabbitai Bot commented Oct 15, 2025 •

Copy link
Copy Markdown

Caution

Review failed

The pull request is closed.

Walkthrough

Adds a new GitHub Actions workflow to export Paraformer models to RKNN across platforms and frameworks, updates a shell script to download additional model assets and run RKNN exports, and adjusts ONNX encoder export code to load checkpoints whether they are raw or nested under "state_dict".

Changes

Cohort / File(s) Summary
CI workflow for Paraformer→RKNN
.github/workflows/export-paraformer-to-rknn.yaml
Introduces a matrix GitHub Actions workflow to export Paraformer models (FunASR and WSChuan-ASR) to RKNN for multiple RK356x targets, installs dependencies/toolkits, converts ONNX (encoder/predictor/decoder) to RKNN, packages artifacts, and conditionally creates releases.
Paraformer RKNN scripts
scripts/paraformer/rknn/download-example-model.sh, scripts/paraformer/rknn/export_encoder_onnx.py
Shell script now downloads additional resources (am.mvn, config.yaml, tokens.json, seg_dict) and runs encoder/predictor/decoder RKNN exports from ONNX artifacts. Python exporter now accepts checkpoints where weights are either the whole object or under a "state_dict" key.

Sequence Diagram(s)

sequenceDiagram
  autonumber
  participant GH as GitHub Actions
  participant Repo as Repository
  participant Py as Python Env
  participant DL as Model Store
  participant RK as RKNN Toolkit
  participant Rel as Release Service

  GH->>Repo: Checkout code
  GH->>Py: Setup Python, install deps (numpy/torch/onnx/onnxruntime/…)
  GH->>RK: Install RKNN toolkit wheel
  rect rgba(200,230,255,0.3)
    note over GH: Matrix over platforms & frameworks
    GH->>DL: Download models & test data (FunASR / WSChuan-ASR)
    GH->>Py: Convert encoder/predictor/decoder to ONNX
    Py->>RK: Build RKNN from ONNX per target
    RK-->>GH: RKNN artifacts
  end
  GH->>GH: Package artifacts (tar.bz2)
  alt Owner: csukuangfj or k2-fsa
    GH->>Rel: Create release and upload tarballs
  else Other owner
    GH->>GH: Skip release
  end
Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Suggested labels

size:L

Poem

I thump my paws—new gears in flight,
Models hop to RKNN’s light.
Scripts fetch snacks (tokens, YAML, tea),
ONNX to chips—swift alchemy.
In CI burrows, artifacts bloom—
Tarballs packed, released—ka-boom!
A happy hare commits with glee. 🐇✨

✨ Finishing touches
  • 📝 Generate docstrings
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Post copyable unit tests in a comment

📜 Recent review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 5de6254 and ea2690c.

📒 Files selected for processing (3)
  • .github/workflows/export-paraformer-to-rknn.yaml (1 hunks)
  • scripts/paraformer/rknn/download-example-model.sh (1 hunks)
  • scripts/paraformer/rknn/export_encoder_onnx.py (1 hunks)

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

@dosubot dosubot Bot added the size:L This PR changes 100-499 lines, ignoring generated files. label Oct 15, 2025
@csukuangfj
csukuangfj merged commit a9985a2 into k2-fsa:master Oct 15, 2025
1 check was pending
@csukuangfj
csukuangfj deleted the ci-paraformer-rknn branch October 15, 2025 08:34
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

size:L This PR changes 100-499 lines, ignoring generated files.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant