diff --git a/nemo_curator/stages/text/utils/text_utils.py b/nemo_curator/stages/text/utils/text_utils.py index 5a7edbdc95..d8130ab772 100644 --- a/nemo_curator/stages/text/utils/text_utils.py +++ b/nemo_curator/stages/text/utils/text_utils.py @@ -92,12 +92,8 @@ def is_paragraph_indices_in_top_or_bottom_only( boilerplate_paragraph_indices: list[int], num_paragraphs: int, ) -> bool: - def _is_contiguous(indices: list[int]) -> bool: - # Indices are sorted in ascending order. - num_indices = len(indices) - 1 - return all(indices[i] + 1 == indices[i + 1] for i in range(num_indices)) - - # See if the indices are contiguous and exclusively at the top/bottom. + # See if the indices form a contiguous run at the top, a contiguous run + # at the bottom, or both, with nothing in between. # Indices are sorted in ascending order. # If num_paragraphs = 11: # Valid indices example : [0, 1, 9, 10] @@ -106,9 +102,11 @@ def _is_contiguous(indices: list[int]) -> bool: # Invalid indices example : [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10] if len(boilerplate_paragraph_indices) == num_paragraphs: return False - return _is_contiguous(boilerplate_paragraph_indices) and ( - boilerplate_paragraph_indices[0] == 0 or boilerplate_paragraph_indices[-1] == num_paragraphs - 1 - ) + num_top = 0 + while num_top < len(boilerplate_paragraph_indices) and boilerplate_paragraph_indices[num_top] == num_top: + num_top += 1 + bottom_indices = boilerplate_paragraph_indices[num_top:] + return bottom_indices == list(range(num_paragraphs - len(bottom_indices), num_paragraphs)) # Node types for processing abstract syntax tree diff --git a/tests/stages/text/modules/test_modifiers.py b/tests/stages/text/modules/test_modifiers.py index 17e097d4ef..a9ec133f2e 100644 --- a/tests/stages/text/modules/test_modifiers.py +++ b/tests/stages/text/modules/test_modifiers.py @@ -18,6 +18,7 @@ from nemo_curator.stages.text.modifiers import DocumentModifier, Modify from nemo_curator.stages.text.modifiers.modifier import _normalize_input_fields, _normalize_output_fields from nemo_curator.stages.text.modifiers.string import ( + BoilerPlateStringModifier, LineRemover, MarkdownRemover, NewlineNormalizer, @@ -268,6 +269,34 @@ def test_dataset_modification(self) -> None: pd.testing.assert_frame_equal(output.data.reset_index(drop=True), expected_df.reset_index(drop=True)) +class TestBoilerPlateStringModifier: + def test_removes_boilerplate_at_top_and_bottom(self) -> None: + text = "We use cookies. See our privacy policy.\n\nFirst paragraph.\n\nSecond paragraph.\n\nTerms of use." + result = BoilerPlateStringModifier().modify_document(text) + assert result == "First paragraph.\n\nSecond paragraph." + + def test_removes_boilerplate_at_bottom_only(self) -> None: + text = "First paragraph.\n\nSecond paragraph.\n\nTerms of use." + result = BoilerPlateStringModifier().modify_document(text) + assert result == "First paragraph.\n\nSecond paragraph." + + def test_keeps_document_with_boilerplate_in_middle(self) -> None: + text = "Terms of use.\n\nFirst paragraph.\n\nPrivacy policy.\n\nSecond paragraph.\n\nTerms of use." + result = BoilerPlateStringModifier().modify_document(text) + assert result == text + + def test_dataset_modification(self) -> None: + docs = [ + "Cookie policy.\n\nKeep me.\n\nTerms of use.", + "Keep me.\n\nAnd me.", + ] + expected_results = ["Keep me.", "Keep me.\n\nAnd me."] + doc_batch = list_to_doc_batch(docs) + output = run_modify(BoilerPlateStringModifier(), doc_batch) + expected_df = pd.DataFrame({"text": expected_results}) + pd.testing.assert_frame_equal(output.data.reset_index(drop=True), expected_df.reset_index(drop=True)) + + class TestSlicer: def test_integer_indices(self) -> None: text = "Hello, world!" diff --git a/tests/stages/text/utils/test_text_utils.py b/tests/stages/text/utils/test_text_utils.py new file mode 100644 index 0000000000..8ecc8a318d --- /dev/null +++ b/tests/stages/text/utils/test_text_utils.py @@ -0,0 +1,41 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. 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 nemo_curator.stages.text.utils.text_utils import is_paragraph_indices_in_top_or_bottom_only + + +class TestIsParagraphIndicesInTopOrBottomOnly: + @pytest.mark.parametrize( + ("indices", "expected"), + [ + ([0], True), + ([0, 1], True), + ([10], True), + ([9, 10], True), + ([0, 1, 9, 10], True), + ([0, 10], True), + ([5], False), + ([1, 2], False), + ([8, 9], False), + ([0, 1, 3, 9, 10], False), + ([0, 1, 3, 5, 6, 9, 10], False), + ([0, 2], False), + ([8, 10], False), + (list(range(11)), False), + ], + ) + def test_indices(self, indices: list[int], expected: bool) -> None: + assert is_paragraph_indices_in_top_or_bottom_only(indices, 11) is expected