diff --git a/rag/nlp/__init__.py b/rag/nlp/__init__.py index 5d84231d28..cbcc381797 100644 --- a/rag/nlp/__init__.py +++ b/rag/nlp/__init__.py @@ -417,7 +417,7 @@ def tokenize_chunks(chunks, doc, eng, pdf_parser=None, child_delimiters_pattern= add_positions(d, [[ii] * 5]) if child_delimiters_pattern: - d["mom_with_weight"] = ck + d["mom_with_weight"] = ck.removeprefix("\n") res.extend(split_with_pattern(d, child_delimiters_pattern, ck, eng, language=language)) continue @@ -440,7 +440,7 @@ def doc_tokenize_chunks_with_images(chunks, doc, eng, child_delimiters_pattern=N if ck.get("ck_type") == "text": if child_delimiters_pattern: - d["mom_with_weight"] = text + d["mom_with_weight"] = text.removeprefix("\n") res.extend(split_with_pattern(d, child_delimiters_pattern, text, eng, language=language)) continue elif ck.get("ck_type") == "image": @@ -463,7 +463,7 @@ def tokenize_chunks_with_images(chunks, doc, eng, images, child_delimiters_patte d["image"] = image add_positions(d, [[ii] * 5]) if child_delimiters_pattern: - d["mom_with_weight"] = ck + d["mom_with_weight"] = ck.removeprefix("\n") res.extend(split_with_pattern(d, child_delimiters_pattern, ck, eng, language=language)) continue tokenize(d, ck, eng, language=language) diff --git a/test/unit_test/rag/nlp/test_tokenize_chunks.py b/test/unit_test/rag/nlp/test_tokenize_chunks.py new file mode 100644 index 0000000000..594305d708 --- /dev/null +++ b/test/unit_test/rag/nlp/test_tokenize_chunks.py @@ -0,0 +1,59 @@ +# +# Copyright 2026 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 pytest + +from rag import nlp + + +@pytest.fixture(autouse=True) +def stub_tokenize(monkeypatch): + def tokenize(doc, text, eng, language="English"): + doc["content_with_weight"] = text + + monkeypatch.setattr(nlp, "tokenize", tokenize) + + +@pytest.mark.p2 +@pytest.mark.parametrize( + ("chunk", "parent"), + [ + ("first\nsecond", "first\nsecond"), + ("\nfirst\nsecond", "first\nsecond"), + ("\n\nfirst\nsecond", "\nfirst\nsecond"), + ], +) +def test_tokenize_chunks_removes_one_leading_newline_from_parent(chunk, parent): + docs = nlp.tokenize_chunks([chunk], {}, True, child_delimiters_pattern="\n") + + assert [doc["mom_with_weight"] for doc in docs] == [parent, parent] + assert [doc["content_with_weight"] for doc in docs] == ["first\n", "second"] + + +@pytest.mark.p2 +def test_doc_tokenize_chunks_with_images_removes_leading_newline_from_parent(): + docs = nlp.doc_tokenize_chunks_with_images([{"text": "\nfirst\nsecond", "ck_type": "text"}], {}, True, child_delimiters_pattern="\n") + + assert [doc["mom_with_weight"] for doc in docs] == ["first\nsecond", "first\nsecond"] + assert [doc["content_with_weight"] for doc in docs] == ["first\n", "second"] + + +@pytest.mark.p2 +def test_tokenize_chunks_with_images_removes_leading_newline_from_parent(): + docs = nlp.tokenize_chunks_with_images(["\nfirst\nsecond"], {}, True, [object()], child_delimiters_pattern="\n") + + assert [doc["mom_with_weight"] for doc in docs] == ["first\nsecond", "first\nsecond"] + assert [doc["content_with_weight"] for doc in docs] == ["first\n", "second"]