134 lines
3.6 KiB
Go
134 lines
3.6 KiB
Go
package embedding
|
||
|
||
import (
|
||
"fmt"
|
||
"math"
|
||
"sort"
|
||
"strings"
|
||
"sync"
|
||
"unicode"
|
||
)
|
||
|
||
// Hit is a single keyword search match: the indexed document id and its
|
||
// relevance score. Higher scores are more relevant.
|
||
type Hit[ID comparable] struct {
|
||
ID ID
|
||
Score float64
|
||
}
|
||
|
||
// Index is a simple in-memory inverted index for keyword full-text search (FTS).
|
||
//
|
||
// It tokenizes each added document into terms and keeps a postings list
|
||
// (term → doc id → term frequency). Search scores each matching document with
|
||
// TF-IDF: the sum over query terms of tf × idf, where idf = ln(1 + N/df). This
|
||
// is the keyword-level retrieval counterpart to the semantic Embedder vectors;
|
||
// both coexist (see DevDocs/design/04-运行时.md §7).
|
||
//
|
||
// The zero value is not ready to use; build one with NewIndex. It is safe for
|
||
// concurrent use.
|
||
type Index[ID comparable] struct {
|
||
mu sync.RWMutex
|
||
postings map[string]map[ID]int // term -> doc id -> term frequency
|
||
docLen map[ID]int // doc id -> term count
|
||
docTerms map[ID][]string // doc id -> indexed terms, for re-indexing
|
||
docCount int // number of distinct documents
|
||
}
|
||
|
||
// NewIndex creates an empty inverted index.
|
||
func NewIndex[ID comparable]() *Index[ID] {
|
||
return &Index[ID]{
|
||
postings: make(map[string]map[ID]int),
|
||
docLen: make(map[ID]int),
|
||
docTerms: make(map[ID][]string),
|
||
}
|
||
}
|
||
|
||
// Add indexes text under id. Re-adding an id replaces its previous term
|
||
// contribution, so the index stays consistent with the latest text for that id.
|
||
func (ix *Index[ID]) Add(text string, id ID) {
|
||
terms := tokenize(text)
|
||
freq := make(map[string]int, len(terms))
|
||
for _, t := range terms {
|
||
freq[t]++
|
||
}
|
||
|
||
ix.mu.Lock()
|
||
defer ix.mu.Unlock()
|
||
|
||
if _, exists := ix.docLen[id]; exists {
|
||
ix.removeDocLocked(id)
|
||
}
|
||
for t, f := range freq {
|
||
if ix.postings[t] == nil {
|
||
ix.postings[t] = make(map[ID]int)
|
||
}
|
||
ix.postings[t][id] = f
|
||
}
|
||
ix.docTerms[id] = terms
|
||
ix.docLen[id] = len(terms)
|
||
ix.docCount++
|
||
}
|
||
|
||
// Search returns every document matching at least one query term, ranked by
|
||
// descending relevance score. Ties are broken by id for determinism. An empty
|
||
// or whitespace-only query returns no hits.
|
||
func (ix *Index[ID]) Search(query string) []Hit[ID] {
|
||
terms := tokenize(query)
|
||
if len(terms) == 0 {
|
||
return []Hit[ID]{}
|
||
}
|
||
|
||
ix.mu.RLock()
|
||
defer ix.mu.RUnlock()
|
||
|
||
scores := make(map[ID]float64)
|
||
n := ix.docCount
|
||
for _, t := range terms {
|
||
df, ok := ix.postings[t]
|
||
if !ok {
|
||
continue
|
||
}
|
||
idf := math.Log(1 + float64(n)/float64(len(df)))
|
||
for id, tf := range df {
|
||
scores[id] += float64(tf) * idf
|
||
}
|
||
}
|
||
|
||
hits := make([]Hit[ID], 0, len(scores))
|
||
for id, s := range scores {
|
||
hits = append(hits, Hit[ID]{ID: id, Score: s})
|
||
}
|
||
sort.Slice(hits, func(i, j int) bool {
|
||
if hits[i].Score != hits[j].Score {
|
||
return hits[i].Score > hits[j].Score
|
||
}
|
||
return fmt.Sprint(hits[i].ID) < fmt.Sprint(hits[j].ID)
|
||
})
|
||
return hits
|
||
}
|
||
|
||
// removeDocLocked removes the previous postings contribution of an existing doc
|
||
// before it is re-indexed. Callers must hold ix.mu.
|
||
func (ix *Index[ID]) removeDocLocked(id ID) {
|
||
for _, t := range ix.docTerms[id] {
|
||
df := ix.postings[t]
|
||
if df != nil {
|
||
delete(df, id)
|
||
if len(df) == 0 {
|
||
delete(ix.postings, t)
|
||
}
|
||
}
|
||
}
|
||
delete(ix.docTerms, id)
|
||
delete(ix.docLen, id)
|
||
ix.docCount--
|
||
}
|
||
|
||
// tokenize splits text into lowercased word terms. Non-letter and non-digit
|
||
// runes act as separators; empty and repeated terms are handled by the caller.
|
||
func tokenize(s string) []string {
|
||
return strings.FieldsFunc(strings.ToLower(s), func(r rune) bool {
|
||
return !unicode.IsLetter(r) && !unicode.IsDigit(r)
|
||
})
|
||
}
|