Add Math pipeline

examples/math/README.md

@@ -0,0 +1,65 @@
+# Math Quality Classifier
+
+This example demonstrates running a model-based math classifier (FineMath)
+
+## Install
+Use uv to create the project environment and install Curator with the text extra:
+
+```bash
+uv sync --extra text
+source .venv/bin/activate
+```
+
+- GPU detection: if `nvidia-smi` shows GPUs but the example logs "No gpus found", install `pynvml` so the backend can discover GPUs:
+
+```bash
+pip install pynvml
+```
+
+## Prerequisites
+- GPU(s) with CUDA for the HF model
+- Python environment with `nemo-curator` installed (uv sync above)
+- Lynx system dependency for HTML rendering to text:
+  - Ubuntu/Debian: `sudo apt-get update && sudo apt-get install -y lynx`
+  - RHEL/Fedora: `sudo dnf install -y lynx` (or `sudo yum install -y lynx`)
+  - Conda: `conda install -c conda-forge lynx`
+
+## Text preprocessing (decode → type-detect → extract)
+
+```bash
+python examples/math/run_text_preprocess.py \
+  --input "examples/math/data/*.parquet" \
+  --output /tmp/math_mock/preprocessed_jsonl
+
+# Optional: Add --report-stats to see extraction statistics
+python examples/math/run_text_preprocess.py \
+  --input "examples/math/data/*.parquet" \
+  --output /tmp/math_mock/preprocessed_jsonl \
+  --report-stats
+```
+
+- Parquet files include columns: `binary_content` (bytes), `url`, `mime_type`.
+- Output JSONL will include `text`, `url`, and `type`.
+
+## Run the classifier pipeline
+Run the pipeline that reads JSONL, classifies with the FineMath model, and writes JSONL outputs:
+
+```bash
+python examples/math/run_quality_classifier.py \
+  --input "/tmp/math_mock/preprocessed_jsonl/*.jsonl" \
+  --output /tmp/math_mock/classified_output
+```
+
+Outputs will be written as JSONL files under `/tmp/math_mock/classified_output/` with columns:
+- `finemath_scores`: float scores (0..5)
+- `finemath_int_scores`: integer scores (0..5)
+
+Output
+```
+{"id":0,"text":"The derivative of x^2 is 2x.","finemath_scores":1.6865234375,"finemath_int_scores":2}
+{"id":1,"text":"This is plain English without math.","finemath_scores":0.9130859375,"finemath_int_scores":1}
+{"id":2,"text":"Let $f(x)=x^2$. Then $f'(x)=2x.","finemath_scores":2.291015625,"finemath_int_scores":2}
+{"id":3,"text":"We have $$\\int_0^1 x^2 dx = 1\/3.$$.","finemath_scores":1.9150390625,"finemath_int_scores":2}
+{"id":4,"text":"Using \\(a^2+b^2=c^2\\) we derive the relation.","finemath_scores":1.93359375,"finemath_int_scores":2}
+{"id":5,"text":"Consider the set A \\subseteq B and A \\in \\mathbb{R}^n.","finemath_scores":1.5458984375,"finemath_int_scores":2}
+```

examples/math/run_quality_classifier.py

@@ -0,0 +1,86 @@
+# Copyright (c) 2025, 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 argparse
+
+from loguru import logger
+
+from nemo_curator.backends.xenna import XennaExecutor
+from nemo_curator.core.client import RayClient
+from nemo_curator.pipeline import Pipeline
+from nemo_curator.stages.math import FineMathClassifier
+from nemo_curator.stages.resources import Resources
+from nemo_curator.stages.text.io.reader import JsonlReader
+from nemo_curator.stages.text.io.writer import JsonlWriter
+
+
+def build_pipeline(input_glob: str, output_dir: str) -> Pipeline:
+    p = Pipeline(name="math_quality_classifier", description="...")
+
+    # Reader (composite): tune both file partitioning and jsonl_reader stages
+    p.add_stage(
+        JsonlReader(file_paths=input_glob).with_(
+            {
+                "file_partitioning": {"resources": Resources(cpus=0.5)},
+                "jsonl_reader": {"resources": Resources(cpus=0.5)},
+            }
+        )
+    )
+
+    # Classifier (composite): tune tokenizer and model sub-stages
+    p.add_stage(
+        FineMathClassifier().with_(
+            {
+                "finemath_classifier_tokenizer": {"resources": Resources(cpus=0.5)},
+                "finemath_classifier_model": {"resources": Resources(cpus=1, gpus=1)},
+            }
+        )
+    )
+
+    # Writer (single stage)
+    p.add_stage(JsonlWriter(path=output_dir).with_(resources=Resources(cpus=0.5)))
+
+    return p
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="Run Math (FineMath) classifier on JSONL files",
+    )
+    parser.add_argument(
+        "--input",
+        required=True,
+        help="Glob or directory for JSONL input files",
+    )
+    parser.add_argument(
+        "--output",
+        required=True,
+        help="Output directory for JSONL results",
+    )
+    args = parser.parse_args()
+
+    ray_client = RayClient()
+    ray_client.start()
+
+    pipeline = build_pipeline(args.input, args.output)
+    logger.info(pipeline.describe())
+
+    executor = XennaExecutor()
+    pipeline.run(executor)
+
+    ray_client.stop()
+
+
+if __name__ == "__main__":
+    main()

