Fork me on GitHub

利用java加载bert模型进行加速推理

这里利用java加载onnx模型,并进行推理。

项目地址:https://github.com/jiangnanboy/onnx-java

步骤

1.利用java加载onnx模型,并进行推理预测。这里采用roberta模型的onnx版。

2.pytorch模型从这里 下载。

3.pytorch模型转为onnx见这里 。

使用

1.sy/BertMask

String text = "中国的首都是[MASK]京。";
Triple<BertTokenizer, Map<String, OnnxTensor>, Integer> triple = null;
try {
    triple = parseInputText(text);
} catch (Exception e) {
    e.printStackTrace();
}
var maskPredictions = predMask(triple);
System.out.println(maskPredictions);

2.result

String text = "中国的首都是[MASK]京。";

tokens -> [[CLS], 中, 国, 的, 首, 都, 是, [MASK], 京, 。, [SEP]]
[MASK] predictions -> [北, 南, 东, 燕, 望]

String text = "我家后面有一[MASK]大树。";

tokens -> [[CLS], 我, 家, 后, 面, 有, 一, [MASK], 大, 树, 。, [SEP]]
[MASK] predictions -> [棵, 个, 株, 只, 颗]
posted @   石头木  阅读(1407)  评论(0编辑  收藏  举报
相关博文:
阅读排行:
· 震惊!C++程序真的从main开始吗?99%的程序员都答错了
· 【硬核科普】Trae如何「偷看」你的代码?零基础破解AI编程运行原理
· 单元测试从入门到精通
· 上周热点回顾(3.3-3.9)
· winform 绘制太阳,地球,月球 运作规律
点击右上角即可分享
微信分享提示