Java中加载语义模型
Java中加载语义模型
一、以BAAI/bge-large-zh-v1.5模型为例
1、模型简介与背景
-
BGE 系列 :BGE 全称 BAAI General Embedding,是北京智源研究院推出的一系列语义嵌入模型,用于将文本映射为稠密向量表示
-
v1 和 v1.5 的差异 :v1.5 系列是在 2023 年 9 月发布的,主要优化了相似度分布,使得检索能力更合理、更稳定
-
参数规模:bge-large-zh-v1.5 属于 large 级别,参数量约为 326M,模型大小约为 1.3 GB
2、功能特性
-
文本嵌入(Embedding):将中文文本转换为固定长度(如 1024 维)的向量表示,可用于检索、相似度计算、分类、聚类等任务
-
支持指令微调:在检索任务中,你可以在查询文本前加上指令以改善表达效果,但 v1.5 提升后无需指令也能表现良好
-
输入长度 :支持最长约 512 tokens,不适用于处理极长文本(如 8192 tokens)的任务。
3、性能表现
-
中文排行榜表现优异 :在中文大规模文本嵌入基准 (C-MTEB) 中,bge-large-zh-v1.5 排名第一,语义检索、分类、相似度等任务上表现突出。
-
在多语言场景中具备一定表现 :虽然该模型主要针对中文,但其架构设计支持多语言场景,对部分其他语言也具备一定能力。
二、在Java中如何调用
1、模型准备
下载地址:https://huggingface.co/BAAI/bge-large-zh-v1.5
1.1 下载有几种方式
-
直接点解下载图标进行下载
-
通过http或者SSH的方式拉取下载,如下图所示
2、在Java的SpringBoot中如何使用
2.1 需要的依赖清单
<!-- DJL 核心 --><dependency><groupId>ai.djl</groupId><artifactId>api</artifactId><version>0.34.0</version></dependency>
<!-- (可选)SentencePiece:只有当 tokenizer 用到了 SentencePiece 才需要 --><dependency><groupId>ai.djl.sentencepiece</groupId><artifactId>sentencepiece</artifactId><version>0.34.0</version></dependency>
<!-- PyTorch 原生库(CPU) --><dependency><groupId>ai.djl.pytorch</groupId><artifactId>pytorch-engine</artifactId><version>0.34.0</version></dependency>
<!-- PyTorch 原生库(CPU) --><dependency><groupId>ai.djl.pytorch</groupId><artifactId>pytorch-native-cpu</artifactId><version>2.7.1</version></dependency>
<!-- PyTorch JNI(版本段:libtorch-djl) --><dependency><groupId>ai.djl.pytorch</groupId><artifactId>pytorch-jni</artifactId><version>2.7.1-0.34.0</version><scope>runtime</scope></dependency>
<!-- HuggingFace 高速分词器 --><dependency><groupId>ai.djl.huggingface</groupId><artifactId>tokenizers</artifactId><version>0.28.0</version></dependency>
2.2 在添加依赖可以避免的坑
注意: 现在模型的使用给我们开发带了很大的便利,但是有时候模型可能存在幻读情况,给出的东西有问题。
例如有些模型在给依赖的时候,有时候名称
或者版本号
有问题,实际在Maven仓库中就不存在,这时候就需要进行验证,甄别一下。
建议:有些依赖包比较大,直接拉取可能比较慢,可以尝试直接下载的方式
可以到Maven仓库进行验证,是否有对应的依赖存在,以及版本号有哪些,如下图所示:
2.3 加载模型的主体代码
Criteria<String, float[]> criteria = Criteria.builder().setTypes(String.class, float[].class).optApplication(Application.NLP.TEXT_EMBEDDING).optEngine("PyTorch").optModelPath(Paths.get("D:/model/bge-large-zh-v1.5")).optTranslatorFactory(new TextEmbeddingTranslatorFactory()).build();model = ModelZoo.loadModel(criteria);
predictor = model.newPredictor();
① setTypes(String.class, float[].class)
告诉 DJL:输入是 String,输出是 float[]。这会指导它选择/校验合适的 Translator。
② optApplication(TEXT_EMBEDDING)
指定“文本嵌入”任务(不是分类/问答等),影响默认 Translator/前后处理
③ optEngine(“PyTorch”)
固定使用 PyTorch 引擎(需要 pytorch-engine + 对应原生库)
④ optModelPath(…)
从本地目录加载模型(离线)
⑤ optTranslatorFactory(new TextEmbeddingTranslatorFactory())
Translator = 前后处理逻辑
⑥ ModelZoo.loadModel(criteria)
读取配置,初始化引擎、载入权重、构建分词器与 Translator(如缺文件/库不匹配会在这里报错)
⑦ model.newPredictor()
基于模型创建推理器。Predictor 内部持有 Translator;你用 predict("文本") 就能拿到向量
2.4 避坑指南
PyTorch 官方保存的模型:
常见是 .pt
或 .pth
,这类文件通常用 torch.save()
保存的 state_dict
或整个模型。DJL 如果你直接用 optModelPath
指向文件,它会走这种直接加载路径
HuggingFace Transformers 模型
大多是.bin(
如pytorch_model.bin
或分片文件 pytorch_model-00001-of-00002.bin),里面是 PyTorch 权重的二进制序列化,但配合config.json
、tokenizer.json
等元数据。DJL 在 NLP 场景下(比如你用了 TextEmbeddingTranslatorFactory
),会按 HuggingFace 格式去解析——此时加载.bin
是完全支持的
2.4.1 踩坑情况一
本人在加载.bin
相关模型的时候,总是会提示XXXX.pt文件不错,想了如下办法转换生成.pt文件 ,转换实现的代码如下:
from transformers import AutoModel
import torchclass SentenceEmbeddingModel(torch.nn.Module):def __init__(self, model):super().__init__()self.model = modeldef forward(self, input_ids, attention_mask):outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)return outputs[0] # 只返回 last_hidden_state(tensor)# 加载 HuggingFace 模型
model = AutoModel.from_pretrained("BAAI/bge-large-zh-v1.5")
wrapped_model = SentenceEmbeddingModel(model)# 生成假输入
example_input_ids = torch.ones(1, 10, dtype=torch.long)
example_attention_mask = torch.ones(1, 10, dtype=torch.long)# Trace
traced = torch.jit.trace(wrapped_model, (example_input_ids, example_attention_mask))
traced.save("model.pt")
最后讲转换生成的文件,放大模型文件下
2.4.2 踩坑情况一
在使用的过程中及时转换有了.pt
文件,有时候还是会报如下的错误
Exception in thread "main" ai.djl.engine.EngineException:
Unknown builtin op: aten::scaled_dot_product_attention.
Here are some suggestions: aten::_scaled_dot_product_attentionThe original call is:File "code/__torch__/transformers/models/bert/modeling_bert.py", line 217x1 = torch.view(_52, [_54, int(_55), 16, 64])value_layer = torch.permute(x1, [0, 2, 1, 3])attn_output = torch.scaled_dot_product_attention(query_layer, key_layer, value_layer, attention_mask)~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HEREattn_output0 = torch.transpose(attn_output, 1, 2)input = torch.reshape(attn_output0, [_42, _43, 1024])at ai.djl.pytorch.jni.PyTorchLibrary.moduleLoad(Native Method)at ai.djl.pytorch.jni.JniUtils.loadModule(JniUtils.java:1795)at ai.djl.pytorch.engine.PtModel.load(PtModel.java:99)at ai.djl.repository.zoo.BaseModelLoader.loadModel(BaseModelLoader.java:176)at ai.djl.repository.zoo.Criteria.loadModel(Criteria.java:151)at ai.djl.repository.zoo.ModelZoo.loadModel(ModelZoo.java:182)at com.shenlan.mcpAiTool.TestAppOptimize.main(TestAppOptimize.java:40)Process finished with exit code 1
找了半天,得到解释一直是:
错误原因
-
aten::scaled_dot_product_attention 是 PyTorch 2.0 引入的新算子 ,DJL 可能依赖的 LibTorch 版本低于 2.0。
-
DJL 的 PyTorch 引擎(尤其是 Java 绑定)没有实现这个新的 builtin op。
-
因此在加载 TorchScript 模块(.pt 或 torchscript 导出的模型)时就报 Unknown builtin op。
解决办法
最后在Maven仓库找到相关的依赖包,更新相关的版本就可以了。
2.5 遗留的问题
值得注意的是,在初次运行程序的时候可以看到,会自动的下载相关的.dll相关的动态链接库。
对于此如何跨平台,以及在离线环境下 可以继续深入研究