From 01c440be060ee3315fc98669f164f73eaef6a238 Mon Sep 17 00:00:00 2001 From: ccyyy1023 <2637034749@qq.com> Date: Thu, 30 Jul 2026 10:02:34 +0800 Subject: [PATCH 1/3] Support GlotLID language and script labels Signed-off-by: ccyyy1023 <2637034749@qq.com> --- .../language-management/language.mdx | 30 ++++++--- .../text/filters/fasttext/fasttext_filters.py | 24 ++++++-- tests/stages/text/filters/__init__.py | 13 ++++ .../stages/text/filters/fasttext/__init__.py | 13 ++++ .../filters/fasttext/test_fasttext_filters.py | 61 +++++++++++++++++++ 5 files changed, 128 insertions(+), 13 deletions(-) create mode 100644 tests/stages/text/filters/__init__.py create mode 100644 tests/stages/text/filters/fasttext/__init__.py create mode 100644 tests/stages/text/filters/fasttext/test_fasttext_filters.py diff --git a/fern/versions/main/pages/curate-text/process-data/language-management/language.mdx b/fern/versions/main/pages/curate-text/process-data/language-management/language.mdx index 8a4714870a..ed5be56ef1 100644 --- a/fern/versions/main/pages/curate-text/process-data/language-management/language.mdx +++ b/fern/versions/main/pages/curate-text/process-data/language-management/language.mdx @@ -1,7 +1,7 @@ --- -description: "Identify document languages accurately using FastText models supporting 176 languages for multilingual text processing" +description: "Identify document languages and scripts using FastText and GlotLID models for multilingual text processing" categories: ["how-to-guides"] -tags: ["language-identification", "fasttext", "multilingual", "176-languages", "detection", "classification"] +tags: ["language-identification", "fasttext", "glotlid", "multilingual", "detection", "classification"] personas: ["data-scientist-focused", "mle-focused"] difficulty: "intermediate" content_type: "how-to" @@ -18,9 +18,9 @@ NeMo Curator's language identification system works through a three-step process 1. **Text Preprocessing**: For FastText classification, normalize input text by stripping whitespace and converting newlines to spaces. -2. **FastText Language Detection**: The pre-trained FastText language identification model ([`lid.176.bin`](https://fasttext.cc/docs/en/language-identification.html)) analyzes the preprocessed text and returns: +2. **FastText Language Detection**: A FastText-compatible language identification model, such as [`lid.176.bin`](https://fasttext.cc/docs/en/language-identification.html) or [GlotLID](https://huggingface.co/cis-lmu/glotlid), analyzes the preprocessed text and returns: - A confidence score (0.0 to 1.0) indicating certainty of the prediction - - A language code (for example, "EN", "ES", "FR") in FastText's two-letter uppercase format + - A language code (for example, `en`) or a GlotLID language-script code (for example, `eng_Latn`) 3. **Filtering and Scoring**: The pipeline filters documents based on a configurable confidence threshold (`min_langid_score`) and stores both the confidence score and language code as metadata. @@ -30,11 +30,11 @@ The `FastTextLangId` filter implements this workflow by: - Loading the FastText language identification model on worker initialization - Processing text through `model.predict()` with `k=1` to get the top language prediction -- Extracting the language code from FastText labels (for example, `__label__en` becomes "EN") +- Removing the `__label__` prefix while preserving the complete model label (for example, `__label__en` becomes `en` and `__label__eng_Latn` becomes `eng_Latn`) - Comparing confidence scores against the threshold to determine document retention - Returning results as `[confidence_score, language_code]` for downstream processing -This approach supports **176 languages** with high accuracy, making it suitable for large-scale multilingual dataset curation where language-specific processing and monolingual dataset creation are critical. +The standard FastText model supports **176 languages**. GlotLID can be used when broader language coverage or script identification is required. ## Usage @@ -110,13 +110,29 @@ if __name__ == "__main__": +### Using GlotLID + +Download the GlotLID FastText model from [Hugging Face](https://huggingface.co/cis-lmu/glotlid), then pass its local path to `FastTextLangId` in the same way as the standard FastText model: + +```python +glotlid_model_path = "/path/to/glotlid/model.bin" + +# Keep every script predicted for English. +english_filter = FastTextLangId(model_path=glotlid_model_path, lang="eng") + +# Keep only English written in the Latin script. +english_latin_filter = FastTextLangId(model_path=glotlid_model_path, lang="eng_Latn") +``` + +Language matching is case-insensitive. A filter without an underscore matches the language portion of a GlotLID label, while a filter containing an underscore matches the complete language-script label. + ## Understanding Results The language identification process adds a score field to each document batch: 1. **`language` field**: Contains the FastText language identification results as a string representation of a list with two elements (for backend compatibility): - Element 0: The confidence score (between 0 and 1) - - Element 1: The language code in FastText format (for example, "EN" for English, "ES" for Spanish) + - Element 1: The complete model label without the `__label__` prefix (for example, `en` or `eng_Latn`) 2. **Task-based processing**: Curator processes documents in batches (tasks), and results are available through the task's Pandas DataFrame: diff --git a/nemo_curator/stages/text/filters/fasttext/fasttext_filters.py b/nemo_curator/stages/text/filters/fasttext/fasttext_filters.py index 81e8987621..6e6619fdef 100644 --- a/nemo_curator/stages/text/filters/fasttext/fasttext_filters.py +++ b/nemo_curator/stages/text/filters/fasttext/fasttext_filters.py @@ -56,12 +56,20 @@ def keep_document(self, score: float) -> bool: class FastTextLangId(DocumentFilter): + """Identify and optionally filter languages predicted by a FastText model. + + Language labels may be simple codes such as ``en`` or language-script + combinations such as ``eng_Latn``. A language-only filter matches every + script for that language, while a language-script filter requires an exact + match. Matching is case-insensitive. + """ + def __init__(self, model_path: str | None = None, min_langid_score: float = 0.3, lang: str | None = None): if model_path is None: msg = "Must provide a valid path to a FastText model to identify languages with this filter" raise ValueError(msg) self._model_path = model_path - self._lang_code = lang.upper() if lang else None + self._lang_code = lang.casefold() if lang else None self._cutoff = min_langid_score self._name = "lang_id" @@ -73,14 +81,14 @@ def model_check_or_download(self) -> None: def load_model(self) -> None: self._fasttext_langid_model = fasttext.load_model(self._model_path) - def score_document(self, text: str) -> list[float | str]: + def score_document(self, text: str) -> str: # See setup() function in modules/filter.py model = self._fasttext_langid_model pp = text.strip().replace("\n", " ") label, score = model.predict([pp], k=1) score = score[0][0].item() - lang_code = label[0][0][-2:].upper() + lang_code = label[0][0].removeprefix("__label__") # Need to convert it to a string to allow backend conversions return str([score, lang_code]) @@ -89,10 +97,14 @@ def keep_document(self, score: float | str) -> bool: if isinstance(score, str): score_lang = eval(score) # noqa: S307 score = score_lang[0] - lang = score_lang[1].upper() - else : + lang = score_lang[1].casefold() + else: msg = "score must be a string convertible to list" raise TypeError(msg) if self._lang_code: - return score >= self._cutoff and lang == self._lang_code + if "_" in self._lang_code: + language_matches = lang == self._lang_code + else: + language_matches = lang.split("_", maxsplit=1)[0] == self._lang_code + return score >= self._cutoff and language_matches return score >= self._cutoff diff --git a/tests/stages/text/filters/__init__.py b/tests/stages/text/filters/__init__.py new file mode 100644 index 0000000000..4fc25d0d3c --- /dev/null +++ b/tests/stages/text/filters/__init__.py @@ -0,0 +1,13 @@ +# 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. diff --git a/tests/stages/text/filters/fasttext/__init__.py b/tests/stages/text/filters/fasttext/__init__.py new file mode 100644 index 0000000000..4fc25d0d3c --- /dev/null +++ b/tests/stages/text/filters/fasttext/__init__.py @@ -0,0 +1,13 @@ +# 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. diff --git a/tests/stages/text/filters/fasttext/test_fasttext_filters.py b/tests/stages/text/filters/fasttext/test_fasttext_filters.py new file mode 100644 index 0000000000..ae39d0e06e --- /dev/null +++ b/tests/stages/text/filters/fasttext/test_fasttext_filters.py @@ -0,0 +1,61 @@ +# 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. + +from unittest.mock import Mock + +import numpy as np +import pytest + +from nemo_curator.stages.text.filters.fasttext import FastTextLangId + + +@pytest.mark.parametrize( + ("label", "expected_language"), + [ + ("__label__en", "en"), + ("__label__eng_Latn", "eng_Latn"), + ], +) +def test_score_document_preserves_complete_fasttext_label(label: str, expected_language: str) -> None: + lang_id = FastTextLangId(model_path="model.bin") + lang_id._fasttext_langid_model = Mock() + lang_id._fasttext_langid_model.predict.return_value = ([[label]], [np.array([0.9])]) + + assert lang_id.score_document("Hello, world!") == str([0.9, expected_language]) + + +@pytest.mark.parametrize( + ("language_filter", "prediction", "expected"), + [ + ("en", "en", True), + ("EN", "en", True), + ("eng", "eng_Latn", True), + ("eng_Latn", "eng_Latn", True), + ("ENG_LATN", "eng_Latn", True), + ("eng_Cyrl", "eng_Latn", False), + ("deu", "eng_Latn", False), + ], +) +def test_keep_document_filters_language_or_language_script( + language_filter: str, prediction: str, expected: bool +) -> None: + lang_id = FastTextLangId(model_path="model.bin", lang=language_filter) + + assert lang_id.keep_document(str([0.9, prediction])) is expected + + +def test_keep_document_applies_score_cutoff_with_glotlid_label() -> None: + lang_id = FastTextLangId(model_path="model.bin", min_langid_score=0.8, lang="eng") + + assert not lang_id.keep_document(str([0.7, "eng_Latn"])) From d545f4502365ab7f69170e22d9bf2d4a9fc726d7 Mon Sep 17 00:00:00 2001 From: ccyyy1023 <2637034749@qq.com> Date: Thu, 30 Jul 2026 10:23:27 +0800 Subject: [PATCH 2/3] Preserve standard FastText label casing Signed-off-by: ccyyy1023 <2637034749@qq.com> --- .../process-data/language-management/language.mdx | 6 +++--- .../stages/text/filters/fasttext/fasttext_filters.py | 4 +++- tests/stages/text/filters/fasttext/test_fasttext_filters.py | 3 ++- 3 files changed, 8 insertions(+), 5 deletions(-) diff --git a/fern/versions/main/pages/curate-text/process-data/language-management/language.mdx b/fern/versions/main/pages/curate-text/process-data/language-management/language.mdx index ed5be56ef1..918d2b8f1f 100644 --- a/fern/versions/main/pages/curate-text/process-data/language-management/language.mdx +++ b/fern/versions/main/pages/curate-text/process-data/language-management/language.mdx @@ -20,7 +20,7 @@ NeMo Curator's language identification system works through a three-step process 2. **FastText Language Detection**: A FastText-compatible language identification model, such as [`lid.176.bin`](https://fasttext.cc/docs/en/language-identification.html) or [GlotLID](https://huggingface.co/cis-lmu/glotlid), analyzes the preprocessed text and returns: - A confidence score (0.0 to 1.0) indicating certainty of the prediction - - A language code (for example, `en`) or a GlotLID language-script code (for example, `eng_Latn`) + - An uppercase language code (for example, `EN`) or a GlotLID language-script code (for example, `eng_Latn`) 3. **Filtering and Scoring**: The pipeline filters documents based on a configurable confidence threshold (`min_langid_score`) and stores both the confidence score and language code as metadata. @@ -30,7 +30,7 @@ The `FastTextLangId` filter implements this workflow by: - Loading the FastText language identification model on worker initialization - Processing text through `model.predict()` with `k=1` to get the top language prediction -- Removing the `__label__` prefix while preserving the complete model label (for example, `__label__en` becomes `en` and `__label__eng_Latn` becomes `eng_Latn`) +- Removing the `__label__` prefix while preserving standard FastText's uppercase output and complete GlotLID labels (for example, `__label__en` becomes `EN` and `__label__eng_Latn` becomes `eng_Latn`) - Comparing confidence scores against the threshold to determine document retention - Returning results as `[confidence_score, language_code]` for downstream processing @@ -132,7 +132,7 @@ The language identification process adds a score field to each document batch: 1. **`language` field**: Contains the FastText language identification results as a string representation of a list with two elements (for backend compatibility): - Element 0: The confidence score (between 0 and 1) - - Element 1: The complete model label without the `__label__` prefix (for example, `en` or `eng_Latn`) + - Element 1: The standard uppercase language code or complete GlotLID label without the `__label__` prefix (for example, `EN` or `eng_Latn`) 2. **Task-based processing**: Curator processes documents in batches (tasks), and results are available through the task's Pandas DataFrame: diff --git a/nemo_curator/stages/text/filters/fasttext/fasttext_filters.py b/nemo_curator/stages/text/filters/fasttext/fasttext_filters.py index 6e6619fdef..9e678adb6e 100644 --- a/nemo_curator/stages/text/filters/fasttext/fasttext_filters.py +++ b/nemo_curator/stages/text/filters/fasttext/fasttext_filters.py @@ -58,7 +58,7 @@ def keep_document(self, score: float) -> bool: class FastTextLangId(DocumentFilter): """Identify and optionally filter languages predicted by a FastText model. - Language labels may be simple codes such as ``en`` or language-script + Language labels may be simple codes such as ``EN`` or language-script combinations such as ``eng_Latn``. A language-only filter matches every script for that language, while a language-script filter requires an exact match. Matching is case-insensitive. @@ -89,6 +89,8 @@ def score_document(self, text: str) -> str: label, score = model.predict([pp], k=1) score = score[0][0].item() lang_code = label[0][0].removeprefix("__label__") + if "_" not in lang_code: + lang_code = lang_code.upper() # Need to convert it to a string to allow backend conversions return str([score, lang_code]) diff --git a/tests/stages/text/filters/fasttext/test_fasttext_filters.py b/tests/stages/text/filters/fasttext/test_fasttext_filters.py index ae39d0e06e..48f3942b13 100644 --- a/tests/stages/text/filters/fasttext/test_fasttext_filters.py +++ b/tests/stages/text/filters/fasttext/test_fasttext_filters.py @@ -23,7 +23,7 @@ @pytest.mark.parametrize( ("label", "expected_language"), [ - ("__label__en", "en"), + ("__label__en", "EN"), ("__label__eng_Latn", "eng_Latn"), ], ) @@ -39,6 +39,7 @@ def test_score_document_preserves_complete_fasttext_label(label: str, expected_l ("language_filter", "prediction", "expected"), [ ("en", "en", True), + ("en", "EN", True), ("EN", "en", True), ("eng", "eng_Latn", True), ("eng_Latn", "eng_Latn", True), From fa849dbd5074f2d6851082c6ab81eed22f8f91d3 Mon Sep 17 00:00:00 2001 From: ccyyy1023 <2637034749@qq.com> Date: Tue, 4 Aug 2026 10:52:08 +0800 Subject: [PATCH 3/3] Address FastText language label review feedback Signed-off-by: ccyyy1023 <2637034749@qq.com> --- .../language-management/language.mdx | 6 +- .../text/filters/fasttext/fasttext_filters.py | 11 ++- tests/stages/text/filters/__init__.py | 13 --- .../stages/text/filters/fasttext/__init__.py | 13 --- .../filters/fasttext/test_fasttext_filters.py | 89 +++++++++++++++++- tests/stages/text/modules/test_filters.py | 90 ++----------------- 6 files changed, 102 insertions(+), 120 deletions(-) diff --git a/fern/versions/main/pages/curate-text/process-data/language-management/language.mdx b/fern/versions/main/pages/curate-text/process-data/language-management/language.mdx index 918d2b8f1f..aa5e911afc 100644 --- a/fern/versions/main/pages/curate-text/process-data/language-management/language.mdx +++ b/fern/versions/main/pages/curate-text/process-data/language-management/language.mdx @@ -20,7 +20,7 @@ NeMo Curator's language identification system works through a three-step process 2. **FastText Language Detection**: A FastText-compatible language identification model, such as [`lid.176.bin`](https://fasttext.cc/docs/en/language-identification.html) or [GlotLID](https://huggingface.co/cis-lmu/glotlid), analyzes the preprocessed text and returns: - A confidence score (0.0 to 1.0) indicating certainty of the prediction - - An uppercase language code (for example, `EN`) or a GlotLID language-script code (for example, `eng_Latn`) + - The model's language label (for example, `en`) or language-script label (for example, `eng_Latn`) 3. **Filtering and Scoring**: The pipeline filters documents based on a configurable confidence threshold (`min_langid_score`) and stores both the confidence score and language code as metadata. @@ -30,7 +30,7 @@ The `FastTextLangId` filter implements this workflow by: - Loading the FastText language identification model on worker initialization - Processing text through `model.predict()` with `k=1` to get the top language prediction -- Removing the `__label__` prefix while preserving standard FastText's uppercase output and complete GlotLID labels (for example, `__label__en` becomes `EN` and `__label__eng_Latn` becomes `eng_Latn`) +- Removing the `__label__` prefix while preserving the model label's original casing (for example, `__label__en` becomes `en` and `__label__eng_Latn` becomes `eng_Latn`) - Comparing confidence scores against the threshold to determine document retention - Returning results as `[confidence_score, language_code]` for downstream processing @@ -132,7 +132,7 @@ The language identification process adds a score field to each document batch: 1. **`language` field**: Contains the FastText language identification results as a string representation of a list with two elements (for backend compatibility): - Element 0: The confidence score (between 0 and 1) - - Element 1: The standard uppercase language code or complete GlotLID label without the `__label__` prefix (for example, `EN` or `eng_Latn`) + - Element 1: The model label without the `__label__` prefix and with its original casing preserved (for example, `en` or `eng_Latn`) 2. **Task-based processing**: Curator processes documents in batches (tasks), and results are available through the task's Pandas DataFrame: diff --git a/nemo_curator/stages/text/filters/fasttext/fasttext_filters.py b/nemo_curator/stages/text/filters/fasttext/fasttext_filters.py index 9e678adb6e..8bee989416 100644 --- a/nemo_curator/stages/text/filters/fasttext/fasttext_filters.py +++ b/nemo_curator/stages/text/filters/fasttext/fasttext_filters.py @@ -69,7 +69,7 @@ def __init__(self, model_path: str | None = None, min_langid_score: float = 0.3, msg = "Must provide a valid path to a FastText model to identify languages with this filter" raise ValueError(msg) self._model_path = model_path - self._lang_code = lang.casefold() if lang else None + self._lang_code = lang or None self._cutoff = min_langid_score self._name = "lang_id" @@ -89,8 +89,6 @@ def score_document(self, text: str) -> str: label, score = model.predict([pp], k=1) score = score[0][0].item() lang_code = label[0][0].removeprefix("__label__") - if "_" not in lang_code: - lang_code = lang_code.upper() # Need to convert it to a string to allow backend conversions return str([score, lang_code]) @@ -104,9 +102,10 @@ def keep_document(self, score: float | str) -> bool: msg = "score must be a string convertible to list" raise TypeError(msg) if self._lang_code: - if "_" in self._lang_code: - language_matches = lang == self._lang_code + lang_filter = self._lang_code.casefold() + if "_" in lang_filter: + language_matches = lang == lang_filter else: - language_matches = lang.split("_", maxsplit=1)[0] == self._lang_code + language_matches = lang.split("_", maxsplit=1)[0] == lang_filter return score >= self._cutoff and language_matches return score >= self._cutoff diff --git a/tests/stages/text/filters/__init__.py b/tests/stages/text/filters/__init__.py index 4fc25d0d3c..e69de29bb2 100644 --- a/tests/stages/text/filters/__init__.py +++ b/tests/stages/text/filters/__init__.py @@ -1,13 +0,0 @@ -# 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. diff --git a/tests/stages/text/filters/fasttext/__init__.py b/tests/stages/text/filters/fasttext/__init__.py index 4fc25d0d3c..e69de29bb2 100644 --- a/tests/stages/text/filters/fasttext/__init__.py +++ b/tests/stages/text/filters/fasttext/__init__.py @@ -1,13 +0,0 @@ -# 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. diff --git a/tests/stages/text/filters/fasttext/test_fasttext_filters.py b/tests/stages/text/filters/fasttext/test_fasttext_filters.py index 48f3942b13..7803eb6da8 100644 --- a/tests/stages/text/filters/fasttext/test_fasttext_filters.py +++ b/tests/stages/text/filters/fasttext/test_fasttext_filters.py @@ -15,15 +15,84 @@ from unittest.mock import Mock import numpy as np +import pandas as pd import pytest +from nemo_curator.stages.text.filters import DocumentFilter, ScoreFilter from nemo_curator.stages.text.filters.fasttext import FastTextLangId +from nemo_curator.tasks import DocumentBatch + + +class FakeQualityFilter(DocumentFilter): + """Emulate ``FastTextQualityFilter`` without loading a model.""" + + def __init__(self, alpha: float = 3, seed: int = 42): + super().__init__() + self._alpha = alpha + self._seed = np.random.seed(seed) # noqa: NPY002 + + def load_model(self) -> None: + pass + + def score_document(self, text: str) -> float: + scores = {"a": 0.00, "b": 0.25, "c": 0.50, "d": 0.75} + try: + return scores[text] + except KeyError: + msg = f"Unexpected text: {text}" + raise ValueError(msg) from None + + def keep_document(self, score: float) -> bool: + return np.random.pareto(self._alpha) > 1 - score # noqa: NPY002 + + +class FakeLangId(DocumentFilter): + """Emulate ``FastTextLangId`` without loading a model.""" + + def __init__(self, min_langid_score: float = 0.3): + super().__init__() + self._cutoff = min_langid_score + + def load_model(self) -> None: + pass + + def score_document(self, text: str) -> str: + scores = { + "a": [0.5, "EN"], + "b": [0.7, "HI"], + "c": [0.2, "PT"], + "d": [0.5, "EN"], + } + try: + return str(scores[text]) + except KeyError: + msg = f"Unexpected text: {text}" + raise ValueError(msg) from None + + def keep_document(self, score: float | str) -> bool: + if isinstance(score, str): + score = eval(score) # noqa: S307 + + return score[0] >= self._cutoff + + +def list_to_dataset(documents: list[str]) -> DocumentBatch: + return DocumentBatch(data=pd.DataFrame({"text": documents}), dataset_name="test_1") + + +def assert_datasets_equal(expected: DocumentBatch, actual: DocumentBatch) -> None: + pd.testing.assert_frame_equal( + expected.to_pandas().reset_index(drop=True), + actual.to_pandas().reset_index(drop=True), + ) + assert actual.dataset_name == expected.dataset_name @pytest.mark.parametrize( ("label", "expected_language"), [ - ("__label__en", "EN"), + ("__label__en", "en"), + ("__label__EN", "EN"), ("__label__eng_Latn", "eng_Latn"), ], ) @@ -60,3 +129,21 @@ def test_keep_document_applies_score_cutoff_with_glotlid_label() -> None: lang_id = FastTextLangId(model_path="model.bin", min_langid_score=0.8, lang="eng") assert not lang_id.keep_document(str([0.7, "eng_Latn"])) + + +def test_fake_quality_filter_pipeline() -> None: + dataset = list_to_dataset(["a", "b", "c", "d"]) + + filtered_data = ScoreFilter(FakeQualityFilter()).process(dataset) + + expected_data = list_to_dataset(["b", "c", "d"]) + assert_datasets_equal(expected_data, filtered_data) + + +def test_fake_langid_filter_pipeline() -> None: + dataset = list_to_dataset(["a", "b", "c", "d"]) + + filtered_data = ScoreFilter(FakeLangId()).process(dataset) + + expected_data = list_to_dataset(["a", "b", "d"]) + assert_datasets_equal(expected_data, filtered_data) diff --git a/tests/stages/text/modules/test_filters.py b/tests/stages/text/modules/test_filters.py index 486b43706f..ed985cd7a0 100644 --- a/tests/stages/text/modules/test_filters.py +++ b/tests/stages/text/modules/test_filters.py @@ -15,7 +15,6 @@ import os import re -import numpy as np import pandas as pd import pytest @@ -87,64 +86,17 @@ def encode(self, text: str) -> list[str]: return text.split() -class FakeQualityFilter(DocumentFilter): - """ - Emulates FastTextQualityFilter without a model - """ - - def __init__(self, alpha: float = 3, seed: int = 42): - super().__init__() - self._alpha = alpha - self._seed = np.random.seed(seed) # noqa: NPY002 +class FakeModelFilter(DocumentFilter): + """Minimal model-backed filter used to test actor-stage detection.""" def load_model(self) -> None: pass def score_document(self, text: str) -> float: - if text == "a": - return 0.00 - elif text == "b": - return 0.25 - elif text == "c": - return 0.50 - elif text == "d": - return 0.75 - else: - msg = f"Unexpected text: {text}" - raise ValueError(msg) + return float(bool(text)) def keep_document(self, score: float) -> bool: - return np.random.pareto(self._alpha) > 1 - score # noqa: NPY002 - - -class FakeLangId(DocumentFilter): - """ - Emulates FastTextLangId without a model - """ - - def __init__(self, min_langid_score: float = 0.3): - super().__init__() - self._cutoff = min_langid_score - - def load_model(self) -> None: - pass - - def score_document(self, text: str) -> str: - if text in ["a", "d"]: - return str([0.5, "EN"]) - if text == "b": - return str([0.7, "HI"]) - if text == "c": - return str([0.2, "PT"]) - else: - msg = f"Unexpected text: {text}" - raise ValueError(msg) - - def keep_document(self, score: float | str) -> bool: - if isinstance(score, str): - score = eval(score) # noqa: S307 - - return score[0] >= self._cutoff + return bool(score) def all_equal(left_dataset: DocumentBatch, right_dataset: DocumentBatch) -> bool: @@ -451,13 +403,9 @@ def test_ray_stage_spec(self) -> None: assert test_filter.ray_stage_spec() == {"is_actor_stage": False} # Has load_model - test_filter = ScoreFilter(FakeQualityFilter(), text_field="documents") - assert test_filter.ray_stage_spec() == {"is_actor_stage": True} - test_filter = Score(FakeQualityFilter(), text_field="documents", score_field="score") + test_filter = ScoreFilter(FakeModelFilter(), text_field="documents") assert test_filter.ray_stage_spec() == {"is_actor_stage": True} - test_filter = ScoreFilter(FakeLangId(), text_field="documents") - assert test_filter.ray_stage_spec() == {"is_actor_stage": True} - test_filter = Score(FakeLangId(), text_field="documents", score_field="score") + test_filter = Score(FakeModelFilter(), text_field="documents", score_field="score") assert test_filter.ray_stage_spec() == {"is_actor_stage": True} # Has load_tokenizer @@ -1288,29 +1236,3 @@ def test_line_statistics( ) -> None: line_statistics = per_extension_filter._line_statistics(content) assert line_statistics == expected, f"Expected {expected} but got {line_statistics}" - - -class TestClassifierFilters: - def test_fake_quality_filter(self) -> None: - dataset = list_to_dataset(["a", "b", "c", "d"]) - filters = ScoreFilter(FakeQualityFilter()) - - filtered_data = filters.process(dataset) - - expected_data = DocumentBatch( - data=pd.DataFrame({"text": ["b", "c", "d"]}), - dataset_name="test_1", - ) - assert all_equal(expected_data, filtered_data), f"Expected {expected_data} but got {filtered_data}" - - def test_fake_langid_filter(self) -> None: - dataset = list_to_dataset(["a", "b", "c", "d"]) - filters = ScoreFilter(FakeLangId()) - - filtered_data = filters.process(dataset) - - expected_data = DocumentBatch( - data=pd.DataFrame({"text": ["a", "b", "d"]}), - dataset_name="test_1", - ) - assert all_equal(expected_data, filtered_data), f"Expected {expected_data} but got {filtered_data}"