mirror of
https://gitcode.com/JianFeeeee/HomeAgent.git
synced 2026-09-21 09:28:14 +00:00
refactor(memory): 核心不再适配具体模型——公共 embedding provider SPI + 注册表
问题:cmd/homed 里 `case "onnx": qwen.New(modelDir)` 把模型适配写进了核心, `type=onnx` 名义上是格式、实际写死了一个模型家族;2117 行 Qwen 专属代码 (BPE、chat template、M-RoPE、Vision_gN 命名)住在内核树里,还带着一对 `//go:build onnxruntime` 的 stub。加任何新模型都要改内核。 现在核心只认一个模型无关的公共契约(pkg/embedding): - 输入是不透明的 Data+MIME,解码/预处理/时序分组全归 provider - 能力是数据(Info.Modalities),不是接口方法——新增模态无需改核心接口 - 不支持的模态返回 embedding.ErrUnsupportedModality(可 errors.Is 识别) - 按名字注册,重复注册 panic;Options 是 provider 私有命名空间,核心不解释 改动: - 新增 pkg/embedding:Modality/Purpose/Input/Info/Provider/Config + 注册表 (Open 校验 Info,ValidateVector 在入库前拦下维度错与非有限值) - providers/qwen3vl:Qwen 实现整体移出内核(git mv),实现公共 SPI 并自注册 - internal/memory/vector:新增 ProviderAdapter(公共 SPI → 内部小接口); ErrModalityUnsupported 改为公共哨兵别名;删除 VideoEmbedder 可选接口 (那正是「核心为每个新模态长方法」的坏味道) - http embedder 也变成普通 provider(注册名 http) - cmd/homed:删除 qwen import 与 onnx/http 分支,改为按 provider 名打开 + 透传 options.*;provider 打开失败只警告并禁用多模态检索,不影响启动 - config:multimodal_space.type/onnx./http.* → provider + options.* - 删除 internal/memory/qwen(整体搬迁) 测试: - pkg/embedding:注册表隔离/未知名字/非法 Info 自动关闭/ValidateVector - vector:适配器原样透传字节与 MIME、维度错被拦、Close 幂等且停止使用、 两个哨兵 errors.Is 互通 - providers/qwen3vl:新增公共 SPI 全链路集成测试(Open→Info→Embed→ 未知模态哨兵),并明确断言 Info 不声明 video 已知未完成(不得当作已验证): - 视频冻结回归 TestEmbedderVideoMatchesONNXReference **显式跳过**:Go 侧 video 模板缺少 processor 按时间组插入的字面时间戳文本 (<0.0 seconds>/<1.0 seconds>),同一输入 Python seq=1190(1152+38)、 Go 只有 22 个文本 token。时间戳也占 M-RoPE 位置,故现有 M-RoPE 自洽断言 通过不能证明与官方实现一致。修复属 provider 内部工作。 - 视觉侧三档已导出并逐档校验通过(cos 1.000000119/1.000000119/1.000000000) 验证:go build ./... ;go vet -tags onnxruntime ./... ; go test -short ./internal/memory/... ./internal/agent/core/... ./internal/sdk/... ./pkg/... ;onnxruntime 下 providers/qwen3vl 全绿(视频为显式 skip)
This commit is contained in:
@ -1,440 +0,0 @@
|
||||
//go:build onnxruntime
|
||||
|
||||
// Package qwen 提供 Qwen3-VL-Embedding 的完整图文共享 ONNX 编码器。
|
||||
// 文本和图像共用 token embedding、28 层 Transformer、last-token 池化与
|
||||
// fingerprint;Vision.onnx 只产生注入 Transformer 的中间特征。
|
||||
package qwen
|
||||
|
||||
import (
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"math"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
ort "github.com/yalue/onnxruntime_go"
|
||||
|
||||
"gitcode.com/JianFeeeee/HomeAgent/internal/memory/vector"
|
||||
)
|
||||
|
||||
type embedConfig struct {
|
||||
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"`
|
||||
}
|
||||
|
||||
type Embedder struct {
|
||||
mu sync.RWMutex
|
||||
|
||||
loaded bool
|
||||
config embedConfig
|
||||
tok *Tokenizer
|
||||
token *ort.DynamicAdvancedSession
|
||||
transform *ort.DynamicAdvancedSession
|
||||
vision *ort.DynamicAdvancedSession
|
||||
fp string
|
||||
close sync.Once
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
var cfg embedConfig
|
||||
if err := json.Unmarshal(cfgRaw, &cfg); err != nil {
|
||||
return nil, fmt.Errorf("parse embed_config.json: %w", err)
|
||||
}
|
||||
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.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.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
|
||||
}
|
||||
if !ort.IsInitialized() {
|
||||
if lib := findOnnxLib(); lib != "" {
|
||||
ort.SetSharedLibraryPath(lib)
|
||||
}
|
||||
if err := ort.InitializeEnvironment(); err != nil {
|
||||
return nil, fmt.Errorf("init onnx env: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
token, err := ort.NewDynamicAdvancedSession(
|
||||
filepath.Join(modelDir, "TokenEmbedding.onnx"),
|
||||
[]string{"input_ids"}, []string{"hidden"}, nil,
|
||||
)
|
||||
if err != nil {
|
||||
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,
|
||||
token: token, transform: transform, vision: vision,
|
||||
fp: computeFingerprint(modelDir),
|
||||
}, nil
|
||||
}
|
||||
|
||||
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, _, 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))
|
||||
}
|
||||
|
||||
// EmbedImageDense 把一张图片编码到统一空间。
|
||||
//
|
||||
// mime 决定这个模态是否在本空间的原生覆盖范围内:Qwen3-VL 能原生编码文本与
|
||||
// 图像,但**不原生支持音频**。音频(以及未抽帧的视频文件)必须返回
|
||||
// ErrModalityUnsupported,而不是拿视觉塔去编码——那会往统一空间里灌入
|
||||
// 语义错误的坐标,而错误是静默的。视频请由上层抽帧后逐帧当作图像编码。
|
||||
func (e *Embedder) EmbedImageDense(raw []byte, mime string) ([]float64, error) {
|
||||
e.mu.RLock()
|
||||
defer e.mu.RUnlock()
|
||||
if !e.loaded {
|
||||
return nil, fmt.Errorf("qwen embedder not loaded")
|
||||
}
|
||||
switch {
|
||||
case strings.HasPrefix(mime, "audio/"):
|
||||
return nil, fmt.Errorf("%w: audio (%s) 需由真正的统一音频模型扩展", vector.ErrModalityUnsupported, mime)
|
||||
case strings.HasPrefix(mime, "video/"):
|
||||
return nil, fmt.Errorf("%w: 视频文件请先抽帧,逐帧按图像编码 (%s)", vector.ErrModalityUnsupported, mime)
|
||||
}
|
||||
pixels, err := preprocessImage(raw)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
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)
|
||||
}
|
||||
in, err := ort.NewTensor(ort.Shape{1, int64(len(ids))}, inputIDs)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("qwen token input: %w", err)
|
||||
}
|
||||
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
|
||||
}
|
||||
|
||||
func (e *Embedder) runVision(pixels []float32) ([][]float32, error) {
|
||||
in, err := ort.NewTensor(ort.Shape{qwenImagePatches, qwenPatchVectorSize}, pixels)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("qwen vision input: %w", err)
|
||||
}
|
||||
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
|
||||
}
|
||||
|
||||
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("qwen hidden tensor: %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)
|
||||
|
||||
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 {
|
||||
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
|
||||
}
|
||||
|
||||
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.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
|
||||
})
|
||||
}
|
||||
|
||||
func computeFingerprint(modelDir string) string {
|
||||
h := sha256.New()
|
||||
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})
|
||||
h.Write(data)
|
||||
h.Write([]byte{0})
|
||||
}
|
||||
}
|
||||
entries, _ := os.ReadDir(modelDir)
|
||||
var names []string
|
||||
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 {
|
||||
if info, err := os.Stat(filepath.Join(modelDir, n)); err == nil {
|
||||
fmt.Fprintf(h, "%s:%d\n", n, info.Size())
|
||||
}
|
||||
}
|
||||
return hex.EncodeToString(h.Sum(nil))
|
||||
}
|
||||
|
||||
func findOnnxLib() string {
|
||||
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
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
@ -1,277 +0,0 @@
|
||||
//go:build onnxruntime
|
||||
|
||||
package qwen
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"image"
|
||||
"image/png"
|
||||
"math"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"gitcode.com/JianFeeeee/HomeAgent/internal/memory/vector"
|
||||
)
|
||||
|
||||
// onnxModelDir 返回三段式 Qwen 多模态 ONNX 产物目录。
|
||||
//
|
||||
// 产物约 8GB(含外部权重),不进仓库;由 scripts/export_qwen3vl_embedding_onnx.py
|
||||
// 自动拉取模型并导出。可通过 QWEN_ONNX_MODEL_DIR 指向别处;目录不存在时相关
|
||||
// 测试跳过,而不是失败——CI 与本机开发者都不一定有这份产物。
|
||||
func onnxModelDir() string {
|
||||
if v := os.Getenv("QWEN_ONNX_MODEL_DIR"); v != "" {
|
||||
return v
|
||||
}
|
||||
return "/home/newqqagent/models/qwen3-vl-embed-multimodal-onnx"
|
||||
}
|
||||
|
||||
func requireONNXArtifacts(t *testing.T) string {
|
||||
t.Helper()
|
||||
dir := onnxModelDir()
|
||||
for _, name := range []string{"TokenEmbedding.onnx", "Transformer.onnx", "Vision.onnx", "embed_config.json", "tokenizer.json"} {
|
||||
if _, err := os.Stat(dir + "/" + name); err != nil {
|
||||
t.Skipf("ONNX 产物不完整(%s: %v),跳过;用 scripts/export_qwen3vl_embedding_onnx.py 导出", name, err)
|
||||
}
|
||||
}
|
||||
return dir
|
||||
}
|
||||
|
||||
// onnxReference 是导出脚本 `--emit-reference` 写出的冻结参考。
|
||||
//
|
||||
// 刻意不把浮点常量硬编码在测试里:参考值必须能追溯到「哪个模型、哪次导出、
|
||||
// 什么输入」,而不是一组无人知道出处的数字。参考文件里的 RGB/尺寸同时用来
|
||||
// 构造测试图片,保证输入与参考按构造一致,不会因测试改动而静默错位。
|
||||
type onnxReference struct {
|
||||
Text string `json:"text"`
|
||||
TextVectorPrefix []float64 `json:"text_vector_prefix"`
|
||||
ImageRGB []int `json:"image_rgb"`
|
||||
ImageSize int `json:"image_size"`
|
||||
ImageVectorPrefix []float64 `json:"image_vector_prefix"`
|
||||
Dim int `json:"dim"`
|
||||
}
|
||||
|
||||
func loadReference(t *testing.T, dir string) *onnxReference {
|
||||
t.Helper()
|
||||
path := os.Getenv("QWEN_ONNX_REFERENCE")
|
||||
if path == "" {
|
||||
path = filepath.Join(dir, "qwen_reference.json")
|
||||
}
|
||||
raw, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
t.Skipf("缺少冻结参考 %s(由 scripts/export_qwen3vl_embedding_onnx.py --emit-reference 生成): %v", path, err)
|
||||
}
|
||||
var ref onnxReference
|
||||
if err := json.Unmarshal(raw, &ref); err != nil {
|
||||
t.Fatalf("解析参考 %s: %v", path, err)
|
||||
}
|
||||
if ref.Text == "" || len(ref.TextVectorPrefix) == 0 || len(ref.ImageRGB) != 3 || ref.ImageSize <= 0 {
|
||||
t.Fatalf("参考 %s 不完整: %+v", path, ref)
|
||||
}
|
||||
return &ref
|
||||
}
|
||||
|
||||
// solidPNG 生成一张 size×size 纯色 PNG,供跨语言冻结向量回归。
|
||||
//
|
||||
// 刻意用纯色且尺寸与视觉塔一致:Go 侧预处理对已是 768×768 的输入不做插值、
|
||||
// 不补边,于是 patch 张量只由布局决定。一旦 patch 排列写错(内层循环顺序、
|
||||
// merge 分组顺序、通道顺序),冻结向量立刻不匹配——而那类错误在人工看图时
|
||||
// 几乎发现不了。
|
||||
func solidPNG(t *testing.T, size int, r, g, b uint8) []byte {
|
||||
t.Helper()
|
||||
img := image.NewNRGBA(image.Rect(0, 0, size, size))
|
||||
for y := 0; y < size; y++ {
|
||||
for x := 0; x < size; x++ {
|
||||
i := img.PixOffset(x, y)
|
||||
img.Pix[i], img.Pix[i+1], img.Pix[i+2], img.Pix[i+3] = r, g, b, 255
|
||||
}
|
||||
}
|
||||
var buf bytes.Buffer
|
||||
if err := png.Encode(&buf, img); err != nil {
|
||||
t.Fatalf("encode png: %v", err)
|
||||
}
|
||||
return buf.Bytes()
|
||||
}
|
||||
|
||||
func newTestEmbedder(t *testing.T, dir string) *Embedder {
|
||||
t.Helper()
|
||||
e, err := New(dir)
|
||||
if err != nil {
|
||||
t.Fatalf("New: %v", err)
|
||||
}
|
||||
t.Cleanup(e.Close)
|
||||
if !e.Loaded() || e.Dim() != 2048 || e.Fingerprint() == "" {
|
||||
t.Fatalf("元数据异常: loaded=%v dim=%d fingerprint=%q", e.Loaded(), e.Dim(), e.Fingerprint())
|
||||
}
|
||||
return e
|
||||
}
|
||||
|
||||
func assertNormalized(t *testing.T, name string, got []float64, dim int) {
|
||||
t.Helper()
|
||||
if dim > 0 && len(got) != dim {
|
||||
t.Fatalf("%s 维度 = %d,期望 %d", name, len(got), dim)
|
||||
}
|
||||
var norm float64
|
||||
for _, v := range got {
|
||||
norm += v * v
|
||||
}
|
||||
if diff := math.Abs(math.Sqrt(norm) - 1); diff > 1e-6 {
|
||||
t.Errorf("%s L2 norm = %.9f,期望 1", name, math.Sqrt(norm))
|
||||
}
|
||||
}
|
||||
|
||||
func assertFrozenPrefix(t *testing.T, name string, got, want []float64) {
|
||||
t.Helper()
|
||||
if len(got) < len(want) {
|
||||
t.Fatalf("%s 向量过短: %d", name, len(got))
|
||||
}
|
||||
for i := range want {
|
||||
if diff := math.Abs(got[i] - want[i]); diff > 2e-5 {
|
||||
t.Errorf("%s 维度 %d = %.10g,参考 %.10g,差 %.3g", name, i, got[i], want[i], diff)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestEmbedderMatchesONNXReference 逐维对比导出脚本写出的冻结参考向量,
|
||||
// 验证完整 Go 路径:模板渲染 → BPE → TokenEmbedding → Transformer →
|
||||
// last-token 池化 → L2 normalize。
|
||||
//
|
||||
// 只覆盖前若干维不是因为放宽正确性(脚本侧的 PyTorch↔ONNX 校验是逐维的),
|
||||
// 而是避免把 2048 个浮点常量塞进仓库;这里负责捕获 Go 张量形状、输入名、
|
||||
// 输出名、池化/归一化或模板接线错误。
|
||||
func TestEmbedderMatchesONNXReference(t *testing.T) {
|
||||
dir := requireONNXArtifacts(t)
|
||||
e := newTestEmbedder(t, dir)
|
||||
ref := loadReference(t, dir)
|
||||
|
||||
ids, _, _, _, err := e.tok.textModelInput(e.config.Instruction, ref.Text, e.config.MaxLength)
|
||||
if err != nil {
|
||||
t.Fatalf("textModelInput: %v", err)
|
||||
}
|
||||
postID, ok := e.tok.SpecialID("<|endoftext|>")
|
||||
if !ok || ids[len(ids)-1] != postID {
|
||||
t.Fatalf("模型输入 post-processor 异常: len=%d tail=%v postID=%d ok=%v", len(ids), ids[len(ids)-1:], postID, ok)
|
||||
}
|
||||
|
||||
got, err := e.VectorizeDense(ref.Text)
|
||||
if err != nil {
|
||||
t.Fatalf("VectorizeDense: %v", err)
|
||||
}
|
||||
assertNormalized(t, "text", got, ref.Dim)
|
||||
assertFrozenPrefix(t, "text", got, ref.TextVectorPrefix)
|
||||
}
|
||||
|
||||
// TestEmbedderImageMatchesONNXReference 冻结一张纯色图的参考向量,
|
||||
// 验证 Go 侧的视觉预处理 + patch 排列 + 视觉注入 + 语言模型与 Python 参考一致。
|
||||
func TestEmbedderImageMatchesONNXReference(t *testing.T) {
|
||||
dir := requireONNXArtifacts(t)
|
||||
e := newTestEmbedder(t, dir)
|
||||
ref := loadReference(t, dir)
|
||||
|
||||
img := solidPNG(t, ref.ImageSize, uint8(ref.ImageRGB[0]), uint8(ref.ImageRGB[1]), uint8(ref.ImageRGB[2]))
|
||||
got, err := e.EmbedImageDense(img, "image/png")
|
||||
if err != nil {
|
||||
t.Fatalf("EmbedImageDense: %v", err)
|
||||
}
|
||||
assertNormalized(t, "image", got, ref.Dim)
|
||||
assertFrozenPrefix(t, "image", got, ref.ImageVectorPrefix)
|
||||
}
|
||||
|
||||
// TestEmbedderTextIsSensitiveToInput 阴性对照:冻结向量必须真的随输入变化。
|
||||
//
|
||||
// 没有这条对照,一个「永远返回同一向量」的错误实现也能通过上面的冻结回归
|
||||
//(只要那个常量恰好等于参考值)。这里验证不同文本给出不同向量,且相似文本
|
||||
// 的余弦高于无关文本——即嵌入确实携带语义,而不是常量。
|
||||
func TestEmbedderTextIsSensitiveToInput(t *testing.T) {
|
||||
dir := requireONNXArtifacts(t)
|
||||
e := newTestEmbedder(t, dir)
|
||||
|
||||
base, err := e.VectorizeDense("今天天气怎么样")
|
||||
if err != nil {
|
||||
t.Fatalf("VectorizeDense: %v", err)
|
||||
}
|
||||
same, err := e.VectorizeDense("今天天气怎么样")
|
||||
if err != nil {
|
||||
t.Fatalf("VectorizeDense: %v", err)
|
||||
}
|
||||
if cosine(base, same) < 0.999999 {
|
||||
t.Errorf("同一输入两次嵌入不一致: cos=%.9f(ONNX 会话被并发复用或存在非确定性)", cosine(base, same))
|
||||
}
|
||||
|
||||
other, err := e.VectorizeDense("数据库索引的选择性是怎么计算的")
|
||||
if err != nil {
|
||||
t.Fatalf("VectorizeDense: %v", err)
|
||||
}
|
||||
if cosine(base, other) > 0.999 {
|
||||
t.Errorf("无关文本的余弦高达 %.6f,嵌入可能是常量", cosine(base, other))
|
||||
}
|
||||
}
|
||||
|
||||
// TestEmbedderImageMatchesONNXReference 的替代:不依赖冻结参考的不变量检查。
|
||||
//
|
||||
// 即使参考文件缺失(没有导出产物)或未重新生成,这些不变量也应成立:
|
||||
// 图像路径必须真的走了视觉塔,且不同图片给出不同坐标。
|
||||
func TestEmbedderImageDiffersFromText(t *testing.T) {
|
||||
dir := requireONNXArtifacts(t)
|
||||
e := newTestEmbedder(t, dir)
|
||||
|
||||
imgVec, err := e.EmbedImageDense(solidPNG(t, qwenImageSize, 200, 30, 30), "image/png")
|
||||
if err != nil {
|
||||
t.Fatalf("EmbedImageDense: %v", err)
|
||||
}
|
||||
txtVec, err := e.VectorizeDense(DefaultInstruction)
|
||||
if err != nil {
|
||||
t.Fatalf("VectorizeDense: %v", err)
|
||||
}
|
||||
if cosine(imgVec, txtVec) > 0.999 {
|
||||
t.Error("图像向量与文本向量几乎相同,视觉塔可能没被真正执行")
|
||||
}
|
||||
|
||||
// 不同颜色的图必须给出不同向量
|
||||
blue, err := e.EmbedImageDense(solidPNG(t, qwenImageSize, 30, 150, 220), "image/png")
|
||||
if err != nil {
|
||||
t.Fatalf("EmbedImageDense: %v", err)
|
||||
}
|
||||
if cosine(imgVec, blue) > 0.999999 {
|
||||
t.Error("不同图片给出相同向量,视觉路径未生效")
|
||||
}
|
||||
}
|
||||
|
||||
// TestEmbedderRejectsUnsupportedModalities 音频与视频文件必须显式报「不在本空间」,
|
||||
// 而不是拿视觉塔硬编码一个语义错误的坐标。
|
||||
//
|
||||
// Qwen3-VL 原生支持文本与图像;音频需要未来接入真正的统一音频模型。
|
||||
// 若这里退化成普通错误,调用方会把它当「本次失败、下次重试」,
|
||||
// 于是每轮启动都重试一批永远不可能成功的条目。
|
||||
func TestEmbedderRejectsUnsupportedModalities(t *testing.T) {
|
||||
dir := requireONNXArtifacts(t)
|
||||
e := newTestEmbedder(t, dir)
|
||||
|
||||
for _, mime := range []string{"audio/wav", "audio/mpeg", "video/mp4", "video/quicktime"} {
|
||||
_, err := e.EmbedImageDense([]byte("not-a-real-media"), mime)
|
||||
if err == nil {
|
||||
t.Fatalf("%s 应返回错误而不是造出向量", mime)
|
||||
}
|
||||
if !errors.Is(err, vector.ErrModalityUnsupported) {
|
||||
t.Errorf("%s 错误应为 ErrModalityUnsupported,实际: %v", mime, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func cosine(a, b []float64) float64 {
|
||||
if len(a) != len(b) || len(a) == 0 {
|
||||
return 0
|
||||
}
|
||||
var dot, na, nb float64
|
||||
for i := range a {
|
||||
dot += a[i] * b[i]
|
||||
na += a[i] * a[i]
|
||||
nb += b[i] * b[i]
|
||||
}
|
||||
if na == 0 || nb == 0 {
|
||||
return 0
|
||||
}
|
||||
return dot / (math.Sqrt(na) * math.Sqrt(nb))
|
||||
}
|
||||
@ -1,28 +0,0 @@
|
||||
//go:build !onnxruntime
|
||||
|
||||
package qwen
|
||||
|
||||
import "fmt"
|
||||
|
||||
// Embedder 在未启用 onnxruntime 时为 no-op 实现。
|
||||
// 默认构建不链接 onnxruntime,保持原有 fastText/TF-IDF 行为不变。
|
||||
type Embedder struct {
|
||||
loaded bool
|
||||
}
|
||||
|
||||
func New(_ string) (*Embedder, error) {
|
||||
return nil, fmt.Errorf("qwen 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("qwen embedder not available")
|
||||
}
|
||||
|
||||
func (e *Embedder) EmbedImageDense(_ []byte, _ string) ([]float64, error) {
|
||||
return nil, fmt.Errorf("qwen embedder not available")
|
||||
}
|
||||
@ -1,157 +0,0 @@
|
||||
//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)))))
|
||||
}
|
||||
@ -1,86 +0,0 @@
|
||||
//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
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@ -1,503 +0,0 @@
|
||||
// Package qwen 实现 Qwen3-VL-Embedding 的字节级 BPE 分词器。
|
||||
//
|
||||
// 为什么不复用 clip 的 tokenizer:CLIP 用的是「小写化 + 空白规整 + 词表 BPE」,
|
||||
// 而千问是 **GPT-2 式字节级 BPE**——先把输入按字节映射到一组可见 unicode,
|
||||
// 再对映射后的字符串做 BPE 合并。两者的预处理不可互换,硬套会在中文和
|
||||
// 空白较多的输入上产出完全不同的 token。
|
||||
//
|
||||
// 与上游(HuggingFace tokenizer.json 的 Rust 实现)对齐时的两处坑:
|
||||
//
|
||||
// 1. pre_tokenizer 正则里的 `\s+(?!\S)` 是**负向前瞻**,Go 的 RE2 不支持
|
||||
// lookaround。该分支只在「空白一直延伸到串尾」时命中,而此时贪婪的
|
||||
// `\s+` 会匹配完全相同的区间,所以直接删掉该分支即为等价改写。
|
||||
// 2. Go 的 `\s` 只覆盖 ASCII,而 Rust regex 的 `\s` 是 Unicode
|
||||
// `\p{White_Space}`。不换成 \p{White_Space} 的话,全角空格、NBSP、
|
||||
// 行分隔符等的切分点会与上游不一致。
|
||||
package qwen
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"strings"
|
||||
"unicode"
|
||||
"unicode/utf8"
|
||||
)
|
||||
|
||||
// 空白判定统一用 unicode.IsSpace(Unicode White_Space 属性)。
|
||||
//
|
||||
// 不能用 Go 正则里的 \s——那只覆盖 ASCII;也不能写 \p{White_Space}——Go 的
|
||||
// regexp 只支持 script/category,不支持二进制属性(会报 invalid character
|
||||
// class range)。上游 Rust regex 的 \s 正是 White_Space,所以这里以
|
||||
// unicode.IsSpace 为准。
|
||||
|
||||
// specialToken 是一个 AddedToken:以整体形式优先匹配,不参与 BPE 拆分。
|
||||
type specialToken struct {
|
||||
content string
|
||||
id int
|
||||
}
|
||||
|
||||
// Tokenizer 是千问的字节级 BPE 分词器。
|
||||
type Tokenizer struct {
|
||||
vocab map[string]int
|
||||
ranks map[string]int
|
||||
|
||||
// byteEnc 是 GPT-2 的 byte→unicode 映射:把 0..255 每个字节映到一个
|
||||
// 「可见且不会与正常文本冲突」的 unicode 码点。因为 BPE 词表基于文本构建,
|
||||
// 直接放原始字节会与合法 UTF-8 冲突。
|
||||
byteEnc map[byte]rune
|
||||
|
||||
// specials 按 content 长度降序,保证「最长优先」——
|
||||
// 否则 `<|im_start|>` 可能被 `<|im_` 之类的短 token 先切走。
|
||||
specials []specialToken
|
||||
|
||||
// MaxLen 是嵌入用途的截断上限(与导出脚本的 MAX_LENGTH 一致)。
|
||||
MaxLen int
|
||||
}
|
||||
|
||||
// tokenizerJSON 只取我们需要的部分。
|
||||
type tokenizerJSON struct {
|
||||
Model struct {
|
||||
Vocab map[string]int `json:"vocab"`
|
||||
Merges []interface{} `json:"merges"`
|
||||
} `json:"model"`
|
||||
AddedTokens []struct {
|
||||
ID int `json:"id"`
|
||||
Content string `json:"content"`
|
||||
Special bool `json:"special"`
|
||||
} `json:"added_tokens"`
|
||||
}
|
||||
|
||||
// LoadTokenizer 从模型目录加载 tokenizer.json。
|
||||
func LoadTokenizer(modelDir string) (*Tokenizer, error) {
|
||||
raw, err := os.ReadFile(filepath.Join(modelDir, "tokenizer.json"))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("read tokenizer.json: %w", err)
|
||||
}
|
||||
var tj tokenizerJSON
|
||||
if err := json.Unmarshal(raw, &tj); err != nil {
|
||||
return nil, fmt.Errorf("parse tokenizer.json: %w", err)
|
||||
}
|
||||
if len(tj.Model.Vocab) == 0 {
|
||||
return nil, fmt.Errorf("tokenizer.json 的 model.vocab 为空")
|
||||
}
|
||||
|
||||
ranks := make(map[string]int, len(tj.Model.Merges))
|
||||
for i, m := range tj.Model.Merges {
|
||||
// merges 有两种形态:字符串 "a b",或数组 ["a","b"]。
|
||||
var pair string
|
||||
switch v := m.(type) {
|
||||
case string:
|
||||
pair = v
|
||||
case []interface{}:
|
||||
if len(v) == 2 {
|
||||
a, _ := v[0].(string)
|
||||
b, _ := v[1].(string)
|
||||
pair = a + " " + b
|
||||
}
|
||||
}
|
||||
if pair != "" {
|
||||
if _, seen := ranks[pair]; !seen {
|
||||
ranks[pair] = i
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
t := &Tokenizer{
|
||||
vocab: tj.Model.Vocab,
|
||||
ranks: ranks,
|
||||
byteEnc: bytesToUnicode(),
|
||||
MaxLen: 512,
|
||||
}
|
||||
for _, at := range tj.AddedTokens {
|
||||
if at.Special && at.Content != "" {
|
||||
t.specials = append(t.specials, specialToken{content: at.Content, id: at.ID})
|
||||
}
|
||||
}
|
||||
// 最长优先,避免短 token 抢走长 token 的前缀。
|
||||
sort.Slice(t.specials, func(i, j int) bool {
|
||||
return len(t.specials[i].content) > len(t.specials[j].content)
|
||||
})
|
||||
return t, nil
|
||||
}
|
||||
|
||||
// VocabSize 返回词表大小(诊断用)。
|
||||
func (t *Tokenizer) VocabSize() int { return len(t.vocab) }
|
||||
|
||||
// SpecialID 返回特殊 token 的 id;不存在时 ok=false。
|
||||
func (t *Tokenizer) SpecialID(content string) (int, bool) {
|
||||
for _, s := range t.specials {
|
||||
if s.content == content {
|
||||
return s.id, true
|
||||
}
|
||||
}
|
||||
return 0, false
|
||||
}
|
||||
|
||||
// DefaultInstruction 是导出脚本随 embed_config.json 写入的默认指令。
|
||||
const DefaultInstruction = "Represent the user's input."
|
||||
|
||||
// renderInstructionInput 按模型自带的对话模板拼输入(无构建标签,便于测试)。
|
||||
//
|
||||
// 必须与 HuggingFace processor 的 apply_chat_template(add_generation_prompt=True)
|
||||
// 产出完全一致:指令放 system、正文放 user、以 assistant 起始符结尾。差一个
|
||||
// 特殊 token,池化取到的「最后一个有效 token」位置就变了,嵌入也就不同——
|
||||
// 而且不会报错。参考数据集里有该模板串的用例,能逐 token 对齐验证。
|
||||
func renderInstructionInput(instruction, text string) string {
|
||||
if instruction == "" {
|
||||
instruction = DefaultInstruction
|
||||
}
|
||||
return "<|im_start|>system\n" + instruction +
|
||||
"<|im_end|>\n<|im_start|>user\n" + text +
|
||||
"<|im_end|>\n<|im_start|>assistant\n"
|
||||
}
|
||||
|
||||
// Encode 把文本编码为 token id 序列(识别输入中已有的特殊 token,
|
||||
// 但不执行 tokenizer.json 的 post_processor,也不做截断)。
|
||||
func (t *Tokenizer) Encode(text string) []int {
|
||||
var ids []int
|
||||
for _, seg := range t.splitSpecials(text) {
|
||||
if seg.specialID >= 0 {
|
||||
ids = append(ids, seg.specialID)
|
||||
continue
|
||||
}
|
||||
ids = append(ids, t.encodeOrdinary(seg.text)...)
|
||||
}
|
||||
return ids
|
||||
}
|
||||
|
||||
// encodeModelInput 执行 TextTower 输入所需的 tokenizer post_processor。
|
||||
//
|
||||
// tokenizer.json 的 TemplateProcessing 规则是 `$A <|endoftext|>`;HuggingFace
|
||||
// 在 truncation=true 时先把 A 截到 maxLen-1,再保留末尾 post token。漏掉它不会
|
||||
// 触发 ONNX 错误,却会改变池化位置和整条嵌入向量,因此不能直接用 Encode 的结果。
|
||||
func (t *Tokenizer) encodeModelInput(text string, maxLen int) ([]int, error) {
|
||||
postID, ok := t.SpecialID("<|endoftext|>")
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("tokenizer.json 缺少 post token <|endoftext|>")
|
||||
}
|
||||
if maxLen <= 0 {
|
||||
return nil, fmt.Errorf("maxLen 必须大于 0")
|
||||
}
|
||||
|
||||
ids := t.Encode(text)
|
||||
if len(ids) >= maxLen {
|
||||
ids = ids[:maxLen-1]
|
||||
}
|
||||
return append(ids, postID), nil
|
||||
}
|
||||
|
||||
// seg 是「普通文本」或「已识别的特殊 token」二选一。
|
||||
type seg struct {
|
||||
text string
|
||||
specialID int // -1 表示普通文本
|
||||
}
|
||||
|
||||
// splitSpecials 把输入切成普通片段与特殊 token 片段。
|
||||
//
|
||||
// 为什么必须先切:`<|im_start|>` 在词表里是一个整体 id(151644),若走 BPE
|
||||
// 会被拆成若干子 token,编码结果与上游不一致,模型看到的输入也就变了。
|
||||
func (t *Tokenizer) splitSpecials(text string) []seg {
|
||||
if len(t.specials) == 0 || text == "" {
|
||||
return []seg{{text: text, specialID: -1}}
|
||||
}
|
||||
var out []seg
|
||||
for len(text) > 0 {
|
||||
// 找最靠前的特殊 token 出现位置(同位置取最长)。
|
||||
bestIdx, bestLen, bestID := -1, 0, -1
|
||||
for _, s := range t.specials {
|
||||
i := strings.Index(text, s.content)
|
||||
if i < 0 {
|
||||
continue
|
||||
}
|
||||
if bestIdx == -1 || i < bestIdx || (i == bestIdx && len(s.content) > bestLen) {
|
||||
bestIdx, bestLen, bestID = i, len(s.content), s.id
|
||||
}
|
||||
}
|
||||
if bestIdx == -1 {
|
||||
out = append(out, seg{text: text, specialID: -1})
|
||||
break
|
||||
}
|
||||
if bestIdx > 0 {
|
||||
out = append(out, seg{text: text[:bestIdx], specialID: -1})
|
||||
}
|
||||
out = append(out, seg{specialID: bestID})
|
||||
text = text[bestIdx+bestLen:]
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// encodeOrdinary 对普通文本做「切分 → 字节映射 → BPE 合并」。
|
||||
func (t *Tokenizer) encodeOrdinary(text string) []int {
|
||||
if text == "" {
|
||||
return nil
|
||||
}
|
||||
var ids []int
|
||||
for _, piece := range t.preTokenize(text) {
|
||||
// 字节级映射:先把 piece 的 UTF-8 字节逐个映射成 unicode 字符。
|
||||
var sb strings.Builder
|
||||
for _, b := range []byte(piece) {
|
||||
sb.WriteRune(t.byteEnc[b])
|
||||
}
|
||||
for _, tok := range t.bpe(sb.String()) {
|
||||
if id, ok := t.vocab[tok]; ok {
|
||||
ids = append(ids, id)
|
||||
}
|
||||
// 词表里找不到的片段直接丢弃:正常情况不会发生
|
||||
//(词表覆盖全部 256 个字节级字符),发生即数据有问题。
|
||||
}
|
||||
}
|
||||
return ids
|
||||
}
|
||||
|
||||
// ---- pre_tokenizer ----
|
||||
//
|
||||
// 上游是一条正则(tokenizer.json 的 pre_tokenizer.pretokenizers[0].pattern):
|
||||
//
|
||||
// (?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}|
|
||||
// ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+
|
||||
//
|
||||
// **为什么不用一个 Go 正则**:末两个分支里的 `\s+(?!\S)` 是负向前瞻,RE2 不
|
||||
// 支持 lookaround;而且它的真实语义依赖**回溯**——`\s+` 先贪婪吃完整段空白,
|
||||
// 发现后面是非空白导致 `(?!\S)` 失败,于是回退一个字符,正好留下末尾一个
|
||||
// 空白给前面那些以 ` ?` / `[^…]?` 开头的分支合并。这个“留一个”直接决定
|
||||
// 切分点(`" leading"` 会切成 `" "` + `" leading"` 而不是 `" "` + `"leading"`),
|
||||
// 近似改写必然对不上,所以按分支顺序显式实现。
|
||||
func (t *Tokenizer) preTokenize(text string) []string {
|
||||
var out []string
|
||||
for len(text) > 0 {
|
||||
switch {
|
||||
case matchApostrophe(text) > 0:
|
||||
n := matchApostrophe(text)
|
||||
out = append(out, text[:n])
|
||||
text = text[n:]
|
||||
case matchWord(text) > 0:
|
||||
n := matchWord(text)
|
||||
out = append(out, text[:n])
|
||||
text = text[n:]
|
||||
case matchDigit(text) > 0:
|
||||
n := matchDigit(text)
|
||||
out = append(out, text[:n])
|
||||
text = text[n:]
|
||||
case matchPunct(text) > 0:
|
||||
n := matchPunct(text)
|
||||
out = append(out, text[:n])
|
||||
text = text[n:]
|
||||
case matchNewline(text) > 0:
|
||||
n := matchNewline(text)
|
||||
out = append(out, text[:n])
|
||||
text = text[n:]
|
||||
default:
|
||||
// `\s+(?!\S)|\s+` 合一:空白段。
|
||||
total, lastStart := wsRun(text)
|
||||
if total == 0 {
|
||||
// 兜底:不应到达(分支覆盖全部字符),防御性前进一个 rune。
|
||||
_, size := utf8.DecodeRuneInString(text)
|
||||
out = append(out, text[:size])
|
||||
text = text[size:]
|
||||
continue
|
||||
}
|
||||
n := total
|
||||
if total < len(text) && lastStart > 0 {
|
||||
n = lastStart // 后面还有非空白 → 回退掉末尾那一个空白
|
||||
}
|
||||
out = append(out, text[:n])
|
||||
text = text[n:]
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func runeAt(s string) (rune, int) { return utf8.DecodeRuneInString(s) }
|
||||
|
||||
func isLetter(r rune) bool { return unicode.IsLetter(r) }
|
||||
func isNumber(r rune) bool { return unicode.IsNumber(r) }
|
||||
func isWS(r rune) bool { return unicode.IsSpace(r) }
|
||||
|
||||
// wsRun 返回开头连续空白段的字节长度,以及最后一个空白 rune 的起始字节位置。
|
||||
func wsRun(s string) (total, lastStart int) {
|
||||
lastStart = -1
|
||||
i := 0
|
||||
for i < len(s) {
|
||||
r, size := runeAt(s[i:])
|
||||
if !isWS(r) {
|
||||
break
|
||||
}
|
||||
lastStart = i
|
||||
i += size
|
||||
}
|
||||
return i, lastStart
|
||||
}
|
||||
|
||||
// matchApostrophe:`(?i:'s|'t|'re|'ve|'m|'ll|'d)`
|
||||
func matchApostrophe(s string) int {
|
||||
if len(s) == 0 || s[0] != '\'' {
|
||||
return 0
|
||||
}
|
||||
rest := s[1:]
|
||||
// 各后缀互为前缀关系(re/ve/ll/s/t/m/d),所以先试长的。
|
||||
for _, suf := range []string{"re", "ve", "ll", "s", "t", "m", "d"} {
|
||||
if len(rest) >= len(suf) && strings.EqualFold(rest[:len(suf)], suf) {
|
||||
return 1 + len(suf)
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
// matchWord:`[^\r\n\p{L}\p{N}]?\p{L}+`
|
||||
//
|
||||
// 注意可选字符**排除** \r \n;若吃了可选字符却没有字母跟上,整个分支失败
|
||||
// (与正则的“该分支不匹配”一致,不能把可选字符当已消耗)。
|
||||
func matchWord(s string) int {
|
||||
i := 0
|
||||
if r, size := runeAt(s); r != '\r' && r != '\n' && !isLetter(r) && !isNumber(r) {
|
||||
i = size
|
||||
}
|
||||
r, size := runeAt(s[i:])
|
||||
if !isLetter(r) {
|
||||
return 0
|
||||
}
|
||||
i += size
|
||||
for i < len(s) {
|
||||
r, size := runeAt(s[i:])
|
||||
if !isLetter(r) {
|
||||
break
|
||||
}
|
||||
i += size
|
||||
}
|
||||
return i
|
||||
}
|
||||
|
||||
// matchDigit:`\p{N}` —— 只吃**一个**数字。
|
||||
func matchDigit(s string) int {
|
||||
if r, size := runeAt(s); isNumber(r) {
|
||||
return size
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
// matchPunct:` ?[^\s\p{L}\p{N}]+[\r\n]*`
|
||||
//
|
||||
// 开头是**字面空格**(不是 \s),所以只可能吃掉一个 U+0020。
|
||||
func matchPunct(s string) int {
|
||||
i := 0
|
||||
if strings.HasPrefix(s, " ") {
|
||||
i = 1
|
||||
}
|
||||
n := 0
|
||||
for i+n < len(s) {
|
||||
r, size := runeAt(s[i+n:])
|
||||
if isWS(r) || isLetter(r) || isNumber(r) {
|
||||
break
|
||||
}
|
||||
n += size
|
||||
}
|
||||
if n == 0 {
|
||||
return 0
|
||||
}
|
||||
i += n
|
||||
for i < len(s) && (s[i] == '\r' || s[i] == '\n') {
|
||||
i++
|
||||
}
|
||||
return i
|
||||
}
|
||||
|
||||
// matchNewline:`\s*[\r\n]+`
|
||||
//
|
||||
// 贪婪+回溯的真实语义:`\s*` 先吃完整段空白,`[\r\n]+` 无可匹配而回退,
|
||||
// 最终停在段内**最后一个** \r 或 \n 之前,再把它之后的连续 \r\n 吃掉。
|
||||
func matchNewline(s string) int {
|
||||
total, _ := wsRun(s)
|
||||
if total == 0 {
|
||||
return 0
|
||||
}
|
||||
last := -1
|
||||
for j := total - 1; j >= 0; j-- {
|
||||
if s[j] == '\r' || s[j] == '\n' {
|
||||
last = j
|
||||
break
|
||||
}
|
||||
}
|
||||
if last < 0 {
|
||||
return 0
|
||||
}
|
||||
end := last
|
||||
for end < len(s) && (s[end] == '\r' || s[end] == '\n') {
|
||||
end++
|
||||
}
|
||||
return end
|
||||
}
|
||||
|
||||
// bpe 是标准字节级 BPE:反复合并 rank 最小的相邻对,直到无可合并。
|
||||
func (t *Tokenizer) bpe(word string) []string {
|
||||
symbols := make([]string, 0, len(word))
|
||||
for _, r := range word {
|
||||
symbols = append(symbols, string(r))
|
||||
}
|
||||
if len(symbols) < 2 {
|
||||
return symbols
|
||||
}
|
||||
|
||||
for {
|
||||
bestRank, bestIdx := -1, -1
|
||||
for i := 0; i+1 < len(symbols); i++ {
|
||||
r, ok := t.ranks[symbols[i]+" "+symbols[i+1]]
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
if bestRank == -1 || r < bestRank {
|
||||
bestRank, bestIdx = r, i
|
||||
}
|
||||
}
|
||||
if bestIdx == -1 {
|
||||
return symbols
|
||||
}
|
||||
merged := symbols[bestIdx] + symbols[bestIdx+1]
|
||||
symbols = append(symbols[:bestIdx], append([]string{merged}, symbols[bestIdx+2:]...)...)
|
||||
if len(symbols) < 2 {
|
||||
return symbols
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// bytesToUnicode 是 GPT-2 的字节↔unicode 映射表。
|
||||
//
|
||||
// 让每个字节都有一个「安全」的可见码点表示,避免原始控制字节混进 BPE 词表。
|
||||
// 可打印 ASCII 与拉丁补充区保持原样,其余字节映射到 256 之后的码点。
|
||||
func bytesToUnicode() map[byte]rune {
|
||||
bs := make([]int, 0, 256)
|
||||
for b := int('!'); b <= int('~'); b++ {
|
||||
bs = append(bs, b)
|
||||
}
|
||||
for b := 0xA1; b <= 0xAC; b++ {
|
||||
bs = append(bs, b)
|
||||
}
|
||||
for b := 0xAE; b <= 0xFF; b++ {
|
||||
bs = append(bs, b)
|
||||
}
|
||||
|
||||
inBS := make(map[int]bool, len(bs))
|
||||
for _, b := range bs {
|
||||
inBS[b] = true
|
||||
}
|
||||
|
||||
cs := make([]int, len(bs))
|
||||
copy(cs, bs)
|
||||
n := 0
|
||||
for b := 0; b < 256; b++ {
|
||||
if inBS[b] {
|
||||
continue
|
||||
}
|
||||
bs = append(bs, b)
|
||||
cs = append(cs, 256+n)
|
||||
n++
|
||||
}
|
||||
|
||||
out := make(map[byte]rune, 256)
|
||||
for i, b := range bs {
|
||||
out[byte(b)] = rune(cs[i])
|
||||
}
|
||||
return out
|
||||
}
|
||||
@ -1,221 +0,0 @@
|
||||
package qwen
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// modelDir 是本地千问模型目录。不存在则跳过——参考数据已固化在 testdata,
|
||||
// 但分词器本身要从 tokenizer.json 加载词表与 merges(11MB,不入库)。
|
||||
const modelDir = "/home/newqqagent/models/models/qwen--Qwen3-VL-Embedding-2B/snapshots/master"
|
||||
|
||||
type tokenizerRef struct {
|
||||
VocabSize int `json:"vocab_size"`
|
||||
Cases []struct {
|
||||
Text string `json:"text"`
|
||||
IDs []int `json:"ids"`
|
||||
Tokens []string `json:"tokens"`
|
||||
} `json:"cases"`
|
||||
AddedTokens []struct {
|
||||
Content string `json:"content"`
|
||||
ID int `json:"id"`
|
||||
Special bool `json:"special"`
|
||||
} `json:"added_tokens"`
|
||||
}
|
||||
|
||||
func loadRef(t *testing.T) *tokenizerRef {
|
||||
t.Helper()
|
||||
raw, err := os.ReadFile(filepath.Join("testdata", "qwen_tokenizer_reference.json"))
|
||||
if err != nil {
|
||||
t.Fatalf("读取参考数据: %v", err)
|
||||
}
|
||||
var ref tokenizerRef
|
||||
if err := json.Unmarshal(raw, &ref); err != nil {
|
||||
t.Fatalf("解析参考数据: %v", err)
|
||||
}
|
||||
return &ref
|
||||
}
|
||||
|
||||
func loadTokenizer(t *testing.T) *Tokenizer {
|
||||
t.Helper()
|
||||
if _, err := os.Stat(filepath.Join(modelDir, "tokenizer.json")); err != nil {
|
||||
t.Skipf("模型目录不可用,跳过: %v", err)
|
||||
}
|
||||
tok, err := LoadTokenizer(modelDir)
|
||||
if err != nil {
|
||||
t.Fatalf("LoadTokenizer: %v", err)
|
||||
}
|
||||
return tok
|
||||
}
|
||||
|
||||
// 与 HuggingFace 的真实 tokenizer 逐条对齐。
|
||||
//
|
||||
// 这是本包唯一的正确性判据:字节级 BPE 的失败模式是「看起来能跑但 token 不同」,
|
||||
// 而 token 不同会让模型收到完全不同的输入,嵌入自然也就错了——不会报任何错。
|
||||
// 所以必须拿真实输出对照,不能靠读代码断言。
|
||||
func TestTokenizerMatchesReference(t *testing.T) {
|
||||
ref := loadRef(t)
|
||||
tok := loadTokenizer(t)
|
||||
|
||||
if got := tok.VocabSize(); got != ref.VocabSize {
|
||||
t.Errorf("词表大小 = %d,参考 %d", got, ref.VocabSize)
|
||||
}
|
||||
|
||||
failed := 0
|
||||
for _, c := range ref.Cases {
|
||||
got := tok.Encode(c.Text)
|
||||
if !sameIDs(got, c.IDs) {
|
||||
failed++
|
||||
t.Errorf("不一致 text=%q\n got %v\n want %v", c.Text, got, c.IDs)
|
||||
}
|
||||
}
|
||||
if failed > 0 {
|
||||
t.Fatalf("%d/%d 条用例不一致", failed, len(ref.Cases))
|
||||
}
|
||||
}
|
||||
|
||||
func sameIDs(a, b []int) bool {
|
||||
if len(a) != len(b) {
|
||||
return false
|
||||
}
|
||||
for i := range a {
|
||||
if a[i] != b[i] {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// 特殊 token 必须整体匹配:走 BPE 会被拆成子 token,模型看到的输入就变了。
|
||||
func TestSpecialTokensMatchWhole(t *testing.T) {
|
||||
ref := loadRef(t)
|
||||
tok := loadTokenizer(t)
|
||||
|
||||
for _, at := range ref.AddedTokens {
|
||||
if !at.Special {
|
||||
continue
|
||||
}
|
||||
got, ok := tok.SpecialID(at.Content)
|
||||
if !ok {
|
||||
t.Errorf("特殊 token %q 未从 tokenizer.json 载入", at.Content)
|
||||
continue
|
||||
}
|
||||
if got != at.ID {
|
||||
t.Errorf("特殊 token %q id=%d,参考 %d", at.Content, got, at.ID)
|
||||
}
|
||||
|
||||
// 单独出现时必须编码成恰好一个 id。
|
||||
ids := tok.Encode(at.Content)
|
||||
if len(ids) != 1 || ids[0] != at.ID {
|
||||
t.Errorf("特殊 token %q 应整体编码为 [%d],实际 %v", at.Content, at.ID, ids)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 最长优先:`<|im_start|>` 不能被更短的 `<|im_end|>` 之类前缀抢走。
|
||||
func TestSpecialTokenLongestFirst(t *testing.T) {
|
||||
tok := loadTokenizer(t)
|
||||
text := "<|im_start|>user\n你好<|im_end|>"
|
||||
|
||||
ids := tok.Encode(text)
|
||||
startID, _ := tok.SpecialID("<|im_start|>")
|
||||
endID, _ := tok.SpecialID("<|im_end|>")
|
||||
|
||||
if len(ids) == 0 || ids[0] != startID {
|
||||
t.Fatalf("应以 <|im_start|>(%d) 开头,实际 %v", startID, ids)
|
||||
}
|
||||
if last := ids[len(ids)-1]; last != endID {
|
||||
t.Fatalf("应以 <|im_end|>(%d) 结尾,实际 %v", endID, ids)
|
||||
}
|
||||
}
|
||||
|
||||
// 空串与单字符边界。
|
||||
func TestTokenizerEdgeCases(t *testing.T) {
|
||||
tok := loadTokenizer(t)
|
||||
if got := tok.Encode(""); len(got) != 0 {
|
||||
t.Errorf("空串应产出 0 个 token,实际 %v", got)
|
||||
}
|
||||
for _, s := range []string{"a", "中", "1", " "} {
|
||||
if got := tok.Encode(s); len(got) == 0 {
|
||||
t.Errorf("%q 应至少产出 1 个 token", s)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 模板渲染必须与参考数据里的整串完全一致,且逐 token 对齐。
|
||||
//
|
||||
// 这是嵌入正确性的前提:模板差一个字符,池化取到的「最后一个有效 token」
|
||||
// 位置就变了,向量也就不同——而且不会报错。
|
||||
func TestRenderInstructionInputMatchesTemplate(t *testing.T) {
|
||||
ref := loadRef(t)
|
||||
tok := loadTokenizer(t)
|
||||
|
||||
const want = "<|im_start|>system\nRepresent the user's input.<|im_end|>\n" +
|
||||
"<|im_start|>user\n你好<|im_end|>\n<|im_start|>assistant\n"
|
||||
|
||||
got := renderInstructionInput("", "你好")
|
||||
if got != want {
|
||||
t.Fatalf("模板渲染不一致:\n got %q\n want %q", got, want)
|
||||
}
|
||||
|
||||
for _, c := range ref.Cases {
|
||||
if c.Text != want {
|
||||
continue
|
||||
}
|
||||
if ids := tok.Encode(got); !sameIDs(ids, c.IDs) {
|
||||
t.Fatalf("模板串 token 不一致:\n got %v\n want %v", ids, c.IDs)
|
||||
}
|
||||
return
|
||||
}
|
||||
t.Fatal("参考数据里缺少该模板串用例")
|
||||
}
|
||||
|
||||
// 模型输入还要执行 tokenizer.json 的 TemplateProcessing:末尾追加
|
||||
// <|endoftext|>;超长输入先给正文留 maxLen-1 个位置,再保留 post token。
|
||||
func TestEncodeModelInputPostProcessor(t *testing.T) {
|
||||
tok := loadTokenizer(t)
|
||||
postID, ok := tok.SpecialID("<|endoftext|>")
|
||||
if !ok {
|
||||
t.Fatal("tokenizer 缺少 <|endoftext|>")
|
||||
}
|
||||
|
||||
shortRaw := tok.Encode("你好")
|
||||
short, err := tok.encodeModelInput("你好", 512)
|
||||
if err != nil {
|
||||
t.Fatalf("短文本 encodeModelInput: %v", err)
|
||||
}
|
||||
if len(short) != len(shortRaw)+1 || short[len(short)-1] != postID {
|
||||
t.Fatalf("短文本 post-processor 异常: raw=%v model=%v", shortRaw, short)
|
||||
}
|
||||
|
||||
longRaw := tok.Encode(strings.Repeat("记忆", 600))
|
||||
long, err := tok.encodeModelInput(strings.Repeat("记忆", 600), 512)
|
||||
if err != nil {
|
||||
t.Fatalf("长文本 encodeModelInput: %v", err)
|
||||
}
|
||||
if len(long) != 512 || long[511] != postID {
|
||||
t.Fatalf("长文本截断异常: len=%d tail=%v", len(long), long[len(long)-1:])
|
||||
}
|
||||
if !sameIDs(long[:511], longRaw[:511]) {
|
||||
t.Fatal("长文本正文未按 maxLen-1 截断")
|
||||
}
|
||||
}
|
||||
|
||||
// byteEnc 必须是双射:256 个字节映射到 256 个互不相同的码点。
|
||||
// 有碰撞就会让不同字节编成同一个 token,静默产生错误输入。
|
||||
func TestBytesToUnicodeBijective(t *testing.T) {
|
||||
m := bytesToUnicode()
|
||||
if len(m) != 256 {
|
||||
t.Fatalf("映射应覆盖 256 个字节,实际 %d", len(m))
|
||||
}
|
||||
seen := map[rune]byte{}
|
||||
for b, r := range m {
|
||||
if prev, dup := seen[r]; dup {
|
||||
t.Fatalf("码点冲突:字节 %d 与 %d 都映射到 %q", prev, b, r)
|
||||
}
|
||||
seen[r] = b
|
||||
}
|
||||
}
|
||||
@ -2,6 +2,7 @@ package vector
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
@ -10,14 +11,50 @@ import (
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"gitcode.com/JianFeeeee/HomeAgent/pkg/embedding"
|
||||
)
|
||||
|
||||
func init() {
|
||||
// http 也是一个普通 provider:核心只按名字打开它,不知道它背后是云 API、
|
||||
// 自建服务还是别的语言写的模型。
|
||||
embedding.Register("http", func(cfg embedding.Config) (embedding.Provider, error) {
|
||||
return NewHTTPEmbedder(HTTPEmbedderConfig{
|
||||
Endpoint: cfg.Options["endpoint"],
|
||||
APIKey: cfg.Options["api_key"],
|
||||
Model: cfg.Options["model"],
|
||||
Fingerprint: cfg.Options["fingerprint"],
|
||||
Dimension: atoiOrZero(cfg.Options["dimension"]),
|
||||
Timeout: durationOrZero(cfg.Options["timeout"]),
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
func atoiOrZero(s string) int {
|
||||
n := 0
|
||||
for _, r := range strings.TrimSpace(s) {
|
||||
if r < '0' || r > '9' {
|
||||
return 0
|
||||
}
|
||||
n = n*10 + int(r-'0')
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
func durationOrZero(s string) time.Duration {
|
||||
d, err := time.ParseDuration(strings.TrimSpace(s))
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
return d
|
||||
}
|
||||
|
||||
// HTTPEmbedderConfig 配置一个外部多模态向量服务。
|
||||
// 服务契约刻意很小:POST Endpoint,输入 modality/data/mime/side,返回 embedding。
|
||||
// 任何云 API 或自建服务只需适配这一个协议,即可复用内核全部向量存储与检索链路。
|
||||
type HTTPEmbedderConfig struct {
|
||||
Endpoint string
|
||||
APIKey string
|
||||
APIKey string
|
||||
Model string
|
||||
Dimension int
|
||||
Timeout time.Duration
|
||||
@ -65,14 +102,14 @@ func NewHTTPEmbedder(cfg HTTPEmbedderConfig) (*HTTPEmbedder, error) {
|
||||
}
|
||||
|
||||
func (e *HTTPEmbedder) VectorizeDense(text string) ([]float64, error) {
|
||||
return e.embed(httpEmbedRequest{Model: e.cfg.Model, Modality: string(ModalityText), Side: "query", Text: text})
|
||||
return e.embed(context.Background(), httpEmbedRequest{Model: e.cfg.Model, Modality: string(ModalityText), Side: "query", Text: text})
|
||||
}
|
||||
|
||||
func (e *HTTPEmbedder) EmbedImageDense(img []byte, mime string) ([]float64, error) {
|
||||
return e.embed(httpEmbedRequest{Model: e.cfg.Model, Modality: string(ModalityImage), Side: "document", Data: base64.StdEncoding.EncodeToString(img), MIME: mime})
|
||||
return e.embed(context.Background(), httpEmbedRequest{Model: e.cfg.Model, Modality: string(ModalityImage), Side: "document", Data: base64.StdEncoding.EncodeToString(img), MIME: mime})
|
||||
}
|
||||
|
||||
func (e *HTTPEmbedder) embed(payload httpEmbedRequest) ([]float64, error) {
|
||||
func (e *HTTPEmbedder) embed(ctx context.Context, payload httpEmbedRequest) ([]float64, error) {
|
||||
e.mu.Lock()
|
||||
closed := e.closed
|
||||
e.mu.Unlock()
|
||||
@ -83,7 +120,7 @@ func (e *HTTPEmbedder) embed(payload httpEmbedRequest) ([]float64, error) {
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req, err := http.NewRequest(http.MethodPost, e.cfg.Endpoint, bytes.NewReader(body))
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, e.cfg.Endpoint, bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@ -119,6 +156,32 @@ func (e *HTTPEmbedder) embed(payload httpEmbedRequest) ([]float64, error) {
|
||||
|
||||
func (e *HTTPEmbedder) Fingerprint() string { return e.cfg.Fingerprint }
|
||||
func (e *HTTPEmbedder) Dim() int { return e.cfg.Dimension }
|
||||
|
||||
// Embed 实现公共 provider 契约:核心只传模态与不透明字节,本实现负责把它
|
||||
// 翻译成外部服务的协议。
|
||||
func (e *HTTPEmbedder) Embed(ctx context.Context, in embedding.Input) ([]float64, error) {
|
||||
req := httpEmbedRequest{
|
||||
Model: e.cfg.Model,
|
||||
Modality: string(in.Modality),
|
||||
Side: string(in.Purpose),
|
||||
Text: in.Text,
|
||||
MIME: in.MIME,
|
||||
}
|
||||
if in.Modality != embedding.ModalityText {
|
||||
req.Data = base64.StdEncoding.EncodeToString(in.Data)
|
||||
}
|
||||
return e.embed(ctx, req)
|
||||
}
|
||||
|
||||
// Info 声明本 provider 的向量空间身份。外部服务的支持模态无法在本地探测,
|
||||
// 因此只声明 text/image 这两条内核真正会走到的路径。
|
||||
func (e *HTTPEmbedder) Info() embedding.Info {
|
||||
return embedding.Info{
|
||||
Dimension: e.cfg.Dimension,
|
||||
Fingerprint: e.cfg.Fingerprint,
|
||||
Modalities: []embedding.Modality{embedding.ModalityText, embedding.ModalityImage},
|
||||
}
|
||||
}
|
||||
func (e *HTTPEmbedder) Loaded() bool {
|
||||
e.mu.Lock()
|
||||
defer e.mu.Unlock()
|
||||
|
||||
83
internal/memory/vector/provider_adapter.go
Normal file
83
internal/memory/vector/provider_adapter.go
Normal file
@ -0,0 +1,83 @@
|
||||
package vector
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
|
||||
"gitcode.com/JianFeeeee/HomeAgent/pkg/embedding"
|
||||
)
|
||||
|
||||
// ProviderAdapter translates the public model-neutral embedding.Provider SPI
|
||||
// to the small internal interface used by the existing memory consumers.
|
||||
// Model selection, media decoding, preprocessing, and runtime details remain
|
||||
// entirely inside the selected provider.
|
||||
type ProviderAdapter struct {
|
||||
provider embedding.Provider
|
||||
info embedding.Info
|
||||
|
||||
mu sync.RWMutex
|
||||
closed bool
|
||||
}
|
||||
|
||||
// AdaptProvider validates and wraps a public provider for internal memory use.
|
||||
func AdaptProvider(provider embedding.Provider) (*ProviderAdapter, error) {
|
||||
info := provider.Info()
|
||||
if err := embedding.ValidateInfo(info); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &ProviderAdapter{provider: provider, info: info}, nil
|
||||
}
|
||||
|
||||
func (a *ProviderAdapter) VectorizeDense(text string) ([]float64, error) {
|
||||
return a.embed(embedding.Input{
|
||||
Modality: embedding.ModalityText,
|
||||
Purpose: embedding.PurposeQuery,
|
||||
Text: text,
|
||||
})
|
||||
}
|
||||
|
||||
func (a *ProviderAdapter) EmbedImageDense(data []byte, mime string) ([]float64, error) {
|
||||
return a.embed(embedding.Input{
|
||||
Modality: embedding.ModalityImage,
|
||||
Purpose: embedding.PurposeDocument,
|
||||
Data: data,
|
||||
MIME: mime,
|
||||
})
|
||||
}
|
||||
|
||||
func (a *ProviderAdapter) embed(input embedding.Input) ([]float64, error) {
|
||||
a.mu.RLock()
|
||||
closed := a.closed
|
||||
a.mu.RUnlock()
|
||||
if closed {
|
||||
return nil, context.Canceled
|
||||
}
|
||||
vec, err := a.provider.Embed(context.Background(), input)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := embedding.ValidateVector(vec, a.info.Dimension); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return vec, nil
|
||||
}
|
||||
|
||||
func (a *ProviderAdapter) Fingerprint() string { return a.info.Fingerprint }
|
||||
func (a *ProviderAdapter) Dim() int { return a.info.Dimension }
|
||||
|
||||
func (a *ProviderAdapter) Loaded() bool {
|
||||
a.mu.RLock()
|
||||
defer a.mu.RUnlock()
|
||||
return !a.closed
|
||||
}
|
||||
|
||||
func (a *ProviderAdapter) Close() {
|
||||
a.mu.Lock()
|
||||
if a.closed {
|
||||
a.mu.Unlock()
|
||||
return
|
||||
}
|
||||
a.closed = true
|
||||
a.mu.Unlock()
|
||||
a.provider.Close()
|
||||
}
|
||||
115
internal/memory/vector/provider_adapter_test.go
Normal file
115
internal/memory/vector/provider_adapter_test.go
Normal file
@ -0,0 +1,115 @@
|
||||
package vector
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"testing"
|
||||
|
||||
"gitcode.com/JianFeeeee/HomeAgent/pkg/embedding"
|
||||
)
|
||||
|
||||
// recordingProvider 记录核心传给 provider 的原始请求,用来断言
|
||||
// 「核心不解释内容、只搬字节」这一契约。
|
||||
type recordingProvider struct {
|
||||
got []embedding.Input
|
||||
dim int
|
||||
closed bool
|
||||
}
|
||||
|
||||
func (p *recordingProvider) Embed(_ context.Context, in embedding.Input) ([]float64, error) {
|
||||
p.got = append(p.got, in)
|
||||
return make([]float64, p.dim), nil
|
||||
}
|
||||
|
||||
func (p *recordingProvider) Info() embedding.Info {
|
||||
return embedding.Info{Dimension: p.dim, Fingerprint: "recording:1"}
|
||||
}
|
||||
|
||||
func (p *recordingProvider) Close() { p.closed = true }
|
||||
|
||||
func TestProviderAdapterPassesOpaqueDataUnchanged(t *testing.T) {
|
||||
inner := &recordingProvider{dim: 3}
|
||||
adapted, err := AdaptProvider(inner)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer adapted.Close()
|
||||
|
||||
// 核心把媒体当作不透明字节搬运:既不解码也不改字节。
|
||||
raw := []byte{0x89, 'P', 'N', 'G', 0x00, 0xff}
|
||||
if _, err := adapted.EmbedImageDense(raw, "image/png"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got := inner.got[0]
|
||||
if string(got.Data) != string(raw) {
|
||||
t.Fatalf("provider 收到的字节被改动: %v", got.Data)
|
||||
}
|
||||
if got.Modality != embedding.ModalityImage || got.MIME != "image/png" {
|
||||
t.Fatalf("模态/MIME 未原样传递: %+v", got)
|
||||
}
|
||||
if got.Purpose != embedding.PurposeDocument {
|
||||
t.Fatalf("用途应为 document: %q", got.Purpose)
|
||||
}
|
||||
|
||||
if _, err := adapted.VectorizeDense("hello"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if inner.got[1].Modality != embedding.ModalityText || inner.got[1].Text != "hello" {
|
||||
t.Fatalf("文本请求不正确: %+v", inner.got[1])
|
||||
}
|
||||
}
|
||||
|
||||
func TestProviderAdapterRejectsWrongDimensionFromProvider(t *testing.T) {
|
||||
// provider 声明 3 维却返回 2 维:必须在进入存储前被拦下,
|
||||
// 否则一个维度错的向量会污染整个余弦检索。
|
||||
bad := &badDimProvider{}
|
||||
adapted, err := AdaptProvider(bad)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer adapted.Close()
|
||||
if _, err := adapted.VectorizeDense("x"); err == nil {
|
||||
t.Fatal("维度不符时应返回错误")
|
||||
}
|
||||
}
|
||||
|
||||
type badDimProvider struct{}
|
||||
|
||||
func (badDimProvider) Embed(context.Context, embedding.Input) ([]float64, error) {
|
||||
return []float64{1, 2}, nil
|
||||
}
|
||||
func (badDimProvider) Info() embedding.Info {
|
||||
return embedding.Info{Dimension: 3, Fingerprint: "bad:1"}
|
||||
}
|
||||
func (badDimProvider) Close() {}
|
||||
|
||||
func TestProviderAdapterCloseIsIdempotentAndStopsUse(t *testing.T) {
|
||||
inner := &recordingProvider{dim: 2}
|
||||
adapted, err := AdaptProvider(inner)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
adapted.Close()
|
||||
adapted.Close() // 重复关闭不应 panic 或二次 Close provider
|
||||
if !inner.closed {
|
||||
t.Fatal("Close 未传递到 provider")
|
||||
}
|
||||
if _, err := adapted.VectorizeDense("x"); err == nil {
|
||||
t.Fatal("关闭后应拒绝调用")
|
||||
}
|
||||
if adapted.Loaded() {
|
||||
t.Fatal("关闭后 Loaded() 应为 false")
|
||||
}
|
||||
}
|
||||
|
||||
func TestModalityUnsupportedSentinelIsShared(t *testing.T) {
|
||||
// 内核侧的哨兵与公共契约的哨兵必须是同一个:provider 返回公共哨兵时,
|
||||
// 内核仍能用自己原有的名字识别。
|
||||
if !errors.Is(ErrModalityUnsupported, embedding.ErrUnsupportedModality) {
|
||||
t.Fatal("vector.ErrModalityUnsupported 与 embedding.ErrUnsupportedModality 未打通")
|
||||
}
|
||||
wrapped := errors.Join(embedding.ErrUnsupportedModality, errors.New("audio/wav"))
|
||||
if !errors.Is(wrapped, ErrModalityUnsupported) {
|
||||
t.Fatal("包装后的错误无法用内核哨兵识别")
|
||||
}
|
||||
}
|
||||
@ -6,6 +6,8 @@ import (
|
||||
"sort"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
"gitcode.com/JianFeeeee/HomeAgent/pkg/embedding"
|
||||
)
|
||||
|
||||
// Vectorizer 接口:将文本转为向量
|
||||
@ -56,8 +58,15 @@ var ErrNotSupported = fmt.Errorf("vectorizer does not support image embedding")
|
||||
// 而不是「这次失败了、下次重试」。绝不能拿另一个模型的向量顶替——那会把
|
||||
// 两套坐标系混进同一空间,检索出来的相似度没有任何意义。
|
||||
//
|
||||
// 例:Qwen3-VL 能原生编码文本/图像,音频需要未来接入真正的统一音频模型。
|
||||
var ErrModalityUnsupported = fmt.Errorf("modality not supported by this embedding space")
|
||||
// 它是公共 provider 契约里那个哨兵值的别名,两者 errors.Is 互通:
|
||||
// provider 在自己的包内返回 embedding.ErrUnsupportedModality 即可,
|
||||
// 内核侧的判断无需改变。
|
||||
var ErrModalityUnsupported = embedding.ErrUnsupportedModality
|
||||
|
||||
// 注:曾经这里还有一个可选的 VideoEmbedder 接口(用类型断言探测视频能力)。
|
||||
// 已删除:那让核心为每一个新模态长出一套模型专属方法,正是“核心适配模型”的
|
||||
// 坏味道。模态能力现在是数据(embedding.Info.Modalities),输入是不透明的
|
||||
// Data+MIME(见 pkg/embedding)。
|
||||
|
||||
// Vector 是带权特征映射:feature → weight
|
||||
type Vector map[string]float64
|
||||
|
||||
Reference in New Issue
Block a user