mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-08-11 09:51:44 +08:00
4b4a6e72f0
Unifies the Go TokenChunker merge path on a single `mergeUnits` core and fixes coordinate-tag drift in the Python JSON merge at `overlap > 0`. Rebased on top of #17979 (delimiter_mode convergence).
502 lines
21 KiB
Python
502 lines
21 KiB
Python
#
|
||
# Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
|
||
#
|
||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||
# you may not use this file except in compliance with the License.
|
||
# You may obtain a copy of the License at
|
||
#
|
||
# http://www.apache.org/licenses/LICENSE-2.0
|
||
#
|
||
# Unless required by applicable law or agreed to in writing, software
|
||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||
# See the License for the specific language governing permissions and
|
||
# limitations under the License.
|
||
import random
|
||
import re
|
||
from copy import deepcopy
|
||
|
||
from common.float_utils import normalize_overlapped_percent
|
||
from common.token_utils import num_tokens_from_string
|
||
from rag.flow.base import ProcessBase, ProcessParamBase
|
||
from rag.flow.chunker.schema import TokenChunkerFromUpstream
|
||
from rag.flow.parser.pdf_chunk_metadata import (
|
||
PDF_POSITIONS_KEY,
|
||
extract_pdf_positions,
|
||
finalize_pdf_chunk,
|
||
restore_pdf_text_previews,
|
||
)
|
||
from rag.nlp import naive_merge
|
||
|
||
# _TAG_RE matches parser-emitted coordinate tags of the form
|
||
# ``@@<page>\t<left>\t<right>\t<top>\t<bottom>##``. Mirrors Go's
|
||
# posTagRemove (internal/ingestion/component/chunker/group.go) so the two
|
||
# languages strip tags identically.
|
||
_TAG_RE = re.compile(r"@@[\t0-9.-]+?##")
|
||
|
||
|
||
def remove_tag(text):
|
||
"""Strip ``@@...##`` coordinate tags from text.
|
||
|
||
Used both when measuring the overlap prefix (so the cut lands on the
|
||
tag-free visible text, matching Go's computeOverlapPrefix) and on the
|
||
final chunk text (so coordinate tags never reach embedding/index).
|
||
"""
|
||
return _TAG_RE.sub("", text or "")
|
||
|
||
|
||
class TokenChunkerParam(ProcessParamBase):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.delimiter_mode = "delimiter"
|
||
self.chunk_token_size = 512
|
||
self.delimiters = ["\n"]
|
||
self.overlapped_percent = 0
|
||
self.children_delimiters = []
|
||
self.table_context_size = 0
|
||
self.image_context_size = 0
|
||
|
||
def check(self):
|
||
# Backward-compat: "token_size" was removed but is behaviorally identical
|
||
# to "delimiter" at runtime (both route through the same code path), so
|
||
# accept and coerce it instead of rejecting legacy configs / pre-fix
|
||
# frontends. Only genuinely unknown values are rejected.
|
||
if self.delimiter_mode == "token_size":
|
||
self.delimiter_mode = "delimiter"
|
||
self.check_valid_value(self.delimiter_mode, "Delimiter mode abnormal.", ["delimiter", "one"])
|
||
if self.delimiters is None:
|
||
self.delimiters = []
|
||
elif isinstance(self.delimiters, str):
|
||
self.delimiters = [self.delimiters]
|
||
else:
|
||
self.delimiters = [d for d in self.delimiters if isinstance(d, str)]
|
||
self.delimiters = [d for d in self.delimiters if d]
|
||
|
||
if self.children_delimiters is None:
|
||
self.children_delimiters = []
|
||
elif isinstance(self.children_delimiters, str):
|
||
self.children_delimiters = [self.children_delimiters]
|
||
else:
|
||
self.children_delimiters = [d for d in self.children_delimiters if isinstance(d, str)]
|
||
self.children_delimiters = [d for d in self.children_delimiters if d]
|
||
|
||
self.check_positive_integer(self.chunk_token_size, "Chunk token size.")
|
||
self.check_decimal_float(self.overlapped_percent, "Overlapped percentage: [0, 1)")
|
||
self.check_nonnegative_number(self.table_context_size, "Table context size.")
|
||
self.check_nonnegative_number(self.image_context_size, "Image context size.")
|
||
|
||
def get_input_form(self) -> dict[str, dict]:
|
||
return {}
|
||
|
||
|
||
def _compile_delimiter_pattern(delimiters):
|
||
# Build the primary delimiter regex from active delimiters wrapped by backticks.
|
||
raw_delimiters = "".join(delimiter for delimiter in (delimiters or []) if delimiter)
|
||
custom_delimiters = [m.group(1) for m in re.finditer(r"`([^`]+)`", raw_delimiters)]
|
||
if not custom_delimiters:
|
||
return ""
|
||
return "|".join(re.escape(text) for text in sorted(set(custom_delimiters), key=len, reverse=True))
|
||
|
||
|
||
def _split_text_by_pattern(text, pattern):
|
||
# Split text by the compiled delimiter pattern and discard delimiters.
|
||
# No atom-split is performed; empty segments between consecutive delimiters
|
||
# are dropped but whitespace-only segments are preserved (the delimiter is
|
||
# the boundary, not stripped away).
|
||
if not pattern:
|
||
return [text or ""]
|
||
|
||
split_texts = re.split(r"(%s)" % pattern, text or "", flags=re.DOTALL)
|
||
chunks = []
|
||
for i in range(0, len(split_texts), 2):
|
||
chunk = split_texts[i]
|
||
if chunk:
|
||
chunks.append(chunk)
|
||
return chunks
|
||
|
||
|
||
def _build_json_chunks(json_result, delimiter_pattern):
|
||
# Convert upstream JSON items into internal working chunks.
|
||
chunks = []
|
||
for item in json_result:
|
||
doc_type = str(item.get("doc_type_kwd") or "").strip().lower()
|
||
if doc_type == "table":
|
||
ck_type = "table"
|
||
elif doc_type == "image":
|
||
ck_type = "image"
|
||
else:
|
||
ck_type = "text"
|
||
|
||
text = item.get("text")
|
||
if not isinstance(text, str):
|
||
text = item.get("content_with_weight")
|
||
if not isinstance(text, str):
|
||
text = ""
|
||
|
||
# Keep PDF coordinates as an internal preview field until the final
|
||
# output is assembled. This avoids leaking two public coordinate
|
||
# formats downstream.
|
||
preview_positions = extract_pdf_positions(item)
|
||
img_id = item.get("img_id")
|
||
|
||
if ck_type == "text":
|
||
text_segments = _split_text_by_pattern(text, delimiter_pattern) if delimiter_pattern else [text]
|
||
for segment in text_segments:
|
||
if not segment or not segment.strip():
|
||
continue
|
||
chunks.append(
|
||
{
|
||
"text": segment,
|
||
"doc_type_kwd": "text",
|
||
"ck_type": "text",
|
||
PDF_POSITIONS_KEY: deepcopy(preview_positions),
|
||
"tk_nums": num_tokens_from_string(segment),
|
||
}
|
||
)
|
||
continue
|
||
|
||
chunks.append(
|
||
{
|
||
"text": text or "",
|
||
"doc_type_kwd": ck_type,
|
||
"ck_type": ck_type,
|
||
"img_id": img_id,
|
||
PDF_POSITIONS_KEY: deepcopy(preview_positions),
|
||
"tk_nums": num_tokens_from_string(text or ""),
|
||
"context_above": "",
|
||
"context_below": "",
|
||
}
|
||
)
|
||
|
||
return chunks
|
||
|
||
|
||
def _take_sentences(text, need_tokens, from_end=False):
|
||
# Take text from one side until the target token budget is reached.
|
||
split_pat = r"([。!?？；！\n]|\. )"
|
||
texts = re.split(split_pat, text or "", flags=re.DOTALL)
|
||
sentences = []
|
||
for i in range(0, len(texts), 2):
|
||
sentences.append(texts[i] + (texts[i + 1] if i + 1 < len(texts) else ""))
|
||
iterator = reversed(sentences) if from_end else sentences
|
||
collected = ""
|
||
for sentence in iterator:
|
||
collected = sentence + collected if from_end else collected + sentence
|
||
if num_tokens_from_string(collected) >= need_tokens:
|
||
break
|
||
return collected
|
||
|
||
|
||
def _attach_context_to_media_chunks(chunks, table_context_size, image_context_size):
|
||
# Add surrounding text to table/image chunks when context windows are enabled.
|
||
for i, chunk in enumerate(chunks):
|
||
if chunk["ck_type"] not in {"table", "image"}:
|
||
continue
|
||
|
||
context_size = image_context_size if chunk["ck_type"] == "image" else table_context_size
|
||
if context_size <= 0:
|
||
continue
|
||
|
||
remain_above = context_size
|
||
remain_below = context_size
|
||
parts_above = []
|
||
parts_below = []
|
||
|
||
prev = i - 1
|
||
while prev >= 0 and remain_above > 0:
|
||
prev_chunk = chunks[prev]
|
||
if prev_chunk["ck_type"] == "text":
|
||
if prev_chunk["tk_nums"] >= remain_above:
|
||
parts_above.insert(0, _take_sentences(prev_chunk["text"], remain_above, from_end=True))
|
||
remain_above = 0
|
||
break
|
||
parts_above.insert(0, prev_chunk["text"])
|
||
remain_above -= prev_chunk["tk_nums"]
|
||
prev -= 1
|
||
|
||
after = i + 1
|
||
while after < len(chunks) and remain_below > 0:
|
||
after_chunk = chunks[after]
|
||
if after_chunk["ck_type"] == "text":
|
||
if after_chunk["tk_nums"] >= remain_below:
|
||
parts_below.append(_take_sentences(after_chunk["text"], remain_below))
|
||
remain_below = 0
|
||
break
|
||
parts_below.append(after_chunk["text"])
|
||
remain_below -= after_chunk["tk_nums"]
|
||
after += 1
|
||
|
||
chunk["context_above"] = "".join(parts_above)
|
||
chunk["context_below"] = "".join(parts_below)
|
||
|
||
|
||
def _merge_text_chunks_by_token_size(chunks, chunk_token_size, overlapped_percent):
|
||
# Merge adjacent text chunks when delimiter-based splitting is not active.
|
||
merged = []
|
||
prev_text_idx = -1
|
||
threshold = chunk_token_size * (100 - overlapped_percent) / 100.0
|
||
|
||
for chunk in chunks:
|
||
if chunk["ck_type"] != "text":
|
||
merged.append(deepcopy(chunk))
|
||
prev_text_idx = -1
|
||
continue
|
||
|
||
current = deepcopy(chunk)
|
||
should_start_new = prev_text_idx < 0 or merged[prev_text_idx]["tk_nums"] > threshold
|
||
# #17799: an over-budget unit stands alone — never merged into the
|
||
# previous chunk. This matches Python naive_merge and the Go
|
||
# TokenChunker (all three paths stand the over-budget unit alone), so
|
||
# the Python JSON path, Python text path, and Go TokenChunker share one
|
||
# contract.
|
||
if current["tk_nums"] > chunk_token_size:
|
||
should_start_new = True
|
||
if should_start_new:
|
||
if prev_text_idx >= 0 and overlapped_percent > 0 and merged[prev_text_idx]["text"]:
|
||
# Mirror Go computeOverlapPrefix: measure the overlap cut on the
|
||
# tag-free *visible* text, never on the raw text that still
|
||
# carries @@...## coordinate tags. This keeps the overlap prefix
|
||
# aligned with Go and prevents a partial tag from leaking into
|
||
# the next chunk when the cut would land inside a tag.
|
||
visible = remove_tag(merged[prev_text_idx]["text"])
|
||
overlap_start = int(len(visible) * (100 - overlapped_percent) / 100.0)
|
||
if 0 <= overlap_start < len(visible):
|
||
overlap_text = visible[overlap_start:]
|
||
else:
|
||
overlap_text = ""
|
||
current["text"] = overlap_text + current["text"]
|
||
current["tk_nums"] = num_tokens_from_string(current["text"])
|
||
merged.append(current)
|
||
prev_text_idx = len(merged) - 1
|
||
continue
|
||
|
||
if merged[prev_text_idx]["text"] and current["text"]:
|
||
merged[prev_text_idx]["text"] += "\n" + current["text"]
|
||
else:
|
||
merged[prev_text_idx]["text"] += current["text"]
|
||
merged[prev_text_idx][PDF_POSITIONS_KEY].extend(current.get(PDF_POSITIONS_KEY) or [])
|
||
merged[prev_text_idx]["tk_nums"] += current["tk_nums"]
|
||
|
||
return merged
|
||
|
||
|
||
def _finalize_json_chunks(chunks):
|
||
# Convert internal chunks into the final token chunker output format.
|
||
docs = []
|
||
for chunk in chunks:
|
||
# Strip parser coordinate tags from the final text so they never
|
||
# reach embedding/index (coordinates already live in the structured
|
||
# PDF_POSITIONS_KEY field). Mirrors Go's removeTag at the chunker
|
||
# output boundary (token.go:544).
|
||
text = remove_tag((chunk.get("context_above") or "") + (chunk.get("text") or "") + (chunk.get("context_below") or ""))
|
||
if not text.strip():
|
||
continue
|
||
|
||
# The internal preview coordinates are converted exactly once into the
|
||
# indexed fields consumed downstream.
|
||
doc = {
|
||
"text": text,
|
||
"doc_type_kwd": chunk.get("doc_type_kwd", "text"),
|
||
}
|
||
if chunk.get(PDF_POSITIONS_KEY):
|
||
doc[PDF_POSITIONS_KEY] = deepcopy(chunk[PDF_POSITIONS_KEY])
|
||
if chunk.get("mom"):
|
||
doc["mom"] = chunk["mom"]
|
||
if chunk.get("img_id"):
|
||
doc["img_id"] = chunk["img_id"]
|
||
docs.append(finalize_pdf_chunk(doc))
|
||
|
||
return docs
|
||
|
||
|
||
def _split_chunk_docs_by_children(chunks, pattern):
|
||
# Apply the secondary children_delimiters split to text chunks only.
|
||
if not pattern:
|
||
return chunks
|
||
|
||
docs = []
|
||
for chunk in chunks:
|
||
if chunk.get("doc_type_kwd", "text") != "text":
|
||
docs.append(chunk)
|
||
continue
|
||
|
||
split_texts = _split_text_by_pattern(chunk.get("text", ""), pattern)
|
||
|
||
mom = chunk.get("text", "")
|
||
for text in split_texts:
|
||
if not text.strip():
|
||
continue
|
||
child = deepcopy(chunk)
|
||
child["mom"] = mom
|
||
child["text"] = text
|
||
docs.append(child)
|
||
|
||
return docs
|
||
|
||
|
||
class TokenChunker(ProcessBase):
|
||
component_name = "TokenChunker"
|
||
|
||
async def _invoke(self, **kwargs):
|
||
try:
|
||
from_upstream = TokenChunkerFromUpstream.model_validate(kwargs)
|
||
except Exception as e:
|
||
self.set_output("_ERROR", f"Input error: {str(e)}")
|
||
return
|
||
|
||
# Build the primary delimiter regex. If no active custom delimiter exists,
|
||
# the token chunker falls back to token-size based merging.
|
||
delimiter_pattern = _compile_delimiter_pattern(self._param.delimiters)
|
||
custom_pattern = "|".join(re.escape(t) for t in sorted(set(self._param.children_delimiters), key=len, reverse=True))
|
||
|
||
self.set_output("output_format", "chunks")
|
||
self.callback(random.randint(1, 5) / 100.0, "Start to split into chunks.")
|
||
overlapped_percent = normalize_overlapped_percent(self._param.overlapped_percent)
|
||
if from_upstream.output_format in ["markdown", "text", "html"]:
|
||
payload = getattr(from_upstream, f"{from_upstream.output_format}_result") or ""
|
||
if self._param.delimiter_mode == "one":
|
||
# Strip parser coordinate tags so they never reach
|
||
# embedding/index (consistent with the JSON merge path).
|
||
self.set_output("chunks", [{"text": remove_tag(payload)}] if payload.strip() else [])
|
||
self.callback(1, "Done.")
|
||
return
|
||
if delimiter_pattern:
|
||
cks = _split_text_by_pattern(payload, delimiter_pattern)
|
||
else:
|
||
cks = naive_merge(
|
||
payload,
|
||
self._param.chunk_token_size,
|
||
"".join(self._param.delimiters),
|
||
overlapped_percent,
|
||
)
|
||
if custom_pattern:
|
||
docs = []
|
||
for c in cks:
|
||
if not c.strip():
|
||
continue
|
||
for text in _split_text_by_pattern(c, custom_pattern):
|
||
if not text.strip():
|
||
continue
|
||
docs.append({"text": text, "mom": c})
|
||
self.set_output("chunks", docs)
|
||
else:
|
||
self.set_output("chunks", [{"text": c.strip()} for c in cks if c.strip()])
|
||
|
||
self.callback(1, "Done.")
|
||
return
|
||
|
||
# json
|
||
json_result = (from_upstream.chunks if from_upstream.output_format == "chunks" else from_upstream.json_result) or []
|
||
if self._param.delimiter_mode == "one":
|
||
sections = []
|
||
for item in json_result:
|
||
text = item.get("text")
|
||
if not isinstance(text, str):
|
||
text = item.get("content_with_weight")
|
||
if isinstance(text, str) and text.strip():
|
||
# Strip parser coordinate tags so they never reach
|
||
# embedding/index (consistent with the JSON merge path).
|
||
sections.append(remove_tag(text))
|
||
merged_text = "\n".join(sections)
|
||
self.set_output("chunks", [{"text": merged_text}] if merged_text.strip() else [])
|
||
self.callback(1, "Done.")
|
||
return
|
||
|
||
# Both branches start from per-item chunks (no pre-split by the
|
||
# delimiter pattern). The delimiter branch splits the buffered text
|
||
# stream while preserving per-segment PDF positions; the no-delimiter
|
||
# branch merges adjacent text items to chunk_token_size (the removed
|
||
# "token_size" behaviour, and a parity match with the Go JSON path).
|
||
text_chunks = _build_json_chunks(json_result, "")
|
||
|
||
if delimiter_pattern:
|
||
chunks = []
|
||
text_buffer = []
|
||
text_buffer_pos = []
|
||
|
||
def flush_text_buffer():
|
||
if not text_buffer:
|
||
return
|
||
# Join buffered text items with "\n" so adjacent item text is not
|
||
# glued together (e.g. "hello" + "world" must not become "helloworld").
|
||
# The delimiter is then applied to the combined text; a segment may
|
||
# span across item boundaries (the "\n" glue is not itself a
|
||
# delimiter), so each segment carries only the PDF positions of the
|
||
# item(s) that contributed to it -- never the union of every item
|
||
# (which previously leaked page-N coordinates into page-M chunks and
|
||
# made all segments share one preview image).
|
||
parts = []
|
||
item_ranges = [] # (start, end) of each buffered item in combined_text
|
||
offset = 0
|
||
for text in text_buffer:
|
||
start = offset
|
||
parts.append(text)
|
||
offset += len(text)
|
||
item_ranges.append((start, offset))
|
||
parts.append("\n")
|
||
offset += 1
|
||
combined_text = "".join(parts[:-1]) # drop the trailing glue
|
||
|
||
raw = re.split(r"(%s)" % delimiter_pattern, combined_text, flags=re.DOTALL)
|
||
segments = [] # (text, start, end) within combined_text
|
||
pos = 0
|
||
for i in range(0, len(raw), 2):
|
||
seg = raw[i]
|
||
seg_start = pos
|
||
seg_end = pos + len(seg)
|
||
if seg:
|
||
segments.append((seg, seg_start, seg_end))
|
||
pos = seg_end
|
||
if i + 1 < len(raw):
|
||
pos += len(raw[i + 1])
|
||
|
||
for text, seg_start, seg_end in segments:
|
||
if not text.strip():
|
||
continue
|
||
seg_pos = []
|
||
for (istart, iend), item_pos in zip(item_ranges, text_buffer_pos, strict=True):
|
||
# A segment overlaps an item when their character ranges
|
||
# intersect; collect that item's coordinates.
|
||
if seg_start < iend and istart < seg_end:
|
||
seg_pos.extend(item_pos or [])
|
||
chunks.append(
|
||
{
|
||
"text": text,
|
||
"doc_type_kwd": "text",
|
||
"ck_type": "text",
|
||
PDF_POSITIONS_KEY: deepcopy(seg_pos),
|
||
"tk_nums": num_tokens_from_string(text),
|
||
}
|
||
)
|
||
text_buffer.clear()
|
||
text_buffer_pos.clear()
|
||
|
||
for chunk in text_chunks:
|
||
if chunk["ck_type"] == "text":
|
||
text_buffer.append(chunk["text"])
|
||
text_buffer_pos.append(chunk.get(PDF_POSITIONS_KEY))
|
||
else:
|
||
flush_text_buffer()
|
||
chunks.append(chunk)
|
||
flush_text_buffer()
|
||
# Apply children_delimiters (secondary split) before finalizing.
|
||
if custom_pattern:
|
||
chunks = _split_chunk_docs_by_children(chunks, custom_pattern)
|
||
_attach_context_to_media_chunks(chunks, self._param.table_context_size, self._param.image_context_size)
|
||
else:
|
||
# No active delimiter: merge adjacent text items to chunk_token_size.
|
||
# This runs on the per-item chunks (NOT a single concatenated chunk),
|
||
# so the token cap is actually enforced -- matching the previous
|
||
# "token_size" mode and the Go JSON path. Media chunks break the merge.
|
||
# Media context is attached on the per-item chunks before merging, as
|
||
# the removed "token_size" branch did, to preserve context windows.
|
||
_attach_context_to_media_chunks(text_chunks, self._param.table_context_size, self._param.image_context_size)
|
||
chunks = _merge_text_chunks_by_token_size(text_chunks, self._param.chunk_token_size, overlapped_percent)
|
||
if custom_pattern:
|
||
chunks = _split_chunk_docs_by_children(chunks, custom_pattern)
|
||
|
||
await restore_pdf_text_previews(chunks, from_upstream, self._canvas)
|
||
self.set_output("chunks", _finalize_json_chunks(chunks))
|
||
self.callback(1, "Done.")
|
||
return
|