kyleandornith-1.5-35b-a3b
7595ad8519
SplitPages is Split's budget/line splitting plus page awareness: a "--- page N ---" line sets the current page and a chunk that starts after the first marker but not on a marker line is prefixed "--- page N (cont.) ---\n\n" for the page its first line belongs to, so every chunk names its page. BM25 (k1=1.2, b=0.75) ranks chunks lexically with counted term frequencies and ln(1+(N-df+0.5)/(df+0.5)) IDF. Co-Authored-By: ornith-1.5-35b-a3b <ornith-1.5-35b-a3b@llama.cpp.invalid>
107 lines
2.2 KiB
Go
107 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 || 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
|
|
}
|