Rewrite and simplify the section chunking algorithm
This commit is contained in:
@@ -29,11 +29,12 @@ type ChunkingConfig struct {
|
||||
// MaxSectionTokens is the maximum allowed tokens per section
|
||||
MaxSectionTokens int
|
||||
|
||||
// MinSectionTokens is the minimum desired tokens per section (for balancing)
|
||||
// MinSectionTokens is a validated soft lower-bound setting retained for
|
||||
// configuration/reporting compatibility.
|
||||
MinSectionTokens int
|
||||
|
||||
// TargetSections is an optional target number of sections
|
||||
// If nil, chunking will optimize for max token bounds only
|
||||
// If nil, section count is derived from total/max token budgeting.
|
||||
TargetSections *int
|
||||
}
|
||||
|
||||
@@ -59,8 +60,18 @@ func NewChunkerWithEstimator(config ChunkingConfig, estimator TokenEstimator) *C
|
||||
}
|
||||
}
|
||||
|
||||
// ChunkTranscript divides a normalized transcript into contiguous token-bounded sections.
|
||||
// Returns an error if the target section count is impossible under the constraints.
|
||||
// ChunkTranscript divides a normalized transcript into contiguous token-bounded
|
||||
// sections using a deterministic balanced forward pass.
|
||||
//
|
||||
// Behavior:
|
||||
// - preserve segment order and never split segments;
|
||||
// - estimate per-segment tokens once, then compute total;
|
||||
// - derive desired section count from ceil(total/max_section_tokens), unless
|
||||
// target_sections is explicitly set;
|
||||
// - prefer section sizes near ceil(total/section_count) while never exceeding
|
||||
// max_section_tokens unless a section consists of a single oversized segment.
|
||||
//
|
||||
// Returns an error if explicit target_sections is impossible under constraints.
|
||||
// The input transcript is never mutated.
|
||||
func (c *Chunker) ChunkTranscript(transcript *schema.Transcript) ([]Section, error) {
|
||||
if transcript == nil || len(transcript.Segments) == 0 {
|
||||
@@ -81,31 +92,51 @@ func (c *Chunker) ChunkTranscript(transcript *schema.Transcript) ([]Section, err
|
||||
c.config.MinSectionTokens, c.config.MaxSectionTokens)
|
||||
}
|
||||
|
||||
// Calculate token counts for each segment (deterministic)
|
||||
// Calculate token counts for each segment (deterministic).
|
||||
segmentTokens := make([]int, len(transcript.Segments))
|
||||
totalTokens := 0
|
||||
for i, seg := range transcript.Segments {
|
||||
segmentTokens[i] = c.estimator.EstimateTokens(seg.Text)
|
||||
totalTokens += segmentTokens[i]
|
||||
}
|
||||
|
||||
// If target sections specified, use target-aware chunking
|
||||
minPossibleSections := c.calculateMinPossibleSections(segmentTokens)
|
||||
|
||||
var desiredSections int
|
||||
useExplicitTarget := false
|
||||
if c.config.TargetSections != nil {
|
||||
target := *c.config.TargetSections
|
||||
|
||||
if target <= 0 {
|
||||
return nil, fmt.Errorf("target_sections must be positive, got %d", target)
|
||||
desiredSections = *c.config.TargetSections
|
||||
useExplicitTarget = true
|
||||
if desiredSections <= 0 {
|
||||
return nil, fmt.Errorf("target_sections must be positive, got %d", desiredSections)
|
||||
}
|
||||
|
||||
// Validate target is achievable
|
||||
if err := c.validateTargetSections(target, segmentTokens); err != nil {
|
||||
if err := c.validateTargetSections(desiredSections, segmentTokens); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Use target-aware chunking strategy
|
||||
return c.chunkWithTarget(transcript.Segments, segmentTokens, target)
|
||||
} else {
|
||||
desiredSections = ceilDiv(totalTokens, c.config.MaxSectionTokens)
|
||||
if desiredSections < minPossibleSections {
|
||||
desiredSections = minPossibleSections
|
||||
}
|
||||
if desiredSections < 1 {
|
||||
desiredSections = 1
|
||||
}
|
||||
if desiredSections > len(transcript.Segments) {
|
||||
desiredSections = len(transcript.Segments)
|
||||
}
|
||||
}
|
||||
|
||||
// No target: just respect max token bounds
|
||||
return c.createSectionsUnderMax(transcript.Segments, segmentTokens)
|
||||
targetTokensPerSection := ceilDiv(totalTokens, desiredSections)
|
||||
if useExplicitTarget {
|
||||
return c.buildSectionsWithExplicitTarget(
|
||||
transcript.Segments,
|
||||
segmentTokens,
|
||||
desiredSections,
|
||||
targetTokensPerSection,
|
||||
)
|
||||
}
|
||||
|
||||
return c.buildSectionsBalanced(transcript.Segments, segmentTokens, targetTokensPerSection), nil
|
||||
}
|
||||
|
||||
// validateTargetSections checks if the target section count is achievable.
|
||||
@@ -131,154 +162,8 @@ func (c *Chunker) validateTargetSections(target int, segmentTokens []int) error
|
||||
return nil
|
||||
}
|
||||
|
||||
// chunkWithTarget creates sections aiming for a specific target count.
|
||||
func (c *Chunker) chunkWithTarget(segments []schema.Segment, segmentTokens []int, target int) ([]Section, error) {
|
||||
// Strategy:
|
||||
// 1. First, create sections respecting max bounds (greedy packing)
|
||||
// 2. If we have more sections than target, try to merge adjacent sections
|
||||
// 3. If we have fewer sections than target, split sections where possible
|
||||
|
||||
// Step 1: Create initial sections respecting max bounds
|
||||
sections, err := c.createSectionsUnderMax(segments, segmentTokens)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Step 2: Adjust to reach target
|
||||
if len(sections) > target {
|
||||
// Need to merge sections
|
||||
sections = c.mergeSectionsToTarget(sections, target)
|
||||
} else if len(sections) < target {
|
||||
// Need to split sections - split the largest sections first
|
||||
sections = c.splitSectionsToTarget(sections, target, segmentTokens)
|
||||
}
|
||||
|
||||
// Verify we achieved the target
|
||||
if len(sections) != target {
|
||||
return nil, fmt.Errorf(
|
||||
"target_sections (%d) is impossible under current constraints (got %d sections)",
|
||||
target, len(sections))
|
||||
}
|
||||
|
||||
// Reindex sections
|
||||
for i := range sections {
|
||||
sections[i].Index = i
|
||||
}
|
||||
|
||||
return sections, nil
|
||||
}
|
||||
|
||||
// mergeSectionsToTarget merges adjacent sections to reduce count to target.
|
||||
func (c *Chunker) mergeSectionsToTarget(sections []Section, target int) []Section {
|
||||
for len(sections) > target {
|
||||
// Find the best pair to merge (smallest combined section that stays under max)
|
||||
bestIdx := -1
|
||||
bestTokens := -1
|
||||
|
||||
for i := 0; i < len(sections)-1; i++ {
|
||||
combinedTokens := sections[i].EstimatedTokens + sections[i+1].EstimatedTokens
|
||||
|
||||
// Can only merge if combined is under max
|
||||
if combinedTokens <= c.config.MaxSectionTokens {
|
||||
// Prefer smaller combined sections to keep sections balanced
|
||||
if bestIdx == -1 || combinedTokens < bestTokens {
|
||||
bestIdx = i
|
||||
bestTokens = combinedTokens
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// No valid merge found
|
||||
if bestIdx == -1 {
|
||||
break
|
||||
}
|
||||
|
||||
// Merge sections[bestIdx] and sections[bestIdx+1]
|
||||
merged := c.mergeTwoSections(sections[bestIdx], sections[bestIdx+1])
|
||||
|
||||
// Replace the two sections with the merged one
|
||||
newSections := make([]Section, 0, len(sections)-1)
|
||||
newSections = append(newSections, sections[:bestIdx]...)
|
||||
newSections = append(newSections, merged)
|
||||
if bestIdx+2 < len(sections) {
|
||||
newSections = append(newSections, sections[bestIdx+2:]...)
|
||||
}
|
||||
sections = newSections
|
||||
}
|
||||
|
||||
return sections
|
||||
}
|
||||
|
||||
// splitSectionsToTarget splits sections to increase count to target.
|
||||
func (c *Chunker) splitSectionsToTarget(sections []Section, target int, segmentTokens []int) []Section {
|
||||
// We can only split sections that have multiple segments
|
||||
// Split the largest sections first
|
||||
|
||||
for len(sections) < target {
|
||||
// Find the best section to split (largest section with multiple segments)
|
||||
bestIdx := -1
|
||||
bestTokenCount := -1
|
||||
|
||||
for i, sec := range sections {
|
||||
if len(sec.Segments) > 1 {
|
||||
if sec.EstimatedTokens > bestTokenCount {
|
||||
bestIdx = i
|
||||
bestTokenCount = sec.EstimatedTokens
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// No splittable section found
|
||||
if bestIdx == -1 {
|
||||
break
|
||||
}
|
||||
|
||||
// Split the section into two
|
||||
sec := sections[bestIdx]
|
||||
splitPoint := len(sec.Segments) / 2
|
||||
|
||||
// Ensure each half has at least one segment
|
||||
if splitPoint < 1 || splitPoint >= len(sec.Segments) {
|
||||
break
|
||||
}
|
||||
|
||||
// Calculate token counts for each half
|
||||
leftTokens := 0
|
||||
for i := 0; i < splitPoint; i++ {
|
||||
leftTokens += c.estimator.EstimateTokens(sec.Segments[i].Text)
|
||||
}
|
||||
rightTokens := sec.EstimatedTokens - leftTokens
|
||||
|
||||
// Create two new sections
|
||||
leftSection := Section{
|
||||
StartSegmentID: sec.Segments[0].ID,
|
||||
EndSegmentID: sec.Segments[splitPoint-1].ID,
|
||||
EstimatedTokens: leftTokens,
|
||||
Segments: sec.Segments[:splitPoint],
|
||||
}
|
||||
|
||||
rightSection := Section{
|
||||
StartSegmentID: sec.Segments[splitPoint].ID,
|
||||
EndSegmentID: sec.Segments[len(sec.Segments)-1].ID,
|
||||
EstimatedTokens: rightTokens,
|
||||
Segments: sec.Segments[splitPoint:],
|
||||
}
|
||||
|
||||
// Replace old section with two new ones
|
||||
newSections := make([]Section, 0, len(sections)+1)
|
||||
newSections = append(newSections, sections[:bestIdx]...)
|
||||
newSections = append(newSections, leftSection, rightSection)
|
||||
if bestIdx+1 < len(sections) {
|
||||
newSections = append(newSections, sections[bestIdx+1:]...)
|
||||
}
|
||||
sections = newSections
|
||||
}
|
||||
|
||||
return sections
|
||||
}
|
||||
|
||||
// createSectionsUnderMax creates initial sections respecting max token bounds.
|
||||
func (c *Chunker) createSectionsUnderMax(segments []schema.Segment, segmentTokens []int) ([]Section, error) {
|
||||
// buildSectionsBalanced creates sections with a deterministic single-pass policy.
|
||||
func (c *Chunker) buildSectionsBalanced(segments []schema.Segment, segmentTokens []int, targetTokensPerSection int) []Section {
|
||||
var sections []Section
|
||||
var currentSegments []schema.Segment
|
||||
currentTokens := 0
|
||||
@@ -286,38 +171,38 @@ func (c *Chunker) createSectionsUnderMax(segments []schema.Segment, segmentToken
|
||||
for i, seg := range segments {
|
||||
tokens := segmentTokens[i]
|
||||
|
||||
// If a single segment exceeds max tokens, it must be its own section
|
||||
if tokens > c.config.MaxSectionTokens {
|
||||
// Flush current section if any
|
||||
if len(currentSegments) > 0 {
|
||||
sections = append(sections, c.buildSection(len(sections), currentSegments, currentTokens))
|
||||
currentSegments = nil
|
||||
currentTokens = 0
|
||||
}
|
||||
// Add the oversized segment as its own section
|
||||
sections = append(sections, c.buildSection(len(sections), []schema.Segment{seg}, tokens))
|
||||
// Empty section: always accept the next segment, including oversized.
|
||||
if len(currentSegments) == 0 {
|
||||
currentSegments = append(currentSegments, seg)
|
||||
currentTokens = tokens
|
||||
continue
|
||||
}
|
||||
|
||||
// Check if adding this segment would exceed max
|
||||
if currentTokens+tokens > c.config.MaxSectionTokens && len(currentSegments) > 0 {
|
||||
// Finish current section
|
||||
// If adding next segment would exceed max, close current section.
|
||||
if currentTokens+tokens > c.config.MaxSectionTokens {
|
||||
sections = append(sections, c.buildSection(len(sections), currentSegments, currentTokens))
|
||||
currentSegments = []schema.Segment{seg}
|
||||
currentTokens = tokens
|
||||
} else {
|
||||
// Add to current section
|
||||
continue
|
||||
}
|
||||
|
||||
// Prefer staying near target tokens per section.
|
||||
if targetTokensPerSection == 0 || currentTokens < targetTokensPerSection {
|
||||
currentSegments = append(currentSegments, seg)
|
||||
currentTokens += tokens
|
||||
continue
|
||||
}
|
||||
|
||||
sections = append(sections, c.buildSection(len(sections), currentSegments, currentTokens))
|
||||
currentSegments = []schema.Segment{seg}
|
||||
currentTokens = tokens
|
||||
}
|
||||
|
||||
// Don't forget the last section
|
||||
if len(currentSegments) > 0 {
|
||||
sections = append(sections, c.buildSection(len(sections), currentSegments, currentTokens))
|
||||
}
|
||||
|
||||
return sections, nil
|
||||
return sections
|
||||
}
|
||||
|
||||
// buildSection creates a Section from segments.
|
||||
@@ -360,17 +245,68 @@ func (c *Chunker) calculateMinPossibleSections(segmentTokens []int) int {
|
||||
return sections
|
||||
}
|
||||
|
||||
// mergeTwoSections combines two adjacent sections.
|
||||
func (c *Chunker) mergeTwoSections(a, b Section) Section {
|
||||
mergedSegments := make([]schema.Segment, 0, len(a.Segments)+len(b.Segments))
|
||||
mergedSegments = append(mergedSegments, a.Segments...)
|
||||
mergedSegments = append(mergedSegments, b.Segments...)
|
||||
// buildSectionsWithExplicitTarget builds exactly desiredSections when feasible.
|
||||
func (c *Chunker) buildSectionsWithExplicitTarget(
|
||||
segments []schema.Segment,
|
||||
segmentTokens []int,
|
||||
desiredSections int,
|
||||
targetTokensPerSection int,
|
||||
) ([]Section, error) {
|
||||
n := len(segments)
|
||||
cursor := 0
|
||||
sections := make([]Section, 0, desiredSections)
|
||||
|
||||
return Section{
|
||||
Index: a.Index, // Will be reassigned
|
||||
StartSegmentID: a.StartSegmentID,
|
||||
EndSegmentID: b.EndSegmentID,
|
||||
EstimatedTokens: a.EstimatedTokens + b.EstimatedTokens,
|
||||
Segments: mergedSegments,
|
||||
for sectionIdx := 0; sectionIdx < desiredSections; sectionIdx++ {
|
||||
if cursor >= n {
|
||||
break
|
||||
}
|
||||
|
||||
remainingSectionsAfter := desiredSections - sectionIdx - 1
|
||||
currentSegments := []schema.Segment{segments[cursor]}
|
||||
currentTokens := segmentTokens[cursor]
|
||||
cursor++
|
||||
|
||||
for cursor < n {
|
||||
remainingSegments := n - cursor
|
||||
|
||||
// Reserve one segment per future section to avoid empty sections.
|
||||
if remainingSegments == remainingSectionsAfter {
|
||||
break
|
||||
}
|
||||
|
||||
nextTokens := segmentTokens[cursor]
|
||||
if currentTokens+nextTokens > c.config.MaxSectionTokens {
|
||||
break
|
||||
}
|
||||
if targetTokensPerSection == 0 || currentTokens < targetTokensPerSection {
|
||||
currentSegments = append(currentSegments, segments[cursor])
|
||||
currentTokens += nextTokens
|
||||
cursor++
|
||||
continue
|
||||
}
|
||||
break
|
||||
}
|
||||
|
||||
sections = append(sections, c.buildSection(len(sections), currentSegments, currentTokens))
|
||||
}
|
||||
|
||||
if cursor != n || len(sections) != desiredSections {
|
||||
return nil, fmt.Errorf(
|
||||
"target_sections (%d) is impossible under current constraints (got %d sections)",
|
||||
desiredSections,
|
||||
len(sections),
|
||||
)
|
||||
}
|
||||
|
||||
return sections, nil
|
||||
}
|
||||
|
||||
func ceilDiv(numerator int, denominator int) int {
|
||||
if denominator <= 0 {
|
||||
return 0
|
||||
}
|
||||
if numerator <= 0 {
|
||||
return 0
|
||||
}
|
||||
return (numerator + denominator - 1) / denominator
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user