1. 为什么要在本地跑嵌入模型1.1 从一个真实的需求场景说起去年年底我接手了一个内部知识库项目需求很朴素把团队积累的几百份技术文档、会议纪要、故障复盘做成一个能语义搜索的东西。一开始想得很简单直接调云端嵌入接口不就行了结果一算账就发现问题了——文档切片之后大概有十几万个 chunk每次全量重建索引的 token 消耗量相当可观而且后续每次增量更新都要走一遍网络请求。更麻烦的是有些内部资料压根不适合往外发数据出域这条红线直接卡死了云端方案。这时候本地嵌入就成了唯一的选择。所谓本地嵌入说白了就是在一个不依赖外部服务的环境里把一段文本转换成一串固定长度的浮点数数组也就是我们常说的向量或者坐标。这串数字代表了这段文本在语义空间里的位置语义相近的文本它们的坐标距离就近。检索的时候把用户的问题也转成坐标然后去找距离最近的几个文档坐标就完成了语义搜索。而ONNX在这里扮演的角色是一个跨框架的模型推理格式。你可以把它理解成模型的“通用集装箱”——不管这个嵌入模型原本是用 PyTorch 还是 TensorFlow 训练的导出成 ONNX 之后就能用一个统一的运行时来加载和执行。对于 Java 技术栈来说这一点尤其重要因为 JVM 生态里原生跑深度学习模型的选择本来就不多ONNX Runtime 提供了相当成熟的 Java binding这就让 Spring 应用直接加载嵌入模型变成了现实。1.2 本地 ONNX 嵌入到底解决了什么问题我把本地 ONNX 嵌入的价值归纳成三条都是实际踩过之后才体会到的。第一是成本归零。云端嵌入按 token 计费量小的时候感觉不到量一大就是持续失血。本地模型跑在自己的 CPU 上除了电费和机器折旧边际成本几乎为零。对于知识库这种“建好之后还要反复重建、反复调参”的场景成本优势非常明显。第二是数据不出域。所有文本的向量化过程都在本地完成原始文本和向量都不需要经过任何外部网络。对于有合规要求的内部系统这一条基本是刚需。第三是延迟可控。本地推理没有网络往返虽然单次推理速度可能不如云端的大模型但胜在稳定。尤其是批量建索引的时候可以充分利用多核 CPU 并行整体吞吐反而更可观。当然本地模型也有它的代价——模型体积、内存占用、推理速度都需要权衡。这就引出了下一个问题模型怎么选。1.3 模型选型的几个关键考量选嵌入模型我一般看四个维度维度、体积、语言支持、推理速度。维度决定了向量占用的存储空间和检索时的计算量。常见的嵌入维度有 384、512、768、1024 等。维度越高表达能力通常越强但存储和计算成本也线性上升。对于中文知识库384 到 768 维基本够用。体积直接关系到部署难度。一个几百 MB 的模型塞进容器镜像里构建和分发都会变慢。我倾向于选择量化后的 ONNX 模型体积能压到几十 MB 甚至更小。语言支持是中文场景的命门。很多英文嵌入模型在中文上表现平平必须选那些在中文语料上训练过的或者至少是多语言的。推理速度则和模型结构、量化方式有关。ONNX Runtime 支持多种图优化和量化实际速度差异可能有好几倍。综合下来我最终选了一个多语言的小型嵌入模型导出为 ONNX 并做了动态量化维度 384模型文件不到 50MB。这个选择在后续的实测中被证明是合理的——单条文本嵌入在普通笔记本 CPU 上大约 10 到 30 毫秒批量处理时吞吐能到每秒几百条。2. Spring AI 与 ONNX 的整合思路2.1 Spring AI 的嵌入抽象层Spring AI 提供了一套相对清晰的抽象核心接口是EmbeddingModel。这个接口定义了两个主要方法一个是把单条文本转成向量一个是批量转换。围绕这个接口Spring AI 还提供了向量存储VectorStore的抽象负责向量的持久化和相似度检索。理解这套抽象的意义在于嵌入模型和向量存储是解耦的。你可以用本地的 ONNX 模型做嵌入同时用内存向量库、关系型数据库或者专门的向量数据库做存储。这种解耦让技术选型变得灵活也让我们可以专注于把本地嵌入这一环做扎实。Spring AI 本身提供了一些云端嵌入模型的实现但本地 ONNX 这一块需要我们自己动手。好在EmbeddingModel是个接口实现起来并不复杂核心工作就是加载 ONNX 模型、做分词、跑推理、归一化输出。2.2 为什么是 ONNX 而不是其他格式在 JVM 上跑模型可选的路子其实有几条。一是用 DJLDeep Java Library它封装得比较好但引入的依赖比较重二是直接用 ONNX Runtime 的 Java API轻量且可控三是通过某种进程间通信调用 Python 服务但这又引入了额外的部署复杂度。我选 ONNX Runtime 直接集成理由有三。第一依赖干净就一个onnxruntime的 jar 包加上它的原生库没有一堆传递依赖。第二控制力强从输入张量的构造到输出张量的解析每一步都在自己手里出问题好排查。第三性能可预期ONNX Runtime 的图优化和线程管理都是成熟的调优空间明确。相比之下DJL 虽然 API 更友好但它对模型格式和分词器的封装有时候会带来“黑盒感”遇到分词不一致或者输出异常时排查起来反而更费劲。对于嵌入这种相对简单的模型直接上 ONNX Runtime 是性价比最高的选择。2.3 分词本地嵌入最容易被忽视的一环很多人以为加载了 ONNX 模型就万事大吉结果跑出来的向量和预期对不上问题十有八九出在分词上。嵌入模型的输入不是原始文本而是 token id 序列。不同的模型用不同的分词器比如 BERT 系用 WordPiece有些模型用 SentencePiece还有些用 BPE。分词器必须和模型训练时用的完全一致否则 token id 对不上向量自然就是错的。在 Python 生态里分词器通常跟着模型一起提供用transformers库加载就行。但在 Java 里没有现成的等价物。我的做法是在 Python 侧用官方分词器把一批文本的 token id 导出来然后在 Java 侧实现一个等价的分词逻辑用这批数据做回归测试。只要 Java 分词的结果和 Python 完全一致就可以放心使用。分词器本身通常是一个词表文件加一套规则。词表可以打包进资源目录规则部分主要是文本规范化比如转小写、去重音、按子词切分、加特殊 token。这部分代码量不大但必须仔细一个字符的差异都可能导致结果偏移。3. 从零搭建本地嵌入的完整实操3.1 环境准备与依赖引入先说依赖。核心就两个ONNX Runtime 的 Java 包以及 Spring AI 的嵌入抽象。Maven 里大概是这样dependency groupIdcom.microsoft.onnxruntime/groupId artifactIdonnxruntime/artifactId version1.17.0/version /dependency dependency groupIdorg.springframework.ai/groupId artifactIdspring-ai-core/artifactId version1.0.0-M2/version /dependencyONNX Runtime 的版本要和模型导出时的 opset 兼容一般选比较新的稳定版就行。Spring AI 的版本迭代比较快接口可能有变化建议锁定一个版本不要频繁升级。模型文件、词表文件放到src/main/resources/models/下面。模型建议用动态量化版本体积小、速度也快精度损失在嵌入任务上通常可以接受。3.2 模型加载与推理会话管理ONNX Runtime 的推理入口是OrtSession。这个对象创建成本较高而且不是线程安全的所以正确的做法是应用启动时创建一个OrtSession然后用一个对象池或者ThreadLocal来管理并发访问。我实测下来单个OrtSession在多线程下直接共享会出问题表现为结果错乱或者直接抛异常。用ThreadLocalOrtSession每个线程持有一个会话内存开销可以接受并发性能也稳定。如果并发量很大可以考虑用对象池但要注意会话的创建和销毁成本。加载模型时OrtSession.SessionOptions里可以配置线程数。对于嵌入这种计算密集但单次计算量不大的任务把 intra-op 线程数设成 CPU 核心数inter-op 线程数设成 1通常是比较优的配置。这个参数需要根据实际机器压测调整没有万能值。3.3 分词器的 Java 实现要点分词器的实现我拆成三步规范化、预分词、子词切分。规范化包括转小写、去除多余空白、处理全角半角。这一步看似简单但中文文本里全角标点、零宽字符很常见不处理干净会影响分词结果。预分词是把文本切成初步的单元。英文按空格和标点切中文按字切。这一步决定了后续子词切分的粒度。子词切分是核心。以 WordPiece 为例它采用贪心最长匹配从当前位置开始尝试匹配词表里最长的子词匹配不到就回退一个字符直到匹配成功或者只剩单字符。匹配不到的单字符用[UNK]代替。每个子词前面加##前缀表示它是词的延续部分。实现的时候有个坑词表的加载顺序和查找方式会影响性能。词表可能有几万个条目用HashMap查找是 O(1)但最长匹配需要多次查找。我的做法是把词表按长度分组先试长词再试短词减少查找次数。3.4 输入张量的构造与输出解析嵌入模型的输入通常有三个张量input_ids、attention_mask、token_type_ids。input_ids是 token 的 id 序列attention_mask标记哪些位置是真实 token1哪些是填充0token_type_ids在单句任务里通常全 0。张量的形状是[batch_size, sequence_length]。sequence_length要固定通常是 128 或 256。短于这个长度的用 0 填充长于这个长度的要截断。截断策略有从头截、从尾截、头尾各截一半一般选头尾各截一半保留首尾信息。输出通常是一个[batch_size, sequence_length, hidden_size]的张量或者已经池化好的[batch_size, hidden_size]。如果是前者需要自己做池化——常见的是取[CLS]位置的向量或者对所有 token 的向量做平均池化。具体用哪种取决于模型训练时的设定这个信息一般在模型的文档或者配置里能找到。拿到向量之后还有一步归一化。把向量除以其 L2 范数使其长度为 1。归一化之后余弦相似度就等价于点积检索时计算更快。这一步很多教程会漏掉但不做归一化相似度排序可能会出问题。4. 向量存储与检索的落地细节4.1 向量存储的选型向量存哪里取决于数据量和检索要求。数据量小几万条以内直接用内存里的一个Listfloat[]加暴力检索就够了简单可靠。数据量大就得上专门的向量索引。我试过几种方案。内存暴力检索在十万条以内单次查询几毫秒到几十毫秒完全可用。再往上暴力检索的延迟就线性增长了这时候需要近似最近邻ANN索引。常见的 ANN 算法有 HNSW、IVF 等各有取舍HNSW 查询快、内存占用高IVF 内存省、但需要训练且召回率略低。对于 Spring AI 项目如果已经在用某个向量数据库的 starter直接用它提供的VectorStore实现就行。如果没有自己基于内存实现一个也不难核心就是余弦相似度排序。4.2 相似度计算与 Top-K 检索相似度计算用余弦相似度。因为向量已经归一化余弦相似度就是点积float dot(float[] a, float[] b) { float sum 0; for (int i 0; i a.length; i) { sum a[i] * b[i]; } return sum; }Top-K 检索就是遍历所有向量算相似度维护一个大小为 K 的最小堆。K 一般取 3 到 10看具体场景。检索出来的结果还要带上原始文本和元数据方便展示和溯源。这里有个优化点批量查询。如果一次要查多个问题可以把查询向量组成矩阵和文档矩阵做矩阵乘法比逐条算快得多。ONNX Runtime 本身也支持批量推理建索引的时候把文本分批送进去吞吐能提升好几倍。4.3 索引构建与增量更新建索引是个批处理任务。流程是读文档、切片、批量嵌入、写入向量存储。切片策略很关键切得太碎会丢失上下文切得太大会稀释语义。我一般按段落切每段控制在 200 到 500 字段间可以留一点重叠避免边界信息丢失。增量更新要处理两个问题新文档的加入和旧文档的删除。加入好办嵌入之后追加就行。删除需要维护一个 id 到向量位置的映射删除时把对应位置标记为无效定期做一次压缩整理。如果用的是支持删除的向量数据库这部分就交给数据库处理。重建索引的时候建议先建到临时索引建完再原子替换避免重建期间检索不可用。5. 常见问题与排查技巧实录5.1 向量结果和 Python 对不上怎么办这是最常见的问题排查顺序如下。先查分词。把同一段文本在 Python 和 Java 里分别分词打印 token id 序列逐位对比。只要有一位不同就顺着分词逻辑找差异。常见差异点大小写处理、全角半角、特殊字符、截断长度。再查张量构造。确认attention_mask是否正确填充位置是否真的填了 0token_type_ids是否全 0。这些细节错了输出会整体偏移。最后查池化和归一化。确认用的是[CLS]还是平均池化归一化是否做了。这几步任何一步不一致向量就会有差异。5.2 推理速度慢的优化方向速度慢先看是不是每次推理都重新创建了OrtSession。如果是改成复用速度能提升一个数量级。再看线程配置。intra-op线程数设成 CPU 核心数inter-op设成 1通常是最优的。如果设反了性能会明显下降。还可以考虑量化。动态量化把权重从 float32 压到 int8速度能提升 2 到 4 倍精度损失在嵌入任务上一般小于 1%。如果对精度要求极高可以用静态量化但需要校准数据集。批量推理也是提速的关键。单条推理的固定开销占比高批量处理能摊薄这个开销。实测批量大小设成 16 或 32吞吐提升明显。5.3 内存占用过高的处理内存占用主要来自三块模型本身、OrtSession、向量数据。模型和会话的内存是固定的量化能显著降低。向量数据的内存和数量、维度成正比十万条 384 维的 float 向量大约 150MB可以接受。如果维度更高或者数量更大可以考虑用 float16 存储内存减半精度损失很小。还有一个容易被忽视的点JVM 堆外内存。ONNX Runtime 的原生内存不在 JVM 堆里-Xmx管不到它。如果发现进程内存远超-Xmx设定大概率是原生内存占用。可以通过SessionOptions限制内存 arena 的大小或者定期重建会话释放内存。5.4 常见问题速查表现象可能原因排查方向向量结果与 Python 不一致分词差异、张量构造错误、池化方式不同逐位对比 token id检查 mask 和池化推理速度慢会话未复用、线程配置不当、未量化复用会话调整线程数启用量化内存占用高模型未量化、向量用 float32、原生内存泄漏量化模型float16 存储限制 arena检索结果不相关未归一化、切片过大、模型不适配中文检查归一化调整切片换中文模型并发下结果错乱OrtSession共享用 ThreadLocal 或对象池启动报 opset 不兼容ONNX Runtime 版本与模型 opset 不匹配升级运行时或重新导出模型6. 一些实操心得与扩展思路6.1 我踩过的几个坑第一个坑是分词器的特殊 token。有些模型在句首加[CLS]句尾加[SEP]有些只加[CLS]。加错了向量会偏。这个必须对着模型文档确认不能想当然。第二个坑是截断长度。我一开始设了 512结果发现大部分文本根本用不到这么长反而拖慢了推理。后来改成 256速度提升明显检索质量几乎没变化。截断长度要根据实际文本长度分布来定不是越长越好。第三个坑是归一化的时机。我一开始在池化之后忘了归一化检索结果排序乱七八糟。后来加上归一化排序立刻正常了。这个细节很容易漏但影响很大。第四个坑是并发下的会话共享。我一开始图省事全局共享一个OrtSession单线程测试没问题一上并发就出各种诡异结果。改成ThreadLocal之后彻底解决。6.2 后续可以怎么扩展本地嵌入搭好之后扩展方向其实不少。一是混合检索。向量检索擅长语义匹配但对精确关键词的匹配不如传统倒排索引。把两者结合用加权或者重排序的方式融合结果检索质量能再上一个台阶。二是重排序。先用向量检索召回一批候选再用一个更精细的交叉编码器模型对候选做重排序。交叉编码器精度高但速度慢只用在少量候选上整体性价比很高。三是多模态扩展。同样的 ONNX 思路可以加载图像嵌入模型把图片也转成向量实现图文混合检索。技术路径是一样的只是输入预处理不同。四是模型热更新。把模型文件放在外部目录应用启动时加载需要更新时替换文件并触发重新加载。这样换模型不用重新打包部署运维上方便很多。6.3 给准备上手的朋友几句实在话如果你正准备在 Spring 项目里搞本地嵌入我的建议是先把分词这一环做扎实。整个链路里分词是最容易出错、也最难排查的部分。花时间把 Java 分词和 Python 分词对齐后面会省很多事。模型选型不要贪大。嵌入任务和生成任务不一样小模型往往就够用而且速度快、部署简单。先用小模型跑通全流程有需要再换大的。向量存储先别急着上向量数据库。内存暴力检索在中小规模下完全够用等数据量真的上来了再迁移迁移成本也不高。最后建索引和检索要分开压测。这两个环节的性能瓶颈不一样混在一起测容易互相干扰定位不到真正的问题。这套方案我在几个内部项目里都用过从几百份文档到几万份文档整体表现稳定。本地 ONNX 嵌入这条路对于有数据合规要求、又想控制成本的团队来说确实是个值得投入的方向。