digest: SplitPages keeps every chunk on its page; BM25 scorer

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>

internal/digest/bm25.go

@@ -0,0 +1,106 @@
+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
+}

internal/digest/bm25_test.go

@@ -0,0 +1,89 @@
+package digest
+
+// P5b task 01: BM25, the lexical ranking used while a long document's embeddings are still
+// being built (and inside the hybrid score after).
+
+import (
+	"math"
+	"strings"
+	"testing"
+)
+
+func TestBM25RanksTheRareTerm(t *testing.T) {
+	docs := []string{
+		"the ship sailed from helsinki to tallinn in the morning",
+		"the ship stayed in port all day",
+		"a zanzibarite sample was found near the harbour",
+		"the harbour master logged the ship",
+	}
+	s := NewBM25(docs).Scores("where was the zanzibarite found")
+	if len(s) != len(docs) {
+		t.Fatalf("%d scores for %d docs", len(s), len(docs))
+	}
+	best := 0
+	for i := range s {
+		if s[i] > s[best] {
+			best = i
+		}
+	}
+	if best != 2 {
+		t.Errorf("best = doc %d (%v), want doc 2, the only one with the rare term", best, s)
+	}
+}
+
+func TestBM25Properties(t *testing.T) {
+	filler := strings.Repeat("word ", 50)
+	docs := []string{
+		"alpha " + filler,                     // 0: tf 1
+		strings.Repeat("alpha ", 10) + filler, // 1: tf 10, same filler
+		"alpha common",                        // 2: tf 1, short
+		"common " + filler,                    // 3: no alpha
+	}
+	idx := NewBM25(docs)
+	s := idx.Scores("alpha")
+	if !(s[1] > s[0]) {
+		t.Errorf("more occurrences must score higher: tf10 %v vs tf1 %v", s[1], s[0])
+	}
+	if !(s[1] < 10*s[0]) {
+		t.Errorf("term frequency must saturate: tf10 %v, 10 × tf1 %v", s[1], 10*s[0])
+	}
+	if !(s[2] > s[0]) {
+		t.Errorf("the same term in a shorter document must score higher: short %v, long %v", s[2], s[0])
+	}
+	if s[3] != 0 {
+		t.Errorf("a document without the term scores %v, want 0", s[3])
+	}
+	// IDF: a term in every document is worth less than a rare one.
+	all := NewBM25([]string{"xx yy", "xx zz", "xx ww"}).Scores("xx")
+	rare := NewBM25([]string{"xx yy", "qq zz", "qq ww"}).Scores("xx")
+	if !(rare[0] > all[0]) {
+		t.Errorf("idf: a term in 1 of 3 docs must outscore one in 3 of 3: %v vs %v", rare[0], all[0])
+	}
+}
+
+func TestBM25Degenerate(t *testing.T) {
+	for _, docs := range [][]string{nil, {}, {""}, {"", "a b"}} {
+		s := NewBM25(docs).Scores("a")
+		if len(s) != len(docs) {
+			t.Errorf("%q: %d scores", docs, len(s))
+		}
+		for _, v := range s {
+			if math.IsNaN(v) || math.IsInf(v, 0) || v < 0 {
+				t.Errorf("%q: score %v", docs, v)
+			}
+		}
+	}
+	for _, v := range NewBM25([]string{"a b", "c d"}).Scores("") {
+		if v != 0 {
+			t.Errorf("empty query scored %v", v)
+		}
+	}
+	// Deterministic.
+	a := NewBM25([]string{"alpha beta", "beta gamma", "gamma alpha"}).Scores("alpha gamma")
+	b := NewBM25([]string{"alpha beta", "beta gamma", "gamma alpha"}).Scores("alpha gamma")
+	for i := range a {
+		if a[i] != b[i] {
+			t.Fatal("scores differ between identical runs")
+		}
+	}
+}

internal/digest/splitpages.go

