mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-30 20:49:21 +08:00
fix(ingestion): clamp unconfigured embedding limit; document chunker title parity (comment-only) (#17539)
## Summary Fixes embedding inputs being emptied and documents chunker heading-detection parity. - `component/tokenizer.go`: `truncateForEmbedding` now mirrors Python `common/token_utils.py:183-185` `truncate(string, max_len)` (keep the first `max_len` tokens). For an unconfigured embedder reporting `maxTokens <= 0`, instead of mirroring Python's `truncate` (which returns `""` for `max_len <= 0` and would make the embeddings API reject the batch with "inputs cannot be empty"), Go **clamps the limit to a safe default** (`defaultEmbeddingTokenLimit = 8192`) so truncation stays active and never produces empty inputs. The previous code returned `""` for `maxTokens <= 10`, which emptied non-empty chunks into the batch and triggered the API error. - A **10-token safety margin** is reserved (mirroring Python's embedding path `rag/svr/task_executor.py` which uses `mdl.max_length - 10`) so Go does not over-send tokens to hard-capped embedding APIs; it is only applied when the limit `> 10` so small limits remain non-empty. - The clamp is applied **centrally inside `truncateForEmbedding`**, covering both Builtin and generic callers — no separate `model_service.go` change is required. - `component/chunker/title.go`: comment-only — documents the already-ported PDF-outline detection (omission 1.5 / Gap C) and the hardcoded `BULLET_PATTERN` fallback (omission 1.7 / Gap C), porting `common.py:_outline_similarity` and `rag/nlp` `BULLET_PATTERN`. - `tokenizer_unit_test.go`: update truncation expectations to the new positive-keeps-tokens / non-positive-clamps-to-default behaviour. ## Test plan - `tokenizer_unit_test.go` updated for the new truncation behaviour (`TestTruncateForEmbedding_SmallMaxTokens`, `TestTruncateForEmbedding_UnconfiguredClampsToDefault`). - `./build.sh --test ./internal/ingestion/component/` passes. 🤖 Generated with [CodeBuddy Code](https://cnb.cool/codebuddy)
This commit is contained in:
@@ -37,15 +37,21 @@
|
||||
// a non-regex text record whose layout flags it as a
|
||||
// section/title/head and whose text passes not_title is promoted to
|
||||
// fallback_level = len(selected_group) + 1. Strategy (1) —
|
||||
// PDF-outline detection — still requires deepdoc/parser binary
|
||||
// access and remains the only parity gap.
|
||||
// PDF-outline detection — is also ported: newLevelContext
|
||||
// (title.go) reads the Parser-supplied file.outline via
|
||||
// outlineFromInputs / resolveOutlineLevels / outlineSimilarity
|
||||
// (port of common.py:_outline_similarity) before falling back to
|
||||
// the regex/layout branch, wired into group.go / hierarchy.go
|
||||
// (Chunker omission 1.5, Gap C closed).
|
||||
//
|
||||
// - The Go port has NO hardcoded BULLET_PATTERN fallback. Heading
|
||||
// detection relies entirely on the user-supplied `levels` param
|
||||
// (which templates carry as comprehensive regex families) paired
|
||||
// with the layout-hint fallback. A PDF without matching levels
|
||||
// produces BODY_LEVEL-only records — Python's BULLET_PATTERN-based
|
||||
// tree_merge / hierarchical_merge would still find structure.
|
||||
// - The Go port SHIPS a hardcoded BULLET_PATTERN fallback
|
||||
// (Chunker omission 1.7, Gap C closed): when the regex/layout
|
||||
// branch yields BODY_LEVEL-only records, resolveTitleLevels
|
||||
// (title.go) calls bulletsCategory over bulletPatterns (title.go,
|
||||
// mirroring Python BULLET_PATTERN at rag/nlp/__init__.py:258) to
|
||||
// recover structure from numbered/bulleted list entries, before
|
||||
// the layout-title fallback. It never overrides an existing level
|
||||
// assignment.
|
||||
//
|
||||
// - GROUP-TITLE and HIERARCHY-TITLE are separate Go files
|
||||
// (`group.go`, `hierarchy.go`); they share the resolve_levels
|
||||
|
||||
@@ -403,13 +403,14 @@ func (c *TokenizerComponent) embedChunks(ctx context.Context, tenantID, kbID, em
|
||||
texts := make([]string, 0, len(chunks))
|
||||
pairs := make([]int, 0, len(chunks))
|
||||
for i, ck := range chunks {
|
||||
txt := concatFields(ck, c.param.Fields)
|
||||
txt = htmlTableRE.ReplaceAllString(txt, " ")
|
||||
raw := concatFields(ck, c.param.Fields)
|
||||
txt := htmlTableRE.ReplaceAllString(raw, " ")
|
||||
txt = strings.TrimSpace(txt)
|
||||
trunc := truncateForEmbedding(txt, embedder.MaxTokens())
|
||||
if txt == "" {
|
||||
continue
|
||||
}
|
||||
texts = append(texts, truncateForEmbedding(txt, embedder.MaxTokens()))
|
||||
texts = append(texts, trunc)
|
||||
pairs = append(pairs, i)
|
||||
}
|
||||
if len(texts) == 0 {
|
||||
@@ -476,15 +477,35 @@ func (c *TokenizerComponent) embedChunks(ctx context.Context, tenantID, kbID, em
|
||||
return chunks, tokenCount, nil
|
||||
}
|
||||
|
||||
// truncateForEmbedding truncates text to fit within maxTokens,
|
||||
// reserving 10 tokens for special/model overhead. Mirrors Python
|
||||
// tokenizer.py truncation: when maxTokens <= 10, the text is
|
||||
// truncated to empty
|
||||
// defaultEmbeddingTokenLimit is the safe fallback used when an embedder reports
|
||||
// no token limit (maxTokens <= 0). It both prevents empty embedding inputs and
|
||||
// keeps truncation active for every path instead of passing the full text through.
|
||||
const defaultEmbeddingTokenLimit = 8192
|
||||
|
||||
// truncateForEmbedding keeps the first maxTokens tokens of text so it fits the
|
||||
// embedding model's limit.
|
||||
//
|
||||
// For a positive maxTokens it mirrors Python common/token_utils.py:183-185
|
||||
// `truncate(string, max_len)` (keep the first max_len tokens).
|
||||
//
|
||||
// An unconfigured embedder reports maxTokens <= 0. Rather than mirror Python's
|
||||
// behaviour of returning "" (which would make the embeddings API reject the whole
|
||||
// batch with "inputs cannot be empty"), Go clamps the limit to a safe default
|
||||
// (defaultEmbeddingTokenLimit = 8192). This both prevents empty inputs AND keeps
|
||||
// truncation active for every path (Builtin and generic) instead of silently
|
||||
// passing the full, untruncated text when no limit is configured.
|
||||
func truncateForEmbedding(text string, maxTokens int) string {
|
||||
if maxTokens <= 10 {
|
||||
return ""
|
||||
if maxTokens <= 0 {
|
||||
maxTokens = defaultEmbeddingTokenLimit
|
||||
}
|
||||
return tokenizer.TrimContentToTokenLimit(text, maxTokens-10)
|
||||
// Keep a 10-token safety margin, mirroring Python's embedding path
|
||||
// (rag/svr/task_executor.py uses `mdl.max_length - 10`). Only apply it
|
||||
// when the limit is large enough; for small limits (<=10) keep the full
|
||||
// value so the result stays non-empty instead of collapsing to "".
|
||||
if maxTokens > 10 {
|
||||
maxTokens -= 10
|
||||
}
|
||||
return tokenizer.TrimContentToTokenLimit(text, maxTokens)
|
||||
}
|
||||
|
||||
func mergeEmbeddingVectors(titleVec, contentVec []float64, titleWeight float64) ([]float64, error) {
|
||||
|
||||
@@ -91,8 +91,8 @@ func (s *stubEmbedder) Encode(ctx context.Context, texts []string) ([]EmbeddingR
|
||||
}
|
||||
|
||||
// newStubEmbedder returns a stub embedder for instance-level resolver injection.
|
||||
// maxTokens defaults to 2048 so truncateForEmbedding (Diff 14: maxTokens <= 10 -> "")
|
||||
// does not empty the content text; tests that exercise truncation set maxTokens explicitly.
|
||||
// maxTokens defaults to 2048 so truncateForEmbedding truncates to the first
|
||||
// 2048 tokens; tests that exercise truncation set maxTokens explicitly.
|
||||
func newStubEmbedder(dim int) *stubEmbedder {
|
||||
return &stubEmbedder{dim: dim, maxTokens: 2048}
|
||||
}
|
||||
@@ -376,21 +376,58 @@ func TestIsPhantomChunk(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestTruncateForEmbedding_SmallMaxTokens covers Tokenizer Diff-14: when
|
||||
// maxTokens <= 10, Go must truncate to empty, matching Python's
|
||||
// truncation-to-empty behaviour (common/token_utils.py:183-185).
|
||||
// TestTruncateForEmbedding_SmallMaxTokens covers Tokenizer Diff-14. For any
|
||||
// positive maxTokens, truncateForEmbedding keeps the first maxTokens tokens
|
||||
// (non-empty) and the result is strictly shorter than the input.
|
||||
//
|
||||
// The unconfigured case (maxTokens <= 0) is covered separately by
|
||||
// TestTruncateForEmbedding_UnconfiguredClampsToDefault: rather than mirroring
|
||||
// Python's `truncate` (which returns "" for max_len <= 0 and would make the
|
||||
// embeddings API reject the batch), Go clamps the limit to a safe default so
|
||||
// every path still truncates instead of passing the full text through.
|
||||
func TestTruncateForEmbedding_SmallMaxTokens(t *testing.T) {
|
||||
long := strings.Repeat("a", 100)
|
||||
if got := truncateForEmbedding(long, 5); got != "" {
|
||||
t.Errorf("truncateForEmbedding(maxTokens=5) = %q, want %q", got, "")
|
||||
// Long enough to produce well over 50 tokens, so 5/10/50 are all clearly
|
||||
// below the total and truncation is observable.
|
||||
long := strings.Repeat("a", 2000)
|
||||
if got := truncateForEmbedding(long, 5); got == "" {
|
||||
t.Error("truncateForEmbedding(maxTokens=5) returned empty, want first 5 tokens")
|
||||
}
|
||||
if got := truncateForEmbedding(long, 10); got != "" {
|
||||
t.Errorf("truncateForEmbedding(maxTokens=10) = %q, want %q", got, "")
|
||||
if got := truncateForEmbedding(long, 10); got == "" {
|
||||
t.Error("truncateForEmbedding(maxTokens=10) returned empty, want first 10 tokens")
|
||||
}
|
||||
// Normal path: maxTokens > 10 should truncate (not return empty).
|
||||
if got := truncateForEmbedding(long, 50); got == "" {
|
||||
t.Error("truncateForEmbedding(maxTokens=50) returned empty, want truncated text")
|
||||
}
|
||||
if got := truncateForEmbedding(long, 50); len(got) >= len(long) {
|
||||
t.Errorf("truncateForEmbedding(maxTokens=50) len = %d, want strictly shorter than %d", len(got), len(long))
|
||||
}
|
||||
}
|
||||
|
||||
// TestTruncateForEmbedding_UnconfiguredClampsToDefault is the regression gate
|
||||
// for the root cause behind embedding truncation: an embedder that reports no
|
||||
// token limit (maxTokens <= 0) must NOT be passed through verbatim. Before the
|
||||
// central clamp, the generic model_service branch left maxTokens = 0, which
|
||||
// silently disabled truncation for the whole generic path. The fix clamps
|
||||
// maxTokens <= 0 to defaultEmbeddingTokenLimit (8192) inside
|
||||
// truncateForEmbedding itself, so every caller — Builtin and generic alike —
|
||||
// keeps truncation active.
|
||||
//
|
||||
// For a clearly-over-limit input and maxTokens = 0, the result must be non-empty
|
||||
// AND strictly shorter than the input: proof that it was truncated to the
|
||||
// default, not returned verbatim.
|
||||
func TestTruncateForEmbedding_UnconfiguredClampsToDefault(t *testing.T) {
|
||||
// ~10000+ CL100K tokens, comfortably above the 8192 default so truncation
|
||||
// is observable.
|
||||
long := strings.Repeat("hello world ", 5000)
|
||||
got := truncateForEmbedding(long, 0)
|
||||
if got == "" {
|
||||
t.Fatal("truncateForEmbedding(maxTokens=0) returned empty; clamp must keep truncation active")
|
||||
}
|
||||
if len(got) >= len(long) {
|
||||
t.Errorf("truncateForEmbedding(maxTokens=0) len = %d, want strictly shorter than %d; clamp to default must truncate",
|
||||
len(got), len(long))
|
||||
}
|
||||
}
|
||||
|
||||
// TestEmbeddingBatchSizeEnvVar covers Tokenizer Omission-3: the batch size
|
||||
|
||||
Reference in New Issue
Block a user