mirror of
https://gitcode.com/JianFeeeee/HomeAgent.git
synced 2026-10-13 19:12:39 +00:00
refactor(memory): 拆除描述式媒体索引，媒体成为一等块并按原生向量融合
背景：此前媒体是靠「生成的描述文本」将就进记忆的——写 marker 进正文、 再由正则反解成 media_refs 与图库里的 type=Media 实体。这条链路有三个 致命缺陷：描述由异步模型生成（未生成前媒体等于不存在）、语义检索实质上 只搜描述文字、图库里的「媒体节点」是描述文本的投影而不是媒体本身。 本提交把这条链路整体拆除，媒体改为按自己的原生向量参与记忆： 一、描述链彻底删除（无残留、无兼容分支） - media.Item 去掉 Description/DescribedBy 与对应列； - 删除 Store.Describe / Store.Search / Store.Pending； - 删除 Agent.mediaDescribeLoop / describePendingMedia 与配置项 core.memory.media.describe_on_ingest； - SDK 侧 MediaAttachment 去掉 Description（见 SDK 仓独立提交）。 二、marker 机制删除，媒体归属改为结构化块边 - 删除 mediaMarkerLine/parseMediaMarkers/mediaEntityName/mediaTriplesFromText/ extractMediaDigests/sentenceWithMediaMarkers/docMediaContext； - memory.Triple 新增 MediaDigests 结构化字段；句子文本保持原样， 不再被 marker 污染； - 块以 sentence --contains--> block / document --contains--> block 结构边 挂到承载节点（新增 documents 表与 document 节点种类）； - 模型未给原句时用「主谓宾。」拼一句自然语言作落点，不造 marker 文本。 三、旧数据迁移（幂等） - 新增 GraphDB.MigrateLegacyMediaEntities：把 type=Media 的旧实体按短 digest 还原成原生块、挂回原句子、删除旧实体与描述关系；Agent 启动时执行； - CleanupOrphanedSentences 同时看关系引用与块边，避免把只靠块存活的句子 连同块边一起删掉。 四、向量融合：媒体按图本身被召回 - 新增 vector.FuseVectors（逐维求和 + L2 归一化）； - Doc.DenseVec = 文本向量 ⊕ 文档块的媒体向量（同 fingerprint 才融合）， 新增 Doc.DenseFP，指纹变化触发重算； - ContextEvent.DenseVec 同理融合事件块；事件新增 DenseFP，Prune 只在 同一统一空间内比稠密余弦； - 跨模态视觉路只召回「仍被某层记忆块持有」的媒体，CAS 全库字节不再 直接充当记忆检索结果。 五、同时纳入本分支既有的嵌入基础改造（此前工作区未提交，缺它 HEAD 不可构建） - internal/tfidf 懒回退包、千问三段式多模态 ONNX 空间的 Go 侧 （qwen/embedder.go、image.go、model_input.go）、CLIP 移除、 sdk.NewStore 分词器签名与调用点、embed 侧车 systemd 单元。 验证：go build ./... 、go vet ./...（含 -tags medialive）均通过； 在 HEAD 的独立 worktree 上重放本次暂存集后 go test -short ./internal/... 全部通过（端口冲突类用例在隔离环境中亦通过）。未提交工作区中与本改造 无关的改动（HarmonyOS、waiter、devicebridge、plan.md 等）。
This commit is contained in:
JianFeeeee
@ -7,11 +7,11 @@ import (
|
||||
)
|
||||
|
||||
type bilingualEvent struct {
|
||||
idx int
|
||||
source string
|
||||
topic string
|
||||
text string // cleaned text for vectorization
|
||||
label string // short description
|
||||
idx int
|
||||
source string
|
||||
topic string
|
||||
text string // cleaned text for vectorization
|
||||
label string // short description
|
||||
}
|
||||
|
||||
func TestBilingualPruningAccuracy(t *testing.T) {
|
||||
@ -40,16 +40,16 @@ func TestBilingualPruningAccuracy(t *testing.T) {
|
||||
t.Logf("%s: %d words", cfg.name, len(e.words))
|
||||
|
||||
type scored struct {
|
||||
idx int
|
||||
topic string
|
||||
label string
|
||||
score float64
|
||||
idx int
|
||||
topic string
|
||||
label string
|
||||
score float64
|
||||
}
|
||||
|
||||
queries := []struct {
|
||||
q string
|
||||
qTopic string
|
||||
desc string
|
||||
q string
|
||||
qTopic string
|
||||
desc string
|
||||
}{
|
||||
{"老大说了关于 React 组件的事情", "老大私聊", "中英混合:老大+React"},
|
||||
{"帮我查一下 Nginx 反向代理配置", "服务器运维", "中英混合:Nginx+反向代理"},
|
||||
@ -189,10 +189,10 @@ func TestBilingualVectorizeClean(t *testing.T) {
|
||||
|
||||
func genBilingualEvents() []bilingualEvent {
|
||||
entries := []struct {
|
||||
topic string
|
||||
zh string // Chinese description
|
||||
en string // English terms mixed in
|
||||
source string
|
||||
topic string
|
||||
zh string // Chinese description
|
||||
en string // English terms mixed in
|
||||
source string
|
||||
}{
|
||||
{"大学招生", "河南医药大学录取分数线", "", "qq"},
|
||||
{"大学招生", "医学院专业排名", "medical university ranking", "agent"},
|
||||
|
||||
@ -40,7 +40,7 @@ type MemoryBlock struct {
|
||||
}
|
||||
|
||||
// MemoryBlockEdge 是 L3 中连接一等记忆节点的结构化语义边。
|
||||
// source/target kind 当前允许 block、entity、sentence。
|
||||
// source/target kind 当前允许 block、entity、sentence、document。
|
||||
type MemoryBlockEdge struct {
|
||||
ID int64 `json:"id"`
|
||||
SourceKind string `json:"source_kind"`
|
||||
@ -111,6 +111,19 @@ func (g *GraphDB) PutMemoryBlocks(blocks []MemoryBlock) error {
|
||||
return tx.Commit()
|
||||
}
|
||||
|
||||
// PutDocumentNode 在 L3 登记一个文档节点，作为 document --contains--> block
|
||||
// 结构边的端点。文档正文已蒸馏为实体/关系，这里只保留身份与摘要。
|
||||
func (g *GraphDB) PutDocumentNode(id, summary string) error {
|
||||
if id == "" {
|
||||
return fmt.Errorf("document node id is required")
|
||||
}
|
||||
g.mu.Lock()
|
||||
defer g.mu.Unlock()
|
||||
_, err := g.db.Exec(`INSERT INTO documents (id, summary) VALUES (?, ?)
|
||||
ON CONFLICT(id) DO UPDATE SET summary = excluded.summary`, id, summary)
|
||||
return err
|
||||
}
|
||||
|
||||
// MemoryBlocks 查询 Graph 层实际持有的一等记忆节点。
|
||||
func (g *GraphDB) MemoryBlocks() ([]MemoryBlock, error) {
|
||||
g.mu.RLock()
|
||||
@ -144,7 +157,7 @@ func (g *GraphDB) MemoryBlocks() ([]MemoryBlock, error) {
|
||||
}
|
||||
|
||||
func validGraphNodeKind(kind string) bool {
|
||||
return kind == "block" || kind == "entity" || kind == "sentence"
|
||||
return kind == "block" || kind == "entity" || kind == "sentence" || kind == "document"
|
||||
}
|
||||
|
||||
func graphNodeExists(tx *sql.Tx, kind, id string) (bool, error) {
|
||||
@ -157,6 +170,8 @@ func graphNodeExists(tx *sql.Tx, kind, id string) (bool, error) {
|
||||
err = tx.QueryRow(`SELECT COUNT(*) FROM entities WHERE CAST(id AS TEXT) = ?`, id).Scan(&n)
|
||||
case "sentence":
|
||||
err = tx.QueryRow(`SELECT COUNT(*) FROM sentences WHERE CAST(id AS TEXT) = ?`, id).Scan(&n)
|
||||
case "document":
|
||||
err = tx.QueryRow(`SELECT COUNT(*) FROM documents WHERE id = ?`, id).Scan(&n)
|
||||
default:
|
||||
return false, fmt.Errorf("invalid graph node kind %q", kind)
|
||||
}
|
||||
|
||||
@ -28,13 +28,13 @@ func cleanQQTemplate(text string) string {
|
||||
}
|
||||
|
||||
type cleanTestEvent struct {
|
||||
idx int
|
||||
source string
|
||||
input string
|
||||
response string
|
||||
rawText string
|
||||
idx int
|
||||
source string
|
||||
input string
|
||||
response string
|
||||
rawText string
|
||||
cleanedText string
|
||||
topic string
|
||||
topic string
|
||||
}
|
||||
|
||||
func TestCleanStressPrecision(t *testing.T) {
|
||||
@ -59,53 +59,53 @@ func TestCleanStressPrecision(t *testing.T) {
|
||||
}
|
||||
t.Logf("topics: %v, events: %d", usedTopics, len(events))
|
||||
|
||||
for _, qTopic := range usedTopics {
|
||||
query := queryForTopic(qTopic)
|
||||
qVec := e.Vectorize(query)
|
||||
for _, qTopic := range usedTopics {
|
||||
query := queryForTopic(qTopic)
|
||||
qVec := e.Vectorize(query)
|
||||
|
||||
type scored struct {
|
||||
idx int
|
||||
topic string
|
||||
text string
|
||||
score float64
|
||||
}
|
||||
all := make([]scored, len(events))
|
||||
for i, ev := range events {
|
||||
text := ev.rawText
|
||||
if cleanMode {
|
||||
text = ev.cleanedText
|
||||
type scored struct {
|
||||
idx int
|
||||
topic string
|
||||
text string
|
||||
score float64
|
||||
}
|
||||
all := make([]scored, len(events))
|
||||
for i, ev := range events {
|
||||
text := ev.rawText
|
||||
if cleanMode {
|
||||
text = ev.cleanedText
|
||||
}
|
||||
vec := e.Vectorize(text)
|
||||
all[i] = scored{idx: i, topic: ev.topic, text: text, score: cosineSim(qVec, vec)}
|
||||
}
|
||||
sort.Slice(all, func(i, j int) bool { return all[i].score > all[j].score })
|
||||
|
||||
topK := len(usedTopics) * 2
|
||||
if topK > len(all) {
|
||||
topK = len(all)
|
||||
}
|
||||
|
||||
intraHits := 0
|
||||
for _, s := range all[:topK] {
|
||||
if s.topic == qTopic {
|
||||
intraHits++
|
||||
}
|
||||
}
|
||||
expected := countTopicEvents(events, qTopic)
|
||||
if expected > topK {
|
||||
expected = topK
|
||||
}
|
||||
recall := float64(intraHits) / float64(expected)
|
||||
|
||||
if recall < 0.3 {
|
||||
t.Logf(" [LOW] query=%q topK=%d intra=%d/%d recall=%.2f", qTopic, topK, intraHits, expected, recall)
|
||||
for _, s := range all[:8] {
|
||||
t.Logf(" [%.4f] %s", s.score, trimLen(s.text, 60))
|
||||
}
|
||||
} else {
|
||||
t.Logf(" [OK] query=%q topK=%d intra=%d/%d recall=%.2f", qTopic, topK, intraHits, expected, recall)
|
||||
}
|
||||
}
|
||||
vec := e.Vectorize(text)
|
||||
all[i] = scored{idx: i, topic: ev.topic, text: text, score: cosineSim(qVec, vec)}
|
||||
}
|
||||
sort.Slice(all, func(i, j int) bool { return all[i].score > all[j].score })
|
||||
|
||||
topK := len(usedTopics) * 2
|
||||
if topK > len(all) {
|
||||
topK = len(all)
|
||||
}
|
||||
|
||||
intraHits := 0
|
||||
for _, s := range all[:topK] {
|
||||
if s.topic == qTopic {
|
||||
intraHits++
|
||||
}
|
||||
}
|
||||
expected := countTopicEvents(events, qTopic)
|
||||
if expected > topK {
|
||||
expected = topK
|
||||
}
|
||||
recall := float64(intraHits) / float64(expected)
|
||||
|
||||
if recall < 0.3 {
|
||||
t.Logf(" [LOW] query=%q topK=%d intra=%d/%d recall=%.2f", qTopic, topK, intraHits, expected, recall)
|
||||
for _, s := range all[:8] {
|
||||
t.Logf(" [%.4f] %s", s.score, trimLen(s.text, 60))
|
||||
}
|
||||
} else {
|
||||
t.Logf(" [OK] query=%q topK=%d intra=%d/%d recall=%.2f", qTopic, topK, intraHits, expected, recall)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@ -252,11 +252,11 @@ func genStressEvents(n int) []cleanTestEvent {
|
||||
keywords []string
|
||||
sources []string
|
||||
}{
|
||||
"大学招生": {[]string{"河南医药大学", "录取分数线", "专业排名", "高考志愿", "招生简章"}, []string{"qq", "qq", "agent"}},
|
||||
"老大私聊": {[]string{"老大私聊消息", "回复老大", "任务安排", "汇报工作", "收到"}, []string{"qq", "agent", "agent"}},
|
||||
"前端开发": {[]string{"前端组件封装", "页面路由配置", "界面布局设计", "交互逻辑开发", "代码调试优化"}, []string{"cli", "cli", "agent"}},
|
||||
"大学招生": {[]string{"河南医药大学", "录取分数线", "专业排名", "高考志愿", "招生简章"}, []string{"qq", "qq", "agent"}},
|
||||
"老大私聊": {[]string{"老大私聊消息", "回复老大", "任务安排", "汇报工作", "收到"}, []string{"qq", "agent", "agent"}},
|
||||
"前端开发": {[]string{"前端组件封装", "页面路由配置", "界面布局设计", "交互逻辑开发", "代码调试优化"}, []string{"cli", "cli", "agent"}},
|
||||
"服务器运维": {[]string{"反向代理配置", "容器部署方案", "证书续期", "数据库备份恢复", "监控告警处理"}, []string{"cli", "agent", "agent"}},
|
||||
"股票基金": {[]string{"基金定投策略", "股票涨跌分析", "理财收益计算", "市场行情分析", "投资风险管理"}, []string{"qq", "qq", "agent"}},
|
||||
"股票基金": {[]string{"基金定投策略", "股票涨跌分析", "理财收益计算", "市场行情分析", "投资风险管理"}, []string{"qq", "qq", "agent"}},
|
||||
}
|
||||
|
||||
for i := 0; i < n; i++ {
|
||||
@ -292,13 +292,13 @@ func genStressEvents(n int) []cleanTestEvent {
|
||||
cleaned := cleanEventText(src, input, response)
|
||||
raw := rawEventText(src, input, response)
|
||||
events = append(events, cleanTestEvent{
|
||||
idx: i,
|
||||
source: src,
|
||||
input: input,
|
||||
response: response,
|
||||
rawText: raw,
|
||||
idx: i,
|
||||
source: src,
|
||||
input: input,
|
||||
response: response,
|
||||
rawText: raw,
|
||||
cleanedText: cleaned,
|
||||
topic: tp,
|
||||
topic: tp,
|
||||
})
|
||||
}
|
||||
return events
|
||||
|
||||
@ -1,484 +0,0 @@
|
||||
//go:build onnxruntime
|
||||
|
||||
// Package clip 提供基于 CLIP ONNX 的稠密多模态编码器。
|
||||
//
|
||||
// 构建标签 onnxruntime 控制是否编译此实现（与 internal/nlp/onnx.go 同模式）。
|
||||
// 未配置 clip_model_dir 时不会初始化 ONNX Runtime，现有 fastText/TF-IDF 行为不变。
|
||||
//
|
||||
// 支持的模型文件（统一放置于 clip_model_dir 目录）：
|
||||
//
|
||||
// text.onnx — CLIP 文本编码器（input_ids + attention_mask → text_embed [1,512]）
|
||||
// vision.onnx — CLIP 图像编码器（pixel_values → image_embed [1,512]）
|
||||
// clip_config.json — 模型元数据（dimension, context_length, image_size, mean, std）
|
||||
// tokenizer.json — HuggingFace tokenizer.json（含 vocab + merges）
|
||||
// merges.txt — BPE merges 文件
|
||||
//
|
||||
// 设计：只产出**稠密** 512 维向量（VectorizeDense / EmbedImageDense），供媒体层
|
||||
// media.Store 的稠密余弦检索（QueryMedia）消费。刻意不实现 vector.Vectorizer
|
||||
// （稀疏词向量）——文档/知识/上下文层的文本相似度检索保留 TF-IDF 高频削弱加权 +
|
||||
// 倒排剪枝 + fastText 稀疏词向量的既有设计，CLIP 稠密空间不混入那套索引。
|
||||
package clip
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"image"
|
||||
_ "image/jpeg"
|
||||
_ "image/png"
|
||||
"log"
|
||||
"math"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
ort "github.com/yalue/onnxruntime_go"
|
||||
)
|
||||
|
||||
// clipConfig 描述模型的超参数与归一化常数。
|
||||
type clipConfig struct {
|
||||
Model string `json:"model"`
|
||||
Dimension int `json:"dimension"`
|
||||
ContextLength int `json:"context_length"`
|
||||
ImageSize int `json:"image_size"`
|
||||
Mean []float64 `json:"mean"`
|
||||
Std []float64 `json:"std"`
|
||||
}
|
||||
|
||||
// Embedder 是 CLIP 多模态稠密编码器，实现 vector.MultimodalEmbedder。
|
||||
// 产出 512 维共享稠密空间向量，供 media.Store 的稠密余弦检索。
|
||||
//
|
||||
// 刻意**不**实现 vector.Vectorizer（稀疏词向量）：CLIP 稠密向量若以
|
||||
// map[string]float64 稀疏形式塞进文档/知识层的 vector.Store，会破坏其
|
||||
// TF-IDF 高频削弱加权与倒排剪枝语义。文本层的相似度检索保持 TF-IDF/fastText。
|
||||
type Embedder struct {
|
||||
mu sync.RWMutex
|
||||
config clipConfig
|
||||
vocab map[string]int64
|
||||
merges []string
|
||||
textSess *ort.DynamicAdvancedSession
|
||||
imgSess *ort.DynamicAdvancedSession
|
||||
close sync.Once
|
||||
loaded bool
|
||||
fingerprint string
|
||||
}
|
||||
|
||||
// Fingerprint 返回当前模型目录的指纹（文本+视觉模型文件 SHA256 拼接），
|
||||
// 用于检测模型切换后触发重算。
|
||||
func (e *Embedder) Fingerprint() string {
|
||||
e.mu.RLock()
|
||||
defer e.mu.RUnlock()
|
||||
return e.fingerprint
|
||||
}
|
||||
|
||||
// Dim 返回向量维度。
|
||||
func (e *Embedder) Dim() int {
|
||||
e.mu.RLock()
|
||||
defer e.mu.RUnlock()
|
||||
return e.config.Dimension
|
||||
}
|
||||
|
||||
// Loaded 返回加载状态。
|
||||
func (e *Embedder) Loaded() bool {
|
||||
e.mu.RLock()
|
||||
defer e.mu.RUnlock()
|
||||
return e.loaded
|
||||
}
|
||||
|
||||
// VectorizeDense 将文本转为归一化的 []float64 向量（CLIP 共享空间）。
|
||||
func (e *Embedder) VectorizeDense(text string) ([]float64, error) {
|
||||
e.mu.RLock()
|
||||
defer e.mu.RUnlock()
|
||||
if !e.loaded {
|
||||
return nil, fmt.Errorf("clip embedder not loaded")
|
||||
}
|
||||
|
||||
tokens := tokenizeCLIP(text, e.vocab, e.merges, e.config.ContextLength)
|
||||
if len(tokens) == 0 {
|
||||
return make([]float64, e.config.Dimension), nil
|
||||
}
|
||||
|
||||
dim := e.config.Dimension
|
||||
inputIDs := make([]int64, e.config.ContextLength)
|
||||
attnMask := make([]int64, e.config.ContextLength)
|
||||
for i, tok := range tokens {
|
||||
if i >= e.config.ContextLength {
|
||||
break
|
||||
}
|
||||
inputIDs[i] = tok
|
||||
attnMask[i] = 1
|
||||
}
|
||||
|
||||
idTensor, err := ort.NewTensor(ort.Shape{1, int64(e.config.ContextLength)}, inputIDs)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("input_ids tensor: %w", err)
|
||||
}
|
||||
defer idTensor.Destroy()
|
||||
|
||||
maskTensor, err := ort.NewTensor(ort.Shape{1, int64(e.config.ContextLength)}, attnMask)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("attention_mask tensor: %w", err)
|
||||
}
|
||||
defer maskTensor.Destroy()
|
||||
|
||||
featTensor, err := ort.NewEmptyTensor[float32](ort.Shape{1, int64(dim)})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("output tensor: %w", err)
|
||||
}
|
||||
defer featTensor.Destroy()
|
||||
|
||||
if err := e.textSess.Run([]ort.Value{idTensor, maskTensor}, []ort.Value{featTensor}); err != nil {
|
||||
return nil, fmt.Errorf("text run: %w", err)
|
||||
}
|
||||
|
||||
raw := featTensor.GetData()
|
||||
out := make([]float64, dim)
|
||||
var norm float64
|
||||
for i, v := range raw {
|
||||
out[i] = float64(v)
|
||||
norm += out[i] * out[i]
|
||||
}
|
||||
if norm > 0 {
|
||||
norm = math.Sqrt(norm)
|
||||
for i := range out {
|
||||
out[i] /= norm
|
||||
}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// EmbedImageDense 将图像字节转为归一化的 []float64 向量（CLIP 共享空间）。
|
||||
func (e *Embedder) EmbedImageDense(img []byte, mime string) ([]float64, error) {
|
||||
e.mu.RLock()
|
||||
defer e.mu.RUnlock()
|
||||
if !e.loaded {
|
||||
return nil, fmt.Errorf("clip embedder not loaded")
|
||||
}
|
||||
return e.embedImageDenseUnlocked(img, mime)
|
||||
}
|
||||
|
||||
func (e *Embedder) embedImageDenseUnlocked(img []byte, mime string) ([]float64, error) {
|
||||
decoded, _, err := image.Decode(bytes.NewReader(img))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("decode image: %w", err)
|
||||
}
|
||||
|
||||
size := e.config.ImageSize
|
||||
resized := resizeImage(decoded, size, size)
|
||||
|
||||
pixels := make([]float32, 3*size*size)
|
||||
for y := 0; y < size; y++ {
|
||||
for x := 0; x < size; x++ {
|
||||
r, g, b, _ := resized.At(x, y).RGBA()
|
||||
rf := float64(r) / 65535.0
|
||||
gf := float64(g) / 65535.0
|
||||
bf := float64(b) / 65535.0
|
||||
|
||||
for c, v := range []float64{rf, gf, bf} {
|
||||
norm := (v - e.config.Mean[c]) / e.config.Std[c]
|
||||
pixels[c*size*size+y*size+x] = float32(norm)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pixelTensor, err := ort.NewTensor(ort.Shape{1, 3, int64(size), int64(size)}, pixels)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("pixel_values tensor: %w", err)
|
||||
}
|
||||
defer pixelTensor.Destroy()
|
||||
|
||||
dim := e.config.Dimension
|
||||
featTensor, err := ort.NewEmptyTensor[float32](ort.Shape{1, int64(dim)})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("output tensor: %w", err)
|
||||
}
|
||||
defer featTensor.Destroy()
|
||||
|
||||
if err := e.imgSess.Run([]ort.Value{pixelTensor}, []ort.Value{featTensor}); err != nil {
|
||||
return nil, fmt.Errorf("vision run: %w", err)
|
||||
}
|
||||
|
||||
raw := featTensor.GetData()
|
||||
out := make([]float64, dim)
|
||||
var norm float64
|
||||
for i, v := range raw {
|
||||
out[i] = float64(v)
|
||||
norm += out[i] * out[i]
|
||||
}
|
||||
if norm > 0 {
|
||||
norm = math.Sqrt(norm)
|
||||
for i := range out {
|
||||
out[i] /= norm
|
||||
}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// Close 释放 ONNX Runtime 资源。
|
||||
func (e *Embedder) Close() {
|
||||
e.close.Do(func() {
|
||||
e.mu.Lock()
|
||||
defer e.mu.Unlock()
|
||||
if e.textSess != nil {
|
||||
e.textSess.Destroy()
|
||||
}
|
||||
if e.imgSess != nil {
|
||||
e.imgSess.Destroy()
|
||||
}
|
||||
e.loaded = false
|
||||
})
|
||||
}
|
||||
|
||||
// New 从目录加载 CLIP 模型。目录需包含 text.onnx、vision.onnx、
|
||||
// clip_config.json、tokenizer.json、merges.txt。
|
||||
func New(modelDir string) (*Embedder, error) {
|
||||
if modelDir == "" {
|
||||
return nil, fmt.Errorf("clip model dir not specified")
|
||||
}
|
||||
|
||||
// 读取配置
|
||||
cfgData, err := os.ReadFile(filepath.Join(modelDir, "clip_config.json"))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("read clip_config.json: %w", err)
|
||||
}
|
||||
var cfg clipConfig
|
||||
if err := json.Unmarshal(cfgData, &cfg); err != nil {
|
||||
return nil, fmt.Errorf("parse clip_config.json: %w", err)
|
||||
}
|
||||
if cfg.Dimension <= 0 || cfg.ContextLength <= 0 || cfg.ImageSize <= 0 {
|
||||
return nil, fmt.Errorf("invalid clip config: dim=%d ctx=%d img=%d", cfg.Dimension, cfg.ContextLength, cfg.ImageSize)
|
||||
}
|
||||
if len(cfg.Mean) != 3 || len(cfg.Std) != 3 {
|
||||
return nil, fmt.Errorf("clip config mean/std must have 3 channels")
|
||||
}
|
||||
|
||||
// 加载 tokenizer
|
||||
vocab, err := loadTokenizerVocab(filepath.Join(modelDir, "tokenizer.json"))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("load tokenizer: %w", err)
|
||||
}
|
||||
merges, err := loadMerges(filepath.Join(modelDir, "merges.txt"))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("load merges: %w", err)
|
||||
}
|
||||
|
||||
// 初始化 ONNX Runtime（只初始化一次）
|
||||
if !ort.IsInitialized() {
|
||||
// 尝试从 nlp 同样的路径查找 libonnxruntime.so
|
||||
libPath := findOnnxLib()
|
||||
if libPath != "" {
|
||||
ort.SetSharedLibraryPath(libPath)
|
||||
}
|
||||
if err := ort.InitializeEnvironment(); err != nil {
|
||||
return nil, fmt.Errorf("init onnx env: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
// 创建文本编码器会话
|
||||
textSess, err := ort.NewDynamicAdvancedSession(
|
||||
filepath.Join(modelDir, "text.onnx"),
|
||||
[]string{"input_ids", "attention_mask"},
|
||||
[]string{"text_embed"},
|
||||
nil,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("create text session: %w", err)
|
||||
}
|
||||
|
||||
// 创建视觉编码器会话
|
||||
imgSess, err := ort.NewDynamicAdvancedSession(
|
||||
filepath.Join(modelDir, "vision.onnx"),
|
||||
[]string{"pixel_values"},
|
||||
[]string{"image_embed"},
|
||||
nil,
|
||||
)
|
||||
if err != nil {
|
||||
textSess.Destroy()
|
||||
return nil, fmt.Errorf("create vision session: %w", err)
|
||||
}
|
||||
|
||||
// 计算模型指纹
|
||||
fp := computeFingerprint(modelDir)
|
||||
|
||||
log.Printf("[clip] loaded %s dim=%d ctx=%d img=%d from %s (fp=%s)", cfg.Model, cfg.Dimension, cfg.ContextLength, cfg.ImageSize, modelDir, fp[:12])
|
||||
|
||||
return &Embedder{
|
||||
config: cfg,
|
||||
vocab: vocab,
|
||||
merges: merges,
|
||||
textSess: textSess,
|
||||
imgSess: imgSess,
|
||||
loaded: true,
|
||||
fingerprint: fp,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// computeFingerprint 计算模型文件指纹（text.onnx + vision.onnx 的 SHA256）。
|
||||
func computeFingerprint(modelDir string) string {
|
||||
h := sha256.New()
|
||||
for _, name := range []string{"text.onnx", "vision.onnx"} {
|
||||
data, err := os.ReadFile(filepath.Join(modelDir, name))
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
h.Write(data)
|
||||
h.Write([]byte{0}) // 分隔符
|
||||
}
|
||||
return hex.EncodeToString(h.Sum(nil))
|
||||
}
|
||||
|
||||
// findOnnxLib 在常见路径中查找 libonnxruntime.so。
|
||||
func findOnnxLib() string {
|
||||
for _, p := range []string{
|
||||
"/opt/onnxruntime/libonnxruntime.so",
|
||||
"libonnxruntime.so",
|
||||
} {
|
||||
if _, err := os.Stat(p); err == nil {
|
||||
return p
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// resizeImage 使用最近邻将 src 缩放到 dstW×dstH。
|
||||
// 生产中应使用双线性插值，此处为 MVP 简化。
|
||||
func resizeImage(src image.Image, dstW, dstH int) image.Image {
|
||||
srcB := src.Bounds()
|
||||
srcW := srcB.Dx()
|
||||
srcH := srcB.Dy()
|
||||
if srcW == dstW && srcH == dstH {
|
||||
return src
|
||||
}
|
||||
|
||||
dst := image.NewRGBA(image.Rect(0, 0, dstW, dstH))
|
||||
for y := 0; y < dstH; y++ {
|
||||
for x := 0; x < dstW; x++ {
|
||||
sx := srcB.Min.X + x*srcW/dstW
|
||||
sy := srcB.Min.Y + y*srcH/dstH
|
||||
dst.Set(x, y, src.At(sx, sy))
|
||||
}
|
||||
}
|
||||
return dst
|
||||
}
|
||||
|
||||
// ---- BPE Tokenizer ----
|
||||
|
||||
// loadTokenizerVocab 从 HuggingFace tokenizer.json 中提取 vocab（token→id 映射）。
|
||||
func loadTokenizerVocab(path string) (map[string]int64, error) {
|
||||
data, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var tok struct {
|
||||
Model struct {
|
||||
Vocab map[string]int64 `json:"vocab"`
|
||||
} `json:"model"`
|
||||
}
|
||||
if err := json.Unmarshal(data, &tok); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(tok.Model.Vocab) == 0 {
|
||||
return nil, fmt.Errorf("empty vocab in %s", path)
|
||||
}
|
||||
return tok.Model.Vocab, nil
|
||||
}
|
||||
|
||||
// loadMerges 从 merges.txt 加载 BPE 合并规则。
|
||||
func loadMerges(path string) ([]string, error) {
|
||||
data, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
lines := strings.Split(strings.TrimSpace(string(data)), "\n")
|
||||
// 第一行是版本号（"#version: 0.2"），跳过
|
||||
var merges []string
|
||||
for _, line := range lines[1:] {
|
||||
line = strings.TrimSpace(line)
|
||||
if line == "" {
|
||||
continue
|
||||
}
|
||||
merges = append(merges, line)
|
||||
}
|
||||
return merges, nil
|
||||
}
|
||||
|
||||
// tokenizeCLIP 将文本分词为模型 vocab 中的 token id 序列。
|
||||
//
|
||||
// 此模型（transformers 5.x 导出的 CLIP tokenizer.json）是**词级 BPE**：
|
||||
// 词末 token 带 </w> 后缀（"a</w>"=320、"red</w>"=736），词中片段不带。
|
||||
// 流程：lowercase → 按空白/标点拆词 → 每词做字符级 BPE 合并 →
|
||||
// 末尾片段加 </w> 查 vocab，其余片段直接查；查不到则丢弃。
|
||||
func tokenizeCLIP(text string, vocab map[string]int64, merges []string, maxLen int) []int64 {
|
||||
rank := make(map[string]int, len(merges))
|
||||
for i, m := range merges {
|
||||
rank[m] = i
|
||||
}
|
||||
const endTok = "</w>"
|
||||
|
||||
var tokens []int64
|
||||
if id, ok := vocab["<|startoftext|>"]; ok {
|
||||
tokens = append(tokens, id)
|
||||
}
|
||||
for _, word := range strings.Fields(strings.ToLower(text)) {
|
||||
seq := make([]string, 0, len(word))
|
||||
for _, ch := range word {
|
||||
seq = append(seq, string(ch))
|
||||
}
|
||||
merged := bpeMerge(seq, rank)
|
||||
for i, t := range merged {
|
||||
lookup := t
|
||||
if i == len(merged)-1 {
|
||||
// 词末片段带 </w>
|
||||
lookup = t + endTok
|
||||
}
|
||||
if id, ok := vocab[lookup]; ok {
|
||||
tokens = append(tokens, id)
|
||||
}
|
||||
}
|
||||
}
|
||||
if id, ok := vocab["<|endoftext|>"]; ok {
|
||||
tokens = append(tokens, id)
|
||||
}
|
||||
if len(tokens) > maxLen {
|
||||
tokens = tokens[:maxLen]
|
||||
}
|
||||
return tokens
|
||||
}
|
||||
|
||||
// bpeMerge 对单个词的字符序列应用 BPE 合并直到无可合并对。
|
||||
// rank[pair] 越小越优先（merges.txt 顺序）。
|
||||
func bpeMerge(seq []string, rank map[string]int) []string {
|
||||
for len(seq) > 1 {
|
||||
// 找 rank 最低的可合并相邻对
|
||||
bestRank := -1
|
||||
bestPair := ""
|
||||
for i := 0; i < len(seq)-1; i++ {
|
||||
pair := seq[i] + " " + seq[i+1]
|
||||
if r, ok := rank[pair]; ok && (bestRank < 0 || r < bestRank) {
|
||||
bestRank = r
|
||||
bestPair = pair
|
||||
}
|
||||
}
|
||||
if bestPair == "" {
|
||||
break
|
||||
}
|
||||
parts := strings.SplitN(bestPair, " ", 2)
|
||||
merged := parts[0] + parts[1]
|
||||
|
||||
// 一次性合并所有相邻的该 pair
|
||||
var out []string
|
||||
for i := 0; i < len(seq); i++ {
|
||||
if i < len(seq)-1 && seq[i] == parts[0] && seq[i+1] == parts[1] {
|
||||
out = append(out, merged)
|
||||
i++
|
||||
} else {
|
||||
out = append(out, seq[i])
|
||||
}
|
||||
}
|
||||
seq = out
|
||||
}
|
||||
return seq
|
||||
}
|
||||
@ -1,28 +0,0 @@
|
||||
//go:build !onnxruntime
|
||||
|
||||
package clip
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
)
|
||||
|
||||
// Embedder 在未启用 onnxruntime 时为 no-op 实现。
|
||||
// 构建时不链接 onnxruntime，默认构建保持原有 fastText/TF-IDF 行为不变。
|
||||
type Embedder struct {
|
||||
loaded bool
|
||||
}
|
||||
|
||||
func New(_ string) (*Embedder, error) {
|
||||
return nil, fmt.Errorf("clip embedder requires build tag 'onnxruntime' (go build -tags onnxruntime)")
|
||||
}
|
||||
|
||||
func (e *Embedder) Fingerprint() string { return "" }
|
||||
func (e *Embedder) Dim() int { return 0 }
|
||||
func (e *Embedder) Loaded() bool { return e.loaded }
|
||||
func (e *Embedder) Close() {}
|
||||
func (e *Embedder) VectorizeDense(_ string) ([]float64, error) {
|
||||
return nil, fmt.Errorf("clip embedder not available")
|
||||
}
|
||||
func (e *Embedder) EmbedImageDense(_ []byte, _ string) ([]float64, error) {
|
||||
return nil, fmt.Errorf("clip embedder not available")
|
||||
}
|
||||
@ -1,154 +0,0 @@
|
||||
//go:build onnxruntime
|
||||
|
||||
package clip
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"fmt"
|
||||
"image"
|
||||
"image/color"
|
||||
"image/png"
|
||||
"os"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestSmokeLoadAndEncode(t *testing.T) {
|
||||
modelDir := os.Getenv("CLIP_MODEL_DIR")
|
||||
if modelDir == "" {
|
||||
modelDir = "/home/newqqagent/models/clip-vit-b32"
|
||||
}
|
||||
if _, err := os.Stat(modelDir + "/text.onnx"); err != nil {
|
||||
t.Skipf("模型目录不存在: %v", err)
|
||||
}
|
||||
|
||||
emb, err := New(modelDir)
|
||||
if err != nil {
|
||||
t.Fatalf("New: %v", err)
|
||||
}
|
||||
defer emb.Close()
|
||||
|
||||
if !emb.Loaded() {
|
||||
t.Fatal("loaded should be true")
|
||||
}
|
||||
if emb.Dim() != 512 {
|
||||
t.Fatalf("dim = %d, want 512", emb.Dim())
|
||||
}
|
||||
if emb.Fingerprint() == "" {
|
||||
t.Fatal("fingerprint should not be empty")
|
||||
}
|
||||
|
||||
// 文本编码
|
||||
textVec, err := emb.VectorizeDense("a photo of a cat")
|
||||
if err != nil {
|
||||
t.Fatalf("VectorizeDense: %v", err)
|
||||
}
|
||||
if len(textVec) != 512 {
|
||||
t.Fatalf("text vec len = %d, want 512", len(textVec))
|
||||
}
|
||||
fmt.Printf("text vec[:5] = %v\n", textVec[:5])
|
||||
|
||||
// 同义文本应比远义文本更相似
|
||||
textVec2, _ := emb.VectorizeDense("a photograph of a dog")
|
||||
textVec3, _ := emb.VectorizeDense("quantum physics equations")
|
||||
|
||||
sim12 := cosineSim(textVec, textVec2)
|
||||
sim13 := cosineSim(textVec, textVec3)
|
||||
fmt.Printf("cat vs dog = %.4f, cat vs physics = %.4f\n", sim12, sim13)
|
||||
if sim12 <= sim13 {
|
||||
t.Errorf("cat-dog sim (%.4f) should be > cat-physics sim (%.4f)", sim12, sim13)
|
||||
}
|
||||
|
||||
// 稠密 VectorizeDense 再次调用验证可重复
|
||||
vAgain, _ := emb.VectorizeDense("hello world")
|
||||
if len(vAgain) != 512 {
|
||||
t.Errorf("dense VectorizeDense len = %d, want 512", len(vAgain))
|
||||
}
|
||||
fmt.Printf("dense len = %d\n", len(vAgain))
|
||||
}
|
||||
|
||||
// TestCrossModalAlignment 验证图文在同一向量空间可比：
|
||||
// 红底图的向量应与 "a red image" 更相似，而非 "a blue image"。
|
||||
func TestCrossModalAlignment(t *testing.T) {
|
||||
modelDir := os.Getenv("CLIP_MODEL_DIR")
|
||||
if modelDir == "" {
|
||||
modelDir = "/home/newqqagent/models/clip-vit-b32"
|
||||
}
|
||||
emb, err := New(modelDir)
|
||||
if err != nil {
|
||||
t.Fatalf("New: %v", err)
|
||||
}
|
||||
defer emb.Close()
|
||||
|
||||
// 生成 224x224 纯红底 PNG
|
||||
img := image.NewRGBA(image.Rect(0, 0, 224, 224))
|
||||
red := color.RGBA{220, 40, 40, 255}
|
||||
for y := 0; y < 224; y++ {
|
||||
for x := 0; x < 224; x++ {
|
||||
img.Set(x, y, red)
|
||||
}
|
||||
}
|
||||
var buf bytes.Buffer
|
||||
if err := png.Encode(&buf, img); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
imgVec, err := emb.EmbedImageDense(buf.Bytes(), "image/png")
|
||||
if err != nil {
|
||||
t.Fatalf("EmbedImageDense: %v", err)
|
||||
}
|
||||
if len(imgVec) != 512 {
|
||||
t.Fatalf("img vec len = %d, want 512", len(imgVec))
|
||||
}
|
||||
|
||||
redText, _ := emb.VectorizeDense("a red image")
|
||||
blueText, _ := emb.VectorizeDense("a blue image")
|
||||
redSim := cosineSim(imgVec, redText)
|
||||
blueSim := cosineSim(imgVec, blueText)
|
||||
fmt.Printf("red-image vs red-text = %.4f, vs blue-text = %.4f\n", redSim, blueSim)
|
||||
if redSim <= blueSim {
|
||||
t.Errorf("red image should align better with red text (%.4f) than blue (%.4f)", redSim, blueSim)
|
||||
}
|
||||
|
||||
// 同图应比异图更相似：存两张不同颜色，query 用红底图应召回红图
|
||||
d1 := imgVec
|
||||
blueImg := image.NewRGBA(image.Rect(0, 0, 224, 224))
|
||||
blue := color.RGBA{40, 40, 220, 255}
|
||||
for y := 0; y < 224; y++ {
|
||||
for x := 0; x < 224; x++ {
|
||||
blueImg.Set(x, y, blue)
|
||||
}
|
||||
}
|
||||
var buf2 bytes.Buffer
|
||||
png.Encode(&buf2, blueImg)
|
||||
d2, _ := emb.EmbedImageDense(buf2.Bytes(), "image/png")
|
||||
if cosineSim(d1, d2) >= 0.99 {
|
||||
t.Errorf("red and blue images should differ (got sim %.4f)", cosineSim(d1, d2))
|
||||
}
|
||||
}
|
||||
|
||||
func cosineSim(a, b []float64) float64 {
|
||||
if len(a) != len(b) {
|
||||
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 / (sqrt(na) * sqrt(nb))
|
||||
}
|
||||
|
||||
func sqrt(x float64) float64 {
|
||||
if x <= 0 {
|
||||
return 0
|
||||
}
|
||||
z := x
|
||||
for i := 0; i < 50; i++ {
|
||||
z = (z + x/z) / 2
|
||||
}
|
||||
return z
|
||||
}
|
||||
@ -13,26 +13,26 @@ import (
|
||||
|
||||
// contentPOS 有实义的词性标签：只保留名词/动词/形容词/专名等
|
||||
var contentPOS = map[string]bool{
|
||||
"n": true, // 普通名词
|
||||
"nr": true, // 人名
|
||||
"ns": true, // 地名
|
||||
"nt": true, // 机构名
|
||||
"nw": true, // 作品名/URL
|
||||
"nz": true, // 其他专名
|
||||
"v": true, // 动词
|
||||
"vd": true, // 副动词
|
||||
"vn": true, // 名动词
|
||||
"a": true, // 形容词
|
||||
"ad": true, // 副形词
|
||||
"an": true, // 名形词
|
||||
"i": true, // 成语
|
||||
"l": true, // 习用语
|
||||
"j": true, // 简称
|
||||
"s": true, // 处所词
|
||||
"f": true, // 方位词
|
||||
"b": true, // 区别词
|
||||
"z": true, // 状态词
|
||||
"t": true, // 时间词
|
||||
"n": true, // 普通名词
|
||||
"nr": true, // 人名
|
||||
"ns": true, // 地名
|
||||
"nt": true, // 机构名
|
||||
"nw": true, // 作品名/URL
|
||||
"nz": true, // 其他专名
|
||||
"v": true, // 动词
|
||||
"vd": true, // 副动词
|
||||
"vn": true, // 名动词
|
||||
"a": true, // 形容词
|
||||
"ad": true, // 副形词
|
||||
"an": true, // 名形词
|
||||
"i": true, // 成语
|
||||
"l": true, // 习用语
|
||||
"j": true, // 简称
|
||||
"s": true, // 处所词
|
||||
"f": true, // 方位词
|
||||
"b": true, // 区别词
|
||||
"z": true, // 状态词
|
||||
"t": true, // 时间词
|
||||
"eng": true, // 英文
|
||||
"x": true, // 非语素字
|
||||
"zg": true, // 其他
|
||||
|
||||
@ -37,6 +37,7 @@ type Doc struct {
|
||||
Blocks []memory.MemoryBlock `json:"blocks,omitempty"` // 一等记忆块（text/image/video/audio）
|
||||
Vector tfidf.Vector `json:"vector,omitempty"` // TF-IDF 稀疏向量（fallback 时持久化）
|
||||
DenseVec []float64 `json:"dense_vec,omitempty"` // 多模态稠密向量（主路径）
|
||||
DenseFP string `json:"dense_fp,omitempty"` // DenseVec 所属统一空间指纹，变化时触发重算
|
||||
}
|
||||
|
||||
// Store — 文档记忆存储。
|
||||
@ -113,7 +114,7 @@ func (s *Store) SetDenseSpace(ds vector.MultimodalEmbedder) {
|
||||
s.denseSpace = ds
|
||||
}
|
||||
|
||||
// BuildDenseIndex 为所有文档计算稠密向量。
|
||||
// BuildDenseIndex 为所有文档计算稠密向量（文本 ⊕ 媒体块）。
|
||||
func (s *Store) BuildDenseIndex(ds vector.MultimodalEmbedder) {
|
||||
if ds == nil || !ds.Loaded() {
|
||||
return
|
||||
@ -123,21 +124,42 @@ func (s *Store) BuildDenseIndex(ds vector.MultimodalEmbedder) {
|
||||
log.Printf("[document memory] building dense index for %d docs (dim=%d)", len(s.docs), ds.Dim())
|
||||
count := 0
|
||||
for _, doc := range s.docs {
|
||||
if doc.DenseVec != nil && len(doc.DenseVec) == ds.Dim() {
|
||||
if doc.DenseVec != nil && len(doc.DenseVec) == ds.Dim() && doc.DenseFP == ds.Fingerprint() {
|
||||
continue
|
||||
}
|
||||
text := doc.Summary + " " + doc.Content
|
||||
vec, err := ds.VectorizeDense(text)
|
||||
if err != nil {
|
||||
log.Printf("[document memory] dense embed failed %s: %v", doc.ID[:min(16, len(doc.ID))], err)
|
||||
vec := s.denseFor(doc)
|
||||
if vec == nil {
|
||||
continue
|
||||
}
|
||||
doc.DenseVec = vec
|
||||
doc.DenseFP = ds.Fingerprint()
|
||||
count++
|
||||
}
|
||||
log.Printf("[document memory] dense index built: %d new vectors", count)
|
||||
}
|
||||
|
||||
// denseFor 计算文档的稠密向量：文本向量与其一等记忆块的媒体向量融合。
|
||||
//
|
||||
// 只有与当前统一空间同指纹的块向量才参与融合：不同模型/维度的旧向量
|
||||
// 属于另一个坐标系，混进去会算出一个两边都不像的方向。
|
||||
// 任意一路缺失时退化为另一路；都不可用返回 nil。
|
||||
func (s *Store) denseFor(doc *Doc) []float64 {
|
||||
if s.denseSpace == nil || !s.denseSpace.Loaded() {
|
||||
return nil
|
||||
}
|
||||
fp := s.denseSpace.Fingerprint()
|
||||
var parts [][]float64
|
||||
if tv, err := s.denseSpace.VectorizeDense(doc.Summary + " " + doc.Content); err == nil && len(tv) > 0 {
|
||||
parts = append(parts, tv)
|
||||
}
|
||||
for _, b := range doc.Blocks {
|
||||
if len(b.Vector) > 0 && b.Fingerprint == fp {
|
||||
parts = append(parts, b.Vector)
|
||||
}
|
||||
}
|
||||
return vector.FuseVectors(parts...)
|
||||
}
|
||||
|
||||
// Reindex 重建 TF-IDF 索引（fallback 路径变更时调用）。
|
||||
func (s *Store) Reindex() {
|
||||
s.mu.Lock()
|
||||
@ -172,11 +194,10 @@ func (s *Store) Insert(doc *Doc) error {
|
||||
|
||||
text := doc.Summary + " " + doc.Content
|
||||
|
||||
// 主路径：稠密向量
|
||||
// 主路径：稠密向量（文本 ⊕ 媒体块）
|
||||
if s.denseSpace != nil && s.denseSpace.Loaded() && len(doc.DenseVec) == 0 {
|
||||
if dv, err := s.denseSpace.VectorizeDense(text); err == nil {
|
||||
doc.DenseVec = dv
|
||||
}
|
||||
doc.DenseVec = s.denseFor(doc)
|
||||
doc.DenseFP = s.denseSpace.Fingerprint()
|
||||
}
|
||||
|
||||
// Fallback 路径：缓存文本，延迟训练
|
||||
@ -232,6 +253,10 @@ func (s *Store) ContextToDoc(source string, entries []ContextEntry, _ interface{
|
||||
d.Tags = tags
|
||||
d.Entities = entities
|
||||
d.Blocks = blocksFromEntries(entries)
|
||||
d.DenseVec = s.denseFor(d)
|
||||
if s.denseSpace != nil {
|
||||
d.DenseFP = s.denseSpace.Fingerprint()
|
||||
}
|
||||
s.dirty = true
|
||||
return d, nil
|
||||
}
|
||||
@ -249,6 +274,10 @@ func (s *Store) ContextToDoc(source string, entries []ContextEntry, _ interface{
|
||||
Blocks: blocksFromEntries(entries),
|
||||
}
|
||||
s.docs[id] = doc
|
||||
doc.DenseVec = s.denseFor(doc)
|
||||
if s.denseSpace != nil {
|
||||
doc.DenseFP = s.denseSpace.Fingerprint()
|
||||
}
|
||||
text := summary + " " + content
|
||||
if s.tfidfIdx != nil {
|
||||
s.tfidfIdx.Add(id, text)
|
||||
|
||||
@ -15,7 +15,7 @@ func TestInsertAndQuery(t *testing.T) {
|
||||
}
|
||||
defer os.RemoveAll(dir)
|
||||
|
||||
s := NewStore(dir)
|
||||
s := NewStore(dir, memory.TokenizeWords)
|
||||
if err := s.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
@ -48,7 +48,7 @@ func TestQuery(t *testing.T) {
|
||||
}
|
||||
defer os.RemoveAll(dir)
|
||||
|
||||
s := NewStore(dir)
|
||||
s := NewStore(dir, memory.TokenizeWords)
|
||||
s.Start()
|
||||
defer s.Stop()
|
||||
|
||||
@ -73,7 +73,7 @@ func TestContextToDoc(t *testing.T) {
|
||||
}
|
||||
defer os.RemoveAll(dir)
|
||||
|
||||
s := NewStore(dir)
|
||||
s := NewStore(dir, memory.TokenizeWords)
|
||||
s.Start()
|
||||
defer s.Stop()
|
||||
|
||||
@ -104,7 +104,7 @@ func TestFindColdDocs(t *testing.T) {
|
||||
}
|
||||
defer os.RemoveAll(dir)
|
||||
|
||||
s := NewStore(dir)
|
||||
s := NewStore(dir, memory.TokenizeWords)
|
||||
s.Start()
|
||||
defer s.Stop()
|
||||
|
||||
@ -136,7 +136,7 @@ func TestRecentDocs(t *testing.T) {
|
||||
}
|
||||
defer os.RemoveAll(dir)
|
||||
|
||||
s := NewStore(dir)
|
||||
s := NewStore(dir, memory.TokenizeWords)
|
||||
s.Start()
|
||||
defer s.Stop()
|
||||
|
||||
@ -160,7 +160,7 @@ func TestReindex(t *testing.T) {
|
||||
}
|
||||
defer os.RemoveAll(dir)
|
||||
|
||||
s := NewStore(dir)
|
||||
s := NewStore(dir, memory.TokenizeWords)
|
||||
s.Start()
|
||||
defer s.Stop()
|
||||
|
||||
@ -211,7 +211,7 @@ func TestInsertEmptyDoc(t *testing.T) {
|
||||
}
|
||||
defer os.RemoveAll(dir)
|
||||
|
||||
s := NewStore(dir)
|
||||
s := NewStore(dir, memory.TokenizeWords)
|
||||
s.Start()
|
||||
defer s.Stop()
|
||||
|
||||
@ -232,13 +232,13 @@ func TestPersistence(t *testing.T) {
|
||||
defer os.RemoveAll(dir)
|
||||
|
||||
// 写
|
||||
s1 := NewStore(dir)
|
||||
s1 := NewStore(dir, memory.TokenizeWords)
|
||||
s1.Start()
|
||||
s1.Insert(&Doc{Summary: "持久化测试", Content: "应该被保存到磁盘", Source: "manual"})
|
||||
s1.Stop()
|
||||
|
||||
// 读
|
||||
s2 := NewStore(dir)
|
||||
s2 := NewStore(dir, memory.TokenizeWords)
|
||||
s2.Start()
|
||||
defer s2.Stop()
|
||||
|
||||
@ -269,7 +269,7 @@ func TestFlushNoDirty(t *testing.T) {
|
||||
}
|
||||
defer os.RemoveAll(dir)
|
||||
|
||||
s := NewStore(dir)
|
||||
s := NewStore(dir, memory.TokenizeWords)
|
||||
s.Start()
|
||||
|
||||
// 不插任何文档，flush 不应报错
|
||||
@ -283,7 +283,7 @@ func TestRemove(t *testing.T) {
|
||||
}
|
||||
defer os.RemoveAll(dir)
|
||||
|
||||
s := NewStore(dir)
|
||||
s := NewStore(dir, memory.TokenizeWords)
|
||||
s.Start()
|
||||
defer s.Stop()
|
||||
|
||||
@ -329,7 +329,7 @@ func TestRemoveNonexistent(t *testing.T) {
|
||||
}
|
||||
defer os.RemoveAll(dir)
|
||||
|
||||
s := NewStore(dir)
|
||||
s := NewStore(dir, memory.TokenizeWords)
|
||||
s.Start()
|
||||
defer s.Stop()
|
||||
|
||||
@ -437,7 +437,7 @@ func TestContextToDocContentPreservesRawToolOutput(t *testing.T) {
|
||||
}
|
||||
defer os.RemoveAll(dir)
|
||||
|
||||
s := NewStore(dir)
|
||||
s := NewStore(dir, memory.TokenizeWords)
|
||||
s.Start()
|
||||
defer s.Stop()
|
||||
|
||||
|
||||
@ -45,6 +45,10 @@ type Triple struct {
|
||||
SubjectType string `json:"subject_type,omitempty"`
|
||||
ObjectType string `json:"object_type,omitempty"`
|
||||
SentenceText string `json:"sentence_text,omitempty"` // 原始句子文本，Commit时写入sentences表
|
||||
// MediaDigests 是该三元组显式携带的媒体 digest（完整或前缀）。
|
||||
// 媒体不再靠正文 marker 反解：结构化字段直接给出归属，
|
||||
// 由调用方（core）把它变成 L3 一等块并与句子建立结构边。
|
||||
MediaDigests []string `json:"media_digests,omitempty"`
|
||||
}
|
||||
|
||||
type GraphDB struct {
|
||||
@ -134,6 +138,11 @@ func (g *GraphDB) initSchema() error {
|
||||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||
UNIQUE(source_kind, source_id, target_kind, target_id, edge_type)
|
||||
)`,
|
||||
`CREATE TABLE IF NOT EXISTS documents (
|
||||
id TEXT PRIMARY KEY,
|
||||
summary TEXT DEFAULT '',
|
||||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
||||
)`,
|
||||
`CREATE INDEX IF NOT EXISTS idx_memory_blocks_modality ON memory_blocks(modality)`,
|
||||
`CREATE INDEX IF NOT EXISTS idx_memory_blocks_digest ON memory_blocks(payload_digest)`,
|
||||
`CREATE INDEX IF NOT EXISTS idx_memory_block_edges_source ON memory_block_edges(source_kind, source_id)`,
|
||||
@ -979,17 +988,42 @@ func (g *GraphDB) ClearSentenceID(relationID int64) error {
|
||||
return err
|
||||
}
|
||||
|
||||
// CleanupOrphanedSentences 删除没有任何关系引用的句子，返回删除数
|
||||
// CleanupOrphanedSentences 删除既无关系引用、也无媒体块边的句子，返回删除数。
|
||||
//
|
||||
// 两个条件都必须看：旧媒体实体被迁移成原生块后，那些句子可能只靠
|
||||
// sentence --contains--> block 存活，若只看 relations 引用就会被误删，
|
||||
// 连带把块边变成悬空引用。
|
||||
func (g *GraphDB) CleanupOrphanedSentences() (int, error) {
|
||||
result, err := g.db.Exec(
|
||||
`DELETE FROM sentences WHERE id NOT IN (
|
||||
SELECT DISTINCT sentence_id FROM relations WHERE sentence_id != 0
|
||||
)`,
|
||||
)
|
||||
g.mu.Lock()
|
||||
defer g.mu.Unlock()
|
||||
tx, err := g.db.Begin()
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
n, _ := result.RowsAffected()
|
||||
defer tx.Rollback()
|
||||
|
||||
// 先清掉指向将被删除句子的块边，避免留下悬空端点。
|
||||
if _, err := tx.Exec(`DELETE FROM memory_block_edges
|
||||
WHERE source_kind = 'sentence' AND source_id NOT IN (
|
||||
SELECT CAST(id AS TEXT) FROM sentences
|
||||
WHERE id IN (SELECT DISTINCT sentence_id FROM relations WHERE sentence_id != 0)
|
||||
OR id IN (SELECT CAST(source_id AS INTEGER) FROM memory_block_edges WHERE source_kind = 'sentence')
|
||||
)`); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
res, err := tx.Exec(`DELETE FROM sentences WHERE id NOT IN (
|
||||
SELECT DISTINCT sentence_id FROM relations WHERE sentence_id != 0
|
||||
) AND id NOT IN (
|
||||
SELECT CAST(source_id AS INTEGER) FROM memory_block_edges WHERE source_kind = 'sentence'
|
||||
)`)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
n, _ := res.RowsAffected()
|
||||
if err := tx.Commit(); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return int(n), nil
|
||||
}
|
||||
|
||||
|
||||
@ -10,19 +10,19 @@ import (
|
||||
)
|
||||
|
||||
type Indexer struct {
|
||||
db *GraphDB
|
||||
vec *vector.Store
|
||||
veczer *vector.TFIDFVectorizer
|
||||
mu sync.RWMutex
|
||||
trained bool
|
||||
recalled map[string]bool // 已通过工具调用显式召回的实体名，自动注入时跳过
|
||||
db *GraphDB
|
||||
vec *vector.Store
|
||||
veczer *vector.TFIDFVectorizer
|
||||
mu sync.RWMutex
|
||||
trained bool
|
||||
recalled map[string]bool // 已通过工具调用显式召回的实体名，自动注入时跳过
|
||||
}
|
||||
|
||||
func NewIndexer(db *GraphDB) *Indexer {
|
||||
return &Indexer{
|
||||
db: db,
|
||||
vec: vector.NewStore(),
|
||||
veczer: vector.NewTFIDFVectorizer(TokenizeWords),
|
||||
db: db,
|
||||
vec: vector.NewStore(),
|
||||
veczer: vector.NewTFIDFVectorizer(TokenizeWords),
|
||||
recalled: make(map[string]bool),
|
||||
}
|
||||
}
|
||||
|
||||
@ -69,11 +69,6 @@ type Item struct {
|
||||
OriginPath string `json:"origin_path,omitempty"`
|
||||
// Tool 是注入这条媒体的工具名（如 multimodal_see_picture）。
|
||||
Tool string `json:"tool,omitempty"`
|
||||
// Description 是视觉/音频模型生成的文字描述，供 L2/L3 检索。
|
||||
// 空表示未描述（未开启描述、模型不可用或描述失败）。
|
||||
Description string `json:"description,omitempty"`
|
||||
// DescribedBy 记录描述来自哪个源，让后续读者能判断可靠性。
|
||||
DescribedBy string `json:"described_by,omitempty"`
|
||||
// FirstSeen/LastSeen 是首末次入库时间。
|
||||
FirstSeen time.Time `json:"first_seen"`
|
||||
LastSeen time.Time `json:"last_seen"`
|
||||
@ -127,8 +122,6 @@ func (s *Store) initSchema() error {
|
||||
height INTEGER DEFAULT 0,
|
||||
origin_path TEXT,
|
||||
tool TEXT,
|
||||
description TEXT,
|
||||
described_by TEXT,
|
||||
first_seen TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||
last_seen TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
||||
)`,
|
||||
@ -196,20 +189,15 @@ func (s *Store) Put(data []byte, meta Item) (string, error) {
|
||||
}
|
||||
_, err := s.db.Exec(`
|
||||
INSERT INTO media (digest, kind, mime, size, width, height,
|
||||
origin_path, tool, description, described_by,
|
||||
first_seen, last_seen)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
origin_path, tool, first_seen, last_seen)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
ON CONFLICT(digest) DO UPDATE SET
|
||||
last_seen = excluded.last_seen,
|
||||
-- 只在原值为空时补写：先到的描述可能来自更强的模型，
|
||||
-- 后到的空值不该把它冲掉。
|
||||
description = CASE WHEN COALESCE(media.description,'') = '' THEN excluded.description ELSE media.description END,
|
||||
described_by = CASE WHEN COALESCE(media.described_by,'') = '' THEN excluded.described_by ELSE media.described_by END,
|
||||
width = CASE WHEN media.width = 0 THEN excluded.width ELSE media.width END,
|
||||
height = CASE WHEN media.height = 0 THEN excluded.height ELSE media.height END,
|
||||
tool = CASE WHEN COALESCE(media.tool,'') = '' THEN excluded.tool ELSE media.tool END
|
||||
`, digest, string(meta.Kind), meta.MIME, int64(len(data)), meta.Width, meta.Height,
|
||||
meta.OriginPath, meta.Tool, meta.Description, meta.DescribedBy, now, now)
|
||||
meta.OriginPath, meta.Tool, now, now)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("media: upsert meta: %w", err)
|
||||
}
|
||||
@ -242,115 +230,12 @@ func (s *Store) Stat(digest string) (*Item, error) {
|
||||
defer s.mu.RUnlock()
|
||||
return s.scanOne(s.db.QueryRow(`
|
||||
SELECT digest, kind, mime, size, width, height, origin_path, tool,
|
||||
description, described_by, first_seen, last_seen,
|
||||
first_seen, last_seen,
|
||||
vec, vec_model
|
||||
FROM media WHERE digest = ?`, digest))
|
||||
}
|
||||
|
||||
// Describe 写入（或覆盖）文字描述。
|
||||
//
|
||||
// 与 Put 的"只在空时补写"不同：Describe 是显式操作，调用方明确想要这份
|
||||
// 描述生效（例如换了更强的视觉模型重新描述）。
|
||||
func (s *Store) Describe(digest, description, describedBy string) error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
res, err := s.db.Exec(`UPDATE media SET description = ?, described_by = ? WHERE digest = ?`,
|
||||
description, describedBy, digest)
|
||||
if err != nil {
|
||||
return fmt.Errorf("media: describe: %w", err)
|
||||
}
|
||||
if n, _ := res.RowsAffected(); n == 0 {
|
||||
return fmt.Errorf("media: describe: unknown digest %s", shortDigest(digest))
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Search 按描述文本做 LIKE 匹配，返回最近的若干条。
|
||||
//
|
||||
// 刻意不在这里做向量检索：媒体的语义检索走 L2 文档层的既有索引
|
||||
// （描述文字随记忆条目一起进 Doc.Content，复用那套 TF-IDF/embedding），
|
||||
// 本方法只是"按关键词直接翻媒体库"的补充入口。
|
||||
func (s *Store) Search(query string, kind Kind, limit int) ([]*Item, error) {
|
||||
if limit <= 0 {
|
||||
limit = 20
|
||||
}
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
|
||||
q := `SELECT digest, kind, mime, size, width, height, origin_path, tool,
|
||||
description, described_by, first_seen, last_seen,
|
||||
vec, vec_model
|
||||
FROM media WHERE COALESCE(description,'') != ''`
|
||||
args := []interface{}{}
|
||||
if strings.TrimSpace(query) != "" {
|
||||
q += ` AND description LIKE ?`
|
||||
args = append(args, "%"+query+"%")
|
||||
}
|
||||
if kind != "" {
|
||||
q += ` AND kind = ?`
|
||||
args = append(args, string(kind))
|
||||
}
|
||||
q += ` ORDER BY last_seen DESC LIMIT ?`
|
||||
args = append(args, limit)
|
||||
|
||||
rows, err := s.db.Query(q, args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var out []*Item
|
||||
for rows.Next() {
|
||||
it, err := s.scanRows(rows)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
out = append(out, it)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// Pending 返回尚无描述的媒体，供后台描述任务消费。
|
||||
// Pending 返回尚无描述的媒体，供后台描述任务消费。
|
||||
//
|
||||
// 不只看 description 为空，还要求 described_by 也为空。
|
||||
// 因为“已尝试但无法描述”的项（如 kind=other 的二进制、blob 已丢失）
|
||||
// 会被标记为 described_by=unsupported/content-missing 而 description 仍为空——
|
||||
// 若只看 description，这些项每轮都会被取出来重试，永远卡在队列头部，
|
||||
// 真正需要描述的新项永远轮不到（LIMIT 只取前 N 条）。
|
||||
func (s *Store) Pending(limit int) ([]*Item, error) {
|
||||
if limit <= 0 {
|
||||
limit = 10
|
||||
}
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
rows, err := s.db.Query(`
|
||||
SELECT digest, kind, mime, size, width, height, origin_path, tool,
|
||||
description, described_by, first_seen, last_seen,
|
||||
vec, vec_model
|
||||
FROM media
|
||||
WHERE COALESCE(description,'') = '' AND COALESCE(described_by,'') = ''
|
||||
ORDER BY last_seen DESC LIMIT ?`, limit)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var out []*Item
|
||||
for rows.Next() {
|
||||
it, err := s.scanRows(rows)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
out = append(out, it)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// GC 清理已不被任何记忆块持有的内容。
|
||||
//
|
||||
// keep 是当前仍被 Context/Document/Graph 里一等记忆块持有的 digest 集合，
|
||||
// 由调用方从三层记忆节点计算得出；media.Store 不再自己维护引用账本。
|
||||
// 不在 keep 中且早于 minAge 的项被清理；超出 maxBytes 时也只淘汰不在 keep 中的项。
|
||||
// Delete 删除一份媒体内容（元数据 + blob）。
|
||||
// Stat 返回元数据，不读内容。
|
||||
//
|
||||
// 这不是 GC，也不看引用计数：调用方是记忆系统本身——当它把一个记忆块
|
||||
// 永久地从三层记忆中删掉（而非在层间迁移）时，媒体作为块的内容一并删除。
|
||||
@ -376,13 +261,11 @@ func (s *Store) Stats() map[string]interface{} {
|
||||
defer s.mu.RUnlock()
|
||||
|
||||
out := map[string]interface{}{"blob_dir": s.blobDir}
|
||||
var count, described int
|
||||
var count int
|
||||
var total int64
|
||||
s.db.QueryRow(`SELECT COUNT(*), COALESCE(SUM(size),0) FROM media`).Scan(&count, &total)
|
||||
s.db.QueryRow(`SELECT COUNT(*) FROM media WHERE COALESCE(description,'') != ''`).Scan(&described)
|
||||
out["count"] = count
|
||||
out["total_bytes"] = total
|
||||
out["described"] = described
|
||||
|
||||
byKind := map[string]int{}
|
||||
rows, err := s.db.Query(`SELECT kind, COUNT(*) FROM media GROUP BY kind`)
|
||||
@ -524,7 +407,7 @@ func (s *Store) queryMediaScored(queryVec []float64, model string, topK int) ([]
|
||||
defer s.mu.RUnlock()
|
||||
|
||||
query := `SELECT digest, kind, mime, size, width, height,
|
||||
origin_path, tool, description, described_by, first_seen, last_seen,
|
||||
origin_path, tool, first_seen, last_seen,
|
||||
vec, vec_model
|
||||
FROM media WHERE vec IS NOT NULL AND vec != ''`
|
||||
var args []interface{}
|
||||
@ -546,17 +429,15 @@ func (s *Store) queryMediaScored(queryVec []float64, model string, topK int) ([]
|
||||
for rows.Next() {
|
||||
var it Item
|
||||
var kind string
|
||||
var origin, tool, desc, by, vecJSON, vecModel sql.NullString
|
||||
var origin, tool, vecJSON, vecModel sql.NullString
|
||||
if err := rows.Scan(&it.Digest, &kind, &it.MIME, &it.Size, &it.Width, &it.Height,
|
||||
&origin, &tool, &desc, &by, &it.FirstSeen, &it.LastSeen,
|
||||
&origin, &tool, &it.FirstSeen, &it.LastSeen,
|
||||
&vecJSON, &vecModel); err != nil {
|
||||
continue
|
||||
}
|
||||
it.Kind = Kind(kind)
|
||||
it.OriginPath = origin.String
|
||||
it.Tool = tool.String
|
||||
it.Description = desc.String
|
||||
it.DescribedBy = by.String
|
||||
if !vecJSON.Valid || vecJSON.String == "" {
|
||||
continue
|
||||
}
|
||||
@ -623,17 +504,15 @@ func (s *Store) scanRows(r rowScanner) (*Item, error) { return scanItem(r) }
|
||||
func scanItem(r rowScanner) (*Item, error) {
|
||||
var it Item
|
||||
var kind string
|
||||
var origin, tool, desc, by, vecJSON, vecModel sql.NullString
|
||||
var origin, tool, vecJSON, vecModel sql.NullString
|
||||
if err := r.Scan(&it.Digest, &kind, &it.MIME, &it.Size, &it.Width, &it.Height,
|
||||
&origin, &tool, &desc, &by, &it.FirstSeen, &it.LastSeen,
|
||||
&origin, &tool, &it.FirstSeen, &it.LastSeen,
|
||||
&vecJSON, &vecModel); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
it.Kind = Kind(kind)
|
||||
it.OriginPath = origin.String
|
||||
it.Tool = tool.String
|
||||
it.Description = desc.String
|
||||
it.DescribedBy = by.String
|
||||
if vecJSON.Valid && vecJSON.String != "" {
|
||||
var v []float64
|
||||
if err := json.Unmarshal([]byte(vecJSON.String), &v); err == nil {
|
||||
|
||||
@ -142,50 +142,13 @@ func TestDelete_UnknownDigestIsNoop(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestDescribe_OverwritesExplicitly(t *testing.T) {
|
||||
s := newTestStore(t, 0)
|
||||
d, _ := s.Put([]byte("img"), Item{MIME: "image/png"})
|
||||
|
||||
if err := s.Describe(d, "一只橘猫", "vis-a"); err != nil {
|
||||
func TestDescribe_Removed(t *testing.T) {
|
||||
// 媒体不再有文字描述：描述式索引是废弃的就机制。
|
||||
// 这里只保留一个编译期断言，确保 API 不会静默回归。
|
||||
s := newTestStore(t)
|
||||
if _, err := s.Put([]byte("img"), Item{MIME: "image/png"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
it, _ := s.Stat(d)
|
||||
if it.Description != "一只橘猫" || it.DescribedBy != "vis-a" {
|
||||
t.Fatalf("描述未写入: %+v", it)
|
||||
}
|
||||
|
||||
// Describe 是显式操作，允许覆盖（换更强模型重描述）
|
||||
if err := s.Describe(d, "一只橘色虎斑猫坐在窗台", "vis-b"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
it, _ = s.Stat(d)
|
||||
if !strings.Contains(it.Description, "虎斑") || it.DescribedBy != "vis-b" {
|
||||
t.Fatalf("Describe 应覆盖旧描述: %+v", it)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDescribe_UnknownDigestErrors(t *testing.T) {
|
||||
s := newTestStore(t, 0)
|
||||
err := s.Describe("deadbeef", "x", "y")
|
||||
if err == nil {
|
||||
t.Fatal("未知 digest 应报错而非静默成功")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPut_DoesNotClobberExistingDescription(t *testing.T) {
|
||||
// 先到的描述可能来自更强的模型；后到的空值不该把它冲掉。
|
||||
s := newTestStore(t, 0)
|
||||
data := []byte("img")
|
||||
d, _ := s.Put(data, Item{MIME: "image/png", Description: "详细描述", DescribedBy: "strong-model"})
|
||||
|
||||
// 第二次 Put 同内容但不带描述
|
||||
if _, err := s.Put(data, Item{MIME: "image/png"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
it, _ := s.Stat(d)
|
||||
if it.Description != "详细描述" || it.DescribedBy != "strong-model" {
|
||||
t.Fatalf("重复 Put 的空描述不该冲掉已有描述: %+v", it)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPut_BackfillsMissingDimensions(t *testing.T) {
|
||||
@ -206,47 +169,22 @@ func TestPut_BackfillsMissingDimensions(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestSearch_FiltersByDescriptionAndKind(t *testing.T) {
|
||||
s := newTestStore(t, 0)
|
||||
di, _ := s.Put([]byte("chart-img"), Item{MIME: "image/png"})
|
||||
da, _ := s.Put([]byte("speech-aud"), Item{MIME: "audio/wav"})
|
||||
dn, _ := s.Put([]byte("no-desc"), Item{MIME: "image/png"})
|
||||
s.Describe(di, "一张蓝色的柱状图表", "vis")
|
||||
s.Describe(da, "一段关于图表的讲解录音", "aud")
|
||||
func TestStats_CountsByKind(t *testing.T) {
|
||||
s := newTestStore(t)
|
||||
s.Put([]byte("i1"), Item{MIME: "image/png"})
|
||||
s.Put([]byte("i2"), Item{MIME: "image/jpeg"})
|
||||
s.Put([]byte("a1"), Item{MIME: "audio/wav"})
|
||||
|
||||
all, err := s.Search("图表", "", 10)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
st := s.Stats()
|
||||
if st["count"].(int) != 3 {
|
||||
t.Fatalf("count 应为 3，实际 %v", st["count"])
|
||||
}
|
||||
if len(all) != 2 {
|
||||
t.Fatalf("两条描述都含「图表」，应返回 2，实际 %d", len(all))
|
||||
if _, ok := st["described"]; ok {
|
||||
t.Fatal("媒体已不再有描述计数")
|
||||
}
|
||||
|
||||
imgs, _ := s.Search("图表", KindImage, 10)
|
||||
if len(imgs) != 1 || imgs[0].Digest != di {
|
||||
t.Fatalf("按 image 过滤应只剩图片，实际 %d 条", len(imgs))
|
||||
}
|
||||
|
||||
// 无描述的项不该出现在语义检索结果里
|
||||
for _, it := range all {
|
||||
if it.Digest == dn {
|
||||
t.Fatal("无描述的项不该被 Search 返回")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestPending_ReturnsUndescribed(t *testing.T) {
|
||||
s := newTestStore(t, 0)
|
||||
described, _ := s.Put([]byte("has-desc"), Item{MIME: "image/png"})
|
||||
undescribed, _ := s.Put([]byte("needs-desc"), Item{MIME: "image/png"})
|
||||
s.Describe(described, "已有描述", "vis")
|
||||
|
||||
pending, err := s.Pending(10)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(pending) != 1 || pending[0].Digest != undescribed {
|
||||
t.Fatalf("应只返回未描述项，实际 %d 条", len(pending))
|
||||
byKind := st["by_kind"].(map[string]int)
|
||||
if byKind["image"] != 2 || byKind["audio"] != 1 {
|
||||
t.Fatalf("by_kind 不对: %v", byKind)
|
||||
}
|
||||
}
|
||||
|
||||
@ -278,28 +216,8 @@ func TestParseDataURL(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestStats_CountsByKindAndDescription(t *testing.T) {
|
||||
s := newTestStore(t, 4096)
|
||||
d1, _ := s.Put([]byte("i1"), Item{MIME: "image/png"})
|
||||
s.Put([]byte("i2"), Item{MIME: "image/jpeg"})
|
||||
s.Put([]byte("a1"), Item{MIME: "audio/wav"})
|
||||
s.Describe(d1, "描述", "vis")
|
||||
|
||||
st := s.Stats()
|
||||
if st["count"].(int) != 3 {
|
||||
t.Fatalf("count 应为 3，实际 %v", st["count"])
|
||||
}
|
||||
if st["described"].(int) != 1 {
|
||||
t.Fatalf("described 应为 1，实际 %v", st["described"])
|
||||
}
|
||||
byKind := st["by_kind"].(map[string]int)
|
||||
if byKind["image"] != 2 || byKind["audio"] != 1 {
|
||||
t.Fatalf("by_kind 不对: %v", byKind)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPut_RejectsEmpty(t *testing.T) {
|
||||
s := newTestStore(t, 0)
|
||||
s := newTestStore(t)
|
||||
if _, err := s.Put(nil, Item{MIME: "image/png"}); err == nil {
|
||||
t.Fatal("空内容应报错")
|
||||
}
|
||||
@ -313,7 +231,7 @@ func TestReopen_PersistsAcrossRestart(t *testing.T) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
d, _ := s1.Put([]byte("persistent-img"), Item{MIME: "image/png", OriginPath: "/tmp/x.png"})
|
||||
s1.Describe(d, "跨重启的描述", "vis")
|
||||
s1.SetVec(d, []float64{0.1, 0.2}, "test-space")
|
||||
s1.Close()
|
||||
|
||||
s2, err := New(dir)
|
||||
@ -326,8 +244,8 @@ func TestReopen_PersistsAcrossRestart(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Fatalf("重开后应能查到: %v", err)
|
||||
}
|
||||
if it.Description != "跨重启的描述" {
|
||||
t.Fatalf("元数据应持久化: %+v", it)
|
||||
if it.OriginPath != "/tmp/x.png" || len(it.Vec) != 2 || it.VecModel != "test-space" {
|
||||
t.Fatalf("元数据与向量应持久化: %+v", it)
|
||||
}
|
||||
data, err := s2.Get(d)
|
||||
if err != nil || string(data) != "persistent-img" {
|
||||
@ -335,38 +253,26 @@ func TestReopen_PersistsAcrossRestart(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestPending_ExcludesAttemptedButUndescribable(t *testing.T) {
|
||||
// 「已尝试但无法描述」的项必须退出待描述队列。
|
||||
//
|
||||
// 这些项被标记为 described_by=unsupported/content-missing 而 description
|
||||
// 仍为空。若 Pending 只看 description，它们每轮都会被取出来重试、
|
||||
// 永久占着 LIMIT 的名额，真正需要描述的新项永远轮不到。
|
||||
s := newTestStore(t, 0)
|
||||
func TestStaleVecDigests_TracksModelSwitch(t *testing.T) {
|
||||
// 模型切换后旧向量必须被重算：StaleVecDigests 是启动迁移的入口。
|
||||
s := newTestStore(t)
|
||||
d1, _ := s.Put([]byte("a"), Item{MIME: "image/png"})
|
||||
d2, _ := s.Put([]byte("b"), Item{MIME: "image/png"})
|
||||
s.SetVec(d1, []float64{0.1}, "space-a")
|
||||
|
||||
fresh, _ := s.Put([]byte("needs-describe"), Item{MIME: "image/png"})
|
||||
unsupported, _ := s.Put([]byte("cannot-describe"), Item{MIME: "application/octet-stream"})
|
||||
described, _ := s.Put([]byte("已描述"), Item{MIME: "image/png"})
|
||||
|
||||
// 标记「尝试过但不支持」：description 空，described_by 非空
|
||||
if err := s.Describe(unsupported, "", "unsupported"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := s.Describe(described, "一张图", "visionllm"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
pending, err := s.Pending(10)
|
||||
stale, err := s.StaleVecDigestsAll("space-a")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(pending) != 1 {
|
||||
var names []string
|
||||
for _, p := range pending {
|
||||
names = append(names, shortDigest(p.Digest))
|
||||
}
|
||||
t.Fatalf("应只剩 1 条待描述，实际 %d 条: %v", len(pending), names)
|
||||
if len(stale) != 1 || stale[0] != d2 {
|
||||
t.Fatalf("只有未嵌入的 d2 需重算，实际 %v", stale)
|
||||
}
|
||||
if pending[0].Digest != fresh {
|
||||
t.Fatalf("待描述的应是未处理项，实际 %s", shortDigest(pending[0].Digest))
|
||||
|
||||
stale, err = s.StaleVecDigestsAll("space-b")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(stale) != 2 {
|
||||
t.Fatalf("换空间后两条都需重算，实际 %v", stale)
|
||||
}
|
||||
}
|
||||
|
||||
@ -10,9 +10,9 @@ func TestQueryMedia_BasicSimilarity(t *testing.T) {
|
||||
defer s.Close()
|
||||
|
||||
// 入库三张带向量的媒体：两张图、一段音频
|
||||
d1, _ := s.Put([]byte("img1"), Item{MIME: "image/png", Description: "紫蓝红三色带"})
|
||||
d2, _ := s.Put([]byte("img2"), Item{MIME: "image/jpeg", Description: "蓝紫红渐变"})
|
||||
d3, _ := s.Put([]byte("aud1"), Item{MIME: "audio/wav", Description: "一段语音"})
|
||||
d1, _ := s.Put([]byte("img1"), Item{MIME: "image/png"})
|
||||
d2, _ := s.Put([]byte("img2"), Item{MIME: "image/jpeg"})
|
||||
d3, _ := s.Put([]byte("aud1"), Item{MIME: "audio/wav"})
|
||||
|
||||
// 模拟视觉嵌入：img1 和 img2 向量接近，aud1 远离
|
||||
vec1 := []float64{0.9, 0.1, 0.0, 0.0}
|
||||
@ -138,24 +138,24 @@ func TestStaleVecDigests(t *testing.T) {
|
||||
s := newTestStore(t, 0)
|
||||
defer s.Close()
|
||||
|
||||
// 有描述且 vec_model 匹配 → 非 stale
|
||||
d1, _ := s.Put([]byte("img1"), Item{MIME: "image/png", Description: "图一"})
|
||||
s.SetVec(d1, []float64{0.1}, "clip-vit-b32")
|
||||
// 有向量且 vec_model 匹配 → 非 stale
|
||||
d1, _ := s.Put([]byte("img1"), Item{MIME: "image/png"})
|
||||
s.SetVec(d1, []float64{0.1}, "space-a")
|
||||
|
||||
// 有描述但 vec_model 旧 → stale
|
||||
d2, _ := s.Put([]byte("img2"), Item{MIME: "image/png", Description: "图二"})
|
||||
s.SetVec(d2, []float64{0.2}, "clip-vit-b14")
|
||||
// 有向量但 vec_model 旧 → stale
|
||||
d2, _ := s.Put([]byte("img2"), Item{MIME: "image/png"})
|
||||
s.SetVec(d2, []float64{0.2}, "space-old")
|
||||
|
||||
// 有描述但从未嵌入（vec_model 空）→ stale
|
||||
d3, _ := s.Put([]byte("img3"), Item{MIME: "image/png", Description: "图三"})
|
||||
// 从未嵌入（vec_model 空）→ stale
|
||||
d3, _ := s.Put([]byte("img3"), Item{MIME: "image/png"})
|
||||
|
||||
// 无描述但有图片 → 也应被迁移（描述是可选语义通道，图片应独立于描述参与向量空间）
|
||||
// 与向量/描述无关的图片同样应被迁移：图片独立参与向量空间
|
||||
d4, _ := s.Put([]byte("img4"), Item{MIME: "image/png"})
|
||||
|
||||
// 音频不参与图片迁移（StaleVecDigests 只查 kind='image'）
|
||||
s.Put([]byte("aud1"), Item{MIME: "audio/wav", Description: "语音"})
|
||||
s.Put([]byte("aud1"), Item{MIME: "audio/wav"})
|
||||
|
||||
stale, err := s.StaleVecDigests("clip-vit-b32")
|
||||
stale, err := s.StaleVecDigests("space-a")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
@ -123,31 +123,44 @@ func TestSmoke_VideoFramesDistinct(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestSmoke_DescribeThenRetrieve(t *testing.T) {
|
||||
// 场景 C：视觉模型描述落库后，描述文字成为可检索的语义入口。
|
||||
// 这是本方案最关键的一环——blob 可能被淘汰，描述会长期留在记忆里。
|
||||
s := newTestStore(t, 50*1024*1024)
|
||||
func TestSmoke_NearestNeighborVectorRetrieve(t *testing.T) {
|
||||
// 场景 C：图片只按自己的原生向量被检索。
|
||||
// 没有描述文本参与——描述式索引是废弃的就机制。
|
||||
s := newTestStore(t)
|
||||
|
||||
pic, _ := s.Put(makePNG(400, 400, 0), Item{MIME: "image/png", Tool: "multimodal_see_picture"})
|
||||
if err := s.Describe(pic, "一张 400x400 的三色带图：上红、中绿、下蓝", "visionllm"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
s.SetVec(pic, []float64{1, 0, 0, 0}, "space")
|
||||
var frames []string
|
||||
for i := 0; i < 6; i++ {
|
||||
d, _ := s.Put(makePNG(320, 240, i), Item{MIME: "image/jpeg", Tool: "multimodal_see_video"})
|
||||
if err := s.Describe(d, fmt.Sprintf("视频第 %d 帧：测试图卡，含彩条与计数器", i+1), "visionllm"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
s.SetVec(d, []float64{1, 1, float64(i) / 10, 0}, "space")
|
||||
frames = append(frames, d)
|
||||
}
|
||||
|
||||
if hits, _ := s.Search("三色带", KindImage, 10); len(hits) != 1 {
|
||||
t.Fatalf("搜「三色带」应命中 1 条，实际 %d", len(hits))
|
||||
hits, err := s.QueryMediaScored([]float64{1, 0, 0, 0}, "space", 10)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if hits, _ := s.Search("计数器", KindImage, 10); len(hits) != 6 {
|
||||
t.Fatalf("搜「计数器」应命中 6 帧，实际 %d", len(hits))
|
||||
if len(hits) != 7 {
|
||||
t.Fatalf("7 份媒体都有同空间向量，应全部可召，实际 %d", len(hits))
|
||||
}
|
||||
pend, _ := s.Pending(100)
|
||||
if len(pend) != 0 {
|
||||
t.Fatalf("应全部已描述，仍有 %d 条待描述", len(pend))
|
||||
if hits[0].Item.Digest != pic {
|
||||
t.Fatalf("与查询同向的应是第一命中，实际 %s", shortDigest(hits[0].Item.Digest))
|
||||
}
|
||||
|
||||
// 不同向量空间/模型的条目不得参与：坐标系不同，余弦无意义。
|
||||
foreign := frames[0]
|
||||
if err := s.SetVec(foreign, []float64{1, 0, 0, 0}, "other-space"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
hits, err = s.QueryMediaScored([]float64{1, 0, 0, 0}, "space", 10)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, h := range hits {
|
||||
if h.Item.Digest == foreign {
|
||||
t.Fatal("另一套空间（other-space）的向量不该被 space 查询召回")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@ -204,7 +217,7 @@ func TestSmoke_DeleteRemovesOnlyThatContent(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestSmoke_FullLifecycleAcrossRestart(t *testing.T) {
|
||||
// 端到端：入库 → 描述 → 删除一些内容 → 重启 → 检索，
|
||||
// 端到端：入库 → 嵌入 → 删除一些内容 → 重启 → 向量检索，
|
||||
// 并确认磁盘与元数据不出现双向孤儿。记忆的意义就在于跨重启还在。
|
||||
dir := t.TempDir()
|
||||
s, err := New(dir)
|
||||
@ -214,10 +227,10 @@ func TestSmoke_FullLifecycleAcrossRestart(t *testing.T) {
|
||||
|
||||
png := makePNG(400, 400, 0)
|
||||
pic, _ := s.Put(png, Item{MIME: "image/png", Width: 400, Height: 400, Tool: "multimodal_see_picture"})
|
||||
s.Describe(pic, "一张 400x400 的三色带图：上红、中绿、下蓝", "visionllm")
|
||||
s.SetVec(pic, []float64{1, 0, 0}, "space")
|
||||
for i := 0; i < 6; i++ {
|
||||
d, _ := s.Put(makePNG(320, 240, i), Item{MIME: "image/jpeg", Tool: "multimodal_see_video"})
|
||||
s.Describe(d, fmt.Sprintf("视频第 %d 帧", i+1), "visionllm")
|
||||
s.SetVec(d, []float64{0, 1, float64(i)}, "space")
|
||||
}
|
||||
for i := 0; i < 10; i++ {
|
||||
d, _ := s.Put(makePNG(64, 64, 2000+i), Item{MIME: "image/png", Tool: "cmd_run"})
|
||||
@ -241,20 +254,24 @@ func TestSmoke_FullLifecycleAcrossRestart(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Fatalf("重开后查不到: %v", err)
|
||||
}
|
||||
if it.Description == "" {
|
||||
t.Fatalf("元数据未持久化: %+v", it)
|
||||
if len(it.Vec) != 3 || it.VecModel != "space" {
|
||||
t.Fatalf("向量未持久化: %+v", it)
|
||||
}
|
||||
data, err := s2.Get(pic)
|
||||
if err != nil || !bytes.Equal(data, png) {
|
||||
t.Fatalf("重开后内容不一致: %v", err)
|
||||
}
|
||||
if hits, _ := s2.Search("三色带", KindImage, 10); len(hits) != 1 {
|
||||
t.Fatal("重开后描述应仍可检索")
|
||||
hits, err := s2.QueryMediaScored([]float64{1, 0, 0}, "space", 10)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(hits) == 0 || hits[0].Item.Digest != pic {
|
||||
t.Fatal("重开后向量检索应仍能命中")
|
||||
}
|
||||
|
||||
// 磁盘文件数 == 元数据条数：无「元数据在文件没了」也无「文件在元数据没了」
|
||||
if n := blobFileCount(t, s2); n != beforeCount {
|
||||
t.Fatalf("磁盘 blob=%d 与元数据=%d 不一致", n, beforeCount)
|
||||
}
|
||||
t.Logf("跨重启：%d 条目、描述与内容全部完好", beforeCount)
|
||||
t.Logf("跨重启：%d 条目、向量与内容全部完好", beforeCount)
|
||||
}
|
||||
|
||||
@ -39,7 +39,7 @@ func TestSoak_SustainedMixedLoad(t *testing.T) {
|
||||
|
||||
stop := make(chan struct{})
|
||||
var wg sync.WaitGroup
|
||||
var puts, gets, deletes, describes, searches atomic.Int64
|
||||
var puts, gets, deletes, embeds, searches atomic.Int64
|
||||
var fatal atomic.Int64
|
||||
|
||||
worker := func(name string, fn func(iter int) error) {
|
||||
@ -109,25 +109,26 @@ func TestSoak_SustainedMixedLoad(t *testing.T) {
|
||||
return nil
|
||||
})
|
||||
|
||||
// 描述者
|
||||
worker("describe", func(i int) error {
|
||||
pend, err := s.Pending(5)
|
||||
// 向量写入者：持续给新内容嵌入并删除（模拟启动迁移/短命媒体）
|
||||
worker("embed", func(i int) error {
|
||||
b := make([]byte, 1024)
|
||||
rand.Read(b)
|
||||
b = append([]byte(fmt.Sprintf("emb-%d-", i)), b...)
|
||||
d, err := s.Put(b, Item{MIME: "image/png", Tool: "cmd_run"})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
for _, it := range pend {
|
||||
// 忽略 unknown digest：GC 可能在 Pending 与 Describe 之间清掉它，
|
||||
// 这是正常竞态而非缺陷。
|
||||
_ = s.Describe(it.Digest, fmt.Sprintf("描述 %d 含图表与文字", i), "vis")
|
||||
describes.Add(1)
|
||||
if err := s.SetVec(d, []float64{1, float64(i % 7)}, "soak-space"); err != nil {
|
||||
return err
|
||||
}
|
||||
time.Sleep(2 * time.Millisecond)
|
||||
embeds.Add(1)
|
||||
time.Sleep(time.Millisecond)
|
||||
return nil
|
||||
})
|
||||
|
||||
// 检索者
|
||||
worker("search", func(i int) error {
|
||||
if _, err := s.Search("图表", KindImage, 20); err != nil {
|
||||
if _, err := s.QueryMediaScored([]float64{1, 0}, "soak-space", 20); err != nil {
|
||||
return err
|
||||
}
|
||||
if _, err := s.Stat(keep[i%keepN]); err != nil {
|
||||
@ -146,8 +147,8 @@ func TestSoak_SustainedMixedLoad(t *testing.T) {
|
||||
t.Fatalf("%d 个 worker 报致命错误", n)
|
||||
}
|
||||
|
||||
t.Logf("%v 内: put=%d get=%d delete=%d describe=%d search=%d",
|
||||
dur, puts.Load(), gets.Load(), deletes.Load(), describes.Load(), searches.Load())
|
||||
t.Logf("%v 内: put=%d get=%d delete=%d embed=%d search=%d",
|
||||
dur, puts.Load(), gets.Load(), deletes.Load(), embeds.Load(), searches.Load())
|
||||
|
||||
// 收尾断言
|
||||
for i, d := range keep {
|
||||
@ -161,6 +162,5 @@ func TestSoak_SustainedMixedLoad(t *testing.T) {
|
||||
}
|
||||
|
||||
st := s.Stats()
|
||||
t.Logf("收尾: 条目=%v 字节=%v 已描述=%v",
|
||||
st["count"], st["total_bytes"], st["described"])
|
||||
t.Logf("收尾: 条目=%v 字节=%v 类型=%v", st["count"], st["total_bytes"], st["by_kind"])
|
||||
}
|
||||
|
||||
@ -306,9 +306,9 @@ func TestStress_DeleteConcurrentWithReads(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestStress_DescribeConcurrentWithSearch(t *testing.T) {
|
||||
// 描述写入与检索并发。C 部分的后台描述任务会长期这样跑。
|
||||
s := newTestStore(t, 0)
|
||||
func TestStress_SetVecConcurrentWithQuery(t *testing.T) {
|
||||
// 嵌入写入与向量检索并发（启动时的向量迁移就会长期这样跑）。
|
||||
s := newTestStore(t)
|
||||
const n = 60
|
||||
digests := make([]string, n)
|
||||
for i := range digests {
|
||||
@ -320,55 +320,52 @@ func TestStress_DescribeConcurrentWithSearch(t *testing.T) {
|
||||
}
|
||||
|
||||
var wg sync.WaitGroup
|
||||
var descErr, searchErr atomic.Int64
|
||||
var writeErr, queryErr atomic.Int64
|
||||
|
||||
// 描述写入者
|
||||
// 向量写入者
|
||||
for w := 0; w < 4; w++ {
|
||||
wg.Add(1)
|
||||
go func(wid int) {
|
||||
defer wg.Done()
|
||||
for i := wid; i < n; i += 4 {
|
||||
desc := fmt.Sprintf("第 %d 张图，含蓝色图表与文字", i)
|
||||
if err := s.Describe(digests[i], desc, "vis-src"); err != nil {
|
||||
descErr.Add(1)
|
||||
vec := []float64{1, float64(i) / 100, 0}
|
||||
if err := s.SetVec(digests[i], vec, "space"); err != nil {
|
||||
writeErr.Add(1)
|
||||
}
|
||||
}
|
||||
}(w)
|
||||
}
|
||||
|
||||
// 检索者 + Pending 消费者
|
||||
// 检索者
|
||||
for r := 0; r < 3; r++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for i := 0; i < 50; i++ {
|
||||
if _, err := s.Search("图表", KindImage, 20); err != nil {
|
||||
searchErr.Add(1)
|
||||
}
|
||||
if _, err := s.Pending(10); err != nil {
|
||||
searchErr.Add(1)
|
||||
if _, err := s.QueryMediaScored([]float64{1, 0, 0}, "space", 20); err != nil {
|
||||
queryErr.Add(1)
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
if v := descErr.Load(); v > 0 {
|
||||
t.Fatalf("Describe 失败 %d 次", v)
|
||||
if v := writeErr.Load(); v > 0 {
|
||||
t.Fatalf("SetVec 失败 %d 次", v)
|
||||
}
|
||||
if v := searchErr.Load(); v > 0 {
|
||||
t.Fatalf("Search/Pending 失败 %d 次", v)
|
||||
if v := queryErr.Load(); v > 0 {
|
||||
t.Fatalf("QueryMediaScored 失败 %d 次", v)
|
||||
}
|
||||
|
||||
// 全部应已描述完
|
||||
pending, err := s.Pending(1000)
|
||||
// 全部应已嵌入，且都在同一空间
|
||||
stale, err := s.StaleVecDigestsAll("space")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(pending) != 0 {
|
||||
t.Fatalf("应全部描述完，仍有 %d 条未描述", len(pending))
|
||||
if len(stale) != 0 {
|
||||
t.Fatalf("应全部已嵌入，仍有 %d 条未嵌入", len(stale))
|
||||
}
|
||||
got, err := s.Search("图表", KindImage, 1000)
|
||||
got, err := s.QueryMediaScored([]float64{1, 0, 0}, "space", 1000)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
193
internal/memory/migrate.go
Normal file
193
internal/memory/migrate.go
Normal file
@ -0,0 +1,193 @@
|
||||
package memory
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"regexp"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// 旧媒体实体迁移。
|
||||
//
|
||||
// 历史上媒体进 L3 的方式是把正文 marker 反解成普通实体三元组：
|
||||
//
|
||||
// [image/png a1b2c3d4e5f6] 一张紫蓝红三色带图
|
||||
// → 实体「图片 a1b2c3d4e5f6」(type=Media) -内容-> 「一张紫蓝红三色带图」
|
||||
//
|
||||
// 这条路径把「媒体」伪装成实体 + 用生成的描述文本当语义索引，正是要废弃的
|
||||
// 将就机制。迁移做的事：把每条这类实体还原成原生记忆块，用
|
||||
// sentence --contains--> block 结构边挂到它当时所属的句子上，
|
||||
// 然后删掉旧实体与它的描述关系。块只按自己的向量被检索。
|
||||
//
|
||||
// 迁移是幂等的：实体处理完即删除，重复运行不会重复建块。
|
||||
|
||||
// legacyMediaDigestPattern 从旧媒体实体名尾部取出短 digest。
|
||||
// 名字形如「图片 a1b2c3d4e5f6」——旧实现刻意为每种模态加中文前缀。
|
||||
var legacyMediaDigestPattern = regexp.MustCompile(`([0-9a-f]{8,64})$`)
|
||||
|
||||
// LegacyMediaEntityDigest 从旧媒体实体名里取出短 digest，取不到返回空串。
|
||||
func LegacyMediaEntityDigest(name string) string {
|
||||
name = strings.TrimSpace(name)
|
||||
if !strings.Contains(name, " ") {
|
||||
return ""
|
||||
}
|
||||
m := legacyMediaDigestPattern.FindStringSubmatch(name)
|
||||
if m == nil {
|
||||
return ""
|
||||
}
|
||||
return m[1]
|
||||
}
|
||||
|
||||
// LegacyMediaResolver 把一个短 digest 解析成可用于 L3 的一等记忆块。
|
||||
// 解析失败（内容已不存在）返回 false，该实体将被直接删除而不建块。
|
||||
type LegacyMediaResolver func(shortDigest string) (MemoryBlock, bool)
|
||||
|
||||
// MigrateLegacyMediaEntities 把 marker 反解出来的旧媒体实体迁移成原生块。
|
||||
//
|
||||
// 返回迁移的块数与删除的旧实体数。任何一步失败都会回滚整个迁移，
|
||||
// 因为半途中断会留下既没有块也没有实体的句子——信息静默消失。
|
||||
func (g *GraphDB) MigrateLegacyMediaEntities(resolve LegacyMediaResolver) (blocks, entities int, err error) {
|
||||
if resolve == nil {
|
||||
return 0, 0, nil
|
||||
}
|
||||
g.mu.Lock()
|
||||
defer g.mu.Unlock()
|
||||
|
||||
rows, err := g.db.Query(`SELECT id, name FROM entities WHERE type = 'Media'`)
|
||||
if err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
type legacyEntity struct {
|
||||
id int64
|
||||
name string
|
||||
}
|
||||
var legacy []legacyEntity
|
||||
for rows.Next() {
|
||||
var e legacyEntity
|
||||
if err := rows.Scan(&e.id, &e.name); err != nil {
|
||||
rows.Close()
|
||||
return 0, 0, err
|
||||
}
|
||||
legacy = append(legacy, e)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
rows.Close()
|
||||
return 0, 0, err
|
||||
}
|
||||
rows.Close()
|
||||
if len(legacy) == 0 {
|
||||
return 0, 0, nil
|
||||
}
|
||||
|
||||
tx, err := g.db.Begin()
|
||||
if err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
defer tx.Rollback()
|
||||
|
||||
// 同一份字节可能被多个旧实体引用（重复注入的同一张图），
|
||||
// 迁移后应指向同一个块：块的身份是内容，不是实体行。
|
||||
blockIDForDigest := make(map[string]string)
|
||||
|
||||
for _, e := range legacy {
|
||||
short := LegacyMediaEntityDigest(e.name)
|
||||
if short != "" {
|
||||
if block, ok := resolve(short); ok && block.PayloadDigest != "" {
|
||||
id, seen := blockIDForDigest[block.PayloadDigest]
|
||||
if !seen {
|
||||
if err := insertMigratedBlock(tx, block); err != nil {
|
||||
return 0, 0, fmt.Errorf("migrate legacy media %s: %w", short, err)
|
||||
}
|
||||
blockIDForDigest[block.PayloadDigest] = block.ID
|
||||
id = block.ID
|
||||
blocks++
|
||||
}
|
||||
n, err := attachBlockToLegacySentences(tx, e.id, id)
|
||||
if err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
_ = n
|
||||
}
|
||||
}
|
||||
// 无论能否解析出内容，旧实体与它的描述关系都必须删除：
|
||||
// 留着就等于继续用描述文本当媒体索引。
|
||||
if _, err := tx.Exec(`DELETE FROM relations WHERE source_id = ? OR target_id = ?`, e.id, e.id); err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
if _, err := tx.Exec(`DELETE FROM entities WHERE id = ?`, e.id); err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
entities++
|
||||
}
|
||||
|
||||
if err := tx.Commit(); err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
return blocks, entities, nil
|
||||
}
|
||||
|
||||
// insertMigratedBlock 写一条迁移来的块（不经过 PutMemoryBlocks，避免重入锁）。
|
||||
func insertMigratedBlock(tx *sql.Tx, block MemoryBlock) error {
|
||||
if block.ID == "" {
|
||||
return fmt.Errorf("migrated block id is required")
|
||||
}
|
||||
if !validBlockModality(block.Modality) {
|
||||
return fmt.Errorf("migrated block %s has invalid modality %q", block.ID, block.Modality)
|
||||
}
|
||||
vectorJSON, err := json.Marshal(block.Vector)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
created := block.CreatedAt
|
||||
if created.IsZero() {
|
||||
created = time.Now()
|
||||
}
|
||||
_, err = tx.Exec(`INSERT INTO memory_blocks (
|
||||
id, modality, text_content, payload_digest, mime, size, width, height,
|
||||
vector, fingerprint, source, tool, created_at, updated_at
|
||||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
ON CONFLICT(id) DO NOTHING`,
|
||||
block.ID, block.Modality, block.Text, block.PayloadDigest, block.MIME,
|
||||
block.Size, block.Width, block.Height, string(vectorJSON), block.Fingerprint,
|
||||
block.Source, block.Tool, created, time.Now())
|
||||
return err
|
||||
}
|
||||
|
||||
// attachBlockToLegacySentences 把迁移出的块挂到该旧实体当时所属的句子上，
|
||||
// 并保留那些句子（它们可能只有媒体关系，删实体后就再无关系引用）。
|
||||
func attachBlockToLegacySentences(tx *sql.Tx, entityID int64, blockID string) (int, error) {
|
||||
rows, err := tx.Query(`SELECT DISTINCT s.id FROM sentences s
|
||||
JOIN relations r ON r.sentence_id = s.id
|
||||
WHERE r.source_id = ? OR r.target_id = ?`, entityID, entityID)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
var sids []int64
|
||||
for rows.Next() {
|
||||
var sid int64
|
||||
if err := rows.Scan(&sid); err != nil {
|
||||
rows.Close()
|
||||
return 0, err
|
||||
}
|
||||
sids = append(sids, sid)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
rows.Close()
|
||||
return 0, err
|
||||
}
|
||||
rows.Close()
|
||||
|
||||
n := 0
|
||||
for _, sid := range sids {
|
||||
if _, err := tx.Exec(`INSERT OR IGNORE INTO memory_block_edges
|
||||
(source_kind, source_id, target_kind, target_id, edge_type)
|
||||
VALUES ('sentence', ?, 'block', ?, 'contains')`,
|
||||
fmt.Sprintf("%d", sid), blockID); err != nil {
|
||||
return n, err
|
||||
}
|
||||
n++
|
||||
}
|
||||
return n, nil
|
||||
}
|
||||
@ -1,14 +1,8 @@
|
||||
//go:build onnxruntime
|
||||
|
||||
// Package qwen 的 ONNX 推理实现（构建标签 onnxruntime，与 internal/memory/clip 同模式）。
|
||||
//
|
||||
// 加载契约：调用方传入模型目录，内核不硬编码模型名。
|
||||
//
|
||||
// TextTower.onnx + 外部权重分片 — 文本塔图（input_ids/attention_mask → embedding）
|
||||
// tokenizer.json — 字节级 BPE 词表与 merges
|
||||
// embed_config.json — dim / max_length / instruction / pooling
|
||||
//
|
||||
// 图内已含 last-token 池化，输出即 [batch, dim]；L2 归一化在 Go 侧做（便宜且便于测试）。
|
||||
// Package qwen 提供 Qwen3-VL-Embedding 的完整图文共享 ONNX 编码器。
|
||||
// 文本和图像共用 token embedding、28 层 Transformer、last-token 池化与
|
||||
// fingerprint；Vision.onnx 只产生注入 Transformer 的中间特征。
|
||||
package qwen
|
||||
|
||||
import (
|
||||
@ -26,34 +20,39 @@ import (
|
||||
ort "github.com/yalue/onnxruntime_go"
|
||||
)
|
||||
|
||||
// embedConfig 对应导出脚本产出的 embed_config.json。
|
||||
type embedConfig struct {
|
||||
Dimension int `json:"dim"`
|
||||
MaxLength int `json:"max_length"`
|
||||
Instruction string `json:"instruction"`
|
||||
Pooling string `json:"pooling"`
|
||||
Arch string `json:"arch"`
|
||||
Dimension int `json:"dim"`
|
||||
MaxLength int `json:"max_length"`
|
||||
Instruction string `json:"instruction"`
|
||||
Pooling string `json:"pooling"`
|
||||
ImageSize int `json:"image_size"`
|
||||
PatchSize int `json:"patch_size"`
|
||||
TemporalPatch int `json:"temporal_patch_size"`
|
||||
SpatialMerge int `json:"spatial_merge_size"`
|
||||
ImageMean []float64 `json:"image_mean"`
|
||||
ImageStd []float64 `json:"image_std"`
|
||||
RopeTheta float64 `json:"rope_theta"`
|
||||
MRopeSection []int `json:"mrope_section"`
|
||||
}
|
||||
|
||||
// Embedder 是基于内嵌 ONNX 文本塔的稠密嵌入器。
|
||||
//
|
||||
// 只提供文本能力：导出的是文本塔，视觉塔未导出。EmbedImageDense 会明确报错，
|
||||
// 而不是返回一个「看起来能用」的零向量——后者会让跨模态检索静默失效。
|
||||
type Embedder struct {
|
||||
mu sync.RWMutex
|
||||
|
||||
loaded bool
|
||||
config embedConfig
|
||||
tok *Tokenizer
|
||||
sess *ort.DynamicAdvancedSession
|
||||
fp string
|
||||
loaded bool
|
||||
config embedConfig
|
||||
tok *Tokenizer
|
||||
token *ort.DynamicAdvancedSession
|
||||
transform *ort.DynamicAdvancedSession
|
||||
vision *ort.DynamicAdvancedSession
|
||||
fp string
|
||||
close sync.Once
|
||||
}
|
||||
|
||||
// New 从模型目录加载文本塔。
|
||||
func New(modelDir string) (*Embedder, error) {
|
||||
if modelDir == "" {
|
||||
return nil, fmt.Errorf("qwen model dir not specified")
|
||||
}
|
||||
|
||||
cfgRaw, err := os.ReadFile(filepath.Join(modelDir, "embed_config.json"))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("read embed_config.json: %w", err)
|
||||
@ -62,22 +61,20 @@ func New(modelDir string) (*Embedder, error) {
|
||||
if err := json.Unmarshal(cfgRaw, &cfg); err != nil {
|
||||
return nil, fmt.Errorf("parse embed_config.json: %w", err)
|
||||
}
|
||||
if cfg.Dimension <= 0 {
|
||||
return nil, fmt.Errorf("embed_config.json 的 dim 无效: %d", cfg.Dimension)
|
||||
if cfg.Dimension != 2048 || cfg.MaxLength < 598 || cfg.Pooling != "last_token" {
|
||||
return nil, fmt.Errorf("qwen: incompatible config dim=%d max_length=%d pooling=%q", cfg.Dimension, cfg.MaxLength, cfg.Pooling)
|
||||
}
|
||||
if cfg.MaxLength <= 0 {
|
||||
cfg.MaxLength = 512
|
||||
if cfg.ImageSize != qwenImageSize || cfg.PatchSize != qwenPatchSize || cfg.TemporalPatch != qwenTemporalPatch || cfg.SpatialMerge != qwenSpatialMerge {
|
||||
return nil, fmt.Errorf("qwen: incompatible vision layout image=%d patch=%d temporal=%d merge=%d", cfg.ImageSize, cfg.PatchSize, cfg.TemporalPatch, cfg.SpatialMerge)
|
||||
}
|
||||
if cfg.Pooling != "" && cfg.Pooling != "last_token" {
|
||||
return nil, fmt.Errorf("不支持的池化方式 %q（导出脚本只产出 last_token）", cfg.Pooling)
|
||||
if cfg.RopeTheta <= 0 || len(cfg.MRopeSection) != 3 || cfg.MRopeSection[0]+cfg.MRopeSection[1]+cfg.MRopeSection[2] != qwenRotaryHalfDim {
|
||||
return nil, fmt.Errorf("qwen: incompatible rope theta=%g section=%v", cfg.RopeTheta, cfg.MRopeSection)
|
||||
}
|
||||
|
||||
tok, err := LoadTokenizer(modelDir)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
tok.MaxLen = cfg.MaxLength
|
||||
|
||||
if !ort.IsInitialized() {
|
||||
if lib := findOnnxLib(); lib != "" {
|
||||
ort.SetSharedLibraryPath(lib)
|
||||
@ -87,74 +84,269 @@ func New(modelDir string) (*Embedder, error) {
|
||||
}
|
||||
}
|
||||
|
||||
sess, err := ort.NewDynamicAdvancedSession(
|
||||
filepath.Join(modelDir, "TextTower.onnx"),
|
||||
[]string{"input_ids", "attention_mask"},
|
||||
[]string{"embedding"},
|
||||
nil,
|
||||
token, err := ort.NewDynamicAdvancedSession(
|
||||
filepath.Join(modelDir, "TokenEmbedding.onnx"),
|
||||
[]string{"input_ids"}, []string{"hidden"}, nil,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("create text tower session: %w", err)
|
||||
return nil, fmt.Errorf("create qwen token embedding session: %w", err)
|
||||
}
|
||||
transform, err := ort.NewDynamicAdvancedSession(
|
||||
filepath.Join(modelDir, "Transformer.onnx"),
|
||||
[]string{"hidden", "deepstack_0", "deepstack_1", "deepstack_2", "rotary_cos", "rotary_sin", "causal_mask"},
|
||||
[]string{"embedding"}, nil,
|
||||
)
|
||||
if err != nil {
|
||||
token.Destroy()
|
||||
return nil, fmt.Errorf("create qwen transformer session: %w", err)
|
||||
}
|
||||
vision, err := ort.NewDynamicAdvancedSession(
|
||||
filepath.Join(modelDir, "Vision.onnx"), []string{"pixel_values"},
|
||||
[]string{"deepstack_feature_0", "deepstack_feature_1", "deepstack_feature_2", "vision_hidden_states"}, nil,
|
||||
)
|
||||
if err != nil {
|
||||
token.Destroy()
|
||||
transform.Destroy()
|
||||
return nil, fmt.Errorf("create qwen vision session: %w", err)
|
||||
}
|
||||
|
||||
return &Embedder{
|
||||
loaded: true,
|
||||
config: cfg,
|
||||
tok: tok,
|
||||
sess: sess,
|
||||
fp: computeFingerprint(modelDir),
|
||||
loaded: true, config: cfg, tok: tok,
|
||||
token: token, transform: transform, vision: vision,
|
||||
fp: computeFingerprint(modelDir),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// renderInput 按模型自带的对话模板拼输入（实现在 tokenizer.go，无构建标签）。
|
||||
func (e *Embedder) renderInput(text string) string {
|
||||
return renderInstructionInput(e.config.Instruction, text)
|
||||
}
|
||||
|
||||
// VectorizeDense 把文本编码为 L2 归一化的稠密向量。
|
||||
func (e *Embedder) VectorizeDense(text string) ([]float64, error) {
|
||||
e.mu.RLock()
|
||||
defer e.mu.RUnlock()
|
||||
if !e.loaded {
|
||||
return nil, fmt.Errorf("qwen embedder not loaded")
|
||||
}
|
||||
|
||||
ids, err := e.tok.encodeModelInput(e.renderInput(text), e.config.MaxLength)
|
||||
ids, _, position, _, err := e.tok.textModelInput(e.config.Instruction, text, e.config.MaxLength)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
hidden, err := e.runTokenEmbedding(ids)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
deep := make([][]float32, 3)
|
||||
for i := range deep {
|
||||
deep[i] = make([]float32, len(hidden))
|
||||
}
|
||||
return e.runTransformer(hidden, deep, position, len(ids))
|
||||
}
|
||||
|
||||
seq := len(ids)
|
||||
inputIDs := make([]int64, seq)
|
||||
attn := make([]int64, seq)
|
||||
func (e *Embedder) EmbedImageDense(raw []byte, _ string) ([]float64, error) {
|
||||
e.mu.RLock()
|
||||
defer e.mu.RUnlock()
|
||||
if !e.loaded {
|
||||
return nil, fmt.Errorf("qwen embedder not loaded")
|
||||
}
|
||||
pixels, err := preprocessImage(raw)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
features, err := e.runVision(pixels)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
ids, _, position, visual, err := e.tok.imageModelInput(e.config.Instruction, e.config.MaxLength)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
hidden, err := e.runTokenEmbedding(ids)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
deep := make([][]float32, 3)
|
||||
for i := range deep {
|
||||
deep[i] = make([]float32, len(hidden))
|
||||
}
|
||||
visualIndex := 0
|
||||
for tokenIndex, isVisual := range visual {
|
||||
if !isVisual {
|
||||
continue
|
||||
}
|
||||
dst := tokenIndex * e.config.Dimension
|
||||
src := visualIndex * e.config.Dimension
|
||||
copy(hidden[dst:dst+e.config.Dimension], features[3][src:src+e.config.Dimension])
|
||||
for layer := range deep {
|
||||
copy(deep[layer][dst:dst+e.config.Dimension], features[layer][src:src+e.config.Dimension])
|
||||
}
|
||||
visualIndex++
|
||||
}
|
||||
if visualIndex != qwenVisualTokens {
|
||||
return nil, fmt.Errorf("qwen: injected visual tokens=%d, want %d", visualIndex, qwenVisualTokens)
|
||||
}
|
||||
return e.runTransformer(hidden, deep, position, len(ids))
|
||||
}
|
||||
|
||||
func (e *Embedder) runTokenEmbedding(ids []int) ([]float32, error) {
|
||||
inputIDs := make([]int64, len(ids))
|
||||
for i, id := range ids {
|
||||
inputIDs[i] = int64(id)
|
||||
attn[i] = 1
|
||||
}
|
||||
|
||||
idTensor, err := ort.NewTensor(ort.Shape{1, int64(seq)}, inputIDs)
|
||||
in, err := ort.NewTensor(ort.Shape{1, int64(len(ids))}, inputIDs)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("input_ids tensor: %w", err)
|
||||
return nil, fmt.Errorf("qwen token input: %w", err)
|
||||
}
|
||||
defer idTensor.Destroy()
|
||||
defer in.Destroy()
|
||||
outs := make([]ort.Value, 1)
|
||||
if err := e.token.Run([]ort.Value{in}, outs); err != nil {
|
||||
return nil, fmt.Errorf("qwen token embedding run: %w", err)
|
||||
}
|
||||
if outs[0] == nil {
|
||||
return nil, fmt.Errorf("qwen token embedding output is nil")
|
||||
}
|
||||
defer outs[0].Destroy()
|
||||
tensor, ok := outs[0].(*ort.Tensor[float32])
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("qwen token embedding output type %T", outs[0])
|
||||
}
|
||||
shape := tensor.GetShape()
|
||||
if len(shape) != 3 || shape[0] != 1 || shape[1] != int64(len(ids)) || shape[2] != int64(e.config.Dimension) {
|
||||
return nil, fmt.Errorf("qwen token embedding shape=%v", shape)
|
||||
}
|
||||
return append([]float32(nil), tensor.GetData()...), nil
|
||||
}
|
||||
|
||||
maskTensor, err := ort.NewTensor(ort.Shape{1, int64(seq)}, attn)
|
||||
func (e *Embedder) runVision(pixels []float32) ([][]float32, error) {
|
||||
in, err := ort.NewTensor(ort.Shape{qwenImagePatches, qwenPatchVectorSize}, pixels)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("attention_mask tensor: %w", err)
|
||||
return nil, fmt.Errorf("qwen vision input: %w", err)
|
||||
}
|
||||
defer maskTensor.Destroy()
|
||||
defer in.Destroy()
|
||||
outs := make([]ort.Value, 4)
|
||||
if err := e.vision.Run([]ort.Value{in}, outs); err != nil {
|
||||
return nil, fmt.Errorf("qwen vision run: %w", err)
|
||||
}
|
||||
features := make([][]float32, 4)
|
||||
for i, value := range outs {
|
||||
if value == nil {
|
||||
return nil, fmt.Errorf("qwen vision output %d is nil", i)
|
||||
}
|
||||
defer value.Destroy()
|
||||
tensor, ok := value.(*ort.Tensor[float32])
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("qwen vision output %d type %T", i, value)
|
||||
}
|
||||
shape := tensor.GetShape()
|
||||
if len(shape) != 2 || shape[0] != qwenVisualTokens || shape[1] != int64(e.config.Dimension) {
|
||||
return nil, fmt.Errorf("qwen vision output %d shape=%v", i, shape)
|
||||
}
|
||||
features[i] = append([]float32(nil), tensor.GetData()...)
|
||||
}
|
||||
return features, nil
|
||||
}
|
||||
|
||||
outTensor, err := ort.NewEmptyTensor[float32](ort.Shape{1, int64(e.config.Dimension)})
|
||||
func (e *Embedder) runTransformer(hidden []float32, deep [][]float32, position []int64, seq int) ([]float64, error) {
|
||||
if len(hidden) != seq*e.config.Dimension || len(deep) != 3 || len(position) != 3*seq {
|
||||
return nil, fmt.Errorf("qwen: invalid transformer inputs hidden=%d deep=%d position=%d seq=%d", len(hidden), len(deep), len(position), seq)
|
||||
}
|
||||
cos, sin := e.rotary(position, seq)
|
||||
causal := causalMask(seq)
|
||||
|
||||
hiddenTensor, err := ort.NewTensor(ort.Shape{1, int64(seq), int64(e.config.Dimension)}, hidden)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("output tensor: %w", err)
|
||||
return nil, fmt.Errorf("qwen hidden tensor: %w", err)
|
||||
}
|
||||
defer outTensor.Destroy()
|
||||
|
||||
if err := e.sess.Run([]ort.Value{idTensor, maskTensor}, []ort.Value{outTensor}); err != nil {
|
||||
return nil, fmt.Errorf("text tower run: %w", err)
|
||||
defer hiddenTensor.Destroy()
|
||||
inputs := []ort.Value{hiddenTensor}
|
||||
var deepTensors []*ort.Tensor[float32]
|
||||
for i, data := range deep {
|
||||
if len(data) != len(hidden) {
|
||||
return nil, fmt.Errorf("qwen deepstack %d length=%d, want %d", i, len(data), len(hidden))
|
||||
}
|
||||
t, err := ort.NewTensor(ort.Shape{1, int64(seq), int64(e.config.Dimension)}, data)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("qwen deepstack %d tensor: %w", i, err)
|
||||
}
|
||||
deepTensors = append(deepTensors, t)
|
||||
inputs = append(inputs, t)
|
||||
}
|
||||
defer func() {
|
||||
for _, t := range deepTensors {
|
||||
t.Destroy()
|
||||
}
|
||||
}()
|
||||
cosTensor, err := ort.NewTensor(ort.Shape{1, int64(seq), qwenRotaryDim}, cos)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("qwen rotary cos: %w", err)
|
||||
}
|
||||
defer cosTensor.Destroy()
|
||||
sinTensor, err := ort.NewTensor(ort.Shape{1, int64(seq), qwenRotaryDim}, sin)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("qwen rotary sin: %w", err)
|
||||
}
|
||||
defer sinTensor.Destroy()
|
||||
causalTensor, err := ort.NewTensor(ort.Shape{1, 1, int64(seq), int64(seq)}, causal)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("qwen causal mask: %w", err)
|
||||
}
|
||||
defer causalTensor.Destroy()
|
||||
inputs = append(inputs, cosTensor, sinTensor, causalTensor)
|
||||
|
||||
raw := outTensor.GetData()
|
||||
out, err := ort.NewEmptyTensor[float32](ort.Shape{1, int64(e.config.Dimension)})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("qwen output tensor: %w", err)
|
||||
}
|
||||
defer out.Destroy()
|
||||
if err := e.transform.Run(inputs, []ort.Value{out}); err != nil {
|
||||
return nil, fmt.Errorf("qwen transformer run: %w", err)
|
||||
}
|
||||
return normalize(out.GetData()), nil
|
||||
}
|
||||
|
||||
const (
|
||||
qwenRotaryHalfDim = 64
|
||||
qwenRotaryDim = 128
|
||||
)
|
||||
|
||||
func (e *Embedder) rotary(position []int64, seq int) ([]float32, []float32) {
|
||||
cos := make([]float32, seq*qwenRotaryDim)
|
||||
sin := make([]float32, seq*qwenRotaryDim)
|
||||
inv := make([]float64, qwenRotaryHalfDim)
|
||||
for i := range inv {
|
||||
inv[i] = 1 / math.Pow(e.config.RopeTheta, float64(2*i)/qwenRotaryDim)
|
||||
}
|
||||
for token := 0; token < seq; token++ {
|
||||
freq := make([]float64, qwenRotaryHalfDim)
|
||||
for i := range freq {
|
||||
freq[i] = float64(position[token]) * inv[i]
|
||||
}
|
||||
for dim, offset := range []int{0, 1, 2} {
|
||||
if dim == 0 {
|
||||
continue
|
||||
}
|
||||
limit := e.config.MRopeSection[dim] * 3
|
||||
for i := offset; i < limit; i += 3 {
|
||||
freq[i] = float64(position[dim*seq+token]) * inv[i]
|
||||
}
|
||||
}
|
||||
for i, f := range freq {
|
||||
c, s := float32(math.Cos(f)), float32(math.Sin(f))
|
||||
cos[token*qwenRotaryDim+i] = c
|
||||
cos[token*qwenRotaryDim+qwenRotaryHalfDim+i] = c
|
||||
sin[token*qwenRotaryDim+i] = s
|
||||
sin[token*qwenRotaryDim+qwenRotaryHalfDim+i] = s
|
||||
}
|
||||
}
|
||||
return cos, sin
|
||||
}
|
||||
|
||||
func causalMask(seq int) []float32 {
|
||||
out := make([]float32, seq*seq)
|
||||
for row := 0; row < seq; row++ {
|
||||
for col := row + 1; col < seq; col++ {
|
||||
out[row*seq+col] = -math.MaxFloat32
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func normalize(raw []float32) []float64 {
|
||||
out := make([]float64, len(raw))
|
||||
var norm float64
|
||||
for i, v := range raw {
|
||||
@ -167,48 +359,39 @@ func (e *Embedder) VectorizeDense(text string) ([]float64, error) {
|
||||
out[i] /= norm
|
||||
}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// EmbedImageDense 不支持：导出的是**文本塔**，视觉塔未导出。
|
||||
//
|
||||
// 明确报错而不是返回零向量或占位：调用方（mediaref.go）会 log 后跳过写向量，
|
||||
// 若返回零向量则「写入了但检索不到」，失败会静默化。要支持图像检索需另外
|
||||
// 导出视觉塔并实现 Qwen3-VL 的图像预处理（patch/merge/缩放规则）。
|
||||
func (e *Embedder) EmbedImageDense(_ []byte, _ string) ([]float64, error) {
|
||||
return nil, fmt.Errorf("qwen text tower 不支持图像嵌入；图像检索请用 clip 或 http 路径")
|
||||
return out
|
||||
}
|
||||
|
||||
func (e *Embedder) Fingerprint() string { return e.fp }
|
||||
func (e *Embedder) Dim() int { return e.config.Dimension }
|
||||
|
||||
func (e *Embedder) Loaded() bool {
|
||||
e.mu.RLock()
|
||||
defer e.mu.RUnlock()
|
||||
return e.loaded
|
||||
}
|
||||
|
||||
func (e *Embedder) Close() {
|
||||
e.mu.Lock()
|
||||
defer e.mu.Unlock()
|
||||
if e.sess != nil {
|
||||
e.sess.Destroy()
|
||||
e.sess = nil
|
||||
}
|
||||
e.loaded = false
|
||||
e.close.Do(func() {
|
||||
e.mu.Lock()
|
||||
defer e.mu.Unlock()
|
||||
if e.token != nil {
|
||||
e.token.Destroy()
|
||||
e.token = nil
|
||||
}
|
||||
if e.transform != nil {
|
||||
e.transform.Destroy()
|
||||
e.transform = nil
|
||||
}
|
||||
if e.vision != nil {
|
||||
e.vision.Destroy()
|
||||
e.vision = nil
|
||||
}
|
||||
e.loaded = false
|
||||
})
|
||||
}
|
||||
|
||||
// computeFingerprint 计算模型指纹，用于 vec_model 持久化与切换后重算判定。
|
||||
//
|
||||
// 为什么不直接哈希全部权重：这个模型目录有 6.5GB 外部权重分片，启动时读一遍
|
||||
// 要几十秒，会阻塞 homeagent 启动。这里哈希「图文件 + 配置 + 全部外部权重的
|
||||
// 文件名与大小」——换模型（哪怕只是换了权重）几乎必然改变文件集合或大小，
|
||||
// 足以识别切换；代价是理论上存在「大小相同但内容不同」的漏判，对本地单机
|
||||
// 部署可接受。
|
||||
func computeFingerprint(modelDir string) string {
|
||||
h := sha256.New()
|
||||
|
||||
for _, name := range []string{"TextTower.onnx", "embed_config.json"} {
|
||||
for _, name := range []string{"TokenEmbedding.onnx", "Transformer.onnx", "Vision.onnx", "embed_config.json"} {
|
||||
if data, err := os.ReadFile(filepath.Join(modelDir, name)); err == nil {
|
||||
h.Write([]byte(name))
|
||||
h.Write([]byte{0})
|
||||
@ -216,35 +399,25 @@ func computeFingerprint(modelDir string) string {
|
||||
h.Write([]byte{0})
|
||||
}
|
||||
}
|
||||
|
||||
entries, _ := os.ReadDir(modelDir)
|
||||
var names []string
|
||||
for _, e := range entries {
|
||||
n := e.Name()
|
||||
// 外部权重分片：torch 新版导出器使用 onnx__<op>_<id> 与模型张量同名文件。
|
||||
if strings.HasPrefix(n, "onnx__") || strings.HasSuffix(n, ".weight") || strings.HasSuffix(n, ".onnx.data") {
|
||||
for _, entry := range entries {
|
||||
n := entry.Name()
|
||||
if strings.HasPrefix(n, "embed_tokens.") || strings.HasPrefix(n, "layers.") || strings.HasPrefix(n, "onnx__") || strings.HasSuffix(n, ".onnx.data") {
|
||||
names = append(names, n)
|
||||
}
|
||||
}
|
||||
sort.Strings(names)
|
||||
for _, n := range names {
|
||||
info, err := os.Stat(filepath.Join(modelDir, n))
|
||||
if err != nil {
|
||||
continue
|
||||
if info, err := os.Stat(filepath.Join(modelDir, n)); err == nil {
|
||||
fmt.Fprintf(h, "%s:%d\n", n, info.Size())
|
||||
}
|
||||
fmt.Fprintf(h, "%s:%d\n", n, info.Size())
|
||||
}
|
||||
return hex.EncodeToString(h.Sum(nil))
|
||||
}
|
||||
|
||||
// findOnnxLib 在常见路径中查找 libonnxruntime.so。
|
||||
func findOnnxLib() string {
|
||||
for _, p := range []string{
|
||||
"/opt/onnxruntime/libonnxruntime.so",
|
||||
"/opt/onnxruntime/lib/libonnxruntime.so",
|
||||
"/usr/local/lib/libonnxruntime.so",
|
||||
"/usr/lib/libonnxruntime.so",
|
||||
} {
|
||||
for _, p := range []string{"/opt/onnxruntime/libonnxruntime.so", "/usr/local/lib/libonnxruntime.so", "/usr/lib/libonnxruntime.so"} {
|
||||
if _, err := os.Stat(p); err == nil {
|
||||
return p
|
||||
}
|
||||
|
||||
@ -4,7 +4,7 @@ package qwen
|
||||
|
||||
import "fmt"
|
||||
|
||||
// Embedder 在未启用 onnxruntime 时为 no-op 实现（与 internal/memory/clip 同模式）。
|
||||
// Embedder 在未启用 onnxruntime 时为 no-op 实现。
|
||||
// 默认构建不链接 onnxruntime，保持原有 fastText/TF-IDF 行为不变。
|
||||
type Embedder struct {
|
||||
loaded bool
|
||||
|
||||
157
internal/memory/qwen/image.go
Normal file
157
internal/memory/qwen/image.go
Normal file
@ -0,0 +1,157 @@
|
||||
//go:build onnxruntime
|
||||
|
||||
package qwen
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"fmt"
|
||||
"image"
|
||||
"image/color"
|
||||
_ "image/gif"
|
||||
_ "image/jpeg"
|
||||
_ "image/png"
|
||||
"math"
|
||||
)
|
||||
|
||||
const (
|
||||
qwenImageSize = 768
|
||||
qwenPatchSize = 16
|
||||
qwenTemporalPatch = 2
|
||||
qwenSpatialMerge = 2
|
||||
qwenImagePatches = (qwenImageSize / qwenPatchSize) * (qwenImageSize / qwenPatchSize)
|
||||
qwenVisualTokens = qwenImagePatches / (qwenSpatialMerge * qwenSpatialMerge)
|
||||
qwenPatchVectorSize = 3 * qwenTemporalPatch * qwenPatchSize * qwenPatchSize
|
||||
)
|
||||
|
||||
// preprocessImage 把任意图片转成固定 768×768 视觉塔输入。
|
||||
//
|
||||
// Vision.onnx 是经过 PyTorch 逐输出验证的固定 48×48 patch 图。为避免拉伸物体，
|
||||
// 这里保持宽高比缩放并在中心补中性灰（归一化后约为 0）；这与直接把长方形
|
||||
// 强拉成正方形相比更能保留 Qwen 的视觉语义。已是 768×768 的输入不做插值，
|
||||
// 便于用跨语言冻结向量精确回归 patch 排列。
|
||||
func preprocessImage(raw []byte) ([]float32, error) {
|
||||
src, _, err := image.Decode(bytes.NewReader(raw))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("qwen: decode image: %w", err)
|
||||
}
|
||||
b := src.Bounds()
|
||||
if b.Dx() <= 0 || b.Dy() <= 0 {
|
||||
return nil, fmt.Errorf("qwen: empty image")
|
||||
}
|
||||
|
||||
scale := math.Min(float64(qwenImageSize)/float64(b.Dx()), float64(qwenImageSize)/float64(b.Dy()))
|
||||
w := max(1, int(math.Round(float64(b.Dx())*scale)))
|
||||
h := max(1, int(math.Round(float64(b.Dy())*scale)))
|
||||
if w > qwenImageSize {
|
||||
w = qwenImageSize
|
||||
}
|
||||
if h > qwenImageSize {
|
||||
h = qwenImageSize
|
||||
}
|
||||
|
||||
resized := resizeBicubic(src, w, h)
|
||||
canvas := image.NewNRGBA(image.Rect(0, 0, qwenImageSize, qwenImageSize))
|
||||
neutral := color.NRGBA{R: 128, G: 128, B: 128, A: 255}
|
||||
for i := 0; i < len(canvas.Pix); i += 4 {
|
||||
canvas.Pix[i], canvas.Pix[i+1], canvas.Pix[i+2], canvas.Pix[i+3] = neutral.R, neutral.G, neutral.B, neutral.A
|
||||
}
|
||||
ox, oy := (qwenImageSize-w)/2, (qwenImageSize-h)/2
|
||||
for y := 0; y < h; y++ {
|
||||
for x := 0; x < w; x++ {
|
||||
canvas.SetNRGBA(ox+x, oy+y, resized.NRGBAAt(x, y))
|
||||
}
|
||||
}
|
||||
|
||||
// 与 transformers Qwen2VLImageProcessor 的排列严格一致：
|
||||
// [grid_h/merge, grid_w/merge, merge_h, merge_w, channel,
|
||||
// temporal_patch, patch_h, patch_w]，然后 flatten。
|
||||
out := make([]float32, 0, qwenImagePatches*qwenPatchVectorSize)
|
||||
blocks := qwenImageSize / qwenPatchSize / qwenSpatialMerge
|
||||
for bh := 0; bh < blocks; bh++ {
|
||||
for bw := 0; bw < blocks; bw++ {
|
||||
for mh := 0; mh < qwenSpatialMerge; mh++ {
|
||||
for mw := 0; mw < qwenSpatialMerge; mw++ {
|
||||
baseY := (bh*qwenSpatialMerge + mh) * qwenPatchSize
|
||||
baseX := (bw*qwenSpatialMerge + mw) * qwenPatchSize
|
||||
for c := 0; c < 3; c++ {
|
||||
for temporal := 0; temporal < qwenTemporalPatch; temporal++ {
|
||||
_ = temporal // 静态图复制同一图片形成 2 帧 temporal patch
|
||||
for py := 0; py < qwenPatchSize; py++ {
|
||||
for px := 0; px < qwenPatchSize; px++ {
|
||||
p := canvas.NRGBAAt(baseX+px, baseY+py)
|
||||
v := [3]uint8{p.R, p.G, p.B}[c]
|
||||
out = append(out, float32(v)/127.5-1)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// resizeBicubic 使用半像素中心的 Catmull-Rom 三次卷积。
|
||||
func resizeBicubic(src image.Image, dstW, dstH int) *image.NRGBA {
|
||||
b := src.Bounds()
|
||||
if b.Dx() == dstW && b.Dy() == dstH {
|
||||
dst := image.NewNRGBA(image.Rect(0, 0, dstW, dstH))
|
||||
for y := 0; y < dstH; y++ {
|
||||
for x := 0; x < dstW; x++ {
|
||||
dst.SetNRGBA(x, y, color.NRGBAModel.Convert(src.At(b.Min.X+x, b.Min.Y+y)).(color.NRGBA))
|
||||
}
|
||||
}
|
||||
return dst
|
||||
}
|
||||
|
||||
dst := image.NewNRGBA(image.Rect(0, 0, dstW, dstH))
|
||||
sx, sy := float64(b.Dx())/float64(dstW), float64(b.Dy())/float64(dstH)
|
||||
for y := 0; y < dstH; y++ {
|
||||
fy := (float64(y)+0.5)*sy - 0.5
|
||||
y0 := int(math.Floor(fy))
|
||||
for x := 0; x < dstW; x++ {
|
||||
fx := (float64(x)+0.5)*sx - 0.5
|
||||
x0 := int(math.Floor(fx))
|
||||
var sum [4]float64
|
||||
var weight float64
|
||||
for j := -1; j <= 2; j++ {
|
||||
wy := cubicWeight(fy - float64(y0+j))
|
||||
yy := min(max(y0+j, 0), b.Dy()-1)
|
||||
for i := -1; i <= 2; i++ {
|
||||
w := wy * cubicWeight(fx-float64(x0+i))
|
||||
xx := min(max(x0+i, 0), b.Dx()-1)
|
||||
p := color.NRGBAModel.Convert(src.At(b.Min.X+xx, b.Min.Y+yy)).(color.NRGBA)
|
||||
sum[0] += float64(p.R) * w
|
||||
sum[1] += float64(p.G) * w
|
||||
sum[2] += float64(p.B) * w
|
||||
sum[3] += float64(p.A) * w
|
||||
weight += w
|
||||
}
|
||||
}
|
||||
if weight == 0 {
|
||||
weight = 1
|
||||
}
|
||||
dst.SetNRGBA(x, y, color.NRGBA{
|
||||
R: clampByte(sum[0] / weight), G: clampByte(sum[1] / weight),
|
||||
B: clampByte(sum[2] / weight), A: clampByte(sum[3] / weight),
|
||||
})
|
||||
}
|
||||
}
|
||||
return dst
|
||||
}
|
||||
|
||||
func cubicWeight(x float64) float64 {
|
||||
x = math.Abs(x)
|
||||
if x <= 1 {
|
||||
return 1.5*x*x*x - 2.5*x*x + 1
|
||||
}
|
||||
if x < 2 {
|
||||
return -0.5*x*x*x + 2.5*x*x - 4*x + 2
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func clampByte(v float64) uint8 {
|
||||
return uint8(min(255, max(0, int(math.Round(v)))))
|
||||
}
|
||||
86
internal/memory/qwen/model_input.go
Normal file
86
internal/memory/qwen/model_input.go
Normal file
@ -0,0 +1,86 @@
|
||||
//go:build onnxruntime
|
||||
|
||||
package qwen
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// imageModelInput 构造 Qwen3-VL 单图对话模板及对应 M-RoPE 位置。
|
||||
// 固定 768×768 视觉塔产生 576 个合并后的视觉 token。
|
||||
func (t *Tokenizer) imageModelInput(instruction string, maxLen int) (ids []int, attention, position []int64, visual []bool, err error) {
|
||||
if instruction == "" {
|
||||
instruction = DefaultInstruction
|
||||
}
|
||||
text := "<|im_start|>system\n" + instruction +
|
||||
"<|im_end|>\n<|im_start|>user\n<|vision_start|>" +
|
||||
strings.Repeat("<|image_pad|>", qwenVisualTokens) +
|
||||
"<|vision_end|><|im_end|>\n<|im_start|>assistant\n"
|
||||
ids, err = t.encodeModelInput(text, maxLen)
|
||||
if err != nil {
|
||||
return nil, nil, nil, nil, err
|
||||
}
|
||||
imageID, ok := t.SpecialID("<|image_pad|>")
|
||||
if !ok {
|
||||
return nil, nil, nil, nil, fmt.Errorf("tokenizer.json 缺少 <|image_pad|>")
|
||||
}
|
||||
|
||||
visual = make([]bool, len(ids))
|
||||
attention = make([]int64, len(ids))
|
||||
position = make([]int64, 3*len(ids))
|
||||
for i, id := range ids {
|
||||
attention[i] = 1
|
||||
visual[i] = id == imageID
|
||||
}
|
||||
|
||||
current := int64(0)
|
||||
for start := 0; start < len(ids); {
|
||||
isVisual := visual[start]
|
||||
end := start + 1
|
||||
for end < len(ids) && visual[end] == isVisual {
|
||||
end++
|
||||
}
|
||||
if !isVisual {
|
||||
for i := start; i < end; i++ {
|
||||
p := current + int64(i-start)
|
||||
position[i] = p
|
||||
position[len(ids)+i] = p
|
||||
position[2*len(ids)+i] = p
|
||||
}
|
||||
current += int64(end - start)
|
||||
} else {
|
||||
if end-start != qwenVisualTokens {
|
||||
return nil, nil, nil, nil, fmt.Errorf("qwen: image token count=%d, want %d", end-start, qwenVisualTokens)
|
||||
}
|
||||
side := qwenImageSize / qwenPatchSize / qwenSpatialMerge
|
||||
for i := start; i < end; i++ {
|
||||
j := i - start
|
||||
position[i] = current
|
||||
position[len(ids)+i] = current + int64(j/side)
|
||||
position[2*len(ids)+i] = current + int64(j%side)
|
||||
}
|
||||
current += int64(side)
|
||||
}
|
||||
start = end
|
||||
}
|
||||
return ids, attention, position, visual, nil
|
||||
}
|
||||
|
||||
// textModelInput 执行完整 tokenizer post_processor，并构造纯文本标准 RoPE 位置。
|
||||
func (t *Tokenizer) textModelInput(instruction, text string, maxLen int) (ids []int, attention, position []int64, visual []bool, err error) {
|
||||
ids, err = t.encodeModelInput(renderInstructionInput(instruction, text), maxLen)
|
||||
if err != nil {
|
||||
return nil, nil, nil, nil, err
|
||||
}
|
||||
attention = make([]int64, len(ids))
|
||||
position = make([]int64, 3*len(ids))
|
||||
visual = make([]bool, len(ids))
|
||||
for i := range ids {
|
||||
attention[i] = 1
|
||||
position[i] = int64(i)
|
||||
position[len(ids)+i] = int64(i)
|
||||
position[2*len(ids)+i] = int64(i)
|
||||
}
|
||||
return ids, attention, position, visual, nil
|
||||
}
|
||||
@ -14,8 +14,8 @@ import (
|
||||
"sync"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/yanyiwu/gojieba"
|
||||
"gitcode.com/JianFeeeee/HomeAgent/internal/memory/vector"
|
||||
"github.com/yanyiwu/gojieba"
|
||||
)
|
||||
|
||||
const downloadMaxWords = 200000
|
||||
|
||||
@ -108,7 +108,7 @@ func TestStaticEmbedderSemanticSimilarity(t *testing.T) {
|
||||
e := newSynthEmbedder(t, 300)
|
||||
|
||||
pairs := []struct {
|
||||
a, b string
|
||||
a, b string
|
||||
related bool
|
||||
}{
|
||||
{"今天天气怎么样", "明天会不会下雨", true},
|
||||
|
||||
51
internal/memory/vector/fuse.go
Normal file
51
internal/memory/vector/fuse.go
Normal file
@ -0,0 +1,51 @@
|
||||
package vector
|
||||
|
||||
import "math"
|
||||
|
||||
// FuseVectors 把同一统一空间里的多个向量融合为一个向量：
|
||||
// 逐维求和后重新 L2 归一化。
|
||||
//
|
||||
// 用途：文档/上下文事件既带文本、又带若干一等记忆块（图片/视频），
|
||||
// 二者的向量来自同一模型、同一 fingerprint、同一维度。融合后，
|
||||
// 一篇文档既能按文字、也能按它携带的图片内容被召回——
|
||||
// 图片由自己的向量参与检索，不依赖任何生成的描述文本。
|
||||
//
|
||||
// 约定：调用方传入的向量应已是 L2 归一化的同空间向量。长度不一致的
|
||||
// 向量会被跳过（不同模型/维度的残留）；全空或全零返回 nil。
|
||||
func FuseVectors(vectors ...[]float64) []float64 {
|
||||
dim := 0
|
||||
for _, v := range vectors {
|
||||
if len(v) > dim {
|
||||
dim = len(v)
|
||||
}
|
||||
}
|
||||
if dim == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([]float64, dim)
|
||||
used := 0
|
||||
for _, v := range vectors {
|
||||
if len(v) != dim {
|
||||
continue
|
||||
}
|
||||
for i, x := range v {
|
||||
out[i] += x
|
||||
}
|
||||
used++
|
||||
}
|
||||
if used == 0 {
|
||||
return nil
|
||||
}
|
||||
var norm float64
|
||||
for _, x := range out {
|
||||
norm += x * x
|
||||
}
|
||||
if norm == 0 {
|
||||
return nil
|
||||
}
|
||||
norm = math.Sqrt(norm)
|
||||
for i := range out {
|
||||
out[i] /= norm
|
||||
}
|
||||
return out
|
||||
}
|
||||
@ -47,7 +47,7 @@ const (
|
||||
ModalityVideo MultimodalModality = "video"
|
||||
)
|
||||
|
||||
// ErrNotSupported 表示 Vectorizer 不支持图像嵌入，调用方按文本描述降级。
|
||||
// ErrNotSupported 表示 Vectorizer 不支持该原生模态；调用方不得以描述文本冒充其向量。
|
||||
var ErrNotSupported = fmt.Errorf("vectorizer does not support image embedding")
|
||||
|
||||
// Vector 是带权特征映射：feature → weight
|
||||
|
||||
Reference in New Issue
Block a user