kenrunio/transcriber-punctuation-onnx

FullStop multilingual punctuation, ONNX

oliverguhr/fullstop-punctuation-multilang-large (English, German, French, Italian), converted to ONNX so it runs with ONNX Runtime and the tokenizers library alone, without PyTorch or Transformers. It is used by Transcriber to put full stops, commas and capitals back into run-on stretches of speech-to-text output.

For each word the model predicts the punctuation that follows it.

The files

• model_int8.onnx (0.6 GB): the same model with its weights stored as 8-bit integers. This is the one Transcriber downloads now; it runs a third faster and needs a quarter of the space.
• model.onnx + model.onnx.data (2.2 GB together; the second holds the weights, keep both in one folder): full precision
• tokenizer.json: the original tokenizer, for the tokenizers library
• labels.json: class id to label: 0 (none), ., ,, ?, -, :

Inputs ids and mask (int64, shape (batch, tokens)); output logits (float32, shape (batch, tokens, 6)). Take the label of each word's first token:

import json, numpy as np, onnxruntime as ort
from tokenizers import Tokenizer
tok = Tokenizer.from_file("tokenizer.json"); tok.enable_truncation(512)
labels = {int(k): v for k, v in json.load(open("labels.json")).items()}
sess = ort.InferenceSession("model_int8.onnx")   # or "model.onnx"
words = "so today we will look at the results and then we will take your questions".split()
enc = tok.encode(words, is_pretokenized=True)
ids = np.array([enc.ids], dtype=np.int64)
pred = sess.run(None, {"ids": ids, "mask": np.ones_like(ids)})[0].argmax(-1)[0]
first = {}
for pos, w in enumerate(enc.word_ids):
    if w is not None and w not in first:
        first[w] = labels[int(pred[pos])]

How it was made, and checked

Exported with torch.onnx.export(..., dynamo=True, opset_version=18, external_data=True) by tools/export_models.py. On the 10,768 words of a 90-minute transcript the ONNX model gave the same label as the PyTorch original for every word.

model_int8.onnx was made from model.onnx with ONNX Runtime's quantize_dynamic(weight_type=QInt8, per_channel=True). On 12,000 words of two earnings-call transcripts it gave the same label as the full model for 99% of words; 599 of the full model's 615 full stops and question marks fell in the same places, and most other differences were commas.

Licence and credit

MIT, as the original (LICENSE). model.onnx is a format conversion of the original weights, unchanged; model_int8.onnx holds the same weights rounded to 8 bits. The model is by Oliver Guhr, trained on Europarl; see the original model card and FullStop.