pdnhan/xlm-roberta-capu-mobile

xlm-roberta-capu for mobile (ONNX)

dragonSwing/xlm-roberta-capu, a Vietnamese capitalization and punctuation restoration model (XLM-RoBERTa-base, 278M parameters), converted to ONNX for on-device use with ONNX Runtime on Android and iOS, with ready-to-use Kotlin and Swift code for everything around the model: tokenizer, chunking of long text and text editing.

It turns lower-cased, unpunctuated text, such as speech recognition output, into written text:

Input	Output
theo đó thủ tướng dự kiến tiếp bộ trưởng nông nghiệp mỹ tom wilsack bộ trưởng thương mại mỹ gina raimondo	Theo đó, Thủ tướng dự kiến tiếp Bộ trưởng Nông nghiệp Mỹ Tom Wilsack, Bộ trưởng Thương mại Mỹ Gina Raimondo.
những gói cước năm g mobifone sẽ mang đến cho bạn những trải nghiệm mới lạ trên cả tuyệt vời	Những gói cước 5G MobiFone sẽ mang đến cho bạn những trải nghiệm mới lạ trên cả tuyệt vời.

It restores . , : ?, capitalization (including words like MobiFone, YouTube), and spoken letters and digits (năm g → 5G). The weights are those of the original model; only the format differs (and the precision, for the 8-bit file).

Files

File	What	Size	SHA-256
onnx/model_quantized.onnx	8-bit model (for apps)	280.8 MB	d61e9353fa1a5533ec6c9b9e5f218c97acbbdfa0122414a8aaabdb3fd5f32c94
onnx/model.onnx	fp32 model	1,110.1 MB	468ff4473976d63189bf99d331f65f89856ae09c5df5276c89252d86f09e642d
unigram.tsv	tokenizer vocabulary	7.8 MB	d025763e9460955c243491b649f4eec42a16d90269d0cfb848922e6d8f3e382c
labels.txt	edit labels	0.2 kB	a31075cfa185b5d24c3b65c009ad0740636b286090f78419eaba2d34202c2b45
d_tags.txt	error detection classes	0.0 kB	926596d65d7b928a3d4dfb553c0cd2a8189f8f2ca7cb0ba26cfbb935b5c5dfad
number_words.tsv	spoken letters and digits	953.0 kB	788dcd04e4057b36005dfd507cf78b455ea12c6655a0fc7d1f2d67f7abae0048
tokenizer.json	tokenizer for Python	17.1 MB	edd5d5202e1c27bf620d7e03f1dbd14ccf89c4ca62d2d3d4d0a59640f66140e5
config.json	ids and pipeline settings	0.5 kB	890cf3038f717b43e26a6d42f549ffef94b11d398e8b5c044990242e177449e6

Mobile apps download onnx/model_quantized.onnx, unigram.tsv, labels.txt and number_words.tsv (about 290 MB in total) and use the code in android/ or ios/.

• onnx/model_quantized.onnx: 8-bit weights. Linear layers are MatMulNBits (8-bit, symmetric, blocks of 128, accuracy_level 4), the embedding an int8 Gather; activations stay float. Needs ONNX Runtime 1.22 or newer (earlier versions only run 4-bit MatMulNBits). Tested with 1.22.0, 1.23.0 and 1.31.1.
• onnx/model.onnx: the fp32 reference export (opset 17, standard operators only).
• unigram.tsv: XLM-RoBERTa's SentencePiece vocabulary, one piece<TAB>score line per token id, and the model's added $START token as the last id (250002).
• labels.txt, d_tags.txt: the classes of the two outputs, one per line in index order.
• number_words.tsv: word<TAB>tag<TAB>replacement for the $TRANSFORM_VERB_* labels (năm VB_VBN 5, gờ VB_VBC G).
• tokenizer.json: the same tokenizer for Python (tokenizers / transformers), with $START added.

Accuracy

On the 347 unique sentences of the FLEURS Vietnamese test split, lower-cased with punctuation removed, plus 43 paragraphs of 8 consecutive sentences (long enough to be split into chunks), against the original transcriptions. F1 over words, for the four punctuation marks and for capitalized words:

Model	Punctuation F1	Capitalization F1	Same text as the original model
original (PyTorch, gec_model.py)	0.763	0.869
onnx/model.onnx (fp32)	0.763	0.869	392 / 392
onnx/model_quantized.onnx (8-bit)	0.766	0.869	315 / 392

The fp32 export with the code in this repo gives exactly the original pipeline's text. The 8-bit model differs on some texts, mostly on borderline commas, with the same accuracy overall. The Kotlin code (run on the JVM with ONNX Runtime 1.31.1) and the Swift code give the same text as the Python reference on all 392 texts; the Python tokenizer gives the same token ids as tokenizers on all 3,480 distinct words of the test texts.

Speed on an x86 desktop CPU (8 threads, one text at a time): 392 texts (21,514 words) in 13.4 s with the 8-bit model, 17.4 s with fp32. Phones are slower; one model run handles up to 64 words.

