【开源】java版本BertTokenizer以及使用onnxruntime推理
python训练的模型,转换为onnx模型后,用python代码可以方便进行推理,但是java代码如何实现呢?
首先ONNX 推理,可以使用onnxruntime
<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代码:
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));
}
}
}
}
关注作者
作者: JadePeng
出处:https://www.cnblogs.com/xiaoqi/p/java-bert-tokenizer.html
版权:本文采用「署名-非商业性使用-相同方式共享 4.0 国际(欢迎转载,但未经作者同意必须保留此段声明,且在文章页面明显位置给出原文链接) 」知识共享许可协议进行许可。
【推荐】国内首个AI IDE,深度理解中文开发场景,立即下载体验Trae
【推荐】编程新体验,更懂你的AI,立即体验豆包MarsCode编程助手
【推荐】抖音旗下AI助手豆包,你的智能百科全书,全免费不限次数
【推荐】轻量又高性能的 SSH 工具 IShell:AI 加持,快人一步
· 全程不用写代码,我用AI程序员写了一个飞机大战
· DeepSeek 开源周回顾「GitHub 热点速览」
· MongoDB 8.0这个新功能碉堡了,比商业数据库还牛
· 记一次.NET内存居高不下排查解决与启示
· 白话解读 Dapr 1.15:你的「微服务管家」又秀新绝活了
2018-09-04 使用.NET Core+Docker 开发微服务