From 0b643afb98862fd48ecaf324acdde9988e9eb35a Mon Sep 17 00:00:00 2001 From: timkaye11 Date: Sat, 21 Feb 2026 12:04:12 -0800 Subject: [PATCH] chunker fixes --- pkg/termite/lib/chunking/chunker.go | 12 +++- pkg/termite/lib/pipelines/chunking.go | 62 ++++++++++++++++--- pkg/termite/lib/pipelines/pipelines.go | 27 ++++++++ registry/index.json | 9 +++ .../mirth/chonky-modernbert-base-1.json | 39 ++++++++++++ 5 files changed, 136 insertions(+), 13 deletions(-) create mode 100644 registry/manifests/mirth/chonky-modernbert-base-1.json diff --git a/pkg/termite/lib/chunking/chunker.go b/pkg/termite/lib/chunking/chunker.go index 0fb1874..42c3d55 100644 --- a/pkg/termite/lib/chunking/chunker.go +++ b/pkg/termite/lib/chunking/chunker.go @@ -171,7 +171,7 @@ func (p *PooledChunker) BackendType() backends.BackendType { // Chunk splits text using neural token classification. // Thread-safe: uses semaphore to limit concurrent pipeline access. -// Note: per-request options (opts) are currently ignored; pipeline uses config from creation time. +// Per-request options (threshold, target_tokens) override pipeline defaults when non-zero. func (p *PooledChunker) Chunk(ctx context.Context, text string, opts chunking.ChunkOptions) ([]chunking.Chunk, error) { if text == "" { p.logger.Debug("Chunk called with empty text") @@ -199,8 +199,14 @@ func (p *PooledChunker) Chunk(ctx context.Context, text string, opts chunking.Ch zap.Int("text_length", textLen), zap.String("text_preview", textPreview)) - // Delegate to ChunkingPipeline.Chunk - pipelineChunks, err := pipeline.Chunk(ctx, text) + // Build per-request overrides from ChunkOptions + reqOpts := pipelines.ChunkRequestOptions{ + Threshold: opts.Threshold, + TargetTokens: opts.TargetTokens, + } + + // Delegate to ChunkingPipeline with per-request options + pipelineChunks, err := pipeline.ChunkWithOptions(ctx, text, reqOpts) if err != nil { p.logger.Error("Chunking failed", zap.Int("pipelineIndex", idx), diff --git a/pkg/termite/lib/pipelines/chunking.go b/pkg/termite/lib/pipelines/chunking.go index 08e0549..289d762 100644 --- a/pkg/termite/lib/pipelines/chunking.go +++ b/pkg/termite/lib/pipelines/chunking.go @@ -306,9 +306,21 @@ func NewChunkingPipeline( } } +// ChunkRequestOptions holds per-request overrides for chunking parameters. +// Zero values mean "use pipeline defaults". +type ChunkRequestOptions struct { + Threshold float32 + TargetTokens int +} + // Chunk splits a single text into semantic chunks. func (p *ChunkingPipeline) Chunk(ctx context.Context, text string) ([]Chunk, error) { - results, err := p.ChunkBatch(ctx, []string{text}) + return p.ChunkWithOptions(ctx, text, ChunkRequestOptions{}) +} + +// ChunkWithOptions splits a single text into semantic chunks with per-request overrides. +func (p *ChunkingPipeline) ChunkWithOptions(ctx context.Context, text string, opts ChunkRequestOptions) ([]Chunk, error) { + results, err := p.ChunkBatchWithOptions(ctx, []string{text}, opts) if err != nil { return nil, err } @@ -320,10 +332,25 @@ func (p *ChunkingPipeline) Chunk(ctx context.Context, text string) ([]Chunk, err // ChunkBatch splits multiple texts into semantic chunks. func (p *ChunkingPipeline) ChunkBatch(ctx context.Context, texts []string) ([][]Chunk, error) { + return p.ChunkBatchWithOptions(ctx, texts, ChunkRequestOptions{}) +} + +// ChunkBatchWithOptions splits multiple texts into semantic chunks with per-request overrides. +func (p *ChunkingPipeline) ChunkBatchWithOptions(ctx context.Context, texts []string, opts ChunkRequestOptions) ([][]Chunk, error) { if len(texts) == 0 { return nil, nil } + // Resolve effective config: per-request overrides take precedence over pipeline defaults. + threshold := p.Config.Threshold + if opts.Threshold > 0 { + threshold = opts.Threshold + } + targetTokens := p.Config.TargetTokens + if opts.TargetTokens > 0 { + targetTokens = opts.TargetTokens + } + // Encode texts with character offsets encoded, err := p.BasePipeline.EncodeWithSpans(texts) if err != nil { @@ -363,7 +390,7 @@ func (p *ChunkingPipeline) ChunkBatch(ctx context.Context, texts []string) ([][] offsets = encoded.Spans[i] } - chunks := p.parseChunks(text, logits[i], offsets) + chunks := p.parseChunksWithConfig(text, logits[i], offsets, threshold, targetTokens) results[i] = chunks } @@ -393,8 +420,13 @@ func (p *ChunkingPipeline) ClassifyTokens(ctx context.Context, inputs *backends. return output.LastHiddenState, nil } -// parseChunks converts token classification results into chunks. +// parseChunks converts token classification results into chunks using pipeline defaults. func (p *ChunkingPipeline) parseChunks(text string, logits [][]float32, offsets []TokenSpan) []Chunk { + return p.parseChunksWithConfig(text, logits, offsets, p.Config.Threshold, p.Config.TargetTokens) +} + +// parseChunksWithConfig converts token classification results into chunks with explicit threshold and targetTokens. +func (p *ChunkingPipeline) parseChunksWithConfig(text string, logits [][]float32, offsets []TokenSpan, threshold float32, targetTokens int) []Chunk { if len(logits) == 0 || len(text) == 0 { return []Chunk{{ Text: text, @@ -405,7 +437,7 @@ func (p *ChunkingPipeline) parseChunks(text string, logits [][]float32, offsets } // Find separator positions based on predicted labels - separatorPositions := p.findSeparatorPositions(logits, offsets, len(text)) + separatorPositions := p.findSeparatorPositionsWithThreshold(logits, offsets, len(text), threshold) // If no separators found, return whole text as single chunk if len(separatorPositions) == 0 { @@ -471,15 +503,20 @@ func (p *ChunkingPipeline) parseChunks(text string, logits [][]float32, offsets } // Apply target tokens aggregation if configured - if p.Config.TargetTokens > 0 && len(chunks) > 1 { - chunks = p.aggregateByTargetTokens(text, chunks) + if targetTokens > 0 && len(chunks) > 1 { + chunks = p.aggregateByTargetTokensN(text, chunks, targetTokens) } return chunks } -// findSeparatorPositions identifies character positions where separators occur. +// findSeparatorPositions identifies character positions where separators occur using pipeline defaults. func (p *ChunkingPipeline) findSeparatorPositions(logits [][]float32, offsets []TokenSpan, textLen int) []int { + return p.findSeparatorPositionsWithThreshold(logits, offsets, textLen, p.Config.Threshold) +} + +// findSeparatorPositionsWithThreshold identifies character positions where separators occur with an explicit threshold. +func (p *ChunkingPipeline) findSeparatorPositionsWithThreshold(logits [][]float32, offsets []TokenSpan, textLen int, threshold float32) []int { var positions []int for tokenIdx, tokenLogits := range logits { @@ -499,7 +536,7 @@ func (p *ChunkingPipeline) findSeparatorPositions(logits [][]float32, offsets [] // Apply confidence threshold using softmax probability prob := softmaxProb(tokenLogits, labelIdx) - if prob < p.Config.Threshold { + if prob < threshold { continue } @@ -531,8 +568,13 @@ func (p *ChunkingPipeline) isSeparatorLabel(label string) bool { strings.HasPrefix(labelLower, "i-") // BIO format: I-SEP (continuation) } -// aggregateByTargetTokens combines chunks until they reach the target token count. +// aggregateByTargetTokens combines chunks until they reach the target token count using pipeline defaults. func (p *ChunkingPipeline) aggregateByTargetTokens(originalText string, chunks []Chunk) []Chunk { + return p.aggregateByTargetTokensN(originalText, chunks, p.Config.TargetTokens) +} + +// aggregateByTargetTokensN combines chunks until they reach the given target token count. +func (p *ChunkingPipeline) aggregateByTargetTokensN(originalText string, chunks []Chunk, targetTokens int) []Chunk { if len(chunks) == 0 { return chunks } @@ -548,7 +590,7 @@ func (p *ChunkingPipeline) aggregateByTargetTokens(originalText string, chunks [ lastEndPos = chunk.End // If adding this chunk exceeds target, finalize current - if currentTokens > 0 && currentTokens+chunkTokens > p.Config.TargetTokens { + if currentTokens > 0 && currentTokens+chunkTokens > targetTokens { combinedText := strings.Join(currentTexts, "\n\n") aggregated = append(aggregated, Chunk{ Text: combinedText, diff --git a/pkg/termite/lib/pipelines/pipelines.go b/pkg/termite/lib/pipelines/pipelines.go index fe8f628..7d71af2 100644 --- a/pkg/termite/lib/pipelines/pipelines.go +++ b/pkg/termite/lib/pipelines/pipelines.go @@ -181,11 +181,24 @@ func (p *Pipeline) Encode(texts []string) (*EncodedBatch, error) { return &EncodedBatch{}, nil } + // Check if the tokenizer is a pure Go tokenizer (TokenizerWithSpans) + // which does not apply post-processing (BOS/EOS special tokens). + _, needsSpecialTokens := p.Tokenizer.(tokenizers.TokenizerWithSpans) + needsSpecialTokens = needsSpecialTokens && p.Config.AddSpecialTokens + // Tokenize all texts var allTokens [][]int maxLen := 0 for _, text := range texts { tokens := p.Tokenizer.Encode(text) + if needsSpecialTokens { + if bosID, err := p.Tokenizer.SpecialTokenID(tokenizers.TokBeginningOfSentence); err == nil { + tokens = append([]int{bosID}, tokens...) + } + if eosID, err := p.Tokenizer.SpecialTokenID(tokenizers.TokEndOfSentence); err == nil { + tokens = append(tokens, eosID) + } + } allTokens = append(allTokens, tokens) if len(tokens) > maxLen { maxLen = len(tokens) @@ -275,6 +288,20 @@ func (p *Pipeline) EncodeWithSpans(texts []string) (*EncodedBatch, error) { encoded := tokWithSpans.EncodeWithSpans(text) result.ids = encoded.IDs result.spans = encoded.Spans + + // The pure Go tokenizer (hftokenizer) does not apply + // post-processing (BOS/EOS), so we add them here when + // AddSpecialTokens is enabled. + if p.Config.AddSpecialTokens { + if bosID, err := p.Tokenizer.SpecialTokenID(tokenizers.TokBeginningOfSentence); err == nil { + result.ids = append([]int{bosID}, result.ids...) + result.spans = append([]TokenSpan{{}}, result.spans...) + } + if eosID, err := p.Tokenizer.SpecialTokenID(tokenizers.TokEndOfSentence); err == nil { + result.ids = append(result.ids, eosID) + result.spans = append(result.spans, TokenSpan{}) + } + } } else { result.ids = p.Tokenizer.Encode(text) // No spans available - leave nil diff --git a/registry/index.json b/registry/index.json index ea5f183..b2cbdaa 100644 --- a/registry/index.json +++ b/registry/index.json @@ -95,6 +95,15 @@ "i8" ] }, + { + "name": "chonky-modernbert-base-1", + "type": "chunker", + "description": "", + "source": "mirth/chonky_modernbert_base_1", + "size": 602606105, + "variants": [], + "owner": "mirth" + }, { "name": "mxbai-rerank-base-v1", "owner": "mixedbread-ai", diff --git a/registry/manifests/mirth/chonky-modernbert-base-1.json b/registry/manifests/mirth/chonky-modernbert-base-1.json new file mode 100644 index 0000000..823225d --- /dev/null +++ b/registry/manifests/mirth/chonky-modernbert-base-1.json @@ -0,0 +1,39 @@ +{ + "schemaVersion": 2, + "name": "chonky-modernbert-base-1", + "type": "chunker", + "description": "", + "source": "mirth/chonky_modernbert_base_1", + "files": [ + { + "name": "config.json", + "digest": "sha256:3ce1327d6e384227a65bb71bfb8318b00a548a08ababaa15ab9fb1ce39224c37", + "size": 1333 + }, + { + "name": "model.onnx", + "digest": "sha256:a4add1ffd0bffeeedd653204eb32d1ef70d37450d59708f47f3e3a129b6502d3", + "size": 598999802 + }, + { + "name": "special_tokens_map.json", + "digest": "sha256:ea97ecdbcc73713039d8d64dbb05e3689495c96657fbd9a18f5bed381be81049", + "size": 694 + }, + { + "name": "tokenizer.json", + "digest": "sha256:c7a995f78d60cc3c253902f4b5becfe2f9d0b44f78e6e2f81a343a0cb71789e6", + "size": 3583327 + }, + { + "name": "tokenizer_config.json", + "digest": "sha256:30be14db339ea143e64398869142116d0c0ba37db0b21bd0503f4b77e7a1309a", + "size": 20949 + } + ], + "owner": "mirth", + "provenance": { + "downloadedFrom": "https://huggingface.co/mirth/chonky_modernbert_base_1", + "downloadedAt": "2026-02-21T19:06:17.718823+00:00" + } +} \ No newline at end of file