Model interface

	Name	Type	Shape
input	input_ids	int64	[batch, tokens]: $START (id 250002), then each word's sub-word tokens
input	attention_mask	int64	[batch, tokens]: 1, or 0 for padding (pad id 1)
input	input_offsets	int64	[batch, words + 1]: index in input_ids of $START (0) and of each word's first token
output	logits	float32	[batch, words + 1, 15]: edit label scores, classes in labels.txt
output	detect_logits	float32	[batch, words + 1, 4]: error detection scores, classes in d_tags.txt (not needed)

How the pipeline works

The model labels each word with an edit; the text is the result of applying the edits. The code in this repo does all of it (python/capu.py is the readable reference, about 300 lines with no dependency besides numpy and onnxruntime):

1. Words: split the text on whitespace.
2. Chunks: up to 48 words are one chunk; longer text is split into windows of 48 words that overlap by 12.
3. Tokens: each word is NFKC-normalized, prefixed with ▁ and split into SentencePiece pieces by Viterbi search over unigram.tsv (unigram language model; a character no piece covers becomes <unk>, id 3).
4. Labels: run the model on $START + the chunk's first 64 words; take the argmax label of each word.
5. Edits: $APPEND_x inserts the punctuation mark x after the word (after $START: before the first word), $TRANSFORM_CASE_* changes its case, $TRANSFORM_VERB_* replaces it from number_words.tsv, $MERGE_SPACE joins it with the next word.
6. Repeat 4 and 5 on the edited chunk (now with punctuation marks as words) up to 3 times, until it stops changing.
7. Merge consecutive chunks inside their overlap, at the first run of words they agree on at least 6 words in, and remove the space before each punctuation mark.

This is the original repo's GecBERTModel (gec_model.py, utils.py) with the settings of its model card (split_chunk=True), rewritten without PyTorch or transformers.

Android (Kotlin)

Copy android/Capu.kt into the app. It uses only the Kotlin standard library and the ONNX Runtime Java API.

// build.gradle.kts: implementation("com.microsoft.onnxruntime:onnxruntime-android:1.22.0") or newer
// dir: a directory with onnx/model_quantized.onnx, unigram.tsv, labels.txt, number_words.tsv
val capu = Capu.load(dir, quantized = true, threads = 4)   // keep one instance: loading takes a while
val text = capu.restore("theo đó thủ tướng dự kiến tiếp bộ trưởng nông nghiệp mỹ")
// "Theo đó, Thủ tướng dự kiến tiếp Bộ trưởng Nông nghiệp Mỹ."

Call restore off the main thread.

iOS (Swift)

Add the Swift package in ios/ (it depends on onnxruntime-swift-package-manager), or copy ios/Sources/Capu/Capu.swift and OnnxLabelModel.swift into an app that already uses ONNX Runtime (SPM module OnnxRuntimeBindings or CocoaPods onnxruntime-objc, 1.22 or newer).

import Capu

// dir: a directory with onnx/model_quantized.onnx, unigram.tsv, labels.txt, number_words.tsv
let model = try OnnxLabelModel(path: dir.appendingPathComponent("onnx/model_quantized.onnx").path, threads: 4)
let capu = try Capu.load(directory: dir, model: model)
let text = try capu.restore("theo đó thủ tướng dự kiến tiếp bộ trưởng nông nghiệp mỹ")
// "Theo đó, Thủ tướng dự kiến tiếp Bộ trưởng Nông nghiệp Mỹ."

The Swift code was compiled and checked on Linux (Swift 6.0) against a stand-in for the ONNX Runtime module that replays the Python reference's model outputs, not yet on an iPhone.

Python

from huggingface_hub import snapshot_download
import sys

repo = snapshot_download("pdnhan/xlm-roberta-capu-mobile", allow_patterns=["onnx/model_quantized.onnx", "*.tsv", "*.txt", "python/*"])
sys.path.insert(0, f"{repo}/python")
from capu import Capu  # numpy + onnxruntime>=1.22

capu = Capu(repo, quantized=True)
print(capu("theo đó thủ tướng dự kiến tiếp bộ trưởng nông nghiệp mỹ"))

How it was made

The scripts are in scripts/: export.py rebuilds the model from the original pytorch_model.bin (XLM-RoBERTa-base with one added token and two linear heads; the original repo's Python code is not run), exports it with torch.onnx.export (opset 17) and quantizes it with ONNX Runtime's MatMulNBitsQuantizer and quantize_dynamic (for the embedding). reference.py runs the original pipeline, validate.py measures the table above. The Kotlin and Swift checks are in scripts/kotlin and scripts/swift-check.

License and credit

CC BY-SA 4.0, the license of the original model; this repo (model files and code) is shared under the same license. All credit for the model goes to its author, dragonSwing: see the original model card for the training data (OSCAR-2109, 5.6M Vietnamese texts) and per-label accuracy. The base model is XLM-RoBERTa (MIT). FLEURS (CC BY 4.0) was used for evaluation only.