Skip to content

Commit ee3611d

Browse files
csukuangfjcsukuangfj
authored
Export Google MedASR to sherpa-onnx (#2934)
This pull request integrates Google's MedASR models into the "sherpa-onnx" project by providing a complete workflow for exporting these specialized medical speech recognition models to ONNX format. It includes scripts for model conversion, dynamic quantization, and testing, enabling efficient deployment and inference of MedASR within the "sherpa-onnx" environment. The changes ensure proper licensing attribution and metadata embedding for the exported models.
1 parent 6cb5005 commit ee3611d

5 files changed

Lines changed: 497 additions & 0 deletions

File tree

‎.github/workflows/export-medasr-ctc-to-onnx.yaml‎

Lines changed: 161 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,161 @@
1+
name: export-medasr-ctc-to-onnx
2+
3+
on:
4+
push:
5+
branches:
6+
- export-medasr-onnx
7+
workflow_dispatch:
8+
9+
concurrency:
10+
group: export-medasr-ctc-to-onnx-${{ github.ref }}
11+
cancel-in-progress: true
12+
13+
jobs:
14+
export-medasr-ctc-to-onnx:
15+
if: github.repository_owner == 'k2-fsa' || github.repository_owner == 'csukuangfj'
16+
name: export medasr ctc
17+
runs-on: ${{ matrix.os }}
18+
strategy:
19+
fail-fast: false
20+
matrix:
21+
os: [macos-latest]
22+
python-version: ["3.10"]
23+
24+
steps:
25+
- uses: actions/checkout@v4
26+
27+
- name: Setup Python ${{ matrix.python-version }}
28+
uses: actions/setup-python@v5
29+
with:
30+
python-version: ${{ matrix.python-version }}
31+
32+
- name: Run
33+
shell: bash
34+
env:
35+
HF_TOKEN: ${{ secrets.HF_TOKEN }}
36+
run: |
37+
cd scripts/medasr
38+
./run.sh
39+
40+
- name: Download test data
41+
shell: bash
42+
run: |
43+
cd scripts/medasr
44+
45+
curl -SL -O https://huggingface.co/csukuangfj/sherpa-onnx-medasr-ctc-en-int8-2025-12-25/resolve/main/test_wavs/0.wav
46+
47+
curl -SL -O https://huggingface.co/csukuangfj/sherpa-onnx-medasr-ctc-en-int8-2025-12-25/resolve/main/test_wavs/transcript.txt
48+
49+
ls -lh
50+
51+
- name: Test fp32
52+
shell: bash
53+
run: |
54+
cd scripts/medasr
55+
56+
python3 test_onnx.py --model ./model.onnx --tokens ./tokens.txt --wav ./0.wav
57+
58+
cat transcript.txt
59+
60+
- name: Test int8
61+
shell: bash
62+
run: |
63+
cd scripts/medasr
64+
65+
python3 test_onnx.py --model ./model.int8.onnx --tokens ./tokens.txt --wav ./0.wav
66+
67+
cat transcript.txt
68+
69+
- name: Collect fp32 files
70+
shell: bash
71+
run: |
72+
cd scripts/medasr
73+
74+
d=sherpa-onnx-medasr-ctc-en-2025-12-25
75+
mkdir -p $d
76+
mkdir -p $d/test_wavs
77+
78+
cp -v model.onnx $d
79+
cp -v README.md $d
80+
cp -v tokens.txt $d
81+
cp -v *.wav $d/test_wavs/
82+
cp -v transcript.txt $d/test_wavs/
83+
84+
tar cjvf $d.tar.bz2 $d
85+
86+
ls -lh $d
87+
ls -lh *.tar.bz2
88+
89+
mv $d ../..
90+
mv $d.tar.bz2 ../..
91+
92+
- name: Collect int8 files
93+
shell: bash
94+
run: |
95+
cd scripts/medasr
96+
97+
d=sherpa-onnx-medasr-ctc-en-int8-2025-12-25
98+
mkdir -p $d
99+
mkdir -p $d/test_wavs
100+
101+
cp -v model.int8.onnx $d
102+
cp -v README.md $d
103+
cp -v tokens.txt $d
104+
cp -v *.wav $d/test_wavs/
105+
cp -v transcript.txt $d/test_wavs/
106+
107+
tar cjvf $d.tar.bz2 $d
108+
109+
ls -lh $d
110+
ls -lh *.tar.bz2
111+
112+
mv $d ../..
113+
mv $d.tar.bz2 ../..
114+
115+
- name: Release
116+
uses: svenstaro/upload-release-action@v2
117+
with:
118+
file_glob: true
119+
file: ./*.tar.bz2
120+
overwrite: true
121+
repo_name: k2-fsa/sherpa-onnx
122+
repo_token: ${{ secrets.UPLOAD_GH_SHERPA_ONNX_TOKEN }}
123+
tag: asr-models
124+
125+
- name: Publish to huggingface
126+
env:
127+
HF_TOKEN: ${{ secrets.HF_TOKEN }}
128+
uses: nick-fields/retry@v3
129+
with:
130+
max_attempts: 5
131+
timeout_seconds: 200
132+
shell: bash
133+
command: |
134+
git config --global user.email "csukuangfj@gmail.com"
135+
git config --global user.name "Fangjun Kuang"
136+
137+
names=(
138+
sherpa-onnx-medasr-ctc-en-2025-12-25
139+
sherpa-onnx-medasr-ctc-en-int8-2025-12-25
140+
)
141+
for d in ${names[@]}; do
142+
if [ ! -d $d ]; then
143+
echo "$d does not exist - skip it"
144+
continue;
145+
fi
146+
147+
export GIT_LFS_SKIP_SMUDGE=1
148+
export GIT_CLONE_PROTECTION_ACTIVE=false
149+
rm -rf huggingface
150+
git clone https://csukuangfj:$HF_TOKEN@huggingface.co/csukuangfj/$d huggingface
151+
cp -av $d/* ./huggingface
152+
cd huggingface
153+
git lfs track "*.onnx"
154+
git lfs track "*.wav"
155+
git status
156+
git add .
157+
git status
158+
git commit -m "add models"
159+
git push https://csukuangfj:$HF_TOKEN@huggingface.co/csukuangfj/$d main
160+
cd ..
161+
done

‎scripts/medasr/README.md‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
---
2+
license: other
3+
license_name: health-ai-developer-foundations
4+
license_link: https://developers.google.com/health-ai-developer-foundations/terms
5+
language:
6+
- en
7+
pipeline_tag: automatic-speech-recognition
8+
library_name: transformers
9+
tags:
10+
- medical-asr
11+
- radiology
12+
- medical
13+
---
14+
15+
# Introduction
16+
17+
This directory includes models sourced from:
18+
19+
https://github.com/Google-Health/medasr
20+
21+
All model files are governed by the Health AI Developer Foundations Terms of Use.
22+
For full licensing details, please refer to:
23+
24+
https://developers.google.com/health-ai-developer-foundations/terms

‎scripts/medasr/export_onnx.py‎

Lines changed: 128 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,128 @@
1+
#!/usr/bin/env python3
2+
# Copyright 2025 Xiaomi Corp. (authors: Fangjun Kuang)
3+
4+
"""
5+
Make sure you have set the environment variable
6+
7+
export HF_TOKEN=hf_xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
8+
9+
where hf_xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx is your Huggingface access token.
10+
"""
11+
12+
from typing import Any, Dict
13+
14+
import onnx
15+
import torch
16+
from onnxruntime.quantization import QuantType, quantize_dynamic
17+
from transformers import AutoModelForCTC, AutoProcessor
18+
19+
20+
def add_meta_data(filename: str, meta_data: Dict[str, Any]):
21+
"""Add meta data to an ONNX model. It is changed in-place.
22+
23+
Args:
24+
filename:
25+
Filename of the ONNX model to be changed.
26+
meta_data:
27+
Key-value pairs.
28+
"""
29+
model = onnx.load(filename)
30+
model.metadata_props.clear()
31+
32+
for key, value in meta_data.items():
33+
meta = model.metadata_props.add()
34+
meta.key = key
35+
meta.value = str(value)
36+
37+
onnx.save(model, filename)
38+
39+
40+
class Wrapper(torch.nn.Module):
41+
def __init__(self, m):
42+
super().__init__()
43+
self.m = m
44+
45+
def forward(self, x: torch.Tensor, mask: torch.Tensor):
46+
"""
47+
Args:
48+
x: (N, T, C), dtype float32
49+
mask: (N, T), dtype int64. Valid positions are 1. Padding positions are 0.
50+
Returns:
51+
logits: (N, T/4, vocob_size), dtype float32
52+
logits_len: (N,), dtype int64
53+
"""
54+
o = self.m(x, mask.bool())
55+
logits_len = self.m._get_subsampling_output_length(mask.sum(-1)).to(torch.int64)
56+
return o.logits, logits_len
57+
58+
59+
def generate_tokens(tokenizer):
60+
vocab = tokenizer.get_vocab()
61+
id2token = {i: t for t, i in vocab.items()}
62+
63+
with open("tokens.txt", "w", encoding="utf-8") as f:
64+
for i in range(tokenizer.vocab_size):
65+
if i == tokenizer.pad_token_id:
66+
f.write(f"<blk> {i}\n")
67+
else:
68+
f.write(f"{id2token[i]} {i}\n")
69+
print("saved to tokens.txt")
70+
71+
72+
@torch.no_grad()
73+
def main():
74+
model_id = "google/medasr"
75+
processor = AutoProcessor.from_pretrained(model_id)
76+
77+
generate_tokens(processor.tokenizer)
78+
79+
model = AutoModelForCTC.from_pretrained(model_id)
80+
81+
w = Wrapper(model)
82+
w.eval()
83+
84+
filename = "model.onnx"
85+
x = torch.rand(1, 100, 128)
86+
mask = torch.ones(1, x.shape[1], dtype=torch.int64)
87+
torch.onnx.export(
88+
w,
89+
(x, mask),
90+
filename,
91+
input_names=["x", "mask"],
92+
output_names=["logits", "logits_len"],
93+
dynamic_axes={
94+
"x": {0: "N", 1: "T"},
95+
"mask": {0: "N", 1: "T"},
96+
"logits": {0: "N", 1: "T_4"},
97+
"logits_len": {0: "N"},
98+
},
99+
opset_version=14,
100+
external_data=False,
101+
dynamo=False,
102+
)
103+
104+
meta_data = {
105+
"model_type": "medasr_ctc",
106+
"version": "20251225",
107+
"model_author": "google",
108+
"maintainer": "k2-fsa",
109+
"vocab_size": processor.tokenizer.vocab_size,
110+
"url": "https://github.com/Google-Health/medasr",
111+
"license": "https://developers.google.com/health-ai-developer-foundations/terms",
112+
}
113+
add_meta_data(filename=filename, meta_data=meta_data)
114+
115+
filename_int8 = "model.int8.onnx"
116+
quantize_dynamic(
117+
model_input=filename,
118+
model_output=filename_int8,
119+
op_types_to_quantize=["MatMul"],
120+
# Note that we have to use QUInt8 here.
121+
#
122+
# When QInt8 is used, C++ onnxruntime produces incorrect results
123+
weight_type=QuantType.QUInt8,
124+
)
125+
126+
127+
if __name__ == "__main__":
128+
main()

‎scripts/medasr/run.sh‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
#!/usr/bin/env bash
2+
3+
if [ -z $HF_TOKEN ]; then
4+
echo "Please first run export HF_TOKEN=your_huggingface_access_token."
5+
exit 1
6+
fi
7+
8+
pip install \
9+
accelerate \
10+
bitsandbytes \
11+
git+https://github.com/huggingface/transformers.git@65dc261512cbdb1ee72b88ae5b222f2605aad8e5 \
12+
onnx==1.17.0 \
13+
onnxruntime==1.17.1 \
14+
librosa \
15+
onnxscript \
16+
"numpy<2" \
17+
kaldi-native-fbank
18+
19+
./export_onnx.py
20+
21+
ls -lh

0 commit comments

Comments
 (0)