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>
91 lines
2.9 KiB
Go
91 lines
2.9 KiB
Go
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))
|
||
}
|
||
}
|