267 lines
8.2 KiB
Go
267 lines
8.2 KiB
Go
package knowledge
|
||
|
||
import (
|
||
"strings"
|
||
"unicode/utf8"
|
||
)
|
||
|
||
// Chunking: turning a document into the units retrieval ranks.
|
||
//
|
||
// The size is a retrieval decision, not a storage one. Too large and a chunk
|
||
// matches on a paragraph the reader does not want, then spends the model's
|
||
// context on the rest of the page; too small and the sentence that answers the
|
||
// question arrives without the sentence that gives it meaning — "this does not
|
||
// apply to agency staff" is worse than useless detached from what "this" is.
|
||
//
|
||
// Paragraph-first, because a document's own paragraph breaks are the author's
|
||
// judgement about what belongs together, and they are better than any window
|
||
// this code could pick. Windows are the fallback for text with no structure.
|
||
|
||
const (
|
||
// TargetChunkRunes is what a chunk aims for. Roughly 250 words, which sits
|
||
// inside every current embedding model's window with room to spare and is
|
||
// about the size of a section a person would quote.
|
||
TargetChunkRunes = 1400
|
||
|
||
// MaxChunkRunes is the hard cap. A paragraph longer than this is split.
|
||
MaxChunkRunes = 2200
|
||
|
||
// OverlapRunes is how much of the previous chunk a split one repeats.
|
||
//
|
||
// Overlap exists for the boundary problem: the answer to a question often
|
||
// straddles a break, and without overlap neither side retrieves well. The
|
||
// cost is duplicated text in the index and occasionally two near-identical
|
||
// results, which the fusion step deduplicates by document and ordinal.
|
||
OverlapRunes = 180
|
||
|
||
// MinChunkRunes is the floor. A fragment shorter than this — a heading on
|
||
// its own, a stray line — is folded into its neighbour rather than indexed,
|
||
// because it will match on a keyword and then say nothing.
|
||
MinChunkRunes = 80
|
||
)
|
||
|
||
// Chunk is one indexable unit.
|
||
type Chunk struct {
|
||
Ordinal int
|
||
Text string
|
||
|
||
// Heading is the trail of headings above this chunk — "Handbook ›
|
||
// Attendance › Lateness". Weighted above the body in the tsvector, and it
|
||
// is what makes a citation read like a location rather than a row id.
|
||
Heading string
|
||
|
||
TokenEstimate int
|
||
}
|
||
|
||
// Split turns a document into chunks.
|
||
//
|
||
// `title` seeds the heading trail, so every chunk carries at least the document
|
||
// it came from. Markdown ATX headings (`#`, `##`) update the trail as they are
|
||
// passed; anything else is body text.
|
||
func Split(title, body string) []Chunk {
|
||
paragraphs, headings := parse(title, body)
|
||
|
||
var (
|
||
chunks []Chunk
|
||
current strings.Builder
|
||
heading string
|
||
)
|
||
|
||
flush := func() {
|
||
text := strings.TrimSpace(current.String())
|
||
current.Reset()
|
||
if text == "" {
|
||
return
|
||
}
|
||
// Too short to stand alone: fold it into the previous chunk rather than
|
||
// index a fragment that matches and then says nothing.
|
||
if utf8.RuneCountInString(text) < MinChunkRunes && len(chunks) > 0 {
|
||
last := &chunks[len(chunks)-1]
|
||
last.Text += "\n\n" + text
|
||
last.TokenEstimate = estimateTokens(last.Text)
|
||
return
|
||
}
|
||
chunks = append(chunks, Chunk{
|
||
Ordinal: len(chunks), Text: text, Heading: heading,
|
||
TokenEstimate: estimateTokens(text),
|
||
})
|
||
}
|
||
|
||
for i, p := range paragraphs {
|
||
if h := headings[i]; h != "" {
|
||
// A new section starts a new chunk. Carrying text across a heading
|
||
// would put two topics in one unit and give it the wrong label.
|
||
flush()
|
||
heading = h
|
||
continue
|
||
}
|
||
|
||
// A paragraph over the cap is split on its own, with overlap.
|
||
if utf8.RuneCountInString(p) > MaxChunkRunes {
|
||
flush()
|
||
for _, piece := range window(p) {
|
||
chunks = append(chunks, Chunk{
|
||
Ordinal: len(chunks), Text: piece, Heading: heading,
|
||
TokenEstimate: estimateTokens(piece),
|
||
})
|
||
}
|
||
continue
|
||
}
|
||
|
||
if current.Len() > 0 && utf8.RuneCountInString(current.String())+utf8.RuneCountInString(p) > TargetChunkRunes {
|
||
flush()
|
||
}
|
||
if current.Len() > 0 {
|
||
current.WriteString("\n\n")
|
||
}
|
||
current.WriteString(p)
|
||
}
|
||
flush()
|
||
|
||
return chunks
|
||
}
|
||
|
||
// parse splits a body into paragraphs, tracking the heading trail.
|
||
//
|
||
// Returns paragraphs and, in step, the heading each one introduces — empty for
|
||
// ordinary text. Two parallel slices rather than a struct because the caller
|
||
// walks them together exactly once.
|
||
func parse(title, body string) (paragraphs []string, headings []string) {
|
||
trail := []string{}
|
||
if t := strings.TrimSpace(title); t != "" {
|
||
trail = append(trail, t)
|
||
}
|
||
|
||
for _, block := range strings.Split(strings.ReplaceAll(body, "\r\n", "\n"), "\n\n") {
|
||
block = strings.TrimSpace(block)
|
||
if block == "" {
|
||
continue
|
||
}
|
||
|
||
if level, text, ok := atxHeading(block); ok {
|
||
// Trim the trail to this heading's depth, then push. The document
|
||
// title is always element 0, so a level-1 heading sits at index 1.
|
||
depth := level
|
||
if depth > len(trail) {
|
||
depth = len(trail)
|
||
}
|
||
trail = append(trail[:depth], text)
|
||
paragraphs = append(paragraphs, block)
|
||
headings = append(headings, strings.Join(trail, " › "))
|
||
continue
|
||
}
|
||
|
||
paragraphs = append(paragraphs, block)
|
||
headings = append(headings, "")
|
||
}
|
||
return paragraphs, headings
|
||
}
|
||
|
||
// atxHeading recognises a markdown heading line.
|
||
//
|
||
// Stricter than "starts with a hash", and it has to be. `#3 on the rota is the
|
||
// closing shift` is prose, and treating it as a heading splits a paragraph
|
||
// mid-thought and mislabels every chunk after it — a mislabelled chunk then
|
||
// cites wrongly, which is the failure that survives longest because the text is
|
||
// right and only the attribution is wrong.
|
||
//
|
||
// Three conditions, all from CommonMark's ATX rule plus one of our own:
|
||
//
|
||
// - One to six hashes, followed by WHITESPACE. This is the condition that
|
||
// `#3` fails, and it is the one CommonMark actually specifies.
|
||
// - A single line. `# Something` followed by prose in the same block is prose
|
||
// that begins with a hash.
|
||
// - Short. A "heading" the length of a paragraph is a paragraph — the cap is
|
||
// ours, not the spec's, and it exists because a heading becomes a citation
|
||
// label and a 400-character label is unusable.
|
||
func atxHeading(block string) (level int, text string, ok bool) {
|
||
if strings.Contains(block, "\n") {
|
||
return 0, "", false
|
||
}
|
||
trimmed := strings.TrimLeft(block, "#")
|
||
level = len(block) - len(trimmed)
|
||
if level == 0 || level > 6 {
|
||
return 0, "", false
|
||
}
|
||
// CommonMark: the hashes must be followed by a space or the end of line.
|
||
if trimmed != "" && !strings.HasPrefix(trimmed, " ") && !strings.HasPrefix(trimmed, "\t") {
|
||
return 0, "", false
|
||
}
|
||
text = strings.TrimSpace(trimmed)
|
||
if text == "" {
|
||
return 0, "", false
|
||
}
|
||
if utf8.RuneCountInString(text) > MaxHeadingRunes {
|
||
return 0, "", false
|
||
}
|
||
return level, text, true
|
||
}
|
||
|
||
// MaxHeadingRunes is how long a heading may be before it is read as a
|
||
// paragraph. A heading becomes a citation label, and a label the length of a
|
||
// paragraph is not a label.
|
||
const MaxHeadingRunes = 120
|
||
|
||
// window splits an over-long paragraph into overlapping pieces.
|
||
//
|
||
// Break points prefer a sentence end near the target, then a space, then the
|
||
// raw offset. Cutting mid-word produces a token nothing matches and a citation
|
||
// that reads as though it were corrupted.
|
||
func window(p string) []string {
|
||
runes := []rune(p)
|
||
var out []string
|
||
|
||
for start := 0; start < len(runes); {
|
||
end := start + TargetChunkRunes
|
||
if end >= len(runes) {
|
||
out = append(out, strings.TrimSpace(string(runes[start:])))
|
||
break
|
||
}
|
||
end = breakNear(runes, start, end)
|
||
out = append(out, strings.TrimSpace(string(runes[start:end])))
|
||
|
||
next := end - OverlapRunes
|
||
if next <= start {
|
||
// Defensive: a pathological break point must not stall the loop.
|
||
next = end
|
||
}
|
||
start = next
|
||
}
|
||
return out
|
||
}
|
||
|
||
// breakNear finds a readable break at or before `end`.
|
||
func breakNear(runes []rune, start, end int) int {
|
||
const look = 220
|
||
floor := end - look
|
||
if floor <= start {
|
||
floor = start + 1
|
||
}
|
||
for i := end; i > floor; i-- {
|
||
switch runes[i-1] {
|
||
case '.', '!', '?', '\n':
|
||
return i
|
||
}
|
||
}
|
||
for i := end; i > floor; i-- {
|
||
if runes[i-1] == ' ' {
|
||
return i
|
||
}
|
||
}
|
||
return end
|
||
}
|
||
|
||
// estimateTokens is a rough token count.
|
||
//
|
||
// Four characters per token, the usual English approximation. Deliberately an
|
||
// estimate: it is used to budget how much context a retrieval may spend, and
|
||
// paying a tokeniser to be exact about a number that is then compared to a soft
|
||
// budget would be precision nobody spends.
|
||
func estimateTokens(s string) int {
|
||
n := utf8.RuneCountInString(s) / 4
|
||
if n < 1 {
|
||
return 1
|
||
}
|
||
return n
|
||
}
|