examples/math/run_text_preprocess.py

@@ -0,0 +1,96 @@
+# Copyright (c) 2025, 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 argparse
+
+import ray.data
+from loguru import logger
+
+from nemo_curator.backends.xenna import XennaExecutor
+from nemo_curator.core.client import RayClient
+from nemo_curator.pipeline import Pipeline
+from nemo_curator.stages.math import MathContentExtractor
+from nemo_curator.stages.resources import Resources
+from nemo_curator.stages.text.download.base.extract import DocumentExtractStage
+from nemo_curator.stages.text.io.reader import ParquetReader
+from nemo_curator.stages.text.io.writer import JsonlWriter
+
+
+def build_pipeline(input_glob: str, output_dir: str) -> Pipeline:
+    p = Pipeline(name="math_text_preprocess", description="Decode (binary) → type → html via lynx → text")
+
+    p.add_stage(
+        ParquetReader(file_paths=input_glob).with_(
+            {
+                "file_partitioning": {"resources": Resources(cpus=0.1)},
+                "parquet_reader": {"resources": Resources(cpus=0.1)},
+            }
+        )
+    )
+
+    p.add_stage(
+        DocumentExtractStage(extractor=MathContentExtractor(), add_filename_column=False).with_(
+            resources=Resources(cpus=1)
+        )
+    )
+
+    p.add_stage(JsonlWriter(path=output_dir).with_(resources=Resources(cpus=0.1)))
+
+    return p
+
+
+def report_extraction_stats(output_dir: str) -> None:
+    """Optional: Report extraction statistics by reading output with Ray Data."""
+    try:
+        ds = ray.data.read_json(f"{output_dir}/*.jsonl")
+        total = ds.count()
+        html_docs = ds.filter(lambda row: row.get("type") == "html").count()
+        html_failed = ds.filter(
+            lambda row: row.get("type") == "html" and (not row.get("text") or row.get("text").strip() == "")
+        ).count()
+
+        logger.info(
+            f"Extraction stats: {total} total documents, {html_docs} HTML, {html_failed} HTML extraction failures"
+        )
+    except Exception as e:  # noqa: BLE001
+        logger.debug(f"Could not compute stats (optional): {e}")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Run math text preprocessing on Parquet files")
+    parser.add_argument("--input", required=True, help="Glob or directory for Parquet input files")
+    parser.add_argument("--output", required=True, help="Output directory for JSONL results")
+    parser.add_argument("--report-stats", action="store_true", help="Report extraction statistics after processing")
+    args = parser.parse_args()
+
+    ray_client = RayClient()
+    ray_client.start()
+
+    pipeline = build_pipeline(args.input, args.output)
+    logger.info(pipeline.describe())
+
+    executor = XennaExecutor()
+    pipeline.run(executor)
+
+    logger.info("Pipeline completed successfully.")
+
+    # Optional: Report extraction statistics
+    if args.report_stats:
+        report_extraction_stats(args.output)
+
+    ray_client.stop()
+
+
+if __name__ == "__main__":
+    main()

nemo_curator/stages/math/__init__.py

@@ -0,0 +1,21 @@
+# Copyright (c) 2025, 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 .classifiers.finemath import FineMathClassifier
+from .download.extract import MathContentExtractor
+
+__all__ = [
+    "FineMathClassifier",
+    "MathContentExtractor",
+]

nemo_curator/stages/math/classifiers/__init__.py

@@ -0,0 +1,19 @@
+# Copyright (c) 2025, 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 .finemath import FineMathClassifier
+
+__all__ = [
+    "FineMathClassifier",
+]