plan-sess_f2e7ba1e-4b44-4bff-92c2-d69b0de3e486.md 3.8 KB

本地嵌入模块(DJL + ONNX Runtime + bge-m3)—— 独立模块方案

定位(按你的要求调整)

  • 完全独立的新模块,不用 agentscope 任何 API,不接入 EmbeddingModelFactory/agent_model 表/RAG 流程,前端无感知。这些现有代码零改动。
  • 对业务暴露一套简洁的 Java API(注入 EmbeddingService 即用);另附一个极小的调试 HTTP 接口用于人工验证(可随时删)。
  • 配置走 application.yaml(zsjz.embedding.*),与项目 zsjz.duckdb/zsjz.insight 等段风格一致,不进数据库。

模块结构(新包 com.zsjz.ai.module.embedding)

module/embedding/
├── config/EmbeddingProperties.java      # @ConfigurationProperties("zsjz.embedding")
├── core/OnnxEmbeddingEngine.java        # 分词 + ONNX 推理 + CLS池化 + L2归一化(懒加载)
├── core/EmbeddingException.java         # 模块内异常
├── service/EmbeddingService.java        # 业务门面(具体类,风格仿 module/search 的 SearchQueryService)
└── controller/EmbeddingDebugController.java  # 可选调试接口

业务 API(EmbeddingService,业务代码这样用)

float[]        embed(String text);              // 单文本 → 1024 维归一化向量
List<float[]>  embedBatch(List<String> texts);  // 批量(内部按 batch-size 分批推理)
double         cosineSimilarity(float[] a, float[] b);
int            getDimension();                  // 从模型输出自动探测(1024)
boolean        isReady();  void warmUp();       // 懒加载状态 / 主动预热

依赖(ai-server/pom.xml,实施时先校验 Maven Central 最新可用版)

  • ai.djl.huggingface:tokenizers(0.36.0,2026-09-10 发布)—— 直接加载 bge-m3/tokenizer.json,原生 Rust 分词
  • com.microsoft.onnxruntime:onnxruntime(1.30.0 左右,校验 metadata 后定版)—— 各平台 CPU native 内置

核心实现要点

  1. 配置项 zsjz.embedding:model-path(默认 bge-m3,相对 user.dir)、max-seq-len(默认 8192)、batch-size(默认 8)、intra-op-threads、opt-level(默认 BASIC)、enabled
  2. 引擎懒加载单例(Spring bean,volatile + synchronized):首次调用才加载 2.3GB 权重,不阻塞启动;session 从文件路径创建,天然支持 onnx/model.onnx + model.onnx_data 外部权重(找不到 onnx/model.onnx 时回退 model.onnx)
  3. 推理:tokenizer 逐条编码(截断到 max-seq-len)→ 批内手动 padding(input_ids / attention_mask / token_type_ids 按会话实际输入名动态构建,输入输出名从 session 动态读取)→ 取输出 last_hidden_state[:, 0](CLS)→ L2 归一化(bge-m3 官方 dense 用法)
  4. 推理在线程池隔离执行,避免阻塞调用线程;失败抛 EmbeddingException,模型目录缺失给出明确提示

调试接口(验证用,Result 包装,符合项目惯例)

POST /js/a/embedding/debug,body {"texts": ["转账给可疑账户", "向陌生账户转移资金"]} → 返回维度、各向量、两两余弦相似度,用于人工确认语义区分度。

验证步骤

  1. mvn -pl ai-server compile 通过
  2. 启动应用,curl 调试接口:确认返回 1024 维、归一化正确(模长≈1)、相关语句相似度显著高于无关语句、首次调用触发模型加载日志
  3. 业务接入示例写入类 Javadoc

风险与注意事项

  • 内存:权重在堆外(native),进程 RSS 约 +2.5~3GB;首次加载 20~60 秒
  • CPU 速度:表结构/短文本(几百 token)毫秒级~秒级;超长文本可通过 max-seq-len 控制
  • GPU:本期 CPU;后续换 onnxruntime_gpu 依赖 + ExecutionProvider 配置即可,API 不变