kyleandClaude Opus 5.5
64d08a358d
Owner review of p5b (tasks 01-03 by Ornith): - background.go: `seen` was never cleared, so a document whose background embedding failed (Dixie hiccup) or was abandoned could never be queued again: every later question on it stayed BM25 for the life of the process. Dedupe now covers only queued and in-progress documents. - background.go: acquireGate gave up after one retry and dropped the document; the task said to keep retrying once a second. It now waits (background work has no deadline). - retrieval.go/background.go: a vector the server left out is not cached as an empty vector. - bm25.go: Scores on a nil *BM25 dereferenced it; returns nil. Regression tests in review_test.go; both fail on the old code. End-to-end against Dixie (local serve): Karpathy's thesis, 95 pages / 51 chunks: bm25 in 19 s, hybrid (51 embedded) after the background; RP2040 datasheet, 642 pages / 284 chunks (a 422 before P5b): bm25 in 13.5 s, all 284 embedded within 90 s, hybrid answer correct with its quote. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
110 lines
2.2 KiB
Go
110 lines
2.2 KiB
Go
package digest
|
|
|
|
import (
|
|
"math"
|
|
"strings"
|
|
"unicode"
|
|
)
|
|
|
|
const (
|
|
bm25K1 = 1.2
|
|
bm25B = 0.75
|
|
bm25Min = 2
|
|
)
|
|
|
|
// BM25 is a lexical ranker over a fixed set of documents. Terms are the same tokens as
|
|
// terms() — lower-case runs of letters and digits, at least bm25Min characters — counted
|
|
// with frequency rather than kept as a set.
|
|
type BM25 struct {
|
|
nDocs int
|
|
avgDL float64
|
|
df map[string]int
|
|
toks []map[string]int
|
|
len []int
|
|
}
|
|
|
|
// NewBM25 indexes docs with the default BM25 parameters (k1 = 1.2, b = 0.75).
|
|
func NewBM25(docs []string) *BM25 {
|
|
b := &BM25{
|
|
nDocs: len(docs),
|
|
df: map[string]int{},
|
|
toks: make([]map[string]int, len(docs)),
|
|
len: make([]int, len(docs)),
|
|
}
|
|
var total int
|
|
for i, doc := range docs {
|
|
counts := docTerms(doc)
|
|
n := 0
|
|
for _, c := range counts {
|
|
n += c
|
|
}
|
|
b.toks[i] = counts
|
|
b.len[i] = n
|
|
total += n
|
|
for t := range counts {
|
|
b.df[t]++
|
|
}
|
|
}
|
|
if len(docs) > 0 {
|
|
b.avgDL = float64(total) / float64(len(docs))
|
|
}
|
|
return b
|
|
}
|
|
|
|
// Scores returns one score per document, in order, each >= 0 and never NaN or Inf. An empty
|
|
// query, an empty index or a document without any query term scores 0.
|
|
func (b *BM25) Scores(query string) []float64 {
|
|
if b == nil {
|
|
return nil
|
|
}
|
|
if b.nDocs == 0 {
|
|
return make([]float64, len(b.toks))
|
|
}
|
|
out := make([]float64, len(b.toks))
|
|
qterms := docTerms(query)
|
|
if len(qterms) == 0 {
|
|
return out
|
|
}
|
|
avgDL := b.avgDL
|
|
if avgDL == 0 {
|
|
avgDL = 1
|
|
}
|
|
for i, toks := range b.toks {
|
|
var score float64
|
|
for term := range qterms {
|
|
tf, ok := toks[term]
|
|
if !ok {
|
|
continue
|
|
}
|
|
idf := math.Log(1 + (float64(b.nDocs)-float64(b.df[term])+0.5)/(float64(b.df[term])+0.5))
|
|
denom := float64(tf) + bm25K1*(1-bm25B+bm25B*float64(b.len[i])/avgDL)
|
|
score += idf * (float64(tf) * (bm25K1 + 1)) / denom
|
|
}
|
|
out[i] = score
|
|
}
|
|
return out
|
|
}
|
|
|
|
// docTerms counts occurrences of each BM25 term in s, using the same tokenization as terms().
|
|
func docTerms(s string) map[string]int {
|
|
out := map[string]int{}
|
|
var b strings.Builder
|
|
flush := func() {
|
|
t := b.String()
|
|
b.Reset()
|
|
if len(t) < bm25Min {
|
|
return
|
|
}
|
|
out[t]++
|
|
}
|
|
for _, r := range strings.ToLower(s) {
|
|
if unicode.IsLetter(r) || unicode.IsDigit(r) {
|
|
b.WriteRune(r)
|
|
} else {
|
|
flush()
|
|
}
|
|
}
|
|
flush()
|
|
return out
|
|
}
|