diff --git a/internal/memory/static_embedder.go b/internal/memory/static_embedder.go index 8e302ff..6eaaa5e 100644 --- a/internal/memory/static_embedder.go +++ b/internal/memory/static_embedder.go @@ -34,11 +34,15 @@ type StaticEmbedder struct { jieba *gojieba.Jieba stopWords map[string]bool - words map[string][]float64 + // words 是词向量表。**用 float32 存**:源文件(fastText 文本格式)本身就是 float32, + // 用 float64 存等于把 578 万……不,是 57.8 万词 × 300 维的常驻内存凭空翻倍 + // (实测生产:float64 → 1.29GB,float32 → 0.65GB)。相似度计算仍在 float64 里累加, + // 精度不受影响。改回 float64 会被 TestStaticEmbedder_VectorMemIsFloat32 拦住。 + words map[string][]float32 dim int loaded bool - unkVec []float64 + unkVec []float32 unkNorm float64 } @@ -150,7 +154,7 @@ func NewStaticEmbedder(modelPaths ...string) *StaticEmbedder { e := &StaticEmbedder{ jieba: GetJieba(), stopWords: sw, - words: make(map[string][]float64), + words: make(map[string][]float32), } if len(modelPaths) == 0 { @@ -254,15 +258,16 @@ func (e *StaticEmbedder) load(spec string, primary bool) error { continue } - vec := make([]float64, dim) + vec := make([]float32, dim) for i := 0; i < dim; i++ { - v, _ := strconv.ParseFloat(fields[i+1], 64) - vec[i] = v + // 源文件是 float32 精度的文本向量:用 32 位解析,与源数据一致。 + v, _ := strconv.ParseFloat(fields[i+1], 32) + vec[i] = float32(v) } e.words[word] = vec if primary { for i := range vecSum { - vecSum[i] += vec[i] + vecSum[i] += float64(vec[i]) } count++ } @@ -276,11 +281,13 @@ func (e *StaticEmbedder) load(spec string, primary bool) error { for i := range vecSum { vecSum[i] /= float64(count) } - e.unkVec = make([]float64, dim) - copy(e.unkVec, vecSum) + e.unkVec = make([]float32, dim) + for i, v := range vecSum { + e.unkVec[i] = float32(v) + } var normSq float64 for _, v := range e.unkVec { - normSq += v * v + normSq += float64(v) * float64(v) } e.unkNorm = float64(math.Sqrt(normSq)) e.loaded = true @@ -366,11 +373,11 @@ func (e *StaticEmbedder) Vectorize(text string) vector.Vector { if !ok { for i, v := range unkVec { - sum[i] += w * v + sum[i] += w * float64(v) } } else { for i, v := range vec { - sum[i] += w * v + sum[i] += w * float64(v) } } weightSum += w diff --git a/internal/memory/static_embedder_mem_test.go b/internal/memory/static_embedder_mem_test.go new file mode 100644 index 0000000..3462848 --- /dev/null +++ b/internal/memory/static_embedder_mem_test.go @@ -0,0 +1,37 @@ +package memory + +import ( + "testing" + "unsafe" +) + +// 词向量必须用 float32 存。 +// +// 这条判据是拿生产内存换来的:向量本体 = 词数 × 维数 × 每元素字节数。 +// 生产配置加载了 200000(zh) + 378151(en) = 57.8 万词 × 300 维 ⇒ +// float64 = 1.29GB、float32 = 0.65GB(差 0.65GB 常驻)。 +// 源数据(fastText 文本格式)本身就是 float32 精度,用 float64 存没有任何收益。 +// +// 若有人把类型改回 float64,本测试**编译失败**(`var vec []float32` 的类型断言), +// 这正是想要的效果。 +func TestStaticEmbedder_VectorMemIsFloat32(t *testing.T) { + e := newSynthEmbedder(t, 300) + + words := 0 + bytes := 0 + for _, vec := range e.words { + var typed []float32 = vec // 编译期断言:存储必须是 []float32 + if len(typed) != e.dim { + t.Fatalf("维度不符: %d != %d", len(typed), e.dim) + } + words++ + bytes += len(typed) * int(unsafe.Sizeof(typed[0])) + } + if words == 0 { + t.Fatal("合成模型应至少加载一个词") + } + // float32:每词 300×4 = 1200 字节;float64 会是 2400 + if want := words * e.dim * 4; bytes != want { + t.Fatalf("向量本体字节数应 %d(float32),实际 %d", want, bytes) + } +}