Java中使用ONNX部署本地模型实战指南

1. ONNX部署本地模型概述

在Java生态中进行大模型应用开发时，Langchain4j和Spring AI是两个主流选择。虽然它们都提供了大模型集成的能力，但在实现细节上存在显著差异。本章将重点介绍如何通过ONNX格式在Langchain4j中部署和使用本地模型。

提示：本系列使用langchain4j-1.9.1版本和JDK 19，所有示例代码均基于纯Java实现，不依赖Spring框架。完整代码已开源在GitHub仓库。

本地模型部署主要有四种常见方式：

1. Local AI：直接在本地服务器部署原始大模型
2. ONNX：使用开放的神经网络交换格式部署模型
3. Ollama：通过Ollama平台简化本地模型部署
4. HuggingFace：利用HuggingFace的TGI解决方案进行云端部署

其中，ONNX特别适合embedding模型的部署，因为：

• 模型体积相对较小
• 计算量要求不高
• 跨平台兼容性好
• 性能损失在可接受范围内

2. ONNX技术解析

2.1 ONNX核心概念

ONNX（Open Neural Network Exchange）是一种开放的神经网络交换格式，具有以下特点：

1. 跨框架兼容：支持PyTorch、TensorFlow等主流框架的模型转换
2. 跨语言支持：提供多种语言的运行时环境
3. 标准化格式：统一的模型表示方式
4. 优化执行：支持硬件加速和模型优化

在Java生态中，主要通过以下两个库实现ONNX支持：

• Deep Java Library (DJL)
• Microsoft ONNX Java Runtime

2.2 ONNX工作流程

典型的ONNX模型使用流程包括：

1. 原始模型训练（使用PyTorch/TensorFlow等框架）
2. 转换为ONNX格式
3. 在目标环境中加载和推理
4. 处理模型输出

对于embedding模型，这种流程特别高效，因为：

• 不需要复杂的预处理
• 推理过程相对简单
• 输出结构固定

3. 实战：部署BGE-small-zh模型

3.1 环境准备

3.1.1 Python环境配置

首先需要准备Python环境来转换模型：

bash复制# 安装基础依赖
pip install modelscope optimum onnx onnxruntime sentence-transformers

3.1.2 模型下载

使用modelscope下载BAAI/bge-small-zh-v1.5模型：

bash复制modelscope download --model BAAI/bge-small-zh-v1.5 --local_dir BAAI

注意：也可以直接从魔塔官网手动下载模型文件

3.1.3 模型转换

将下载的模型转换为ONNX格式：

bash复制optimum-cli export onnx --task sentence-similarity --model BAAI onnx

转换完成后会生成两个关键文件：

• model.onnx：模型权重文件
• tokenizer.json：分词器配置文件

3.2 Java项目配置

3.2.1 Maven依赖

在pom.xml中添加必要依赖：

xml复制<dependencies>
    <dependency>
        <groupId>dev.langchain4j</groupId>
        <artifactId>langchain4j</artifactId>
        <version>1.9.1</version>
    </dependency>
    <dependency>
        <groupId>dev.langchain4j</groupId>
        <artifactId>langchain4j-embeddings</artifactId>
        <version>1.9.1</version>
    </dependency>
</dependencies>

3.2.2 资源文件准备

将转换好的ONNX模型文件放入resources目录：

code复制src/main/resources/
└── bge-small-zh/
    ├── model.onnx
    └── tokenizer.json

3.3 模型加载与使用

3.3.1 基础使用示例

创建ONNXTest类演示基本用法：

java复制import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.model.embedding.onnx.OnnxEmbeddingModel;
import dev.langchain4j.model.embedding.onnx.PoolingMode;
import dev.langchain4j.model.output.Response;

import java.io.File;
import java.util.Arrays;

public class ONNXTest {
    public static void main(String[] args) {
        OnnxEmbeddingModel embeddingModel = new OnnxEmbeddingModel(
                new File(ClassLoader.getSystemResource("bge-small-zh/model.onnx").getPath()).toPath(),
                new File(ClassLoader.getSystemResource("bge-small-zh/tokenizer.json").getPath()).toPath(),
                PoolingMode.CLS);
        
        Response<Embedding> embeddingResponse = embeddingModel.embed("测试进行嵌入");
        System.out.println("维度= " + embeddingResponse.content().dimension());
        System.out.println(Arrays.toString(embeddingResponse.content().vector()));
    }
}

3.3.2 输出结果分析

运行程序后会输出：

1. 嵌入向量的维度（通常是384或768）
2. 具体的向量数值

