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>
75 lines
1.9 KiB
Go
75 lines
1.9 KiB
Go
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
|
|
}
|