【开源】java版本BertTokenizer以及使用onnxruntime推理

python训练的模型,转换为onnx模型后,用python代码可以方便进行推理,但是java代码如何实现呢?

首先ONNX 推理,可以使用onnxruntime

Copy
<dependency> <groupId>com.microsoft.onnxruntime</groupId> <artifactId>onnxruntime</artifactId> <version>1.15.1</version> </dependency>

另外,训练的模型需要用到bert分词器,将单词和字变成token id, github上有 https://github.com/ankiteciitkgp/bertTokenizer,我们基于这个库简单改造下,来适配bert onnx模型的输入,改造后代码见: https://github.com/jadepeng/bertTokenizer

主要新增了tokenizeOnnxTensor 方法,返回适配bert模型输入的onnx tensor

完整demo代码:

Copy
public class OnnxTests { public static void main(String[] args) throws IOException, OrtException { BertTokenizer bertTokenizer = new BertTokenizer("D:\\model\\vocab.txt"); var env = OrtEnvironment.getEnvironment(); var session = env.createSession("D:\\model\\output\\onnx\\fp16_model.onnx", new OrtSession.SessionOptions()); var inputMap = bertTokenizer.tokenizeOnnxTensor(Arrays.asList("hello world 你好", "肿瘤治疗未来发展趋势")); try (var results = session.run(inputMap)) { System.out.println(results); var embeddings = (float[][])results.get(0).getValue(); for (var embedding : embeddings) { System.out.println(JSON.toJSONString(embedding)); } } } }

关注作者

欢迎关注作者微信公众号, 一起交流软件开发:欢迎关注作者微信公众号

posted @   JadePeng  阅读(1109)  评论(0编辑  收藏  举报
相关博文:
阅读排行:
· 全程不用写代码,我用AI程序员写了一个飞机大战
· DeepSeek 开源周回顾「GitHub 热点速览」
· MongoDB 8.0这个新功能碉堡了,比商业数据库还牛
· 记一次.NET内存居高不下排查解决与启示
· 白话解读 Dapr 1.15:你的「微服务管家」又秀新绝活了
历史上的今天:
2018-09-04 使用.NET Core+Docker 开发微服务
点击右上角即可分享
微信分享提示