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 torch

class SentenceEmbeddingModel(torch.nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model

    def 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_attention

The original call is:
  File "code/__torch__/transformers/models/bert/modeling_bert.py", line 217
    x1 = 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)
                  ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
    attn_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相关的动态链接库。

对于此如何跨平台，以及在离线环境下 可以继续深入研究