examples/add_punctuation/add_punctuation.py

#!/usr/bin/python3
# -*- coding: utf-8 -*-
import argparse

import sherpa_onnx

from project_settings import project_path


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--model_file",
        default=(project_path / "pretrained_models/huggingface/csukuangfj/sherpa-onnx-punct-ct-transformer-zh-en-vocab272727-2024-04-12/model.onnx").as_posix(),
        type=str
    )
    parser.add_argument(
        "--text",
        default="i'm a google virtual assistant recording this call for the person you're trying to reach before i try to connect you can ask what you're calling about",
        type=str
    )
    args = parser.parse_args()
    return args


def main():
    args = get_args()

    config = sherpa_onnx.OfflinePunctuationConfig(
        model=sherpa_onnx.OfflinePunctuationModelConfig(
            ct_transformer=args.model_file
        ),
    )

    punctuation_model = sherpa_onnx.OfflinePunctuation(config)

    text = punctuation_model.add_punctuation(args.text)
    print("text: {}".format(text))
    return


if __name__ == '__main__':
    main()