mirror of
https://gitcode.com/JianFeeeee/HomeAgent.git
synced 2026-09-21 09:28:14 +00:00
307 lines
7.1 KiB
Go
307 lines
7.1 KiB
Go
// Package tfidf 提供 TF-IDF 向量化作为 MultimodalEmbedder 的一个实现。
|
||
//
|
||
// 设计：核心（内侧）只认 vector.MultimodalEmbedder 接口；本包是外层可插拔
|
||
// 的一个具体实现，与 Jina HTTP / ONNX CLIP 并列。启动时由 cmd/homed 按配置
|
||
// 选择注入，核心代码零改动。
|
||
//
|
||
// 稀疏向量 → 稠密桥接：TF-IDF 产出的是稀疏 map[feature]weight，通过特征哈希
|
||
// 投射到固定维度（默认 4096）的 []float64，供统一的 cosine 检索使用。
|
||
// 哈希碰撞在 4096 维下可接受（英文单词 ~50k，碰撞率 ~3%）。
|
||
package tfidf
|
||
|
||
import (
|
||
"fmt"
|
||
"hash/fnv"
|
||
"math"
|
||
"sort"
|
||
"strings"
|
||
"sync"
|
||
)
|
||
|
||
// Tokenizer 将文本拆分为词级 token。
|
||
type Tokenizer func(text string) []string
|
||
|
||
// Vector 是带权特征映射：feature → weight（稀疏表示）。
|
||
type Vector map[string]float64
|
||
|
||
// Embedder 实现 vector.MultimodalEmbedder，将 TF-IDF 稀疏向量投射为固定维度稠密向量。
|
||
type Embedder struct {
|
||
mu sync.RWMutex
|
||
tokenizer Tokenizer
|
||
docFreq map[string]int
|
||
totalDocs int
|
||
dim int
|
||
loaded bool
|
||
fingerprint string
|
||
}
|
||
|
||
// NewEmbedder 创建 TF-IDF 多模态嵌入器。
|
||
// dim 是投射后的稠密维度（默认 4096，哈希空间大小）。
|
||
func NewEmbedder(tokenizer Tokenizer, dim int) *Embedder {
|
||
if dim <= 0 {
|
||
dim = 4096
|
||
}
|
||
return &Embedder{
|
||
tokenizer: tokenizer,
|
||
docFreq: make(map[string]int),
|
||
dim: dim,
|
||
}
|
||
}
|
||
|
||
// Train 用文档集训练 IDF 统计。启动时调用一次。
|
||
func (e *Embedder) Train(docs []string) {
|
||
e.mu.Lock()
|
||
defer e.mu.Unlock()
|
||
|
||
e.totalDocs = len(docs)
|
||
e.docFreq = make(map[string]int)
|
||
|
||
for _, doc := range docs {
|
||
features := e.tokenizer(doc)
|
||
seen := make(map[string]bool)
|
||
for _, f := range features {
|
||
if !seen[f] {
|
||
e.docFreq[f]++
|
||
seen[f] = true
|
||
}
|
||
}
|
||
}
|
||
e.fingerprint = fmt.Sprintf("tfidf:d%d:f%d", e.totalDocs, len(e.docFreq))
|
||
e.loaded = true
|
||
}
|
||
|
||
// vectorize 将文本转为 TF-IDF 稀疏向量（内部方法，调用方已持锁）。
|
||
func (e *Embedder) vectorize(text string) Vector {
|
||
features := e.tokenizer(text)
|
||
tf := make(map[string]float64)
|
||
for _, f := range features {
|
||
tf[f]++
|
||
}
|
||
maxTF := 0.0
|
||
for _, c := range tf {
|
||
if c > maxTF {
|
||
maxTF = c
|
||
}
|
||
}
|
||
|
||
vec := make(Vector)
|
||
for f, count := range tf {
|
||
tfNorm := count / maxTF
|
||
if e.totalDocs < 3 {
|
||
vec[f] = tfNorm
|
||
continue
|
||
}
|
||
df := e.docFreq[f]
|
||
if df <= 0 {
|
||
continue
|
||
}
|
||
idf := math.Log(float64(e.totalDocs+1) / float64(df+1))
|
||
if idf < 0.1 {
|
||
continue
|
||
}
|
||
vec[f] = tfNorm * idf
|
||
}
|
||
return vec
|
||
}
|
||
|
||
// sparseToDense 将稀疏向量投射为固定维度稠密向量（FNV-1a 哈希映射）。
|
||
func (e *Embedder) sparseToDense(sparse Vector) []float64 {
|
||
dense := make([]float64, e.dim)
|
||
var norm float64
|
||
for feature, weight := range sparse {
|
||
idx := fnvHash(feature) % uint32(e.dim)
|
||
dense[idx] += weight
|
||
norm += weight * weight
|
||
}
|
||
// L2 归一化
|
||
if norm > 0 {
|
||
norm = math.Sqrt(norm)
|
||
for i := range dense {
|
||
dense[i] /= norm
|
||
}
|
||
}
|
||
return dense
|
||
}
|
||
|
||
func fnvHash(s string) uint32 {
|
||
h := fnv.New32a()
|
||
h.Write([]byte(s))
|
||
return h.Sum32()
|
||
}
|
||
|
||
// --- vector.MultimodalEmbedder 接口实现 ---
|
||
|
||
func (e *Embedder) VectorizeDense(text string) ([]float64, error) {
|
||
e.mu.RLock()
|
||
defer e.mu.RUnlock()
|
||
if !e.loaded {
|
||
return nil, fmt.Errorf("tfidf: not trained")
|
||
}
|
||
sparse := e.vectorize(text)
|
||
return e.sparseToDense(sparse), nil
|
||
}
|
||
|
||
func (e *Embedder) EmbedImageDense(img []byte, mime string) ([]float64, error) {
|
||
return nil, fmt.Errorf("tfidf: image embedding not supported")
|
||
}
|
||
|
||
func (e *Embedder) Fingerprint() string {
|
||
e.mu.RLock()
|
||
defer e.mu.RUnlock()
|
||
return e.fingerprint
|
||
}
|
||
|
||
func (e *Embedder) Dim() int { return e.dim }
|
||
|
||
func (e *Embedder) Loaded() bool {
|
||
e.mu.RLock()
|
||
defer e.mu.RUnlock()
|
||
return e.loaded
|
||
}
|
||
|
||
func (e *Embedder) Close() {}
|
||
|
||
// --- 检索（供 document.Store 使用，非接口方法）---
|
||
|
||
// DocHit 是一条检索命中。
|
||
type DocHit struct {
|
||
ID string
|
||
Score float64
|
||
}
|
||
|
||
// SearchableIndex 是支持 TF-IDF 倒排检索的索引。
|
||
// document.Store 在 TF-IDF 模式下使用此索引替代 brute-force。
|
||
type SearchableIndex struct {
|
||
mu sync.RWMutex
|
||
docs map[string]Vector // id → tfidf sparse vector
|
||
texts map[string]string // id → 原文
|
||
emb *Embedder
|
||
}
|
||
|
||
// NewSearchableIndex 创建可检索索引。
|
||
func NewSearchableIndex(emb *Embedder) *SearchableIndex {
|
||
return &SearchableIndex{
|
||
docs: make(map[string]Vector),
|
||
texts: make(map[string]string),
|
||
emb: emb,
|
||
}
|
||
}
|
||
|
||
// Add 添加或更新一条文档。
|
||
func (idx *SearchableIndex) Add(id, text string) {
|
||
idx.mu.Lock()
|
||
defer idx.mu.Unlock()
|
||
idx.emb.mu.RLock()
|
||
vec := idx.emb.vectorize(text)
|
||
idx.emb.mu.RUnlock()
|
||
idx.docs[id] = vec
|
||
idx.texts[id] = text
|
||
}
|
||
|
||
// Remove 移除一条文档。
|
||
func (idx *SearchableIndex) Remove(id string) {
|
||
idx.mu.Lock()
|
||
defer idx.mu.Unlock()
|
||
delete(idx.docs, id)
|
||
delete(idx.texts, id)
|
||
}
|
||
|
||
// Search 用查询文本检索 topK 个最相似的文档。
|
||
func (idx *SearchableIndex) Search(query string, topK int) []DocHit {
|
||
idx.mu.RLock()
|
||
defer idx.mu.RUnlock()
|
||
|
||
if len(idx.docs) == 0 {
|
||
return nil
|
||
}
|
||
|
||
idx.emb.mu.RLock()
|
||
qVec := idx.emb.vectorize(query)
|
||
idx.emb.mu.RUnlock()
|
||
|
||
type scored struct {
|
||
id string
|
||
score float64
|
||
}
|
||
var results []scored
|
||
for id, docVec := range idx.docs {
|
||
score := sparseCosine(qVec, docVec)
|
||
if score > 0.01 {
|
||
results = append(results, scored{id, score})
|
||
}
|
||
}
|
||
|
||
sort.Slice(results, func(i, j int) bool { return results[i].score > results[j].score })
|
||
if len(results) > topK {
|
||
results = results[:topK]
|
||
}
|
||
|
||
out := make([]DocHit, len(results))
|
||
for i, r := range results {
|
||
out[i] = DocHit{ID: r.id, Score: r.score}
|
||
}
|
||
return out
|
||
}
|
||
|
||
// sparseCosine 计算两个稀疏向量的余弦相似度。
|
||
func sparseCosine(a, b Vector) float64 {
|
||
var dot, na, nb float64
|
||
for f, va := range a {
|
||
dot += va * b[f]
|
||
na += va * va
|
||
}
|
||
for _, vb := range b {
|
||
nb += vb * vb
|
||
}
|
||
if na == 0 || nb == 0 {
|
||
return 0
|
||
}
|
||
return dot / math.Sqrt(na*nb)
|
||
}
|
||
|
||
// Text 返回文档原文（供调试/展示）。
|
||
func (idx *SearchableIndex) Text(id string) string {
|
||
idx.mu.RLock()
|
||
defer idx.mu.RUnlock()
|
||
return idx.texts[id]
|
||
}
|
||
|
||
// IDs 返回所有文档 ID（供重建索引）。
|
||
func (idx *SearchableIndex) IDs() []string {
|
||
idx.mu.RLock()
|
||
defer idx.mu.RUnlock()
|
||
ids := make([]string, 0, len(idx.docs))
|
||
for id := range idx.docs {
|
||
ids = append(ids, id)
|
||
}
|
||
sort.Strings(ids)
|
||
return ids
|
||
}
|
||
|
||
// Size 返回索引中的文档数。
|
||
func (idx *SearchableIndex) Size() int {
|
||
idx.mu.RLock()
|
||
defer idx.mu.RUnlock()
|
||
return len(idx.docs)
|
||
}
|
||
|
||
// Train 用文档集训练并建立索引。
|
||
func (idx *SearchableIndex) Train(docs map[string]string) {
|
||
idx.mu.Lock()
|
||
defer idx.mu.Unlock()
|
||
for id, text := range docs {
|
||
idx.emb.mu.RLock()
|
||
vec := idx.emb.vectorize(text)
|
||
idx.emb.mu.RUnlock()
|
||
idx.docs[id] = vec
|
||
idx.texts[id] = text
|
||
}
|
||
}
|
||
|
||
// TokenizeWords 是默认的中英文分词器（需外部注入 jieba 分词函数）。
|
||
// 外层 cmd/homed 负责注入，核心不直接依赖 jieba。
|
||
var TokenizeWords Tokenizer = func(text string) []string {
|
||
// 简单 fallback：按空白和标点拆分
|
||
return strings.Fields(text)
|
||
}
|