典型输出示例：

code复制维度= 384
[0.123, -0.456, ..., 0.789]

4. 源码深度解析

4.1 OnnxEmbeddingModel类结构

OnnxEmbeddingModel继承自AbstractInProcessEmbeddingModel，核心功能包括：

1. 模型加载：

   • 支持从JAR包加载
   • 支持从文件系统加载
   • 支持多种路径格式（Path、String等）

2. 嵌入计算：

   • 单文本处理
   • 批量文本处理（并行计算）
   • 支持多种池化模式（CLS、MEAN等）

4.2 关键实现细节

4.2.1 模型加载流程

java复制// AbstractInProcessEmbeddingModel.java
protected static OnnxBertBiEncoder loadFromFileSystem(Path modelPath, Path tokenizerPath) {
    try {
        OrtEnvironment environment = OrtEnvironment.getEnvironment();
        Tokenizer tokenizer = Tokenizer.fromJson(tokenizerPath);
        return new OnnxBertBiEncoder(environment, tokenizer, modelPath);
    } catch (Exception e) {
        throw new RuntimeException(e);
    }
}

4.2.2 嵌入计算实现

java复制// AbstractInProcessEmbeddingModel.java
@Override
public Response<List<Embedding>> embedAll(List<String> texts) {
    if (texts.size() == 1) {
        return Response.from(Collections.singletonList(embed(texts.get(0)).content()));
    }
    
    // 并行处理多个文本
    return Response.from(texts.parallelStream()
            .map(text -> embed(text).content())
            .collect(Collectors.toList()));
}

4.3 OnnxBertBiEncoder核心组件

OnnxBertBiEncoder类包含两个关键成员：

1. environment：ONNX运行时环境
2. tokenizer：文本分词器

实际计算过程：

1. 文本分词
2. 转换为模型输入格式
3. 执行ONNX推理
4. 处理输出结果

5. 性能优化与实践建议

5.1 性能优化技巧

1. 批量处理：尽量使用embedAll方法处理多个文本
2. 线程池配置：调整并行计算的线程数
3. 模型量化：使用量化后的ONNX模型减小内存占用
4. 缓存机制：对频繁使用的文本嵌入结果进行缓存

5.2 常见问题排查

5.2.1 模型加载失败

可能原因：

• 文件路径错误
• ONNX模型版本不兼容
• 缺少必要的依赖库

解决方案：

1. 检查文件路径是否正确
2. 确认ONNX模型版本与运行时兼容
3. 确保所有依赖已正确安装

5.2.2 推理速度慢

优化建议：

1. 使用更小的模型
2. 启用GPU加速（如果可用）
3. 减少批量大小

5.2.3 内存不足

处理方法：

1. 减小批量大小
2. 使用量化模型
3. 增加JVM堆内存

6. 扩展应用场景

6.1 文本相似度计算

利用embedding结果计算文本相似度：

java复制public float cosineSimilarity(float[] vec1, float[] vec2) {
    float dotProduct = 0.0f;
    float norm1 = 0.0f;
    float norm2 = 0.0f;
    
    for (int i = 0; i < vec1.length; i++) {
        dotProduct += vec1[i] * vec2[i];
        norm1 += vec1[i] * vec1[i];
        norm2 += vec2[i] * vec2[i];
    }
    
    return dotProduct / (float)(Math.sqrt(norm1) * Math.sqrt(norm2));
}

6.2 语义搜索实现

构建简单的语义搜索引擎：

1. 为文档集合生成嵌入
2. 存储嵌入向量
3. 查询时计算查询与文档的相似度
4. 返回最相似的文档

6.3 混合模型部署

可以结合多种部署方式：

• 使用ONNX部署embedding模型
• 使用Ollama部署LLM模型
• 通过HuggingFace接入云端模型

这种混合方案可以：

• 降低本地资源需求
• 提高关键组件的性能
• 保持系统的灵活性

7. 进阶主题

7.1 自定义模型支持

要支持自定义ONNX模型，需要：

1. 确保模型输入输出格式符合预期
2. 实现自定义的Tokenizer
3. 扩展AbstractInProcessEmbeddingModel

7.2 模型微调与再训练

虽然ONNX主要用于推理，但可以通过：

1. 在原框架中微调模型
2. 重新转换为ONNX格式
3. 更新Java应用中的模型文件

7.3 性能监控与调优

建议添加：

1. 推理耗时监控
2. 内存使用监控
3. 批量处理效率统计
4. 错误率跟踪

这些指标可以帮助发现性能瓶颈和优化机会。