mirror of
https://gitcode.com/JianFeeeee/HomeAgent.git
synced 2026-10-07 16:14:30 +00:00
feat(memory): 千问三段式 ONNX 嵌入补齐——可复现导出脚本 + Go 侧首次完整验证
此前三段式拆分后 ONNX 路径从未从 Go 侧跑通：embedder_onnx_test.go 仍引用 分段前的 API（e.renderInput、TextTower.onnx、旧目录），go vet -tags onnxruntime 直接编译失败。导出脚本只在 /tmp 且硬编码本机路径、从第三个目录拷贝固定形状的 Vision.onnx，完全不可复现。音频会被视觉塔编码，静默往统一空间灌入错误坐标。 本提交补齐这些缺口： 一、可复现导出脚本（scripts/export_qwen3vl_embedding_onnx.py） - 自动拉取模型（HuggingFace 优先，失败回落 ModelScope，支持 HF_ENDPOINT 镜像）； - 导出 TokenEmbedding + Transformer + Vision 三段图，图文共用同一 token embedding、28 层 Transformer、last-token 池化与 fingerprint； - 双重自检（不可省）：分段 PyTorch vs 完整模型 + 导出后的 ONNX vs 完整模型， cos < 0.999999 即非零退出——「能加载」不等于「算得对」； - 默认把 L2 归一化后的冻结参考向量写入产物目录（qwen_reference.json）—— Go 测试据此做逐维冻结回归，且「该目录是哪次导出的」从文件本身可追溯； - --verify-only 校验既有产物不重新导出，可用来确认线上在用的图没坏。 关键实测结论（已写入 docs/zh/multimodal-space.md 与长期记忆）： 原生多帧视频不可行——Qwen3-VL 视觉塔把 grid_thw 当 Python 值消费 （grid_thw.tolist()），legacy tracer 固化为常量，导出后图中根本没有 grid_thw 输入，换帧数调用直接 Invalid input name: grid_thw。故视觉塔固定 (1,48,48)， 视频由上层抽帧后逐帧按图像编码（同模型/同维度/同 fingerprint），音频明确 unsupported。 二、模态边界（vector.ErrModalityUnsupported） - 新增 vector.ErrModalityUnsupported：表示「该模态不在本统一空间的原生覆盖 范围内」，与普通错误语义不同——调用方应把它当「永远不会有向量」而非 「本次失败、下次重试」； - qwen.EmbedImageDense 按 mime 拒绝 audio/* 与 video/*：此前它会拿视觉塔 去解音频字节，往统一空间灌入语义错误的坐标且静默； - reembedStaleMedia 对 ErrModalityUnsupported 不计失败、不重试、不用别的 模型向量顶替（TestReembedStaleMedia_SkipsUnsupportedWithoutFaking 守住）。 三、Go ONNX 测试首次完整通过 - 重写 embedder_onnx_test.go：修复编译 + 文本冻结回归 + 图像冻结回归 + 两条阴性对照（不同输入必须不同、图像与文本必须不同）+ 不支持模态断言； - 参考值从产物目录的 qwen_reference.json 读取（不在测试里硬编码浮点）； - 用线上部署产物实测全部通过（text cos=0.999999940, image cos=0.999999762）。 四、.gitignore 修复 - /scripts/ 此前被列在「运行时产物」下，但它是作者维护的工具目录 （模型导出、侧车、部署校验），deploy/systemd/embed-sidecar.service 直接 引用 scripts/embed_sidecar.py，忽略它会让那份 unit 在别人的机器上指向 不存在的文件。改为只忽略 __pycache__。 五、文档（docs/zh/multimodal-space.md） - 获取/启用/产物契约/模态边界/验证/资源成本/与现有部署产物的等价性。 验证：go build ./...、go vet ./...、go vet -tags onnxruntime ./...、 go test -short 全部通过；ONNX 标签测试对线上部署产物全部通过。
This commit is contained in:
JianFeeeee
7
.gitignore
vendored
7
.gitignore
vendored
@ -48,7 +48,12 @@ codegraph.json
|
||||
/adapters/
|
||||
/knowledge/
|
||||
/memory/
|
||||
/scripts/
|
||||
# 注：/scripts/ **不**忽略。它是作者维护的工具目录（模型导出、侧车、部署校验），
|
||||
# 不是运行期产物：deploy/systemd/embed-sidecar.service 直接引用
|
||||
# scripts/embed_sidecar.py，忽略它会让那份 unit 在别人的机器上指向不存在的文件。
|
||||
# 只忽略其中的缓存。
|
||||
/scripts/__pycache__/
|
||||
__pycache__/
|
||||
terminal_locked_log.txt
|
||||
|
||||
dist/
|
||||
|
||||
153
docs/zh/multimodal-space.md
Normal file
153
docs/zh/multimodal-space.md
Normal file
@ -0,0 +1,153 @@
|
||||
# 统一多模态向量空间（Qwen3-VL-Embedding-2B）
|
||||
|
||||
文本、图像、**视频帧** 在同一模型、同一 2048 维、同一 fingerprint 空间里被编码。
|
||||
记忆系统用它做三件事：多模态图记忆的跨模态召回、multimodal doc 的向量融合、
|
||||
multimodal context 的相关性裁剪/淘汰。
|
||||
|
||||
统一空间取代了此前「把图片交给视觉模型生成文字描述、再按描述检索」的做法。
|
||||
那条链路有三个致命缺陷：描述是异步生成的（未生成前媒体等于不存在）、语义检索
|
||||
实际上只搜描述文字、图库里的「媒体节点」只是描述文本的投影而不是媒体本身。
|
||||
**不要再引入任何描述式索引。**
|
||||
|
||||
## 一、产物与获取
|
||||
|
||||
产物约 8 GB（含外部权重），**不进仓库**；用导出脚本自动拉取模型并导出：
|
||||
|
||||
```bash
|
||||
# 自动拉取（HuggingFace 优先，失败回落 ModelScope）+ 导出 + 自检
|
||||
python3 scripts/export_qwen3vl_embedding_onnx.py \
|
||||
--out /home/newqqagent/models/qwen3-vl-embed-multimodal-onnx
|
||||
|
||||
# 已下载过模型：跳过拉取
|
||||
python3 scripts/export_qwen3vl_embedding_onnx.py \
|
||||
--model-dir /path/to/Qwen3-VL-Embedding-2B \
|
||||
--out /home/newqqagent/models/qwen3-vl-embed-multimodal-onnx
|
||||
|
||||
# 参考向量默认直接写进产物目录（<out>/qwen_reference.json），无需额外参数
|
||||
python3 scripts/export_qwen3vl_embedding_onnx.py --model-dir ... --out ...
|
||||
```
|
||||
|
||||
国内镜像：导出脚本沿用 `huggingface_hub` 的约定，直接 `export HF_ENDPOINT=https://hf-mirror.com` 即可。
|
||||
依赖：`torch`（CPU 版即可）、`transformers>=4.57`、`onnx`、`onnxruntime`、`pillow`、`numpy`，
|
||||
以及可选的 `huggingface_hub` / `modelscope`。显存不需要，内存建议 ≥ 16 GB（FP32 加载约 8 GB）。
|
||||
|
||||
导出脚本**会清空 --out 目录**后重写，避免旧图/旧外部权重污染 fingerprint
|
||||
（fingerprint 变化会触发一次无意义的全量向量重算）。因此不要直接覆盖线上正在使用的目录，
|
||||
先导出到新目录再切换。
|
||||
|
||||
### 产物契约（Go 侧按此读取）
|
||||
|
||||
| 文件 | 输入 | 输出 |
|
||||
|---|---|---|
|
||||
| `TokenEmbedding.onnx` | `input_ids` int64 `[1,seq]` | `hidden` float `[1,seq,2048]` |
|
||||
| `Transformer.onnx` | `hidden`、`deepstack_0/1/2` `[1,seq,2048]`、`rotary_cos/sin` `[1,seq,128]`、`causal_mask` `[1,1,seq,seq]` | `embedding` `[1,2048]` |
|
||||
| `Vision.onnx(+.data)` | `pixel_values` `[2304,1536]` | `deepstack_feature_0/1/2`、`vision_hidden_states` `[576,2048]` |
|
||||
|
||||
外加 `tokenizer.json`、`tokenizer_config.json`、`chat_template.jinja`、`embed_config.json`。
|
||||
|
||||
三段只是部署形式，不是三个向量空间：图文共用同一 token embedding、同一 28 层
|
||||
Transformer、同一 last-token 池化。RoPE 与视觉特征散射故意留在 Go 计算，
|
||||
因为旧式 tracer 会把 `seq=598 / visual=576` 烘焙进图里——签名上写着 dynamic
|
||||
axis，实际却只能用导出的那个长度运行。
|
||||
|
||||
### 导出脚本自检（不可省）
|
||||
|
||||
脚本内部跑两道校验，任一道 cos < 0.999999 就以非零码退出：
|
||||
|
||||
1. 分段 PyTorch（三段组合）对比完整模型前向；
|
||||
2. 用 onnxruntime 跑**导出后**的三段图，再对比完整模型前向。
|
||||
|
||||
「能加载」不等于「算得对」：形状错、输入名错、池化位置错的图都能正常 load。
|
||||
|
||||
## 二、启用
|
||||
|
||||
```bash
|
||||
# 配置库（config.db）或 WebUI 设置页
|
||||
core.memory.multimodal_space.type = onnx
|
||||
core.memory.multimodal_space.onnx.model_dir = /home/newqqagent/models/qwen3-vl-embed-multimodal-onnx
|
||||
```
|
||||
|
||||
注意事项：
|
||||
|
||||
- `homed` 必须带 `onnxruntime` build tag 构建，且 `libonnxruntime.so` 可被找到
|
||||
（`/opt/onnxruntime/libonnxruntime.so` 等）。未带 tag 时 `qwen` 是 no-op stub。
|
||||
- 改配置后需重启进程生效。
|
||||
- 未配置时优雅降级：文档层退到 TF-IDF 稀疏检索，媒体块仍按结构边关联，只是没有跨模态召回。
|
||||
|
||||
## 三、模态覆盖范围
|
||||
|
||||
| 模态 | 状态 | 说明 |
|
||||
|---|---|---|
|
||||
| 文本 | ✅ 原生 | `VectorizeDense` |
|
||||
| 图像 | ✅ 原生 | `EmbedImageDense`，固定 768×768 视觉塔 |
|
||||
| 视频 | ⚠️ 逐帧 | 上层抽帧后**逐帧按图像编码**，同模型/同维度/同 fingerprint；不做跨帧时序注意力 |
|
||||
| 音频 | ❌ 明确不支持 | 返回 `vector.ErrModalityUnsupported` |
|
||||
|
||||
**音频不得用视觉塔硬编码**，也**不得**拿另一个模型的向量顶替——那会把两套坐标系
|
||||
混进同一空间，检索出的相似度没有任何意义，而且错误是静默的。未来接入真正的统一
|
||||
音频模型后再扩展。
|
||||
|
||||
### 为什么视频不做原生时序（已实测，勿重复尝试）
|
||||
|
||||
Qwen3-VL 视觉塔把 `grid_thw` 当 Python 值消费（源码里是 `grid_thw.tolist()`）。
|
||||
legacy tracer（`dynamo=False`）会把它固化成常量：实测导出后 ONNX 图里**根本没有**
|
||||
`grid_thw` 输入，用别的帧数调用直接报 `Invalid input name: grid_thw`；
|
||||
导出时的 TracerWarning 明确提示 `Converting a tensor to a Python list might cause
|
||||
the trace to be incorrect`。
|
||||
|
||||
因此视觉塔固定 `grid=(1,48,48)`。要做到原生多帧需要换 `torch.export`/dynamo 路径，
|
||||
而该路径此前已产生过「形状看似动态、实际错误」的静默故障（Core.onnx 的
|
||||
`3 by 23 / 3 by 598` 广播错误），在时序维度上重试的收益不足以抵消风险。
|
||||
视频价值由「逐帧进入同一空间」提供：帧是真实媒体块，按自己的向量被召回。
|
||||
|
||||
## 四、验证
|
||||
|
||||
```bash
|
||||
# Go 侧：ONNX 路径（模型目录缺失时自动 skip）
|
||||
QWEN_ONNX_MODEL_DIR=/home/newqqagent/models/qwen3-vl-embed-multimodal-onnx \
|
||||
go test -tags onnxruntime ./internal/memory/qwen/ -v
|
||||
|
||||
# 排除二进制交付问题的替代：先单独验证模型与 CSV 无关的 ONNX 图
|
||||
go vet -tags onnxruntime ./...
|
||||
```
|
||||
|
||||
Go 测试覆盖：冻结参考向量（文本/图像各 12 维）、同输入确定性、不同输入敏感性、
|
||||
图像与文本向量必须不同、以及音频/视频必须返回 `ErrModalityUnsupported`。
|
||||
|
||||
冻结参考向量由导出脚本写入**产物目录本身**（`<out>/qwen_reference.json`），
|
||||
来源可追溯：同一脚本既产出模型，也产出「这个模型对固定输入应有的输出」。
|
||||
重新导出后若参考值变化，说明权重或图结构变了，必须显式更新参考而不是放宽阈值。
|
||||
|
||||
> **参考向量是 L2 归一化后的值。** ONNX 图返回的是 final norm 之后的原始
|
||||
> last hidden（量级约 100），而 Go 侧 `VectorizeDense` / `EmbedImageDense`
|
||||
> 返回归一化向量。写参考时忘归一化，Go 测试会全线不匹配，而现象看起来
|
||||
> 像“模型不对”，实际只是两边对“向量”的定义不同。
|
||||
|
||||
验证既有产物（不重新导出）：
|
||||
|
||||
```bash
|
||||
python3 scripts/export_qwen3vl_embedding_onnx.py --verify-only --model-dir <model> \
|
||||
--out /home/newqqagent/models/qwen3-vl-embed-multimodal-onnx
|
||||
```
|
||||
|
||||
脚本会顺便把归一化后的参考向量写入该目录。
|
||||
|
||||
### 与现有部署产物的等价性
|
||||
|
||||
本仓库脚本对同一源模型导出时，`TokenEmbedding.onnx` 与 `Transformer.onnx` 与
|
||||
线上在用的产物**逐字节相同**（sha256 一致）；`Vision.onnx` 差异仅在打包形式：
|
||||
旧产物把权重量到外部 `Vision.onnx.data`，新脚本内联在图里。两者数值等价。
|
||||
|
||||
注意这会带来一个**操作性**差异：Go 的结构指纹（`computeFingerprint`）把
|
||||
`*.onnx.data` 的文件名与大小算在内，因此「外部权重版 ↔ 内联版」互换会让
|
||||
fingerprint 变化，从而触发一次全量向量重算。重算不会**算错**（数值等价），
|
||||
只是白花一次 CPU；若不想触发，就保持产物打包形式不变。
|
||||
|
||||
## 五、资源成本
|
||||
|
||||
- 产物磁盘约 8 GB；导出过程峰值内存约 10–12 GB（FP32 加载）。
|
||||
- 单次 CPU 推理：文本约几十毫秒量级，图像（2304 patch 过 24 层视觉塔 + 28 层语言模型）
|
||||
明显更重，因此入库时不阻塞对话，靠 `reembedStaleMedia` 在启动时并发迁移
|
||||
（ONNX 路径 4 worker）。
|
||||
- fingerprint 由三段图 + `embed_config.json` + 外部权重文件名/大小共同决定；
|
||||
换模型或重新导出都会让它变化，从而触发历史向量重算——这是预期行为。
|
||||
@ -1,11 +1,13 @@
|
||||
package core
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"log"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
|
||||
"gitcode.com/JianFeeeee/HomeAgent/internal/memory"
|
||||
"gitcode.com/JianFeeeee/HomeAgent/internal/memory/vector"
|
||||
)
|
||||
|
||||
// 媒体与记忆块的生命周期辅助。
|
||||
@ -126,7 +128,7 @@ func (a *Agent) reembedStaleMedia() {
|
||||
len(digests), shortFP(fp), a.multimodalSpace.Dim(), workers)
|
||||
|
||||
jobs := make(chan string, workers*2)
|
||||
var done, failed int64
|
||||
var done, failed, unsupported int64
|
||||
var failedMu sync.Mutex
|
||||
var wg sync.WaitGroup
|
||||
|
||||
@ -135,13 +137,18 @@ func (a *Agent) reembedStaleMedia() {
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for d := range jobs {
|
||||
if err := a.reembedOne(d, fp); err != nil {
|
||||
switch err := a.reembedOne(d, fp); {
|
||||
case err == nil:
|
||||
atomic.AddInt64(&done, 1)
|
||||
case errors.Is(err, vector.ErrModalityUnsupported):
|
||||
// 该模态不在本空间内（如音频）：不重试、不计失败，
|
||||
// 也不拿另一个模型的向量顶替。
|
||||
atomic.AddInt64(&unsupported, 1)
|
||||
default:
|
||||
failedMu.Lock()
|
||||
failed++
|
||||
failedMu.Unlock()
|
||||
continue
|
||||
}
|
||||
atomic.AddInt64(&done, 1)
|
||||
}
|
||||
}()
|
||||
}
|
||||
@ -155,11 +162,14 @@ func (a *Agent) reembedStaleMedia() {
|
||||
}
|
||||
close(jobs)
|
||||
wg.Wait()
|
||||
log.Printf("[media] 向量迁移完成: 成功=%d 失败=%d 总计=%d fp=%s",
|
||||
done, failed, len(digests), shortFP(fp))
|
||||
log.Printf("[media] 向量迁移完成: 成功=%d 失败=%d 不在本空间=%d 总计=%d fp=%s",
|
||||
done, failed, unsupported, len(digests), shortFP(fp))
|
||||
}
|
||||
|
||||
// reembedOne 为单条媒体重新计算向量并写入（stat/get 失败时跳过该条目）。
|
||||
//
|
||||
// 模态不在本空间覆盖范围时返回 ErrModalityUnsupported，调用方据此区分
|
||||
// 「永久无向量」与「本次失败重试」。
|
||||
func (a *Agent) reembedOne(digest, fp string) error {
|
||||
it, err := a.mediaStore.Stat(digest)
|
||||
if err != nil {
|
||||
|
||||
@ -2,12 +2,15 @@ package core
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"gitcode.com/JianFeeeee/HomeAgent/internal/memory"
|
||||
"gitcode.com/JianFeeeee/HomeAgent/internal/memory/document"
|
||||
"gitcode.com/JianFeeeee/HomeAgent/internal/memory/media"
|
||||
"gitcode.com/JianFeeeee/HomeAgent/internal/memory/vector"
|
||||
)
|
||||
|
||||
// 媒体与记忆块的生命周期测试。
|
||||
@ -81,6 +84,53 @@ func TestHeldMediaDigests_CollectsAcrossLayers(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// fakeSpace 是一个只覆盖图像的假统一空间，用来验证「不在本空间」与
|
||||
// 「本次失败」必须被区分对待。
|
||||
type fakeSpace struct{}
|
||||
|
||||
func (fakeSpace) VectorizeDense(string) ([]float64, error) { return []float64{1, 0}, nil }
|
||||
|
||||
func (fakeSpace) EmbedImageDense(_ []byte, mime string) ([]float64, error) {
|
||||
if strings.HasPrefix(mime, "audio/") || strings.HasPrefix(mime, "video/") {
|
||||
return nil, fmt.Errorf("%w: %s", vector.ErrModalityUnsupported, mime)
|
||||
}
|
||||
return []float64{1, 0}, nil
|
||||
}
|
||||
|
||||
func (fakeSpace) Fingerprint() string { return "fake-space" }
|
||||
func (fakeSpace) Dim() int { return 2 }
|
||||
func (fakeSpace) Loaded() bool { return true }
|
||||
func (fakeSpace) Close() {}
|
||||
|
||||
// TestReembedStaleMedia_SkipsUnsupportedWithoutFaking 验证向量迁移不会：
|
||||
// - 把音频当失败反复重试；
|
||||
// - 更不能拿另一个模型的向量顶替音频（那会污染统一空间且静默）。
|
||||
func TestReembedStaleMedia_SkipsUnsupportedWithoutFaking(t *testing.T) {
|
||||
a, ms := newMediaLoopAgent(t)
|
||||
img, _ := ms.Put([]byte("img-bytes"), media.Item{MIME: "image/png"})
|
||||
aud, _ := ms.Put([]byte("aud-bytes"), media.Item{MIME: "audio/wav"})
|
||||
|
||||
a.multimodalSpace = fakeSpace{}
|
||||
a.reembedStaleMedia()
|
||||
|
||||
it, err := ms.Stat(img)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(it.Vec) != 2 || it.VecModel != "fake-space" {
|
||||
t.Fatalf("图像应拿到本空间向量，实际 vec=%v model=%q", it.Vec, it.VecModel)
|
||||
}
|
||||
|
||||
audIt, err := ms.Stat(aud)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(audIt.Vec) != 0 || audIt.VecModel != "" {
|
||||
t.Fatalf("音频不得被写入任何向量（不能用别的模型顶替），实际 vec=%v model=%q",
|
||||
audIt.Vec, audIt.VecModel)
|
||||
}
|
||||
}
|
||||
|
||||
// payloadHeld 是删除前的活查询。
|
||||
func TestPayloadHeld(t *testing.T) {
|
||||
a, ms := newMediaLoopAgent(t)
|
||||
|
||||
@ -18,6 +18,8 @@ import (
|
||||
"sync"
|
||||
|
||||
ort "github.com/yalue/onnxruntime_go"
|
||||
|
||||
"gitcode.com/JianFeeeee/HomeAgent/internal/memory/vector"
|
||||
)
|
||||
|
||||
type embedConfig struct {
|
||||
@ -138,12 +140,24 @@ func (e *Embedder) VectorizeDense(text string) ([]float64, error) {
|
||||
return e.runTransformer(hidden, deep, position, len(ids))
|
||||
}
|
||||
|
||||
func (e *Embedder) EmbedImageDense(raw []byte, _ string) ([]float64, error) {
|
||||
// EmbedImageDense 把一张图片编码到统一空间。
|
||||
//
|
||||
// mime 决定这个模态是否在本空间的原生覆盖范围内：Qwen3-VL 能原生编码文本与
|
||||
// 图像，但**不原生支持音频**。音频（以及未抽帧的视频文件）必须返回
|
||||
// ErrModalityUnsupported，而不是拿视觉塔去编码——那会往统一空间里灌入
|
||||
// 语义错误的坐标，而错误是静默的。视频请由上层抽帧后逐帧当作图像编码。
|
||||
func (e *Embedder) EmbedImageDense(raw []byte, mime string) ([]float64, error) {
|
||||
e.mu.RLock()
|
||||
defer e.mu.RUnlock()
|
||||
if !e.loaded {
|
||||
return nil, fmt.Errorf("qwen embedder not loaded")
|
||||
}
|
||||
switch {
|
||||
case strings.HasPrefix(mime, "audio/"):
|
||||
return nil, fmt.Errorf("%w: audio (%s) 需由真正的统一音频模型扩展", vector.ErrModalityUnsupported, mime)
|
||||
case strings.HasPrefix(mime, "video/"):
|
||||
return nil, fmt.Errorf("%w: 视频文件请先抽帧，逐帧按图像编码 (%s)", vector.ErrModalityUnsupported, mime)
|
||||
}
|
||||
pixels, err := preprocessImage(raw)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
|
||||
@ -3,68 +3,275 @@
|
||||
package qwen
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"image"
|
||||
"image/png"
|
||||
"math"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"gitcode.com/JianFeeeee/HomeAgent/internal/memory/vector"
|
||||
)
|
||||
|
||||
// 本地模型路径与 tokenizer_test.go 共用；产物不在仓库时跳过。
|
||||
const onnxModelDir = "/home/newqqagent/models/qwen3-vl-embed-text-onnx"
|
||||
|
||||
// TestEmbedderMatchesONNXReference 冻结一条由 Python onnxruntime 1.28.0 生成的
|
||||
// FP32 参考向量，验证完整 Go 路径：模板渲染 → BPE → ONNX → L2 normalize。
|
||||
// onnxModelDir 返回三段式 Qwen 多模态 ONNX 产物目录。
|
||||
//
|
||||
// 只校验前 12 维不是为了放宽正确性，而是避免把 2048 个浮点常量塞进仓库；
|
||||
// tokenizer 的全部 token 已由 TestTokenizerMatchesReference 逐条精确校验，图本身
|
||||
// 另有 PyTorch↔ONNX 的多形状验证。这里负责捕获 Go 张量形状、输入名、输出名、
|
||||
// 池化/归一化或模板接线错误。
|
||||
func TestEmbedderMatchesONNXReference(t *testing.T) {
|
||||
if _, err := os.Stat(onnxModelDir + "/TextTower.onnx"); err != nil {
|
||||
t.Skipf("ONNX 产物不可用，跳过: %v", err)
|
||||
// 产物约 8GB（含外部权重），不进仓库；由 scripts/export_qwen3vl_embedding_onnx.py
|
||||
// 自动拉取模型并导出。可通过 QWEN_ONNX_MODEL_DIR 指向别处；目录不存在时相关
|
||||
// 测试跳过，而不是失败——CI 与本机开发者都不一定有这份产物。
|
||||
func onnxModelDir() string {
|
||||
if v := os.Getenv("QWEN_ONNX_MODEL_DIR"); v != "" {
|
||||
return v
|
||||
}
|
||||
return "/home/newqqagent/models/qwen3-vl-embed-multimodal-onnx"
|
||||
}
|
||||
|
||||
e, err := New(onnxModelDir)
|
||||
func requireONNXArtifacts(t *testing.T) string {
|
||||
t.Helper()
|
||||
dir := onnxModelDir()
|
||||
for _, name := range []string{"TokenEmbedding.onnx", "Transformer.onnx", "Vision.onnx", "embed_config.json", "tokenizer.json"} {
|
||||
if _, err := os.Stat(dir + "/" + name); err != nil {
|
||||
t.Skipf("ONNX 产物不完整（%s: %v），跳过；用 scripts/export_qwen3vl_embedding_onnx.py 导出", name, err)
|
||||
}
|
||||
}
|
||||
return dir
|
||||
}
|
||||
|
||||
// onnxReference 是导出脚本 `--emit-reference` 写出的冻结参考。
|
||||
//
|
||||
// 刻意不把浮点常量硬编码在测试里：参考值必须能追溯到「哪个模型、哪次导出、
|
||||
// 什么输入」，而不是一组无人知道出处的数字。参考文件里的 RGB/尺寸同时用来
|
||||
// 构造测试图片，保证输入与参考按构造一致，不会因测试改动而静默错位。
|
||||
type onnxReference struct {
|
||||
Text string `json:"text"`
|
||||
TextVectorPrefix []float64 `json:"text_vector_prefix"`
|
||||
ImageRGB []int `json:"image_rgb"`
|
||||
ImageSize int `json:"image_size"`
|
||||
ImageVectorPrefix []float64 `json:"image_vector_prefix"`
|
||||
Dim int `json:"dim"`
|
||||
}
|
||||
|
||||
func loadReference(t *testing.T, dir string) *onnxReference {
|
||||
t.Helper()
|
||||
path := os.Getenv("QWEN_ONNX_REFERENCE")
|
||||
if path == "" {
|
||||
path = filepath.Join(dir, "qwen_reference.json")
|
||||
}
|
||||
raw, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
t.Skipf("缺少冻结参考 %s（由 scripts/export_qwen3vl_embedding_onnx.py --emit-reference 生成）: %v", path, err)
|
||||
}
|
||||
var ref onnxReference
|
||||
if err := json.Unmarshal(raw, &ref); err != nil {
|
||||
t.Fatalf("解析参考 %s: %v", path, err)
|
||||
}
|
||||
if ref.Text == "" || len(ref.TextVectorPrefix) == 0 || len(ref.ImageRGB) != 3 || ref.ImageSize <= 0 {
|
||||
t.Fatalf("参考 %s 不完整: %+v", path, ref)
|
||||
}
|
||||
return &ref
|
||||
}
|
||||
|
||||
// solidPNG 生成一张 size×size 纯色 PNG，供跨语言冻结向量回归。
|
||||
//
|
||||
// 刻意用纯色且尺寸与视觉塔一致：Go 侧预处理对已是 768×768 的输入不做插值、
|
||||
// 不补边，于是 patch 张量只由布局决定。一旦 patch 排列写错（内层循环顺序、
|
||||
// merge 分组顺序、通道顺序），冻结向量立刻不匹配——而那类错误在人工看图时
|
||||
// 几乎发现不了。
|
||||
func solidPNG(t *testing.T, size int, r, g, b uint8) []byte {
|
||||
t.Helper()
|
||||
img := image.NewNRGBA(image.Rect(0, 0, size, size))
|
||||
for y := 0; y < size; y++ {
|
||||
for x := 0; x < size; x++ {
|
||||
i := img.PixOffset(x, y)
|
||||
img.Pix[i], img.Pix[i+1], img.Pix[i+2], img.Pix[i+3] = r, g, b, 255
|
||||
}
|
||||
}
|
||||
var buf bytes.Buffer
|
||||
if err := png.Encode(&buf, img); err != nil {
|
||||
t.Fatalf("encode png: %v", err)
|
||||
}
|
||||
return buf.Bytes()
|
||||
}
|
||||
|
||||
func newTestEmbedder(t *testing.T, dir string) *Embedder {
|
||||
t.Helper()
|
||||
e, err := New(dir)
|
||||
if err != nil {
|
||||
t.Fatalf("New: %v", err)
|
||||
}
|
||||
defer e.Close()
|
||||
t.Cleanup(e.Close)
|
||||
if !e.Loaded() || e.Dim() != 2048 || e.Fingerprint() == "" {
|
||||
t.Fatalf("元数据异常: loaded=%v dim=%d fingerprint=%q", e.Loaded(), e.Dim(), e.Fingerprint())
|
||||
}
|
||||
return e
|
||||
}
|
||||
|
||||
ids, err := e.tok.encodeModelInput(e.renderInput("今天天气怎么样"), e.config.MaxLength)
|
||||
if err != nil {
|
||||
t.Fatalf("encodeModelInput: %v", err)
|
||||
func assertNormalized(t *testing.T, name string, got []float64, dim int) {
|
||||
t.Helper()
|
||||
if dim > 0 && len(got) != dim {
|
||||
t.Fatalf("%s 维度 = %d，期望 %d", name, len(got), dim)
|
||||
}
|
||||
postID, ok := e.tok.SpecialID("<|endoftext|>")
|
||||
if !ok || len(ids) != 23 || ids[len(ids)-1] != postID {
|
||||
t.Fatalf("模型输入 post-processor 异常: len=%d tail=%v postID=%d ok=%v", len(ids), ids[len(ids)-1:], postID, ok)
|
||||
}
|
||||
|
||||
got, err := e.VectorizeDense("今天天气怎么样")
|
||||
if err != nil {
|
||||
t.Fatalf("VectorizeDense: %v", err)
|
||||
}
|
||||
if len(got) != 2048 {
|
||||
t.Fatalf("向量维度 = %d，期望 2048", len(got))
|
||||
}
|
||||
|
||||
want := []float64{
|
||||
-0.0288955811, 0.0522339381, 0.0360119902, -0.000676361844,
|
||||
-0.0431003496, 0.00342374574, -0.000226021366, 0.0157372113,
|
||||
-0.028889874, 0.0264931992, 0.0117276432, -0.00513622677,
|
||||
}
|
||||
for i := range want {
|
||||
if diff := math.Abs(got[i] - want[i]); diff > 2e-5 {
|
||||
t.Errorf("维度 %d = %.10g，参考 %.10g，差 %.3g", i, got[i], want[i], diff)
|
||||
}
|
||||
}
|
||||
|
||||
var norm float64
|
||||
for _, v := range got {
|
||||
norm += v * v
|
||||
}
|
||||
if diff := math.Abs(math.Sqrt(norm) - 1); diff > 1e-6 {
|
||||
t.Errorf("L2 norm = %.9f，期望 1", math.Sqrt(norm))
|
||||
}
|
||||
if !e.Loaded() || e.Dim() != 2048 || e.Fingerprint() == "" {
|
||||
t.Errorf("元数据异常: loaded=%v dim=%d fingerprint=%q", e.Loaded(), e.Dim(), e.Fingerprint())
|
||||
t.Errorf("%s L2 norm = %.9f，期望 1", name, math.Sqrt(norm))
|
||||
}
|
||||
}
|
||||
|
||||
func assertFrozenPrefix(t *testing.T, name string, got, want []float64) {
|
||||
t.Helper()
|
||||
if len(got) < len(want) {
|
||||
t.Fatalf("%s 向量过短: %d", name, len(got))
|
||||
}
|
||||
for i := range want {
|
||||
if diff := math.Abs(got[i] - want[i]); diff > 2e-5 {
|
||||
t.Errorf("%s 维度 %d = %.10g，参考 %.10g，差 %.3g", name, i, got[i], want[i], diff)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestEmbedderMatchesONNXReference 逐维对比导出脚本写出的冻结参考向量，
|
||||
// 验证完整 Go 路径：模板渲染 → BPE → TokenEmbedding → Transformer →
|
||||
// last-token 池化 → L2 normalize。
|
||||
//
|
||||
// 只覆盖前若干维不是因为放宽正确性（脚本侧的 PyTorch↔ONNX 校验是逐维的），
|
||||
// 而是避免把 2048 个浮点常量塞进仓库；这里负责捕获 Go 张量形状、输入名、
|
||||
// 输出名、池化/归一化或模板接线错误。
|
||||
func TestEmbedderMatchesONNXReference(t *testing.T) {
|
||||
dir := requireONNXArtifacts(t)
|
||||
e := newTestEmbedder(t, dir)
|
||||
ref := loadReference(t, dir)
|
||||
|
||||
ids, _, _, _, err := e.tok.textModelInput(e.config.Instruction, ref.Text, e.config.MaxLength)
|
||||
if err != nil {
|
||||
t.Fatalf("textModelInput: %v", err)
|
||||
}
|
||||
postID, ok := e.tok.SpecialID("<|endoftext|>")
|
||||
if !ok || ids[len(ids)-1] != postID {
|
||||
t.Fatalf("模型输入 post-processor 异常: len=%d tail=%v postID=%d ok=%v", len(ids), ids[len(ids)-1:], postID, ok)
|
||||
}
|
||||
|
||||
got, err := e.VectorizeDense(ref.Text)
|
||||
if err != nil {
|
||||
t.Fatalf("VectorizeDense: %v", err)
|
||||
}
|
||||
assertNormalized(t, "text", got, ref.Dim)
|
||||
assertFrozenPrefix(t, "text", got, ref.TextVectorPrefix)
|
||||
}
|
||||
|
||||
// TestEmbedderImageMatchesONNXReference 冻结一张纯色图的参考向量，
|
||||
// 验证 Go 侧的视觉预处理 + patch 排列 + 视觉注入 + 语言模型与 Python 参考一致。
|
||||
func TestEmbedderImageMatchesONNXReference(t *testing.T) {
|
||||
dir := requireONNXArtifacts(t)
|
||||
e := newTestEmbedder(t, dir)
|
||||
ref := loadReference(t, dir)
|
||||
|
||||
img := solidPNG(t, ref.ImageSize, uint8(ref.ImageRGB[0]), uint8(ref.ImageRGB[1]), uint8(ref.ImageRGB[2]))
|
||||
got, err := e.EmbedImageDense(img, "image/png")
|
||||
if err != nil {
|
||||
t.Fatalf("EmbedImageDense: %v", err)
|
||||
}
|
||||
assertNormalized(t, "image", got, ref.Dim)
|
||||
assertFrozenPrefix(t, "image", got, ref.ImageVectorPrefix)
|
||||
}
|
||||
|
||||
// TestEmbedderTextIsSensitiveToInput 阴性对照：冻结向量必须真的随输入变化。
|
||||
//
|
||||
// 没有这条对照，一个「永远返回同一向量」的错误实现也能通过上面的冻结回归
|
||||
//（只要那个常量恰好等于参考值）。这里验证不同文本给出不同向量，且相似文本
|
||||
// 的余弦高于无关文本——即嵌入确实携带语义，而不是常量。
|
||||
func TestEmbedderTextIsSensitiveToInput(t *testing.T) {
|
||||
dir := requireONNXArtifacts(t)
|
||||
e := newTestEmbedder(t, dir)
|
||||
|
||||
base, err := e.VectorizeDense("今天天气怎么样")
|
||||
if err != nil {
|
||||
t.Fatalf("VectorizeDense: %v", err)
|
||||
}
|
||||
same, err := e.VectorizeDense("今天天气怎么样")
|
||||
if err != nil {
|
||||
t.Fatalf("VectorizeDense: %v", err)
|
||||
}
|
||||
if cosine(base, same) < 0.999999 {
|
||||
t.Errorf("同一输入两次嵌入不一致: cos=%.9f（ONNX 会话被并发复用或存在非确定性）", cosine(base, same))
|
||||
}
|
||||
|
||||
other, err := e.VectorizeDense("数据库索引的选择性是怎么计算的")
|
||||
if err != nil {
|
||||
t.Fatalf("VectorizeDense: %v", err)
|
||||
}
|
||||
if cosine(base, other) > 0.999 {
|
||||
t.Errorf("无关文本的余弦高达 %.6f，嵌入可能是常量", cosine(base, other))
|
||||
}
|
||||
}
|
||||
|
||||
// TestEmbedderImageMatchesONNXReference 的替代：不依赖冻结参考的不变量检查。
|
||||
//
|
||||
// 即使参考文件缺失（没有导出产物）或未重新生成，这些不变量也应成立：
|
||||
// 图像路径必须真的走了视觉塔，且不同图片给出不同坐标。
|
||||
func TestEmbedderImageDiffersFromText(t *testing.T) {
|
||||
dir := requireONNXArtifacts(t)
|
||||
e := newTestEmbedder(t, dir)
|
||||
|
||||
imgVec, err := e.EmbedImageDense(solidPNG(t, qwenImageSize, 200, 30, 30), "image/png")
|
||||
if err != nil {
|
||||
t.Fatalf("EmbedImageDense: %v", err)
|
||||
}
|
||||
txtVec, err := e.VectorizeDense(DefaultInstruction)
|
||||
if err != nil {
|
||||
t.Fatalf("VectorizeDense: %v", err)
|
||||
}
|
||||
if cosine(imgVec, txtVec) > 0.999 {
|
||||
t.Error("图像向量与文本向量几乎相同，视觉塔可能没被真正执行")
|
||||
}
|
||||
|
||||
// 不同颜色的图必须给出不同向量
|
||||
blue, err := e.EmbedImageDense(solidPNG(t, qwenImageSize, 30, 150, 220), "image/png")
|
||||
if err != nil {
|
||||
t.Fatalf("EmbedImageDense: %v", err)
|
||||
}
|
||||
if cosine(imgVec, blue) > 0.999999 {
|
||||
t.Error("不同图片给出相同向量，视觉路径未生效")
|
||||
}
|
||||
}
|
||||
|
||||
// TestEmbedderRejectsUnsupportedModalities 音频与视频文件必须显式报「不在本空间」，
|
||||
// 而不是拿视觉塔硬编码一个语义错误的坐标。
|
||||
//
|
||||
// Qwen3-VL 原生支持文本与图像；音频需要未来接入真正的统一音频模型。
|
||||
// 若这里退化成普通错误，调用方会把它当「本次失败、下次重试」，
|
||||
// 于是每轮启动都重试一批永远不可能成功的条目。
|
||||
func TestEmbedderRejectsUnsupportedModalities(t *testing.T) {
|
||||
dir := requireONNXArtifacts(t)
|
||||
e := newTestEmbedder(t, dir)
|
||||
|
||||
for _, mime := range []string{"audio/wav", "audio/mpeg", "video/mp4", "video/quicktime"} {
|
||||
_, err := e.EmbedImageDense([]byte("not-a-real-media"), mime)
|
||||
if err == nil {
|
||||
t.Fatalf("%s 应返回错误而不是造出向量", mime)
|
||||
}
|
||||
if !errors.Is(err, vector.ErrModalityUnsupported) {
|
||||
t.Errorf("%s 错误应为 ErrModalityUnsupported，实际: %v", mime, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func cosine(a, b []float64) float64 {
|
||||
if len(a) != len(b) || len(a) == 0 {
|
||||
return 0
|
||||
}
|
||||
var dot, na, nb float64
|
||||
for i := range a {
|
||||
dot += a[i] * b[i]
|
||||
na += a[i] * a[i]
|
||||
nb += b[i] * b[i]
|
||||
}
|
||||
if na == 0 || nb == 0 {
|
||||
return 0
|
||||
}
|
||||
return dot / (math.Sqrt(na) * math.Sqrt(nb))
|
||||
}
|
||||
|
||||
@ -50,6 +50,15 @@ const (
|
||||
// ErrNotSupported 表示 Vectorizer 不支持该原生模态；调用方不得以描述文本冒充其向量。
|
||||
var ErrNotSupported = fmt.Errorf("vectorizer does not support image embedding")
|
||||
|
||||
// ErrModalityUnsupported 表示该模态不在本统一向量空间的原生覆盖范围内。
|
||||
//
|
||||
// 它与普通错误语义不同：调用方应把它当作「这条媒体本空间永远不会有向量」
|
||||
// 而不是「这次失败了、下次重试」。绝不能拿另一个模型的向量顶替——那会把
|
||||
// 两套坐标系混进同一空间，检索出来的相似度没有任何意义。
|
||||
//
|
||||
// 例：Qwen3-VL 能原生编码文本/图像，音频需要未来接入真正的统一音频模型。
|
||||
var ErrModalityUnsupported = fmt.Errorf("modality not supported by this embedding space")
|
||||
|
||||
// Vector 是带权特征映射：feature → weight
|
||||
type Vector map[string]float64
|
||||
|
||||
|
||||
203
scripts/embed_sidecar.py
Normal file
203
scripts/embed_sidecar.py
Normal file
@ -0,0 +1,203 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Jina v5-omni-nano embedding sidecar for HomeAgent.
|
||||
|
||||
Minimal HTTP server exposing POST /embed matching the HTTPEmbedder contract:
|
||||
Request: {"modality":"text|image", "side":"query|document", "text":"...", "data":"base64...", "mime":"..."}
|
||||
Response: {"embedding":[float...]}
|
||||
|
||||
Health: GET /health → {"status":"ok","model":"jina-v5-omni-nano","dim":768,"loaded":true}
|
||||
"""
|
||||
|
||||
import base64
|
||||
import io
|
||||
import logging
|
||||
import os
|
||||
import signal
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
from http.server import HTTPServer, BaseHTTPRequestHandler
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
# ─── Config ──────────────────────────────────────────────────────────────────
|
||||
MODEL_DIR = os.environ.get("JINA_MODEL_DIR", "/home/newqqagent/models/jina-v5-omni-nano")
|
||||
PORT = int(os.environ.get("JINA_PORT", "18999"))
|
||||
DIMENSION = int(os.environ.get("JINA_DIMENSION", "768"))
|
||||
MAX_WORKERS = int(os.environ.get("JINA_MAX_WORKERS", "4"))
|
||||
BATCH_SIZE = int(os.environ.get("JINA_BATCH_SIZE", "8"))
|
||||
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="%(asctime)s [embed_sidecar] %(message)s",
|
||||
datefmt="%Y-%m-%d %H:%M:%S",
|
||||
)
|
||||
log = logging.getLogger("embed_sidecar")
|
||||
|
||||
# ─── Model Loading ───────────────────────────────────────────────────────────
|
||||
model = None
|
||||
processor = None
|
||||
model_lock = threading.Lock()
|
||||
ready = False
|
||||
|
||||
|
||||
def load_model():
|
||||
global model, processor, ready
|
||||
log.info("loading model from %s ...", MODEL_DIR)
|
||||
t0 = time.time()
|
||||
|
||||
from transformers import AutoModel, AutoProcessor
|
||||
|
||||
model = AutoModel.from_pretrained(
|
||||
MODEL_DIR,
|
||||
trust_remote_code=True,
|
||||
local_files_only=True,
|
||||
default_task="retrieval",
|
||||
modality="vision",
|
||||
dtype=torch.float32,
|
||||
).eval()
|
||||
|
||||
processor = AutoProcessor.from_pretrained(
|
||||
MODEL_DIR,
|
||||
trust_remote_code=True,
|
||||
local_files_only=True,
|
||||
)
|
||||
|
||||
elapsed = time.time() - t0
|
||||
ready = True
|
||||
log.info("model loaded in %.1fs, dim=%d", elapsed, DIMENSION)
|
||||
|
||||
|
||||
def embed_text(text: str, side: str = "query") -> list[float]:
|
||||
"""Embed text with proper Query:/Document: prefix for retrieval."""
|
||||
prefix = "Query: " if side == "query" else "Document: "
|
||||
inputs = processor(
|
||||
text=[prefix + text],
|
||||
padding=True,
|
||||
truncation=True,
|
||||
max_length=1024,
|
||||
return_tensors="pt",
|
||||
)
|
||||
with torch.inference_mode():
|
||||
vec = model.embed(**inputs)
|
||||
return vec.float().cpu().numpy()[0].tolist()
|
||||
|
||||
|
||||
def embed_image(data_b64: str, mime: str, side: str = "document") -> list[float]:
|
||||
"""Embed image from base64 data."""
|
||||
from PIL import Image
|
||||
|
||||
img_bytes = base64.b64decode(data_b64)
|
||||
img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
|
||||
|
||||
prefix = "Query: " if side == "query" else "Document: "
|
||||
inputs = processor(
|
||||
images=img,
|
||||
text=f"{prefix}<image>",
|
||||
return_tensors="pt",
|
||||
)
|
||||
with torch.inference_mode():
|
||||
vec = model.embed(**inputs)
|
||||
return vec.float().cpu().numpy()[0].tolist()
|
||||
|
||||
|
||||
# ─── HTTP Server ─────────────────────────────────────────────────────────────
|
||||
class EmbedHandler(BaseHTTPRequestHandler):
|
||||
"""Handle /embed and /health endpoints."""
|
||||
|
||||
def log_message(self, fmt, *args):
|
||||
# Suppress default access log for /health
|
||||
if "/health" not in str(args[0]):
|
||||
log.info(fmt, *args)
|
||||
|
||||
def do_GET(self):
|
||||
parsed = urlparse(self.path)
|
||||
if parsed.path == "/health":
|
||||
self._respond(200, {
|
||||
"status": "ok" if ready else "loading",
|
||||
"model": "jina-v5-omni-nano",
|
||||
"dim": DIMENSION,
|
||||
"loaded": ready,
|
||||
})
|
||||
else:
|
||||
self._respond(404, {"error": "not found"})
|
||||
|
||||
def do_POST(self):
|
||||
parsed = urlparse(self.path)
|
||||
if parsed.path != "/embed":
|
||||
self._respond(404, {"error": "not found"})
|
||||
return
|
||||
|
||||
if not ready:
|
||||
self._respond(503, {"error": "model not loaded"})
|
||||
return
|
||||
|
||||
# Read request body
|
||||
try:
|
||||
length = int(self.headers.get("Content-Length", 0))
|
||||
body = self.rfile.read(length)
|
||||
req = __import__("json").loads(body)
|
||||
except Exception as e:
|
||||
self._respond(400, {"error": f"invalid request: {e}"})
|
||||
return
|
||||
|
||||
modality = req.get("modality", "text")
|
||||
side = req.get("side", "query")
|
||||
|
||||
try:
|
||||
with model_lock:
|
||||
if modality == "text":
|
||||
text = req.get("text", "")
|
||||
if not text:
|
||||
self._respond(400, {"error": "missing text field"})
|
||||
return
|
||||
vec = embed_text(text, side)
|
||||
elif modality == "image":
|
||||
data = req.get("data", "")
|
||||
mime = req.get("mime", "image/png")
|
||||
if not data:
|
||||
self._respond(400, {"error": "missing data field"})
|
||||
return
|
||||
vec = embed_image(data, mime, side)
|
||||
else:
|
||||
self._respond(400, {"error": f"unsupported modality: {modality}"})
|
||||
return
|
||||
|
||||
self._respond(200, {"embedding": vec})
|
||||
except Exception as e:
|
||||
log.error("embed error: %s", e, exc_info=True)
|
||||
self._respond(500, {"error": str(e)})
|
||||
|
||||
def _respond(self, status: int, data: dict):
|
||||
import json as json_mod
|
||||
body = json_mod.dumps(data).encode("utf-8")
|
||||
self.send_response(status)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.send_header("Content-Length", str(len(body)))
|
||||
self.end_headers()
|
||||
self.wfile.write(body)
|
||||
|
||||
|
||||
def main():
|
||||
# Load model in background thread so server can start accepting /health
|
||||
threading.Thread(target=load_model, daemon=True).start()
|
||||
|
||||
server = HTTPServer(("0.0.0.0", PORT), EmbedHandler)
|
||||
|
||||
def shutdown(signum, frame):
|
||||
log.info("shutting down...")
|
||||
server.shutdown()
|
||||
sys.exit(0)
|
||||
|
||||
signal.signal(signal.SIGTERM, shutdown)
|
||||
signal.signal(signal.SIGINT, shutdown)
|
||||
|
||||
log.info("listening on :%d", PORT)
|
||||
server.serve_forever()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
516
scripts/export_qwen3vl_embedding_onnx.py
Normal file
516
scripts/export_qwen3vl_embedding_onnx.py
Normal file
@ -0,0 +1,516 @@
|
||||
#!/usr/bin/env python3
|
||||
"""自动拉取 Qwen3-VL-Embedding-2B 并导出 HomeAgent 用的三段式 ONNX 统一向量空间。
|
||||
|
||||
产物（写入 --out 目录，约 8GB）::
|
||||
|
||||
TokenEmbedding.onnx input_ids -> hidden
|
||||
Transformer.onnx hidden + deepstack + RoPE + causal mask -> embedding
|
||||
Vision.onnx(+.data) pixel_values -> 3 层 DeepStack + 主视觉特征
|
||||
tokenizer.json / tokenizer_config.json / chat_template.jinja
|
||||
embed_config.json Go 侧读取的布局与契约常量
|
||||
|
||||
为什么是「三段」而不是一张图
|
||||
--------------------------
|
||||
文本与图像共用同一 token embedding、同一 28 层 Transformer、同一 last-token
|
||||
池化与同一 fingerprint；分段只是部署形式。把 RoPE 与视觉特征散射留在 Go 计算，
|
||||
是为了避开旧式 tracer 把 seq=598 / visual=576 烘焙进图里——那样签名上写着
|
||||
dynamic_axes，实际却只能用导出的那个长度运行。
|
||||
|
||||
Vision 为什么固定 768×768（不支持原生多帧视频）
|
||||
--------------------------------------------
|
||||
Qwen3-VL 的视觉塔把 ``grid_thw`` 当 Python 值消费（``grid_thw.tolist()``），
|
||||
legacy tracer 会把它的内容固化成常量：实测导出后 ONNX 图里根本没有
|
||||
``grid_thw`` 输入，用别的帧数调用会直接报 Invalid input name。因此这里把
|
||||
grid 固定为 (1, 48, 48)，并在导出处做 PyTorch↔ONNX 一致性校验。
|
||||
视频由上层抽帧后逐帧按图像编码——同一模型、同一维度、同一 fingerprint，
|
||||
只是不做跨帧时序注意力；音频不在本空间覆盖范围内（见 vector.ErrModalityUnsupported）。
|
||||
|
||||
自检是不可省的
|
||||
------------
|
||||
导出脚本必须自己证明产物正确，而不是只比较有没有报错：
|
||||
1. 分段 PyTorch（TokenEmbedding+Transformer+Vision 的组合）对比完整模型前向；
|
||||
2. 再用 onnxruntime 跑导出后的三段图，对比完整模型前向。
|
||||
两步都要求 cos ≥ 0.999999，否则以非零码退出——绝不产出一个「能加载但算错」的模型。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import sys
|
||||
import time
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
INSTRUCTION = "Represent the user's input."
|
||||
IMAGE_SIZE = 768
|
||||
PATCH_SIZE = 16
|
||||
TEMPORAL_PATCH = 2
|
||||
SPATIAL_MERGE = 2
|
||||
MAX_LENGTH = 1024 # 768x768 有 576 个视觉 token；512 会截断视觉占位符
|
||||
DEFAULT_MODEL_ID = "Qwen/Qwen3-VL-Embedding-2B"
|
||||
REFERENCE_TEXT = "今天天气怎么样"
|
||||
REFERENCE_IMAGE_RGB = (200, 30, 30)
|
||||
|
||||
|
||||
def log(msg: str) -> None:
|
||||
print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True)
|
||||
|
||||
|
||||
# ─────────────────────────── 模型获取 ───────────────────────────
|
||||
|
||||
def pull_model(model_id: str, store: str) -> str:
|
||||
"""把模型拉到本地并返回快照目录。优先 HuggingFace，失败回落 ModelScope。
|
||||
|
||||
两个源都必须显式给出目标目录：默认的 HF blobs+snapshots 结构会把权重存成
|
||||
符号链接树，导出脚本按固定文件名读取时很容易踩到不存在的路径。
|
||||
HF_ENDPOINT 会被 huggingface_hub 自动识别，因此国内镜像（如
|
||||
https://hf-mirror.com ）无需额外参数。
|
||||
"""
|
||||
os.makedirs(store, exist_ok=True)
|
||||
local = os.path.join(store, model_id.replace("/", "--"))
|
||||
marker = os.path.join(local, "config.json")
|
||||
if os.path.exists(marker):
|
||||
log(f"复用已下载模型: {local}")
|
||||
return local
|
||||
|
||||
errors = []
|
||||
try:
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
log(f"从 HuggingFace 拉取 {model_id} -> {local}（HF_ENDPOINT={os.environ.get('HF_ENDPOINT', '默认')}）")
|
||||
snapshot_download(repo_id=model_id, local_dir=local, max_workers=4)
|
||||
if os.path.exists(marker):
|
||||
return local
|
||||
errors.append("huggingface: 下载完成但缺少 config.json")
|
||||
except Exception as exc: # noqa: BLE001 - 需要回落到 ModelScope
|
||||
errors.append(f"huggingface: {exc}")
|
||||
log(f"HuggingFace 拉取失败：{exc}")
|
||||
|
||||
try:
|
||||
from modelscope import snapshot_download as ms_snapshot
|
||||
|
||||
log(f"从 ModelScope 拉取 {model_id} -> {local}")
|
||||
ms_snapshot(model_id, local_dir=local)
|
||||
if os.path.exists(marker):
|
||||
return local
|
||||
errors.append("modelscope: 下载完成但缺少 config.json")
|
||||
except Exception as exc: # noqa: BLE001
|
||||
errors.append(f"modelscope: {exc}")
|
||||
log(f"ModelScope 拉取失败：{exc}")
|
||||
|
||||
raise SystemExit("模型拉取失败：\n - " + "\n - ".join(errors))
|
||||
|
||||
|
||||
# ─────────────────────────── 模型分段 ───────────────────────────
|
||||
|
||||
class TokenEmbedding(torch.nn.Module):
|
||||
def __init__(self, embed_tokens):
|
||||
super().__init__()
|
||||
self.embed_tokens = embed_tokens
|
||||
|
||||
def forward(self, input_ids):
|
||||
return self.embed_tokens(input_ids)
|
||||
|
||||
|
||||
class Transformer(torch.nn.Module):
|
||||
"""28 层语言模型 + 前置 3 层 DeepStack 相加 + final norm + last-token 池化。
|
||||
|
||||
池化放在图里（而不是 Go）是有意的：last-token 的位置由 attention_mask 决定，
|
||||
一旦 Go 侧算错位置就会静默取到 padding 的 hidden，而向量照样归一化、照样
|
||||
能比余弦——那种错误只能靠与参考向量对比才能发现。
|
||||
"""
|
||||
|
||||
def __init__(self, lm):
|
||||
super().__init__()
|
||||
self.layers = lm.layers
|
||||
self.norm = lm.norm
|
||||
|
||||
def forward(self, hidden, deepstack_0, deepstack_1, deepstack_2,
|
||||
rotary_cos, rotary_sin, causal_mask):
|
||||
deep = (deepstack_0, deepstack_1, deepstack_2)
|
||||
for i, layer in enumerate(self.layers):
|
||||
hidden = layer(
|
||||
hidden_states=hidden,
|
||||
attention_mask=causal_mask,
|
||||
position_embeddings=(rotary_cos, rotary_sin),
|
||||
use_cache=False,
|
||||
)
|
||||
if i < 3:
|
||||
hidden = hidden + deep[i]
|
||||
return self.norm(hidden)[:, -1]
|
||||
|
||||
|
||||
class VisionTower(torch.nn.Module):
|
||||
"""视觉塔：固定 grid 的 patch 张量 -> 主视觉特征 + 3 层 DeepStack。
|
||||
|
||||
grid_thw 作为 buffer 固化。原因见模块 docstring：legacy tracer 无法把
|
||||
grid_thw 保留为运行时输入，写成输入只会得到一个实际不含该输入的图。
|
||||
"""
|
||||
|
||||
def __init__(self, visual, grid_thw):
|
||||
super().__init__()
|
||||
self.visual = visual
|
||||
self.register_buffer("grid_thw", grid_thw)
|
||||
|
||||
def forward(self, pixel_values):
|
||||
out = self.visual(pixel_values, grid_thw=self.grid_thw, return_dict=True)
|
||||
d = out.deepstack_features
|
||||
return d[0], d[1], d[2], out.pooler_output
|
||||
|
||||
|
||||
# ─────────────────────────── 输入构造 ───────────────────────────
|
||||
|
||||
def render_text(instruction: str, text: str) -> str:
|
||||
"""与 Go 侧 tokenizer.renderInstructionInput 逐字符一致。
|
||||
|
||||
指令放 system、正文放 user、以 assistant 起始符结尾。差一个特殊 token，
|
||||
last-token 池化取到的位置就变了，嵌入也就不同——而且不会报错。
|
||||
"""
|
||||
return (f"<|im_start|>system\n{instruction}<|im_end|>\n"
|
||||
f"<|im_start|>user\n{text}<|im_end|>\n<|im_start|>assistant\n")
|
||||
|
||||
|
||||
def text_position_ids(attention_mask):
|
||||
"""无 padding 的单批语义下，三个 RoPE 轴都等于累计可见位置。"""
|
||||
pos = attention_mask.long().cumsum(-1) - 1
|
||||
return pos.clamp(min=0).unsqueeze(0).expand(3, -1, -1)
|
||||
|
||||
|
||||
def text_inputs(processor, lm, text):
|
||||
rendered = render_text(INSTRUCTION, text)
|
||||
x = processor.tokenizer([rendered], return_tensors="pt", truncation=True,
|
||||
max_length=MAX_LENGTH, padding=True)
|
||||
pos = text_position_ids(x["attention_mask"])
|
||||
hidden = lm.embed_tokens(x["input_ids"])
|
||||
cos, sin = lm.rotary_emb(hidden, pos)
|
||||
zero = torch.zeros_like(hidden)
|
||||
return x, hidden, (zero, zero, zero), cos, sin, causal_mask(hidden.shape[1]), pos
|
||||
|
||||
|
||||
def image_inputs(processor, model, image):
|
||||
conv = [{"role": "system", "content": [{"type": "text", "text": INSTRUCTION}]},
|
||||
{"role": "user", "content": [{"type": "image", "image": image}]}]
|
||||
rendered = processor.apply_chat_template([conv], add_generation_prompt=True, tokenize=False)
|
||||
x = processor(text=rendered, images=[image], do_resize=False,
|
||||
return_tensors="pt", truncation=True, max_length=MAX_LENGTH)
|
||||
lm = model.model.language_model
|
||||
with torch.no_grad():
|
||||
vo = model.model.visual(x["pixel_values"], grid_thw=x["image_grid_thw"], return_dict=True)
|
||||
pos, _ = model.model.get_rope_index(
|
||||
x["input_ids"], x["mm_token_type_ids"],
|
||||
image_grid_thw=x["image_grid_thw"], attention_mask=x["attention_mask"])
|
||||
mask = x["mm_token_type_ids"] == 1
|
||||
hidden = lm.embed_tokens(x["input_ids"])
|
||||
hidden = hidden.clone()
|
||||
hidden[mask] = vo.pooler_output
|
||||
deep = []
|
||||
for d in vo.deepstack_features:
|
||||
full = torch.zeros_like(hidden)
|
||||
full[mask] = d
|
||||
deep.append(full)
|
||||
cos, sin = lm.rotary_emb(hidden, pos)
|
||||
return x, hidden, tuple(deep), cos, sin, causal_mask(hidden.shape[1]), pos
|
||||
|
||||
|
||||
def causal_mask(seq: int) -> torch.Tensor:
|
||||
m = torch.full((1, 1, seq, seq), torch.finfo(torch.float32).min)
|
||||
return torch.triu(m, diagonal=1)
|
||||
|
||||
|
||||
def pool_last(hidden, mask):
|
||||
last = mask.shape[1] - mask.flip(1).argmax(1) - 1
|
||||
return hidden[torch.arange(hidden.shape[0], device=hidden.device), last]
|
||||
|
||||
|
||||
# ─────────────────────────── 校验 ───────────────────────────
|
||||
|
||||
def compare(name: str, a, b, floor: float = 0.999999) -> float:
|
||||
a = torch.nn.functional.normalize(torch.as_tensor(a).float(), dim=-1)
|
||||
b = torch.nn.functional.normalize(torch.as_tensor(b).float(), dim=-1)
|
||||
cos = float((a * b).sum())
|
||||
md = float((a - b).abs().max())
|
||||
log(f" {name}: cos={cos:.9f} maxdiff={md:.3e}")
|
||||
if cos < floor:
|
||||
raise SystemExit(f"导出校验失败：{name} 与完整模型不等价 (cos={cos:.9f})")
|
||||
return cos
|
||||
|
||||
|
||||
def verify_split_torch(model, processor, transformer) -> None:
|
||||
log("校验①：分段 PyTorch vs 完整模型")
|
||||
x, h, d, cos, sin, cm, pos = text_inputs(processor, model.model.language_model, REFERENCE_TEXT)
|
||||
with torch.no_grad():
|
||||
got = transformer(h, *d, cos, sin, cm)
|
||||
ref = model.model(input_ids=x["input_ids"], attention_mask=x["attention_mask"],
|
||||
position_ids=pos, use_cache=False).last_hidden_state[:, -1]
|
||||
compare("text/split-torch", got, ref)
|
||||
|
||||
img = reference_image()
|
||||
x, h, d, cos, sin, cm, _ = image_inputs(processor, model, img)
|
||||
with torch.no_grad():
|
||||
got = transformer(h, *d, cos, sin, cm)
|
||||
ref = model.model(input_ids=x["input_ids"], attention_mask=x["attention_mask"],
|
||||
pixel_values=x["pixel_values"], image_grid_thw=x["image_grid_thw"],
|
||||
mm_token_type_ids=x["mm_token_type_ids"],
|
||||
use_cache=False).last_hidden_state[:, -1]
|
||||
compare("image/split-torch", got, ref)
|
||||
log(f" 形状: seq={h.shape[1]} visual={(x['mm_token_type_ids'] == 1).sum().item()}")
|
||||
|
||||
|
||||
def unit(vec):
|
||||
"""L2 归一化。
|
||||
|
||||
必须对**写出的参考向量**归一化：ONNX 图返回的是 final norm 之后的原始 last hidden，
|
||||
而 Go 侧的 public 接口返回的是归一化后的向量。若参考用原始值，Go 测试会全线不匹配——
|
||||
且这个差异看起来像“模型不对”，实际上只是两边对“向量”的定义不同。
|
||||
"""
|
||||
v = np.asarray(vec, dtype=np.float64)
|
||||
n = float(np.linalg.norm(v))
|
||||
return v / n if n > 0 else v
|
||||
|
||||
|
||||
def verify_onnx(out_dir: str, model, processor) -> dict:
|
||||
"""用 onnxruntime 跑导出后的三段图，对比完整模型前向。
|
||||
|
||||
返回参考向量（供 Go 侧测试冻结使用）：Go 必须复现同一套预处理与模板，
|
||||
因此这里把同一输入下的期望向量前若干维导出。
|
||||
"""
|
||||
import onnxruntime as ort
|
||||
|
||||
log("校验②：导出后的 ONNX 三段图 vs 完整模型")
|
||||
ts = ort.InferenceSession(os.path.join(out_dir, "TokenEmbedding.onnx"), providers=["CPUExecutionProvider"])
|
||||
xs = ort.InferenceSession(os.path.join(out_dir, "Transformer.onnx"), providers=["CPUExecutionProvider"])
|
||||
vs = ort.InferenceSession(os.path.join(out_dir, "Vision.onnx"), providers=["CPUExecutionProvider"])
|
||||
lm = model.model.language_model
|
||||
|
||||
reference: dict[str, object] = {}
|
||||
|
||||
def run_transform(hidden, deep, cos, sin):
|
||||
seq = hidden.shape[1]
|
||||
return xs.run(None, {
|
||||
"hidden": hidden.astype(np.float32),
|
||||
"deepstack_0": deep[0].astype(np.float32),
|
||||
"deepstack_1": deep[1].astype(np.float32),
|
||||
"deepstack_2": deep[2].astype(np.float32),
|
||||
"rotary_cos": cos.astype(np.float32),
|
||||
"rotary_sin": sin.astype(np.float32),
|
||||
"causal_mask": causal_mask(seq).numpy(),
|
||||
})[0]
|
||||
|
||||
x, h, d, cos, sin, _cm, pos = text_inputs(processor, lm, REFERENCE_TEXT)
|
||||
with torch.no_grad():
|
||||
ref_text = model.model(input_ids=x["input_ids"], attention_mask=x["attention_mask"],
|
||||
position_ids=pos, use_cache=False).last_hidden_state[:, -1].numpy()
|
||||
h_onnx = ts.run(None, {"input_ids": x["input_ids"].numpy().astype(np.int64)})[0]
|
||||
zero = np.zeros_like(h_onnx)
|
||||
got = run_transform(h_onnx, (zero, zero, zero), cos.numpy(), sin.numpy())
|
||||
compare("text/onnx-vs-full", got, ref_text)
|
||||
reference["text"] = REFERENCE_TEXT
|
||||
reference["text_vector_prefix"] = [float(v) for v in unit(got[0])[:12]]
|
||||
reference["text_norm_raw"] = float(np.linalg.norm(got[0]))
|
||||
|
||||
img = reference_image()
|
||||
x, _h, _d, cos, sin, _cm, _ = image_inputs(processor, model, img)
|
||||
with torch.no_grad():
|
||||
ref_img = model.model(input_ids=x["input_ids"], attention_mask=x["attention_mask"],
|
||||
pixel_values=x["pixel_values"], image_grid_thw=x["image_grid_thw"],
|
||||
mm_token_type_ids=x["mm_token_type_ids"],
|
||||
use_cache=False).last_hidden_state[:, -1].numpy()
|
||||
h_onnx = ts.run(None, {"input_ids": x["input_ids"].numpy().astype(np.int64)})[0]
|
||||
vo = vs.run(None, {"pixel_values": x["pixel_values"].numpy().astype(np.float32)})
|
||||
mask = x["mm_token_type_ids"].numpy() == 1
|
||||
h_onnx = h_onnx.copy()
|
||||
h_onnx[mask] = vo[3]
|
||||
deep = []
|
||||
for d in vo[:3]:
|
||||
full = np.zeros_like(h_onnx)
|
||||
full[mask] = d
|
||||
deep.append(full)
|
||||
got = run_transform(h_onnx, tuple(deep), cos.numpy(), sin.numpy())
|
||||
compare("image/onnx-vs-full", got, ref_img)
|
||||
reference["image_rgb"] = list(REFERENCE_IMAGE_RGB)
|
||||
reference["image_size"] = IMAGE_SIZE
|
||||
reference["image_vector_prefix"] = [float(v) for v in unit(got[0])[:12]]
|
||||
reference["image_norm_raw"] = float(np.linalg.norm(got[0]))
|
||||
reference["dim"] = int(got.shape[1])
|
||||
return reference
|
||||
|
||||
|
||||
def reference_image():
|
||||
from PIL import Image
|
||||
|
||||
return Image.new("RGB", (IMAGE_SIZE, IMAGE_SIZE), REFERENCE_IMAGE_RGB)
|
||||
|
||||
|
||||
# ─────────────────────────── 导出 ───────────────────────────
|
||||
|
||||
def export_graphs(out_dir: str, model, processor, model_dir: str, transformer) -> None:
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
# 清掉旧产物，避免 fingerprint 把死文件算进去（旧图/旧外部权重会让
|
||||
# 空间指纹变化，触发一次毫无意义的全量重算）。
|
||||
for name in os.listdir(out_dir):
|
||||
p = os.path.join(out_dir, name)
|
||||
if os.path.isfile(p):
|
||||
os.remove(p)
|
||||
|
||||
lm = model.model.language_model
|
||||
|
||||
log("导出 TokenEmbedding.onnx")
|
||||
ids = torch.tensor([[151643, 151643]], dtype=torch.long)
|
||||
with torch.no_grad():
|
||||
torch.onnx.export(
|
||||
TokenEmbedding(lm.embed_tokens).eval(), (ids,), os.path.join(out_dir, "TokenEmbedding.onnx"),
|
||||
input_names=["input_ids"], output_names=["hidden"],
|
||||
dynamic_axes={"input_ids": {1: "seq"}, "hidden": {1: "seq"}},
|
||||
opset_version=17, do_constant_folding=True, dynamo=False,
|
||||
)
|
||||
|
||||
log("导出 Transformer.onnx")
|
||||
x, h, d, cos, sin, cm, _ = image_inputs(processor, model, reference_image())
|
||||
with torch.no_grad():
|
||||
torch.onnx.export(
|
||||
transformer, (h, *d, cos, sin, cm), os.path.join(out_dir, "Transformer.onnx"),
|
||||
input_names=["hidden", "deepstack_0", "deepstack_1", "deepstack_2",
|
||||
"rotary_cos", "rotary_sin", "causal_mask"],
|
||||
output_names=["embedding"],
|
||||
dynamic_axes={
|
||||
"hidden": {1: "seq"}, "deepstack_0": {1: "seq"}, "deepstack_1": {1: "seq"},
|
||||
"deepstack_2": {1: "seq"}, "rotary_cos": {1: "seq"}, "rotary_sin": {1: "seq"},
|
||||
"causal_mask": {2: "seq", 3: "seq"},
|
||||
},
|
||||
opset_version=17, do_constant_folding=True, dynamo=False,
|
||||
)
|
||||
|
||||
log("导出 Vision.onnx（固定 grid 1×48×48）")
|
||||
grid = torch.tensor([[1, IMAGE_SIZE // PATCH_SIZE, IMAGE_SIZE // PATCH_SIZE]], dtype=torch.long)
|
||||
pv = x["pixel_values"]
|
||||
with torch.no_grad():
|
||||
torch.onnx.export(
|
||||
VisionTower(model.model.visual, grid).eval(), (pv,), os.path.join(out_dir, "Vision.onnx"),
|
||||
input_names=["pixel_values"],
|
||||
output_names=["deepstack_feature_0", "deepstack_feature_1",
|
||||
"deepstack_feature_2", "vision_hidden_states"],
|
||||
opset_version=17, do_constant_folding=True, dynamo=False,
|
||||
)
|
||||
|
||||
for name in ("tokenizer.json", "tokenizer_config.json", "chat_template.jinja", "added_tokens.json"):
|
||||
src = os.path.join(model_dir, name)
|
||||
if os.path.exists(src):
|
||||
shutil.copy2(src, os.path.join(out_dir, name))
|
||||
|
||||
|
||||
def write_config(out_dir: str, model, processor) -> None:
|
||||
cfg = model.config
|
||||
text_cfg = getattr(cfg, "text_config", cfg)
|
||||
rope_scaling = getattr(text_cfg, "rope_scaling", None) or {}
|
||||
mrope_section = rope_scaling.get("mrope_section") or [24, 20, 20]
|
||||
vision = cfg.vision_config
|
||||
meta = {
|
||||
"arch": "qwen3-vl-embedding-2b-multimodal",
|
||||
"runtime": "homeagent-onnx-three-part",
|
||||
"dim": int(getattr(text_cfg, "hidden_size", 2048)),
|
||||
"max_length": MAX_LENGTH,
|
||||
"instruction": INSTRUCTION,
|
||||
"pooling": "last_token",
|
||||
"normalize": True,
|
||||
"image_size": IMAGE_SIZE,
|
||||
"patch_size": int(vision.patch_size),
|
||||
"temporal_patch_size": int(vision.temporal_patch_size),
|
||||
"spatial_merge_size": int(vision.spatial_merge_size),
|
||||
"image_mean": [0.5, 0.5, 0.5],
|
||||
"image_std": [0.5, 0.5, 0.5],
|
||||
"rope_theta": float(getattr(text_cfg, "rope_theta", 5000000)),
|
||||
"mrope_section": [int(v) for v in mrope_section],
|
||||
"num_layers": int(getattr(text_cfg, "num_hidden_layers", 28)),
|
||||
"supports_native_video": False,
|
||||
"unsupported_modalities": ["audio", "video"],
|
||||
"notes": ("视频由上层抽帧后逐帧按图像编码（同模型/同维度/同 fingerprint）；"
|
||||
"音频需未来接入真正的统一音频模型。grid_thw 被 legacy tracer 固化为常量，"
|
||||
"故视觉塔固定 1×48×48，详见导出脚本 docstring。"),
|
||||
}
|
||||
with open(os.path.join(out_dir, "embed_config.json"), "w") as f:
|
||||
json.dump(meta, f, ensure_ascii=False, indent=2)
|
||||
log(f"写出 embed_config.json dim={meta['dim']} rope_theta={meta['rope_theta']} "
|
||||
f"mrope={meta['mrope_section']}")
|
||||
|
||||
|
||||
def main() -> int:
|
||||
ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
ap.add_argument("--out", required=True, help="ONNX 产物目录（写入约 8GB）")
|
||||
ap.add_argument("--model-dir", default="", help="已下载的模型目录；给出则跳过自动拉取")
|
||||
ap.add_argument("--model-id", default=DEFAULT_MODEL_ID, help=f"模型仓库 id（默认 {DEFAULT_MODEL_ID}）")
|
||||
ap.add_argument("--model-store", default=os.path.expanduser("~/.cache/homeagent-qwen-models"),
|
||||
help="自动拉取时的模型存放目录")
|
||||
ap.add_argument("--no-reference", action="store_true",
|
||||
help="不写 <out>/qwen_reference.json（默认会写；Go 测试靠它做冻结回归）")
|
||||
ap.add_argument("--skip-verify", action="store_true", help="跳过导出后校验（仅调试用，不推荐）")
|
||||
ap.add_argument("--verify-only", action="store_true",
|
||||
help="不重新导出，只校验已存在的 <out> 并（重新）写出参考向量")
|
||||
args = ap.parse_args()
|
||||
|
||||
if args.verify_only:
|
||||
# 校验既有产物目录：既能确认线上在用的图没坏，也能给旧目录补上参考向量。
|
||||
for name in ("TokenEmbedding.onnx", "Transformer.onnx", "Vision.onnx"):
|
||||
if not os.path.exists(os.path.join(args.out, name)):
|
||||
raise SystemExit(f"{args.out} 下缺少 {name}，无法 --verify-only")
|
||||
model_dir = args.model_dir or pull_model(args.model_id, args.model_store)
|
||||
from transformers import AutoProcessor
|
||||
from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLForConditionalGeneration
|
||||
|
||||
log("加载 processor / model（--verify-only）")
|
||||
processor = AutoProcessor.from_pretrained(model_dir, trust_remote_code=True, padding_side="right")
|
||||
model = Qwen3VLForConditionalGeneration.from_pretrained(
|
||||
model_dir, dtype=torch.float32, low_cpu_mem_usage=True).eval()
|
||||
reference = verify_onnx(args.out, model, processor)
|
||||
if not args.no_reference:
|
||||
path = os.path.join(args.out, "qwen_reference.json")
|
||||
with open(path, "w") as f:
|
||||
json.dump(reference, f, ensure_ascii=False, indent=2)
|
||||
log(f"写出冻结参考向量: {path}")
|
||||
log(f"校验完成: {args.out}")
|
||||
return 0
|
||||
|
||||
if args.model_dir:
|
||||
model_dir = args.model_dir
|
||||
if not os.path.exists(os.path.join(model_dir, "config.json")):
|
||||
raise SystemExit(f"--model-dir {model_dir} 下没有 config.json")
|
||||
log(f"使用本地模型: {model_dir}")
|
||||
else:
|
||||
model_dir = pull_model(args.model_id, args.model_store)
|
||||
|
||||
# transformers 只在真正导出时才需要（拉取模型本身只用 huggingface_hub）。
|
||||
from PIL import Image # noqa: F401 确保依赖存在并给出清晰报错
|
||||
from transformers import AutoProcessor
|
||||
from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLForConditionalGeneration
|
||||
|
||||
log("加载 processor / model（FP32，CPU）")
|
||||
processor = AutoProcessor.from_pretrained(model_dir, trust_remote_code=True, padding_side="right")
|
||||
model = Qwen3VLForConditionalGeneration.from_pretrained(
|
||||
model_dir, dtype=torch.float32, low_cpu_mem_usage=True).eval()
|
||||
transformer = Transformer(model.model.language_model).eval()
|
||||
|
||||
verify_split_torch(model, processor, transformer)
|
||||
export_graphs(args.out, model, processor, model_dir, transformer)
|
||||
write_config(args.out, model, processor)
|
||||
|
||||
reference = None if args.skip_verify else verify_onnx(args.out, model, processor)
|
||||
# 参考写进产物目录本身：这样任何一个 ONNX 目录都自带「它应当给出什么输出」，
|
||||
# Go 测试无需额外配置就能找到，也不会出现「模型换了、参考还是旧的」的错配。
|
||||
if reference is not None and not args.no_reference:
|
||||
path = os.path.join(args.out, "qwen_reference.json")
|
||||
with open(path, "w") as f:
|
||||
json.dump(reference, f, ensure_ascii=False, indent=2)
|
||||
log(f"写出冻结参考向量: {path}")
|
||||
|
||||
total = sum(os.path.getsize(os.path.join(args.out, n))
|
||||
for n in os.listdir(args.out) if os.path.isfile(os.path.join(args.out, n)))
|
||||
log(f"完成: {args.out}（{total / 2**30:.2f} GiB）")
|
||||
log("Go 侧用法: core.memory.multimodal_space.type=onnx + "
|
||||
f"core.memory.multimodal_space.onnx.model_dir={args.out}")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
192
scripts/tfidf_sidecar.py
Normal file
192
scripts/tfidf_sidecar.py
Normal file
@ -0,0 +1,192 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
TF-IDF embedding sidecar for HomeAgent.
|
||||
|
||||
HTTP interface:
|
||||
POST /embed → {"text":"...", "side":"query|document"} → {"embedding":[...]}
|
||||
POST /train → {"documents": {"id":"text", ...}} → {"status":"ok", "count":N}
|
||||
POST /search → {"query":"...", "topK":5} → {"results":[{"id":"...", "score":0.3, "text":"..."}]}
|
||||
GET /health → {"status":"ok","count":N,"type":"tfidf"}
|
||||
|
||||
Provides TF-IDF sparse vectors via the same HTTP contract as Jina/ONNX dense vectors,
|
||||
allowing the core to treat all embedding providers uniformly.
|
||||
"""
|
||||
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
from collections import Counter
|
||||
from http.server import HTTPServer, BaseHTTPRequestHandler
|
||||
from urllib.parse import urlparse
|
||||
|
||||
# ─── Chinese Tokenizer (jieba) ──────────────────────────────────────────────
|
||||
try:
|
||||
import jieba
|
||||
def tokenize(text):
|
||||
return [w for w in jieba.cut(text) if w.strip()]
|
||||
except ImportError:
|
||||
# Fallback: simple character/word split
|
||||
def tokenize(text):
|
||||
return re.findall(r'[\w\u4e00-\u9fff]+', text)
|
||||
|
||||
# ─── TF-IDF Engine ──────────────────────────────────────────────────────────
|
||||
class TFIDFEngine:
|
||||
def __init__(self):
|
||||
self.lock = threading.RLock()
|
||||
self.doc_freq = Counter()
|
||||
self.total_docs = 0
|
||||
self.docs = {} # id -> {"text": ..., "vec": ...}
|
||||
|
||||
def train(self, docs: dict):
|
||||
"""Train IDF statistics from a batch of documents."""
|
||||
with self.lock:
|
||||
self.total_docs = len(docs)
|
||||
self.doc_freq = Counter()
|
||||
seen_per_doc = []
|
||||
for text in docs.values():
|
||||
features = tokenize(text)
|
||||
seen = set()
|
||||
for f in features:
|
||||
if f not in seen:
|
||||
self.doc_freq[f] += 1
|
||||
seen.add(f)
|
||||
# Store docs with their vectors
|
||||
for did, text in docs.items():
|
||||
self.docs[did] = {"text": text, "vec": self._vectorize(text)}
|
||||
print(f"[tfidf] trained on {len(docs)} docs, {len(self.doc_freq)} features", flush=True)
|
||||
|
||||
def _vectorize(self, text):
|
||||
features = tokenize(text)
|
||||
tf = Counter(features)
|
||||
max_tf = max(tf.values()) if tf else 1
|
||||
vec = {}
|
||||
for f, count in tf.items():
|
||||
tf_norm = count / max_tf
|
||||
if self.total_docs < 3:
|
||||
vec[f] = tf_norm
|
||||
continue
|
||||
df = self.doc_freq.get(f, 0)
|
||||
if df <= 0:
|
||||
continue
|
||||
idf = math.log((self.total_docs + 1) / (df + 1))
|
||||
if idf < 0.1:
|
||||
continue
|
||||
vec[f] = tf_norm * idf
|
||||
return vec
|
||||
|
||||
def vectorize(self, text):
|
||||
with self.lock:
|
||||
return self._vectorize(text)
|
||||
|
||||
def search(self, query, top_k=5):
|
||||
with self.lock:
|
||||
q_vec = self._vectorize(query)
|
||||
results = []
|
||||
for did, entry in self.docs.items():
|
||||
score = self._cosine(q_vec, entry["vec"])
|
||||
if score > 0.01:
|
||||
results.append({"id": did, "score": score, "text": entry["text"][:200]})
|
||||
results.sort(key=lambda x: -x["score"])
|
||||
return results[:top_k]
|
||||
|
||||
def add(self, did, text):
|
||||
with self.lock:
|
||||
self.docs[did] = {"text": text, "vec": self._vectorize(text)}
|
||||
|
||||
def remove(self, did):
|
||||
with self.lock:
|
||||
self.docs.pop(did, None)
|
||||
|
||||
@staticmethod
|
||||
def _cosine(a, b):
|
||||
dot = sum(a.get(k, 0) * b.get(k, 0) for k in set(a) | set(b))
|
||||
na = math.sqrt(sum(v * v for v in a.values()))
|
||||
nb = math.sqrt(sum(v * v for v in b.values()))
|
||||
if na == 0 or nb == 0:
|
||||
return 0
|
||||
return dot / (na * nb)
|
||||
|
||||
# ─── HTTP Server ─────────────────────────────────────────────────────────────
|
||||
PORT = int(os.environ.get("TFIDF_PORT", "18998"))
|
||||
engine = TFIDFEngine()
|
||||
|
||||
class TFIDFHandler(BaseHTTPRequestHandler):
|
||||
def log_message(self, fmt, *args):
|
||||
if "/health" not in str(args[0]):
|
||||
print(f"[tfidf] {fmt % args}", flush=True)
|
||||
|
||||
def do_GET(self):
|
||||
if urlparse(self.path).path == "/health":
|
||||
self._respond(200, {"status": "ok", "count": len(engine.docs), "type": "tfidf"})
|
||||
else:
|
||||
self._respond(404, {"error": "not found"})
|
||||
|
||||
def do_POST(self:
|
||||
path = urlparse(self.path).path
|
||||
try:
|
||||
length = int(self.headers.get("Content-Length", 0))
|
||||
body = json.loads(self.rfile.read(length))
|
||||
except Exception as e:
|
||||
self._respond(400, {"error": str(e)})
|
||||
return
|
||||
|
||||
if path == "/train":
|
||||
docs = body.get("documents", {})
|
||||
engine.train(docs)
|
||||
self._respond(200, {"status": "ok", "count": len(docs)})
|
||||
|
||||
elif path == "/embed":
|
||||
text = body.get("text", "")
|
||||
vec = engine.vectorize(text)
|
||||
# Convert to list format matching HTTPEmbedder contract
|
||||
# Sparse vector → dense-ish list (feature indices as keys)
|
||||
embedding = [vec.get(k, 0) for k in sorted(vec.keys())] if vec else []
|
||||
self._respond(200, {"embedding": embedding, "sparse": vec})
|
||||
|
||||
elif path == "/search":
|
||||
query = body.get("query", "")
|
||||
top_k = body.get("topK", 5)
|
||||
results = engine.search(query, top_k)
|
||||
self._respond(200, {"results": results})
|
||||
|
||||
elif path == "/add":
|
||||
did = body.get("id", "")
|
||||
text = body.get("text", "")
|
||||
engine.add(did, text)
|
||||
self._respond(200, {"status": "ok"})
|
||||
|
||||
elif path == "/remove":
|
||||
did = body.get("id", "")
|
||||
engine.remove(did)
|
||||
self._respond(200, {"status": "ok"})
|
||||
|
||||
else:
|
||||
self._respond(404, {"error": "not found"})
|
||||
|
||||
def _respond(self, status, data):
|
||||
body = json.dumps(data).encode("utf-8")
|
||||
self.send_response(status)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.send_header("Content-Length", str(len(body)))
|
||||
self.end_headers()
|
||||
self.wfile.write(body)
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Load from disk if available
|
||||
data_file = os.environ.get("TFIDF_DATA", "/home/newqqagent/memory/tfidf_index.json")
|
||||
if os.path.exists(data_file):
|
||||
try:
|
||||
with open(data_file) as f:
|
||||
docs = json.load(f)
|
||||
engine.train(docs)
|
||||
print(f"[tfidf] loaded {len(docs)} docs from {data_file}", flush=True)
|
||||
except Exception as e:
|
||||
print(f"[tfidf] failed to load: {e}", flush=True)
|
||||
|
||||
server = HTTPServer(("0.0.0.0", PORT), TFIDFHandler)
|
||||
print(f"[tfidf] listening on :{PORT}", flush=True)
|
||||
server.serve_forever()
|
||||
106
scripts/verify_deploy.sh
Executable file
106
scripts/verify_deploy.sh
Executable file
@ -0,0 +1,106 @@
|
||||
#!/usr/bin/env bash
|
||||
# HomeAgent 部署验证脚本：检查 healthcheck 残留 / 记忆去重 / archived 残留 / 嵌入规格
|
||||
# 用法: verify_deploy.sh [data_dir] 默认 /home/newqqagent
|
||||
set -u
|
||||
DATA_DIR="${1:-/home/newqqagent}"
|
||||
PASS=0; FAIL=0; SKIP=0
|
||||
ok() { echo "[PASS] $1"; PASS=$((PASS+1)); }
|
||||
bad() { echo "[FAIL] $1"; FAIL=$((FAIL+1)); }
|
||||
skip() { echo "[SKIP] $1"; SKIP=$((SKIP+1)); }
|
||||
|
||||
echo "== HomeAgent 部署验证 =="
|
||||
echo "数据目录: $DATA_DIR"
|
||||
echo "时间: $(date '+%Y-%m-%d %H:%M:%S')"
|
||||
echo "模式: 只读检测，不修改/删除任何数据"
|
||||
echo
|
||||
|
||||
# ---- 0. 数据目录可用性 ----
|
||||
if [ ! -d "$DATA_DIR" ]; then
|
||||
echo "[FATAL] 数据目录不存在: $DATA_DIR"
|
||||
exit 2
|
||||
fi
|
||||
|
||||
# ---- 1. healthcheck 残留检查 ----
|
||||
echo "--- 1. healthcheck 残留 ---"
|
||||
hc_files=$(find "$DATA_DIR/knowledge" "$DATA_DIR/memory" -maxdepth 2 \
|
||||
\( -name '*_hc_*' -o -name 'gotest' -o -name 'luatest' \) 2>/dev/null)
|
||||
if [ -z "$hc_files" ]; then
|
||||
ok "无 _hc_*/gotest/luatest 残留"
|
||||
else
|
||||
echo " 发现残留:"
|
||||
echo "$hc_files" | sed 's/^/ /'
|
||||
bad "存在 healthcheck 残留（上列，需人工确认清理）"
|
||||
fi
|
||||
|
||||
# ---- 2. relations 去重检查 ----
|
||||
echo "--- 2. relations 去重 ---"
|
||||
GRAPH_DB="$DATA_DIR/memory/graph.db"
|
||||
if [ -f "$GRAPH_DB" ] && command -v sqlite3 >/dev/null 2>&1; then
|
||||
dup=$(sqlite3 "$GRAPH_DB" "SELECT count(*) - count(DISTINCT source_id||'|'||target_id||'|'||relation_type||'|'||COALESCE(session_id,'')) FROM relations;" 2>/dev/null)
|
||||
total=$(sqlite3 "$GRAPH_DB" "SELECT count(*) FROM relations;" 2>/dev/null)
|
||||
if [ -n "$dup" ] && [ "$dup" -le 0 ] 2>/dev/null; then
|
||||
ok "relations 去重良好 (共 ${total:-0} 条, 重复 ${dup:-0})"
|
||||
else
|
||||
bad "relations 存在重复: 共 ${total:-?} 条, 重复 ${dup:-?}"
|
||||
fi
|
||||
# UNIQUE 索引存在性：查 relations 表 DDL 是否含 UNIQUE 约束
|
||||
# （注意：UNIQUE 约束创建的 sqlite_autoindex_* 索引其 sql 字段为 NULL，
|
||||
# 用 LIKE 查索引 sql 会漏报——须查表 DDL）
|
||||
ddl=$(sqlite3 "$GRAPH_DB" "SELECT sql FROM sqlite_master WHERE type='table' AND name='relations';" 2>/dev/null)
|
||||
if echo "$ddl" | grep -q 'UNIQUE'; then
|
||||
# 进一步确认唯一约束覆盖的列
|
||||
auto_idx=$(sqlite3 "$GRAPH_DB" "SELECT name FROM sqlite_master WHERE type='index' AND name LIKE 'sqlite_autoindex_relations%' AND tbl_name='relations';" 2>/dev/null | head -1)
|
||||
if [ -n "$auto_idx" ]; then
|
||||
cols=$(sqlite3 "$GRAPH_DB" "SELECT group_concat(name,',') FROM pragma_index_info('$auto_idx');" 2>/dev/null)
|
||||
echo " UNIQUE 自动索引: $auto_idx, 覆盖列: ${cols:-?}"
|
||||
fi
|
||||
ok "relations 复合 UNIQUE 约束存在"
|
||||
else
|
||||
bad "relations 复合 UNIQUE 约束缺失（迁移未生效？）"
|
||||
fi
|
||||
else
|
||||
skip "graph.db 不存在或 sqlite3 不可用 ($GRAPH_DB)"
|
||||
fi
|
||||
|
||||
# ---- 3. archived 模板垃圾检查 ----
|
||||
echo "--- 3. archived 残留 ---"
|
||||
if [ -f "$GRAPH_DB" ] && command -v sqlite3 >/dev/null 2>&1; then
|
||||
# 兜底检测：relation_type 直接为 context_archived，或 target_id 含 context_archived
|
||||
arch=$(sqlite3 "$GRAPH_DB" "SELECT count(*) FROM relations WHERE relation_type='context_archived' OR relation_type='来源' AND target_id LIKE '%context_archived%';" 2>/dev/null)
|
||||
if [ -n "$arch" ] && [ "$arch" -eq 0 ] 2>/dev/null; then
|
||||
ok "无 context_archived 模板垃圾"
|
||||
else
|
||||
bad "context_archived 模板垃圾: ${arch:-?} 条"
|
||||
fi
|
||||
else
|
||||
skip "graph.db 不存在或 sqlite3 不可用"
|
||||
fi
|
||||
|
||||
# ---- 4. 嵌入规格检查 ----
|
||||
echo "--- 4. 嵌入规格 ---"
|
||||
CONFIG_DB="$DATA_DIR/config.db"
|
||||
if [ -f "$CONFIG_DB" ] && command -v sqlite3 >/dev/null 2>&1; then
|
||||
emb=$(sqlite3 "$CONFIG_DB" "SELECT value FROM config WHERE key='core.agent.embedding_model_path';" 2>/dev/null)
|
||||
echo " embedding_model_path = ${emb:-（未配置）}"
|
||||
if [ -n "$emb" ]; then
|
||||
# 模型路径含 #topN 则裁剪生效；否则全量加载（300 维 ≈ 1.5G/模型）
|
||||
if [[ "$emb" == *"#top"* ]]; then
|
||||
ok "嵌入模型已启用 #topN 裁剪（内存预算受控）"
|
||||
else
|
||||
# 逗号分隔的模型个数估算内存
|
||||
n=$(echo "$emb" | awk -F',' '{print NF}')
|
||||
est=$(( n * 1500 ))
|
||||
echo " 全量加载 ${n} 个模型，估算常驻 ≈ ${est}MB（未裁剪）"
|
||||
bad "嵌入模型未启用 #topN 裁剪（内存预算偏高）"
|
||||
fi
|
||||
else
|
||||
skip "embedding_model_path 未配置（可能是默认值）"
|
||||
fi
|
||||
else
|
||||
skip "config.db 不存在或 sqlite3 不可用 ($CONFIG_DB)"
|
||||
fi
|
||||
|
||||
# ---- 汇总 ----
|
||||
echo
|
||||
echo "== 汇总: PASS=$PASS FAIL=$FAIL SKIP=$SKIP =="
|
||||
[ "$FAIL" -eq 0 ] && exit 0 || exit 1
|
||||
Reference in New Issue
Block a user