@@ -0,0 +1,74 @@
+package digest
+
+import (
+	"fmt"
+	"regexp"
+	"strconv"
+	"strings"
+)
+
+// pageMarker matches a page header line like "--- page 7 ---"; the number is group 1.
+var pageMarker = regexp.MustCompile(`^--- page (\d+) ---$`)
+
+// SplitPages splits text into chunks of at most maxTok tokens (EstTokens) on line
+// boundaries, like Split, but page-aware. The text is expected to carry "--- page N ---"
+// markers (as pdfdoc.Full produces). A marker line sets the current page; a chunk that
+// starts after the first marker but not on a marker line is prefixed with
+// "--- page N (cont.) ---\n\n" for the page its first line belongs to, so every chunk
+// names the page it cites. Text with no marker lines yields exactly Split's result.
+func SplitPages(text string, maxTok int) []string {
+	if maxTok <= 0 {
+		maxTok = ChunkTokens
+	}
+	lines := strings.Split(text, "\n")
+	var chunks []string
+	var buf strings.Builder
+	var firstLine string
+	var chunkPage int
+	hasLines, startedAfterMarker, markerSeen := false, false, false
+	currentPage := 0
+
+	flush := func() {
+		s := strings.TrimSpace(buf.String())
+		if s == "" {
+			buf.Reset()
+			hasLines = false
+			return
+		}
+		if pageMarker.MatchString(firstLine) || !startedAfterMarker {
+			chunks = append(chunks, s)
+		} else {
+			chunks = append(chunks, fmt.Sprintf("--- page %d (cont.) ---\n\n", chunkPage)+s)
+		}
+		buf.Reset()
+		hasLines = false
+	}
+
+	for _, line := range lines {
+		if m := pageMarker.FindStringSubmatch(line); m != nil {
+			if n, err := strconv.Atoi(m[1]); err == nil {
+				currentPage = n
+				markerSeen = true
+			}
+		}
+		next := line + "\n"
+		if hasLines {
+			est := EstTokens(buf.String() + next)
+			if startedAfterMarker {
+				est += EstTokens(fmt.Sprintf("--- page %d (cont.) ---\n\n", chunkPage))
+			}
+			if est > maxTok {
+				flush()
+			}
+		}
+		if !hasLines {
+			firstLine = line
+			chunkPage = currentPage
+			startedAfterMarker = markerSeen
+			hasLines = true
+		}
+		buf.WriteString(next)
+	}
+	flush()
+	return chunks
+}

internal/digest/splitpages_test.go

@@ -0,0 +1,90 @@
+package digest
+
+// P5b task 01: SplitPages keeps every chunk of a PDF's text attached to its page, so a digest
+// built from a chunk that starts mid-page can still cite the page.
+
+import (
+	"fmt"
+	"regexp"
+	"strings"
+	"testing"
+)
+
+// pagesText builds n pages of about perPage tokens each, every line naming its page.
+func pagesText(n, perPage int) string {
+	var b strings.Builder
+	for p := 1; p <= n; p++ {
+		if p > 1 {
+			b.WriteString("\n\n")
+		}
+		fmt.Fprintf(&b, "--- page %d ---\n\n", p)
+		for l := 0; l*12 < perPage; l++ {
+			fmt.Fprintf(&b, "page %d line %d alpha bravo charlie delta echo foxtrot golf hotel\n", p, l)
+		}
+	}
+	return b.String()
+}
+
+var markerRE = regexp.MustCompile(`^--- page (\d+)( \(cont\.\))? ---`)
+
+func TestSplitPagesEveryChunkNamesItsPage(t *testing.T) {
+	text := pagesText(4, 3000) // each page ~3000 tokens: several chunks per page at 1200
+	chunks := SplitPages(text, ChunkTokens)
+	if len(chunks) < 8 {
+		t.Fatalf("%d chunks, want at least 8 for 4 × 3000 tokens at %d", len(chunks), ChunkTokens)
+	}
+	cont := 0
+	for i, c := range chunks {
+		m := markerRE.FindStringSubmatch(c)
+		if m == nil {
+			t.Fatalf("chunk %d does not start with a page marker: %.80q", i, c)
+		}
+		if m[2] != "" {
+			cont++
+			// A continuation marker names the page the chunk's first line belongs to.
+			first := strings.SplitN(strings.TrimPrefix(c, m[0]), "\n", 3)
+			line := strings.TrimSpace(strings.Join(first, " "))
+			if !strings.Contains(line, "page "+m[1]+" line ") {
+				t.Errorf("chunk %d marked page %s but starts with %.60q", i, m[1], line)
+			}
+		}
+		if EstTokens(c) > ChunkTokens+20 {
+			t.Errorf("chunk %d is %d tokens, over the %d budget (plus a marker line)", i, EstTokens(c), ChunkTokens)
+		}
+	}
+	if cont == 0 {
+		t.Error("no continuation markers: pages longer than a chunk must produce some")
+	}
+}
+
+// Nothing is lost or duplicated: removing the continuation markers gives back the original lines.
+func TestSplitPagesPreservesContent(t *testing.T) {
+	text := pagesText(3, 2500)
+	var got []string
+	for _, c := range SplitPages(text, ChunkTokens) {
+		for _, line := range strings.Split(c, "\n") {
+			if strings.HasSuffix(line, " (cont.) ---") || strings.TrimSpace(line) == "" {
+				continue
+			}
+			got = append(got, line)
+		}
+	}
+	var want []string
+	for _, line := range strings.Split(text, "\n") {
+		if strings.TrimSpace(line) != "" {
+			want = append(want, line)
+		}
+	}
+	if strings.Join(got, "\n") != strings.Join(want, "\n") {
+		t.Errorf("content changed: %d lines back, %d in", len(got), len(want))
+	}
+}
+
+// Text with no markers splits exactly as Split does.
+func TestSplitPagesWithoutMarkersIsSplit(t *testing.T) {
+	text := strings.Repeat("plain words without any page markers at all\n", 2000)
+	a, b := SplitPages(text, ChunkTokens), Split(text, ChunkTokens)
+	if strings.Join(a, "\x00") != strings.Join(b, "\x00") {
+		t.Errorf("SplitPages differs from Split on unmarked text: %d vs %d chunks", len(a), len(b))
+	}
+}