diff --git a/.github/workflows/cicd-main.yml b/.github/workflows/cicd-main.yml index 75e3ec93c2..d01ae03287 100644 --- a/.github/workflows/cicd-main.yml +++ b/.github/workflows/cicd-main.yml @@ -63,7 +63,7 @@ jobs: matrix: os: [ubuntu-latest] python-version: ["3.10", "3.12"] - folder: ["backends", "config", "core", "models", "pipelines", "stages-audio", "stages-common", "stages-deduplication", "stages-image", "stages-math_stages", "stages-synthetic", "stages-text", "stages-video", "tasks", "utils"] + folder: ["backends", "config", "core", "models", "pipelines", "stages-audio", "stages-common", "stages-deduplication", "stages-image", "stages-interleaved", "stages-math_stages", "stages-synthetic", "stages-text", "stages-video", "tasks", "utils"] needs: [pre-flight, cicd-wait-in-queue] runs-on: ${{ matrix.os }} name: Unit_Test_${{ matrix.folder}}_CPU_python-${{ matrix.python-version }} diff --git a/benchmarking/nightly-benchmark.yaml b/benchmarking/nightly-benchmark.yaml index ab0b66f3e0..4d71c0f352 100644 --- a/benchmarking/nightly-benchmark.yaml +++ b/benchmarking/nightly-benchmark.yaml @@ -947,6 +947,62 @@ entries: num_gpus: 0 enable_object_spilling: false + - name: interleaved_filter_xenna + enabled: false + script: interleaved_filter_benchmark.py + args: >- + --benchmark-results-path={session_entry_dir} + --executor=xenna + --input-path={dataset:multimodal_mint1t,wds_tar_dir} + --output-path={session_entry_dir}/scratch/output + --clip-model-dir={dataset:mscoco_model_weights,files} + --files-per-partition=2 + --mode=overwrite + --no-materialize-on-write + --output-max-batch-bytes=2000000 + timeout_s: 1800 + sink_data: + - name: slack + additional_metrics: + - throughput_rows_per_sec + - num_rows + - num_output_files + - materialize_error_count + ping_on_failure: + - UGFRT88RE + ray: + num_cpus: 64 + num_gpus: 4 + enable_object_spilling: false + + - name: interleaved_filter_xenna_materialize + enabled: false + script: interleaved_filter_benchmark.py + args: >- + --benchmark-results-path={session_entry_dir} + --executor=xenna + --input-path={dataset:multimodal_mint1t,wds_tar_dir} + --output-path={session_entry_dir}/scratch/output + --clip-model-dir={dataset:mscoco_model_weights,files} + --files-per-partition=2 + --mode=overwrite + --materialize-on-write + --output-max-batch-bytes=2000000 + timeout_s: 1800 + sink_data: + - name: slack + additional_metrics: + - throughput_rows_per_sec + - num_rows + - num_output_files + - materialize_error_count + ping_on_failure: + - UGFRT88RE + ray: + num_cpus: 64 + num_gpus: 4 + enable_object_spilling: false + - name: alm_pipeline_xenna enabled: true script: alm_pipeline_benchmark.py diff --git a/benchmarking/scripts/interleaved_filter_benchmark.py b/benchmarking/scripts/interleaved_filter_benchmark.py new file mode 100644 index 0000000000..d7fb80baa9 --- /dev/null +++ b/benchmarking/scripts/interleaved_filter_benchmark.py @@ -0,0 +1,291 @@ +# 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. + +"""Benchmark for interleaved row-wise filters covered by tests/stages/interleaved/filter/.""" + +import argparse +import time +import traceback +from pathlib import Path +from typing import Any + +from loguru import logger +from utils import ( + collect_interleaved_parquet_metrics, + collect_interleaved_wds_metrics, + setup_executor, + write_benchmark_results, +) + +from nemo_curator.pipeline import Pipeline +from nemo_curator.stages.interleaved.filter import ( + InterleavedBlurFilterStage, + InterleavedCLIPScoreFilterStage, + InterleavedImageToTextRatioFilterStage, + InterleavedQRCodeFilterStage, +) +from nemo_curator.stages.interleaved.io import ( + InterleavedParquetReader, + InterleavedParquetWriterStage, + InterleavedWebdatasetReader, + InterleavedWebdatasetWriterStage, +) +from nemo_curator.tasks.utils import TaskPerfUtils + + +def create_pipeline(args: argparse.Namespace) -> Pipeline: + pipeline = Pipeline( + name="interleaved_filter_benchmark", + description=("Benchmark: interleaved reader -> blur, QR, CLIP score, image/text ratio filters -> writer"), + ) + + # ── Reader ──────────────────────────────────────────────────────────────── + if args.reader_type == "wds": + pipeline.add_stage( + InterleavedWebdatasetReader( + file_paths=args.input_path, + files_per_partition=args.files_per_partition, + blocksize=args.input_blocksize, + max_batch_bytes=args.output_max_batch_bytes, + materialize_on_read=args.materialize_on_read, + per_image_fields=tuple(args.per_image_fields) if args.per_image_fields else (), + per_text_fields=tuple(args.per_text_fields) if args.per_text_fields else (), + ) + ) + else: + pipeline.add_stage( + InterleavedParquetReader( + file_paths=args.input_path, + files_per_partition=args.files_per_partition, + blocksize=args.input_blocksize, + max_batch_bytes=args.output_max_batch_bytes, + fields=tuple(args.reader_fields) if args.reader_fields else None, + ) + ) + + # ── Filters ──────────────────────────────────────────────────────────────── + max_ratio = float("inf") if args.image_text_max_ratio is None else args.image_text_max_ratio + pipeline.add_stage( + InterleavedBlurFilterStage( + drop_invalid_rows=args.drop_invalid_rows, + score_threshold=args.blur_score_threshold, + ) + ) + pipeline.add_stage( + InterleavedQRCodeFilterStage( + drop_invalid_rows=args.drop_invalid_rows, + score_threshold=args.qrcode_score_threshold, + ) + ) + pipeline.add_stage( + InterleavedCLIPScoreFilterStage( + drop_invalid_rows=args.drop_invalid_rows, + model_dir=args.clip_model_dir, + min_score=args.clip_min_score, + ) + ) + pipeline.add_stage( + InterleavedImageToTextRatioFilterStage( + drop_invalid_rows=args.drop_invalid_rows, + min_ratio=args.image_text_min_ratio, + max_ratio=max_ratio, + ) + ) + + # ── Writer ──────────────────────────────────────────────────────────────── + if args.writer_format == "wds": + pipeline.add_stage( + InterleavedWebdatasetWriterStage( + path=args.output_path, + materialize_on_write=args.materialize_on_write, + on_materialize_error=args.on_materialize_error, + mode=args.mode, + ) + ) + else: + write_kwargs: dict[str, Any] = {} + if args.parquet_row_group_size is not None: + write_kwargs["row_group_size"] = args.parquet_row_group_size + if args.parquet_compression is not None: + write_kwargs["compression"] = args.parquet_compression + pipeline.add_stage( + InterleavedParquetWriterStage( + path=args.output_path, + materialize_on_write=args.materialize_on_write, + on_materialize_error=args.on_materialize_error, + write_kwargs=write_kwargs, + mode=args.mode, + ) + ) + + return pipeline + + +def run_benchmark(args: argparse.Namespace) -> dict[str, Any]: + executor = setup_executor(args.executor) + input_path = str(Path(args.input_path).absolute()) + output_path = Path(args.output_path).absolute() + output_path.mkdir(parents=True, exist_ok=True) + + input_metrics_start = time.perf_counter() + collect_fn = ( + collect_interleaved_parquet_metrics if args.reader_type == "parquet" else collect_interleaved_wds_metrics + ) + input_metrics = {f"input_{k}": v for k, v in collect_fn(args.input_path).items()} + input_metrics_elapsed = time.perf_counter() - input_metrics_start + logger.info("collect_input_metrics took {:.3f}s", input_metrics_elapsed) + + start = time.perf_counter() + output_tasks = [] + success = False + try: + pipeline = create_pipeline(args) + logger.info("Pipeline:\n{}", pipeline.describe()) + output_tasks = pipeline.run(executor) + success = True + except Exception as e: + logger.error("Benchmark failed: {}", e) + logger.debug(traceback.format_exc()) + + elapsed = time.perf_counter() - start + metrics_start = time.perf_counter() + collect_fn = ( + collect_interleaved_parquet_metrics if args.writer_format == "parquet" else collect_interleaved_wds_metrics + ) + output_metrics = {f"output_{k}": v for k, v in collect_fn(output_path).items()} + metrics_elapsed = time.perf_counter() - metrics_start + logger.info("collect_output_metrics took {:.3f}s", metrics_elapsed) + task_metrics = TaskPerfUtils.aggregate_task_metrics(output_tasks, prefix="task") + writer_stats = {k: v for k, v in task_metrics.items() if "interleaved_" in k and "_writer" in k} + logger.info("Writer stage stats: {}", writer_stats) + + rows = output_metrics["output_num_rows"] + samples = output_metrics["output_num_samples"] + return { + "params": { + "executor": args.executor, + "input_path": input_path, + "output_path": str(output_path), + "reader_type": args.reader_type, + "writer_format": args.writer_format, + "files_per_partition": args.files_per_partition, + "input_blocksize": args.input_blocksize, + "output_max_batch_bytes": args.output_max_batch_bytes, + "materialize_on_read": args.materialize_on_read, + "materialize_on_write": args.materialize_on_write, + "on_materialize_error": args.on_materialize_error, + "reader_fields": list(args.reader_fields), + "per_image_fields": list(args.per_image_fields) if args.per_image_fields else [], + "per_text_fields": list(args.per_text_fields) if args.per_text_fields else [], + "parquet_row_group_size": args.parquet_row_group_size, + "parquet_compression": args.parquet_compression, + "mode": args.mode, + "clip_model_dir": args.clip_model_dir, + "drop_invalid_rows": args.drop_invalid_rows, + "blur_score_threshold": args.blur_score_threshold, + "qrcode_score_threshold": args.qrcode_score_threshold, + "clip_min_score": args.clip_min_score, + "image_text_min_ratio": args.image_text_min_ratio, + "image_text_max_ratio": args.image_text_max_ratio, + }, + "metrics": { + "is_success": success, + "time_taken_s": elapsed, + "throughput_rows_per_sec": (rows / elapsed) if (elapsed > 0 and rows > 0) else 0.0, + "throughput_samples_per_sec": (samples / elapsed) if (elapsed > 0 and samples > 0) else 0.0, + **input_metrics, + **task_metrics, + **output_metrics, + "num_rows": rows, + "num_output_files": output_metrics["output_num_files"], + "materialize_error_count": output_metrics.get("output_materialize_error_count", 0), + }, + "tasks": output_tasks, + } + + +def main() -> int: + parser = argparse.ArgumentParser( + description="Interleaved filter benchmark (blur, QR, CLIP score, image/text ratio)" + ) + parser.add_argument("--benchmark-results-path", type=Path, required=True) + parser.add_argument("--executor", default="ray_data", choices=["xenna", "ray_data"]) + parser.add_argument("--input-path", type=str, required=True) + parser.add_argument("--output-path", type=str, required=True) + parser.add_argument("--clip-model-dir", type=str, required=True) + parser.add_argument("--reader-type", default="wds", choices=["wds", "parquet"]) + parser.add_argument("--writer-format", default="parquet", choices=["parquet", "wds"]) + parser.add_argument("--files-per-partition", type=int, default=1) + parser.add_argument("--input-blocksize", type=str, default=None) + parser.add_argument("--output-max-batch-bytes", type=int, default=None) + parser.add_argument( + "--reader-fields", + nargs="*", + default=[ + "image_metadata", + "url", + "language_id_whole_page_fasttext", + "pdf_name", + "previous_word_count", + "bff_contained_ngram_count_before_dedupe", + ], + ) + parser.add_argument("--materialize-on-read", action="store_true", dest="materialize_on_read") + parser.add_argument("--no-materialize-on-read", action="store_false", dest="materialize_on_read") + parser.add_argument("--materialize-on-write", action="store_true", dest="materialize_on_write") + parser.add_argument("--no-materialize-on-write", action="store_false", dest="materialize_on_write") + parser.add_argument( + "--on-materialize-error", + default="error", + choices=["error", "warn", "drop_row", "drop_sample"], + dest="on_materialize_error", + ) + parser.add_argument("--parquet-row-group-size", type=int, default=None) + parser.add_argument("--parquet-compression", type=str, default=None) + parser.add_argument("--mode", type=str, default="overwrite", choices=["ignore", "overwrite", "append", "error"]) + parser.add_argument("--per-image-fields", nargs="*", default=["image_metadata"]) + parser.add_argument("--per-text-fields", nargs="*", default=[]) + parser.add_argument("--blur-score-threshold", type=float, default=100.0) + parser.add_argument("--qrcode-score-threshold", type=float, default=0.05) + parser.add_argument("--clip-min-score", type=float, default=0.15) + parser.add_argument("--image-text-min-ratio", type=float, default=0.0) + parser.add_argument( + "--image-text-max-ratio", + type=float, + default=None, + help="Upper bound on image/text ratio; omit for no upper limit (infinity).", + ) + parser.add_argument("--drop-invalid-rows", action="store_true", dest="drop_invalid_rows") + parser.add_argument("--no-drop-invalid-rows", action="store_false", dest="drop_invalid_rows") + parser.set_defaults(materialize_on_write=True, materialize_on_read=True, drop_invalid_rows=True) + args = parser.parse_args() + + try: + results = run_benchmark(args) + except Exception as e: + logger.error("Benchmark crashed: {}", e) + logger.debug(traceback.format_exc()) + results = { + "params": vars(args), + "metrics": {"is_success": False}, + "tasks": [], + } + finally: + write_benchmark_results(results, args.benchmark_results_path) + + return 0 if results["metrics"]["is_success"] else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/nemo_curator/models/clip.py b/nemo_curator/models/clip.py index 788019e43f..673d767d57 100644 --- a/nemo_curator/models/clip.py +++ b/nemo_curator/models/clip.py @@ -106,6 +106,24 @@ def __call__(self, images: torch.Tensor | npt.NDArray[np.uint8] | list[np.ndarra # Normalize embeddings return embed / torch.linalg.vector_norm(embed, dim=-1, keepdim=True) # type: ignore[no-any-return] + @torch.no_grad() + def encode_text(self, texts: list[str]) -> torch.Tensor: + """Encode text(s) to normalized CLIP text embeddings. + + Args: + texts: List of strings to encode. + + Returns: + Normalized text embeddings, shape (len(texts), dim). + """ + if not texts: + msg = "encode_text requires at least one text" + raise ValueError(msg) + inputs = self.processor(text=texts, return_tensors="pt", padding=True, truncation=True) + inputs = {k: v.to(self.device) for k, v in inputs.items()} + embed = self.clip.get_text_features(**inputs) + return embed / torch.linalg.vector_norm(embed, dim=-1, keepdim=True) # type: ignore[no-any-return] + @classmethod def download_weights_on_node(cls, model_dir: str) -> None: """Download the weights for the CLIPImageEmbeddings model on the node.""" diff --git a/nemo_curator/stages/image/filters/__init__.py b/nemo_curator/stages/image/filters/__init__.py index d97c52d1d5..aad321081e 100644 --- a/nemo_curator/stages/image/filters/__init__.py +++ b/nemo_curator/stages/image/filters/__init__.py @@ -15,4 +15,7 @@ from .aesthetic_filter import ImageAestheticFilterStage from .nsfw_filter import ImageNSFWFilterStage -__all__ = ["ImageAestheticFilterStage", "ImageNSFWFilterStage"] +__all__ = [ + "ImageAestheticFilterStage", + "ImageNSFWFilterStage", +] diff --git a/nemo_curator/stages/interleaved/filter/__init__.py b/nemo_curator/stages/interleaved/filter/__init__.py new file mode 100644 index 0000000000..0b9d947725 --- /dev/null +++ b/nemo_curator/stages/interleaved/filter/__init__.py @@ -0,0 +1,27 @@ +# 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 nemo_curator.stages.interleaved.filter.blur_filter import InterleavedBlurFilterStage +from nemo_curator.stages.interleaved.filter.clip_score_filter import InterleavedCLIPScoreFilterStage +from nemo_curator.stages.interleaved.filter.image_to_text_ratio_filter import ( + InterleavedImageToTextRatioFilterStage, +) +from nemo_curator.stages.interleaved.filter.qrcode_filter import InterleavedQRCodeFilterStage + +__all__ = [ + "InterleavedBlurFilterStage", + "InterleavedCLIPScoreFilterStage", + "InterleavedImageToTextRatioFilterStage", + "InterleavedQRCodeFilterStage", +] diff --git a/nemo_curator/stages/interleaved/filter/blur_filter.py b/nemo_curator/stages/interleaved/filter/blur_filter.py new file mode 100644 index 0000000000..832d574d10 --- /dev/null +++ b/nemo_curator/stages/interleaved/filter/blur_filter.py @@ -0,0 +1,73 @@ +# 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 __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import cv2 +import pandas as pd +from loguru import logger + +from nemo_curator.stages.interleaved.stages import BaseInterleavedFilterStage +from nemo_curator.stages.interleaved.utils import image_bytes_to_array + +if TYPE_CHECKING: + from collections.abc import Hashable + + import numpy as np + + from nemo_curator.tasks import InterleavedBatch + +DEFAULT_BLUR_SCORE_THRESHOLD: float = 100.0 + + +def _sharpness_score(image: np.ndarray, row_index: Hashable | None = None) -> float: + """Compute Laplacian variance as sharpness score; higher is sharper.""" + try: + return float(cv2.Laplacian(image, cv2.CV_64F).var()) + except cv2.error as e: + logger.debug( + "cv2.Laplacian failed (row_index={} image_shape={}): {}", + row_index, + getattr(image, "shape", None), + e, + ) + return 0.0 + + +@dataclass +class InterleavedBlurFilterStage(BaseInterleavedFilterStage): + """Filter interleaved image rows by sharpness (Laplacian variance); drop blurry images.""" + + score_threshold: float = DEFAULT_BLUR_SCORE_THRESHOLD + name: str = "interleaved_blur_filter" + + def content_keep_mask(self, task: InterleavedBatch, df: pd.DataFrame) -> pd.Series: + keep_mask = pd.Series(True, index=df.index, dtype=bool) + image_mask = df["modality"] == "image" + if not image_mask.any(): + return keep_mask + for idx, image_bytes in self.iter_materialized_bytes(task=task, df=df, row_mask=image_mask): + if image_bytes is None: + keep_mask.loc[idx] = False + continue + image = image_bytes_to_array(image_bytes, row_index=idx) + if image is None: + keep_mask.loc[idx] = False + continue + sharpness = _sharpness_score(image, row_index=idx) + keep_mask.loc[idx] = sharpness >= self.score_threshold + return keep_mask diff --git a/nemo_curator/stages/interleaved/filter/clip_score_filter.py b/nemo_curator/stages/interleaved/filter/clip_score_filter.py new file mode 100644 index 0000000000..1b8753b76c --- /dev/null +++ b/nemo_curator/stages/interleaved/filter/clip_score_filter.py @@ -0,0 +1,116 @@ +# 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 __future__ import annotations + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING + +import pandas as pd + +from nemo_curator.models.clip import CLIPImageEmbeddings +from nemo_curator.stages.interleaved.stages import BaseInterleavedFilterStage +from nemo_curator.stages.interleaved.utils import image_bytes_to_array +from nemo_curator.stages.resources import Resources + +if TYPE_CHECKING: + import numpy as np + + from nemo_curator.backends.base import NodeInfo, WorkerMetadata + from nemo_curator.tasks import InterleavedBatch + +DEFAULT_CLIP_MIN_SCORE: float = 0.15 + + +def _sample_texts_list_from_df(df: pd.DataFrame, sample_id: str) -> list[str]: + """Return list of text_content from all text rows for the given sample_id (non-empty).""" + if "text_content" not in df.columns or "modality" not in df.columns: + return [] + subset = df[(df["sample_id"] == sample_id) & (df["modality"] == "text")] + if subset.empty: + return [] + return [s.strip() for s in subset["text_content"].dropna().astype(str).tolist() if s.strip()] + + +def _indices_and_decoded_images_from_rows( + rows: list[tuple[int, bytes]], keep_mask: pd.Series +) -> tuple[list[int], list[np.ndarray]]: + """Decode image bytes per row; clear keep_mask entries where decode fails.""" + indices: list[int] = [] + images: list[np.ndarray] = [] + for idx, b in rows: + arr = image_bytes_to_array(b, row_index=idx) + if arr is None: + keep_mask.loc[idx] = False + continue + indices.append(idx) + images.append(arr) + return indices, images + + +@dataclass +class InterleavedCLIPScoreFilterStage(BaseInterleavedFilterStage): + """Filter interleaved image rows by CLIP image-text relevance score. + + For each image row, all text rows with the same sample_id form (image, text) + pairs. CLIP similarity is computed for each pair. An image is kept only if at + least one pair has score >= min_score; otherwise it is dropped. + """ + + model_dir: str | None = None + min_score: float = DEFAULT_CLIP_MIN_SCORE + name: str = "interleaved_clip_score_filter" + resources: Resources = field(default_factory=lambda: Resources(gpu_memory_gb=20.0)) + + def setup(self, worker_metadata: WorkerMetadata | None = None) -> None: # noqa: ARG002 + self._model = CLIPImageEmbeddings(self.model_dir) + self._model.setup() + + def setup_on_node(self, node_info: NodeInfo, worker_metadata: WorkerMetadata) -> None: # noqa: ARG002 + """Download the weights for the CLIP model on the node.""" + if self.model_dir is None: + msg = "InterleavedCLIPScoreFilterStage requires model_dir to be set" + raise RuntimeError(msg) + CLIPImageEmbeddings.download_weights_on_node(self.model_dir) + + def content_keep_mask(self, task: InterleavedBatch, df: pd.DataFrame) -> pd.Series: + keep_mask = pd.Series(True, index=df.index, dtype=bool) + image_mask = df["modality"] == "image" + if not image_mask.any(): + return keep_mask + + sample_id_to_rows: dict[str, list[tuple[int, bytes]]] = {} + for idx, image_bytes in self.iter_materialized_bytes(task=task, df=df, row_mask=image_mask): + if image_bytes is None: + keep_mask.loc[idx] = False + continue + sample_id = df.loc[idx, "sample_id"] + sample_id_to_rows.setdefault(sample_id, []).append((idx, image_bytes)) + + for sample_id, rows in sample_id_to_rows.items(): + texts = _sample_texts_list_from_df(df, sample_id) + if not texts: + for idx, _ in rows: + keep_mask.loc[idx] = False + continue + indices, images = _indices_and_decoded_images_from_rows(rows, keep_mask) + if not images: + continue + img_emb = self._model(images) + text_emb = self._model.encode_text(texts) + scores = img_emb @ text_emb.T + for i, idx in enumerate(indices): + keep_mask.loc[idx] = (scores[i].max() >= self.min_score).item() + + return keep_mask diff --git a/nemo_curator/stages/interleaved/filter/image_to_text_ratio_filter.py b/nemo_curator/stages/interleaved/filter/image_to_text_ratio_filter.py new file mode 100644 index 0000000000..0d0cbaf4e4 --- /dev/null +++ b/nemo_curator/stages/interleaved/filter/image_to_text_ratio_filter.py @@ -0,0 +1,71 @@ +# 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 __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import pandas as pd + +from nemo_curator.stages.interleaved.stages import BaseInterleavedFilterStage + +if TYPE_CHECKING: + from nemo_curator.tasks import InterleavedBatch + +DEFAULT_IMAGE_TO_TEXT_MIN_RATIO: float = 0.0 +DEFAULT_IMAGE_TO_TEXT_MAX_RATIO: float = float("inf") + + +def _text_word_count(text: str | None) -> int: + """Count words in text by splitting on whitespace.""" + if text is None or (isinstance(text, float) and pd.isna(text)): + return 0 + return len(str(text).split()) + + +@dataclass +class InterleavedImageToTextRatioFilterStage(BaseInterleavedFilterStage): + """Filter interleaved samples by image-to-text ratio (images per word). + + Groups rows by sample_id. For each sample: + - image_count = number of rows with modality == 'image' + - text_word_count = sum of len(text_content.split()) over text rows + - ratio = image_count / max(text_word_count, 1) + + Samples with ratio outside [min_ratio, max_ratio] are dropped (all their rows). + """ + + min_ratio: float = DEFAULT_IMAGE_TO_TEXT_MIN_RATIO + max_ratio: float = DEFAULT_IMAGE_TO_TEXT_MAX_RATIO + name: str = "interleaved_image_to_text_ratio_filter" + + def content_keep_mask(self, task: InterleavedBatch, df: pd.DataFrame) -> pd.Series: # noqa: ARG002 + keep_mask = pd.Series(True, index=df.index, dtype=bool) + if "sample_id" not in df.columns: + return keep_mask + + sample_keep: dict[str, bool] = {} + for sample_id, group in df.groupby("sample_id"): + image_count = int((group["modality"] == "image").sum()) + text_mask = group["modality"] == "text" + text_word_count = 0 + if text_mask.any() and "text_content" in group.columns: + text_word_count = sum(_text_word_count(t) for t in group.loc[text_mask, "text_content"].tolist()) + ratio = image_count / max(text_word_count, 1) + sample_keep[sample_id] = self.min_ratio <= ratio <= self.max_ratio + + keep_mask = df["sample_id"].map(sample_keep) + keep_mask = keep_mask.fillna(True) + return keep_mask.astype(bool) diff --git a/nemo_curator/stages/interleaved/filter/qrcode_filter.py b/nemo_curator/stages/interleaved/filter/qrcode_filter.py new file mode 100644 index 0000000000..a185e6f9cb --- /dev/null +++ b/nemo_curator/stages/interleaved/filter/qrcode_filter.py @@ -0,0 +1,86 @@ +# 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 __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import cv2 +import numpy as np +import pandas as pd +from loguru import logger + +from nemo_curator.stages.interleaved.stages import BaseInterleavedFilterStage +from nemo_curator.stages.interleaved.utils import image_bytes_to_array + +if TYPE_CHECKING: + from collections.abc import Hashable + + from nemo_curator.tasks import InterleavedBatch + +DEFAULT_QRCODE_SCORE_THRESHOLD: float = 0.05 + + +def _qr_code_ratio(image: np.ndarray, row_index: Hashable | None = None) -> float: + """Return the ratio of image area covered by all detected QR code(s), in [0, 1].""" + img_shape = image.shape + height, width = img_shape[:2] + img_area = float(height * width) + if img_area <= 0: + return 0.0 + try: + detector = cv2.QRCodeDetector() + retval, _decoded_info, points, _ = detector.detectAndDecodeMulti(image) + if not retval or points is None or points.size == 0: + return 0.0 + points = np.asarray(points, dtype=np.float32) + total_qr_area = 0.0 + for i in range(len(points)): + pts = points[i].reshape(-1, 1, 2) + total_qr_area += cv2.contourArea(pts) + return total_qr_area / img_area + except cv2.error as e: + logger.debug( + "cv2 QR code ratio computation failed (row_index={} image_shape={}): {}", + row_index, + img_shape, + e, + ) + return 0.0 + + +@dataclass +class InterleavedQRCodeFilterStage(BaseInterleavedFilterStage): + """Filter interleaved image rows by QR code area ratio; drop images with high QR coverage.""" + + score_threshold: float = DEFAULT_QRCODE_SCORE_THRESHOLD + name: str = "interleaved_qrcode_filter" + + def content_keep_mask(self, task: InterleavedBatch, df: pd.DataFrame) -> pd.Series: + keep_mask = pd.Series(True, index=df.index, dtype=bool) + image_mask = df["modality"] == "image" + if not image_mask.any(): + return keep_mask + for idx, image_bytes in self.iter_materialized_bytes(task=task, df=df, row_mask=image_mask): + if image_bytes is None: + keep_mask.loc[idx] = False + continue + image = image_bytes_to_array(image_bytes, row_index=idx) + if image is None: + keep_mask.loc[idx] = False + continue + qr_ratio = _qr_code_ratio(image, row_index=idx) + keep_mask.loc[idx] = qr_ratio < self.score_threshold + return keep_mask diff --git a/nemo_curator/stages/interleaved/utils/__init__.py b/nemo_curator/stages/interleaved/utils/__init__.py index a7a6f6a1c2..4f40fae27e 100644 --- a/nemo_curator/stages/interleaved/utils/__init__.py +++ b/nemo_curator/stages/interleaved/utils/__init__.py @@ -17,6 +17,9 @@ DEFAULT_JSON_EXTENSIONS, DEFAULT_WEBDATASET_EXTENSIONS, ) +from nemo_curator.stages.interleaved.utils.image_utils import ( + image_bytes_to_array, +) from nemo_curator.stages.interleaved.utils.materialization import ( materialize_task_binary_content, ) @@ -34,6 +37,7 @@ "DEFAULT_JSON_EXTENSIONS", "DEFAULT_WEBDATASET_EXTENSIONS", "align_table", + "image_bytes_to_array", "materialize_task_binary_content", "reconcile_schema", "resolve_storage_options", diff --git a/nemo_curator/stages/interleaved/utils/image_utils.py b/nemo_curator/stages/interleaved/utils/image_utils.py new file mode 100644 index 0000000000..5aab36d5aa --- /dev/null +++ b/nemo_curator/stages/interleaved/utils/image_utils.py @@ -0,0 +1,35 @@ +# 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 __future__ import annotations + +from typing import TYPE_CHECKING + +import cv2 +import numpy as np +from loguru import logger + +if TYPE_CHECKING: + from collections.abc import Hashable + + +def image_bytes_to_array(image_bytes: bytes, *, row_index: Hashable | None = None) -> np.ndarray | None: + """Decode image bytes to RGB numpy array for OpenCV.""" + try: + arr = np.frombuffer(image_bytes, dtype=np.uint8) + image = cv2.imdecode(arr, cv2.IMREAD_COLOR) + return cv2.cvtColor(image, cv2.COLOR_BGR2RGB) + except cv2.error as e: + logger.debug("cv2 image decode failed (row_index={}): {}", row_index, e) + return None diff --git a/tests/models/test_clip.py b/tests/models/test_clip.py index f195378f64..001e45d5c2 100644 --- a/tests/models/test_clip.py +++ b/tests/models/test_clip.py @@ -206,6 +206,53 @@ def test_call_with_list_images_uses_processor_and_normalizes(self) -> None: norms = torch.linalg.vector_norm(out, dim=-1) assert torch.allclose(norms, torch.ones_like(norms), atol=1e-5) + def test_encode_text_empty_raises(self) -> None: + """encode_text must reject an empty text list.""" + self.model.clip = Mock() + self.model.processor = Mock() + + with pytest.raises(ValueError, match="encode_text requires at least one text"): + self.model.encode_text([]) + + self.model.processor.assert_not_called() + self.model.clip.get_text_features.assert_not_called() + + def test_encode_text_processor_clip_path_and_normalizes(self) -> None: + """encode_text should tokenize via processor, call get_text_features, and L2-normalize rows.""" + with patch("nemo_curator.models.clip.torch.cuda.is_available", return_value=False): + model = CLIPImageEmbeddings(model_dir="test_models/clip") + + mock_clip = Mock() + embed = torch.tensor([[3.0, 4.0], [1.0, 2.0]], dtype=torch.float32) + mock_clip.get_text_features.return_value = embed + + mock_processor = Mock() + input_ids = torch.tensor([[1, 2, 3], [4, 5, 0]], dtype=torch.long) + attention_mask = torch.tensor([[1, 1, 1], [1, 1, 0]], dtype=torch.long) + mock_processor.return_value = {"input_ids": input_ids, "attention_mask": attention_mask} + + model.clip = mock_clip + model.processor = mock_processor + + texts = ["a cat", "a dog"] + out = model.encode_text(texts) + + mock_processor.assert_called_once_with( + text=texts, + return_tensors="pt", + padding=True, + truncation=True, + ) + mock_clip.get_text_features.assert_called_once() + call_kwargs = mock_clip.get_text_features.call_args.kwargs + assert torch.equal(call_kwargs["input_ids"], input_ids) + assert torch.equal(call_kwargs["attention_mask"], attention_mask) + + norms = torch.linalg.vector_norm(out, dim=-1) + assert torch.allclose(norms, torch.ones_like(norms), atol=1e-5) + expected = embed / torch.linalg.vector_norm(embed, dim=-1, keepdim=True) + assert torch.allclose(out, expected, atol=1e-5) + class TestCLIPAestheticScorer: """Test cases for CLIPAestheticScorer model class.""" diff --git a/tests/stages/interleaved/filter/__init__.py b/tests/stages/interleaved/filter/__init__.py new file mode 100644 index 0000000000..4fc25d0d3c --- /dev/null +++ b/tests/stages/interleaved/filter/__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/interleaved/filter/conftest.py b/tests/stages/interleaved/filter/conftest.py new file mode 100644 index 0000000000..a55d117f21 --- /dev/null +++ b/tests/stages/interleaved/filter/conftest.py @@ -0,0 +1,39 @@ +# 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 io import BytesIO + +import numpy as np +import pyarrow as pa +from PIL import Image + +from nemo_curator.tasks import InterleavedBatch +from nemo_curator.tasks.interleaved import INTERLEAVED_SCHEMA + + +def interleaved_task(rows: list[dict]) -> InterleavedBatch: + table = pa.Table.from_pylist(rows, schema=INTERLEAVED_SCHEMA) + return InterleavedBatch(task_id="test", dataset_name="d", data=table) + + +def make_jpeg_bytes(width: int = 32, height: int = 32, sharp: bool = True) -> bytes: + buf = BytesIO() + if sharp: + rng = np.random.default_rng() + arr = rng.integers(0, 256, size=(height, width, 3), dtype=np.uint8) + img = Image.fromarray(arr) + else: + img = Image.new("RGB", (width, height), (128, 128, 128)) + img.save(buf, format="JPEG") + return buf.getvalue() diff --git a/tests/stages/interleaved/filter/test_blur_filter.py b/tests/stages/interleaved/filter/test_blur_filter.py new file mode 100644 index 0000000000..55031752ae --- /dev/null +++ b/tests/stages/interleaved/filter/test_blur_filter.py @@ -0,0 +1,274 @@ +# 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 collections.abc import Iterator +from typing import Any +from unittest.mock import patch + +import numpy as np +import pandas as pd + +from nemo_curator.stages.interleaved.filter.blur_filter import ( + InterleavedBlurFilterStage, + _sharpness_score, +) +from nemo_curator.stages.interleaved.utils import image_bytes_to_array + +from .conftest import interleaved_task, make_jpeg_bytes + + +def test_sharpness_score_solid_image_is_low() -> None: + arr = np.full((10, 10, 3), 100, dtype=np.uint8) + assert _sharpness_score(arr) == 0.0 + + +def test_sharpness_score_high_frequency_is_high() -> None: + rng = np.random.default_rng() + arr = rng.integers(0, 256, size=(20, 20, 3), dtype=np.uint8) + score = _sharpness_score(arr) + assert score > 0.0 + + +def test_image_bytes_to_array_valid_jpeg() -> None: + jpeg = make_jpeg_bytes() + arr = image_bytes_to_array(jpeg) + assert arr.shape[-1] == 3 + + +def test_blur_filter_text_only_passthrough() -> None: + rows = [ + { + "sample_id": "s1", + "position": 0, + "modality": "text", + "content_type": "text/plain", + "text_content": "hello", + "binary_content": None, + "source_ref": None, + "materialize_error": None, + }, + { + "sample_id": "s1", + "position": 1, + "modality": "text", + "content_type": "text/plain", + "text_content": "world", + "binary_content": None, + "source_ref": None, + "materialize_error": None, + }, + ] + task = interleaved_task(rows) + stage = InterleavedBlurFilterStage(score_threshold=100.0) + out = stage.process(task) + out_frame = out.to_pandas() + assert len(out_frame) == 2 + + +def test_blur_filter_image_with_binary_content_sharp_kept() -> None: + jpeg = make_jpeg_bytes(sharp=True) + rows = [ + { + "sample_id": "s1", + "position": 0, + "modality": "image", + "content_type": "image/jpeg", + "text_content": None, + "binary_content": jpeg, + "source_ref": None, + "materialize_error": None, + }, + ] + task = interleaved_task(rows) + stage = InterleavedBlurFilterStage(score_threshold=0.0) + out = stage.process(task) + out_frame = out.to_pandas() + assert len(out_frame) == 1 + + +def test_blur_filter_image_with_binary_content_blurry_dropped() -> None: + jpeg = make_jpeg_bytes(sharp=False) + rows = [ + { + "sample_id": "s1", + "position": 0, + "modality": "image", + "content_type": "image/jpeg", + "text_content": None, + "binary_content": jpeg, + "source_ref": None, + "materialize_error": None, + }, + ] + task = interleaved_task(rows) + stage = InterleavedBlurFilterStage(score_threshold=1e6) + out = stage.process(task) + out_frame = out.to_pandas() + assert len(out_frame) == 0 + + +def test_blur_filter_image_bytes_none_drops_row() -> None: + rows = [ + { + "sample_id": "s1", + "position": 0, + "modality": "image", + "content_type": "image/jpeg", + "text_content": None, + "binary_content": b"unused", + "source_ref": None, + "materialize_error": None, + }, + ] + task = interleaved_task(rows) + + def iter_materialized_bytes_none( + self: object, + task: object, + df: pd.DataFrame, + row_mask: pd.Series, + ) -> Iterator[tuple[Any, None]]: + del self, task + for idx in df[row_mask].index: + yield idx, None + + with patch.object(InterleavedBlurFilterStage, "iter_materialized_bytes", iter_materialized_bytes_none): + stage = InterleavedBlurFilterStage(score_threshold=0.0) + out = stage.process(task) + out_frame = out.to_pandas() + assert len(out_frame) == 0 + + +def test_blur_filter_empty_task_unchanged() -> None: + task = interleaved_task([]) + stage = InterleavedBlurFilterStage(score_threshold=100.0) + out = stage.process(task) + assert out.num_items == 0 + + +def test_blur_filter_metadata_row_preserved_with_text() -> None: + rows = [ + { + "sample_id": "s1", + "position": -1, + "modality": "metadata", + "content_type": "application/json", + "text_content": None, + "binary_content": None, + "source_ref": None, + "materialize_error": None, + }, + { + "sample_id": "s1", + "position": 0, + "modality": "text", + "content_type": "text/plain", + "text_content": "hello", + "binary_content": None, + "source_ref": None, + "materialize_error": None, + }, + ] + task = interleaved_task(rows) + stage = InterleavedBlurFilterStage(score_threshold=100.0) + out = stage.process(task) + out_frame = out.to_pandas() + assert len(out_frame) == 2 + assert (out_frame["modality"] == "metadata").sum() == 1 + + +def test_blur_filter_mixed_images_one_dropped_one_kept() -> None: + sharp_jpeg = make_jpeg_bytes(sharp=True) + blur_jpeg = make_jpeg_bytes(sharp=False) + rows = [ + { + "sample_id": "s1", + "position": 0, + "modality": "image", + "content_type": "image/jpeg", + "text_content": None, + "binary_content": sharp_jpeg, + "source_ref": None, + "materialize_error": None, + }, + { + "sample_id": "s1", + "position": 1, + "modality": "image", + "content_type": "image/jpeg", + "text_content": None, + "binary_content": blur_jpeg, + "source_ref": None, + "materialize_error": None, + }, + ] + task = interleaved_task(rows) + stage = InterleavedBlurFilterStage(score_threshold=100.0) + out = stage.process(task) + out_frame = out.to_pandas() + assert len(out_frame) == 1 + assert out_frame.iloc[0]["modality"] == "image" + assert out_frame.iloc[0]["position"] == 0 + + +def test_blur_filter_invalid_modality_dropped_when_drop_invalid_rows() -> None: + jpeg = make_jpeg_bytes(sharp=True) + rows = [ + { + "sample_id": "s1", + "position": 0, + "modality": "audio", + "content_type": "audio/wav", + "text_content": None, + "binary_content": None, + "source_ref": None, + "materialize_error": None, + }, + { + "sample_id": "s1", + "position": 1, + "modality": "image", + "content_type": "image/jpeg", + "text_content": None, + "binary_content": jpeg, + "source_ref": None, + "materialize_error": None, + }, + ] + task = interleaved_task(rows) + stage = InterleavedBlurFilterStage(score_threshold=0.0, drop_invalid_rows=True) + out = stage.process(task) + out_frame = out.to_pandas() + assert len(out_frame) == 1 + assert out_frame.iloc[0]["modality"] == "image" + + +def test_blur_filter_invalid_modality_kept_when_not_drop_invalid_rows() -> None: + rows = [ + { + "sample_id": "s1", + "position": 0, + "modality": "audio", + "content_type": "audio/wav", + "text_content": None, + "binary_content": None, + "source_ref": None, + "materialize_error": None, + }, + ] + task = interleaved_task(rows) + stage = InterleavedBlurFilterStage(score_threshold=100.0, drop_invalid_rows=False) + out = stage.process(task) + out_frame = out.to_pandas() + assert len(out_frame) == 1 diff --git a/tests/stages/interleaved/filter/test_clip_score_filter.py b/tests/stages/interleaved/filter/test_clip_score_filter.py new file mode 100644 index 0000000000..2d527a0dd7 --- /dev/null +++ b/tests/stages/interleaved/filter/test_clip_score_filter.py @@ -0,0 +1,257 @@ +# 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 MagicMock, patch + +import pandas as pd +import pytest +import torch + +from nemo_curator.backends.base import NodeInfo, WorkerMetadata +from nemo_curator.stages.interleaved.filter.clip_score_filter import ( + InterleavedCLIPScoreFilterStage, + _sample_texts_list_from_df, +) + +from .conftest import interleaved_task, make_jpeg_bytes + + +def test_clip_score_filter_requires_model_dir() -> None: + stage = InterleavedCLIPScoreFilterStage(model_dir=None) + with pytest.raises(RuntimeError, match="model_dir"): + stage.setup_on_node(NodeInfo(), WorkerMetadata()) + + +def test_sample_texts_list_from_df_missing_modality_column() -> None: + sample_frame = pd.DataFrame({"sample_id": ["s1"], "text_content": ["hello"]}) + assert _sample_texts_list_from_df(sample_frame, "s1") == [] + + +def test_sample_texts_list_from_df_missing_text_content_column() -> None: + sample_frame = pd.DataFrame({"sample_id": ["s1"], "modality": ["text"]}) + assert _sample_texts_list_from_df(sample_frame, "s1") == [] + + +def test_sample_texts_list_from_df_multiple_text_rows_order() -> None: + sample_frame = pd.DataFrame( + { + "sample_id": ["s1", "s1"], + "modality": ["text", "text"], + "text_content": ["first line", "second line"], + } + ) + assert _sample_texts_list_from_df(sample_frame, "s1") == ["first line", "second line"] + + +def test_sample_texts_list_from_df_strips_and_skips_empty() -> None: + sample_frame = pd.DataFrame( + { + "sample_id": ["s1", "s1", "s1"], + "modality": ["text", "text", "text"], + "text_content": [" padded ", "", " "], + } + ) + assert _sample_texts_list_from_df(sample_frame, "s1") == ["padded"] + + +def test_clip_score_filter_empty_task_unchanged() -> None: + task = interleaved_task([]) + with ( + patch("nemo_curator.stages.interleaved.filter.clip_score_filter.CLIPImageEmbeddings.download_weights_on_node"), + patch("nemo_curator.stages.interleaved.filter.clip_score_filter.CLIPImageEmbeddings") as mock_clip_class, + ): + mock_clip_class.return_value.return_value = torch.zeros(1, 1) + mock_clip_class.return_value.encode_text.return_value = torch.zeros(1, 1) + stage = InterleavedCLIPScoreFilterStage(model_dir="/fake/clip") + stage.setup_on_node(NodeInfo(), WorkerMetadata()) + stage.setup() + out = stage.process(task) + assert out.num_items == 0 + + +@patch("nemo_curator.stages.interleaved.filter.clip_score_filter.CLIPImageEmbeddings.download_weights_on_node") +@patch("nemo_curator.stages.interleaved.filter.clip_score_filter.CLIPImageEmbeddings") +def test_clip_score_filter_drops_image_when_sample_has_no_text( + mock_clip_class: MagicMock, mock_download: MagicMock +) -> None: + mock_download.return_value = None + dim = 512 + mock_model = mock_clip_class.return_value + mock_model.return_value = torch.ones(1, dim) / (dim**0.5) + mock_model.encode_text.return_value = torch.ones(1, dim) / (dim**0.5) + + jpeg = make_jpeg_bytes() + rows = [ + { + "sample_id": "solo", + "position": 0, + "modality": "image", + "content_type": "image/jpeg", + "text_content": None, + "binary_content": jpeg, + "source_ref": None, + "materialize_error": None, + }, + ] + task = interleaved_task(rows) + stage = InterleavedCLIPScoreFilterStage(model_dir="/fake/clip", min_score=0.15) + stage.setup_on_node(NodeInfo(), WorkerMetadata()) + stage.setup() + out = stage.process(task) + assert out.num_items == 0 + mock_model.encode_text.assert_not_called() + + +@patch("nemo_curator.stages.interleaved.filter.clip_score_filter.CLIPImageEmbeddings.download_weights_on_node") +@patch("nemo_curator.stages.interleaved.filter.clip_score_filter.CLIPImageEmbeddings") +def test_clip_score_filter_passes_all_sample_texts_to_encode_text( + mock_clip_class: MagicMock, mock_download: MagicMock +) -> None: + mock_download.return_value = None + dim = 512 + mock_model = mock_clip_class.return_value + mock_model.return_value = torch.ones(1, dim) / (dim**0.5) + mock_model.encode_text.return_value = torch.ones(2, dim) / (dim**0.5) + + jpeg = make_jpeg_bytes() + rows = [ + { + "sample_id": "s1", + "position": 0, + "modality": "text", + "content_type": "text/plain", + "text_content": "alpha", + "binary_content": None, + "source_ref": None, + "materialize_error": None, + }, + { + "sample_id": "s1", + "position": 1, + "modality": "text", + "content_type": "text/plain", + "text_content": "beta", + "binary_content": None, + "source_ref": None, + "materialize_error": None, + }, + { + "sample_id": "s1", + "position": 2, + "modality": "image", + "content_type": "image/jpeg", + "text_content": None, + "binary_content": jpeg, + "source_ref": None, + "materialize_error": None, + }, + ] + task = interleaved_task(rows) + stage = InterleavedCLIPScoreFilterStage(model_dir="/fake/clip", min_score=0.15) + stage.setup_on_node(NodeInfo(), WorkerMetadata()) + stage.setup() + stage.process(task) + mock_model.encode_text.assert_called_once() + texts_arg = mock_model.encode_text.call_args[0][0] + assert texts_arg == ["alpha", "beta"] + + +@patch("nemo_curator.stages.interleaved.filter.clip_score_filter.CLIPImageEmbeddings.download_weights_on_node") +@patch("nemo_curator.stages.interleaved.filter.clip_score_filter.CLIPImageEmbeddings") +def test_clip_score_filter_process_keeps_image_when_score_above_threshold( + mock_clip_class: MagicMock, mock_download: MagicMock +) -> None: + mock_download.return_value = None + dim = 512 + mock_model = mock_clip_class.return_value + mock_model.return_value = torch.ones(1, dim) / (dim**0.5) + mock_model.encode_text.return_value = torch.ones(2, dim) / (dim**0.5) + + jpeg = make_jpeg_bytes() + rows = [ + { + "sample_id": "s1", + "position": 0, + "modality": "text", + "content_type": "text/plain", + "text_content": "a cat", + "binary_content": None, + "source_ref": None, + "materialize_error": None, + }, + { + "sample_id": "s1", + "position": 1, + "modality": "image", + "content_type": "image/jpeg", + "text_content": None, + "binary_content": jpeg, + "source_ref": None, + "materialize_error": None, + }, + ] + task = interleaved_task(rows) + stage = InterleavedCLIPScoreFilterStage(model_dir="/fake/clip", min_score=0.15) + stage.setup_on_node(NodeInfo(), WorkerMetadata()) + stage.setup() + out = stage.process(task) + out_frame = out.to_pandas() + assert len(out_frame) == 2 + assert (out_frame["modality"] == "image").sum() == 1 + assert (out_frame["modality"] == "text").sum() == 1 + + +@patch("nemo_curator.stages.interleaved.filter.clip_score_filter.CLIPImageEmbeddings.download_weights_on_node") +@patch("nemo_curator.stages.interleaved.filter.clip_score_filter.CLIPImageEmbeddings") +def test_clip_score_filter_process_drops_image_when_score_below_threshold( + mock_clip_class: MagicMock, mock_download: MagicMock +) -> None: + mock_download.return_value = None + dim = 512 + mock_model = mock_clip_class.return_value + mock_model.return_value = torch.ones(1, dim) / (dim**0.5) + mock_model.encode_text.return_value = -torch.ones(1, dim) / (dim**0.5) + + jpeg = make_jpeg_bytes() + rows = [ + { + "sample_id": "s1", + "position": 0, + "modality": "text", + "content_type": "text/plain", + "text_content": "unrelated", + "binary_content": None, + "source_ref": None, + "materialize_error": None, + }, + { + "sample_id": "s1", + "position": 1, + "modality": "image", + "content_type": "image/jpeg", + "text_content": None, + "binary_content": jpeg, + "source_ref": None, + "materialize_error": None, + }, + ] + task = interleaved_task(rows) + stage = InterleavedCLIPScoreFilterStage(model_dir="/fake/clip", min_score=0.15) + stage.setup_on_node(NodeInfo(), WorkerMetadata()) + stage.setup() + out = stage.process(task) + out_frame = out.to_pandas() + assert len(out_frame) == 1 + assert (out_frame["modality"] == "text").sum() == 1 + assert (out_frame["modality"] == "image").sum() == 0 diff --git a/tests/stages/interleaved/filter/test_image_to_text_ratio_filter.py b/tests/stages/interleaved/filter/test_image_to_text_ratio_filter.py new file mode 100644 index 0000000000..5aee10c642 --- /dev/null +++ b/tests/stages/interleaved/filter/test_image_to_text_ratio_filter.py @@ -0,0 +1,342 @@ +# 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 nemo_curator.stages.interleaved.filter.image_to_text_ratio_filter import ( + InterleavedImageToTextRatioFilterStage, + _text_word_count, +) + +from .conftest import interleaved_task + + +def test_text_word_count_none_is_zero() -> None: + assert _text_word_count(None) == 0 + + +def test_text_word_count_nan_float_is_zero() -> None: + assert _text_word_count(float("nan")) == 0 + + +def test_text_word_count_splits_on_whitespace() -> None: + assert _text_word_count(" one two three ") == 3 + + +def test_image_to_text_ratio_empty_task_unchanged() -> None: + task = interleaved_task([]) + stage = InterleavedImageToTextRatioFilterStage() + out = stage.process(task) + assert out.num_items == 0 + + +def test_image_to_text_ratio_image_only_uses_denominator_one() -> None: + rows = [ + { + "sample_id": "solo", + "position": 0, + "modality": "image", + "content_type": "image/jpeg", + "text_content": None, + "binary_content": None, + "source_ref": None, + "materialize_error": None, + }, + { + "sample_id": "solo", + "position": 1, + "modality": "image", + "content_type": "image/jpeg", + "text_content": None, + "binary_content": None, + "source_ref": None, + "materialize_error": None, + }, + ] + task = interleaved_task(rows) + stage = InterleavedImageToTextRatioFilterStage(min_ratio=1.5, max_ratio=2.5) + out = stage.process(task) + out_frame = out.to_pandas() + assert len(out_frame) == 2 + + +def test_image_to_text_ratio_boundary_inclusive_min_and_max() -> None: + rows = [ + { + "sample_id": "s1", + "position": 0, + "modality": "text", + "content_type": "text/plain", + "text_content": "one two three four", + "binary_content": None, + "source_ref": None, + "materialize_error": None, + }, + { + "sample_id": "s1", + "position": 1, + "modality": "image", + "content_type": "image/jpeg", + "text_content": None, + "binary_content": None, + "source_ref": None, + "materialize_error": None, + }, + ] + task = interleaved_task(rows) + stage = InterleavedImageToTextRatioFilterStage(min_ratio=0.25, max_ratio=0.25) + out = stage.process(task) + out_frame = out.to_pandas() + assert len(out_frame) == 2 + assert out.num_items == 1 + + +def test_image_to_text_ratio_metadata_only_sample_dropped_as_orphan() -> None: + rows = [ + { + "sample_id": "meta_only", + "position": -1, + "modality": "metadata", + "content_type": "application/json", + "text_content": None, + "binary_content": None, + "source_ref": None, + "materialize_error": None, + }, + ] + task = interleaved_task(rows) + stage = InterleavedImageToTextRatioFilterStage() + out = stage.process(task) + assert out.num_items == 0 + + +def test_image_to_text_ratio_text_rows_contribute_word_count_across_rows() -> None: + rows = [ + { + "sample_id": "s1", + "position": 0, + "modality": "text", + "content_type": "text/plain", + "text_content": "one two", + "binary_content": None, + "source_ref": None, + "materialize_error": None, + }, + { + "sample_id": "s1", + "position": 1, + "modality": "text", + "content_type": "text/plain", + "text_content": "three four", + "binary_content": None, + "source_ref": None, + "materialize_error": None, + }, + { + "sample_id": "s1", + "position": 2, + "modality": "image", + "content_type": "image/jpeg", + "text_content": None, + "binary_content": None, + "source_ref": None, + "materialize_error": None, + }, + ] + task = interleaved_task(rows) + stage = InterleavedImageToTextRatioFilterStage(min_ratio=0.15, max_ratio=0.25) + out = stage.process(task) + out_frame = out.to_pandas() + assert len(out_frame) == 3 + assert out.num_items == 1 + + +def test_image_to_text_ratio_no_sample_id_passthrough() -> None: + # InterleavedBatch requires sample_id; test content_keep_mask with frame missing sample_id. + rows = [ + { + "sample_id": "x", + "position": 0, + "modality": "text", + "content_type": "text/plain", + "text_content": "hello", + "binary_content": None, + "source_ref": None, + "materialize_error": None, + }, + { + "sample_id": "x", + "position": 1, + "modality": "image", + "content_type": "image/jpeg", + "text_content": None, + "binary_content": None, + "source_ref": None, + "materialize_error": None, + }, + ] + task = interleaved_task(rows) + task_frame = task.to_pandas().drop(columns=["sample_id"]) + stage = InterleavedImageToTextRatioFilterStage(min_ratio=0.0, max_ratio=1.0) + keep = stage.content_keep_mask(task, task_frame) + assert keep.all() + + +def test_image_to_text_ratio_ratio_in_range_kept() -> None: + rows = [ + { + "sample_id": "s1", + "position": 0, + "modality": "text", + "content_type": "text/plain", + "text_content": "one two three four", + "binary_content": None, + "source_ref": None, + "materialize_error": None, + }, + { + "sample_id": "s1", + "position": 1, + "modality": "image", + "content_type": "image/jpeg", + "text_content": None, + "binary_content": None, + "source_ref": None, + "materialize_error": None, + }, + ] + task = interleaved_task(rows) + stage = InterleavedImageToTextRatioFilterStage(min_ratio=0.2, max_ratio=1.0) + out = stage.process(task) + out_frame = out.to_pandas() + assert len(out_frame) == 2 + + +def test_image_to_text_ratio_ratio_below_min_dropped() -> None: + rows = [ + { + "sample_id": "s1", + "position": 0, + "modality": "text", + "content_type": "text/plain", + "text_content": "one two three four five", + "binary_content": None, + "source_ref": None, + "materialize_error": None, + }, + { + "sample_id": "s1", + "position": 1, + "modality": "image", + "content_type": "image/jpeg", + "text_content": None, + "binary_content": None, + "source_ref": None, + "materialize_error": None, + }, + ] + task = interleaved_task(rows) + stage = InterleavedImageToTextRatioFilterStage(min_ratio=1.0, max_ratio=2.0) + out = stage.process(task) + out_frame = out.to_pandas() + assert len(out_frame) == 0 + + +def test_image_to_text_ratio_ratio_above_max_dropped() -> None: + rows = [ + { + "sample_id": "s1", + "position": 0, + "modality": "text", + "content_type": "text/plain", + "text_content": "x", + "binary_content": None, + "source_ref": None, + "materialize_error": None, + }, + { + "sample_id": "s1", + "position": 1, + "modality": "image", + "content_type": "image/jpeg", + "text_content": None, + "binary_content": None, + "source_ref": None, + "materialize_error": None, + }, + { + "sample_id": "s1", + "position": 2, + "modality": "image", + "content_type": "image/jpeg", + "text_content": None, + "binary_content": None, + "source_ref": None, + "materialize_error": None, + }, + ] + task = interleaved_task(rows) + stage = InterleavedImageToTextRatioFilterStage(min_ratio=0.0, max_ratio=1.0) + out = stage.process(task) + out_frame = out.to_pandas() + assert len(out_frame) == 0 + + +def test_image_to_text_ratio_multiple_samples_one_dropped() -> None: + rows = [ + { + "sample_id": "keep", + "position": 0, + "modality": "text", + "content_type": "text/plain", + "text_content": "a b", + "binary_content": None, + "source_ref": None, + "materialize_error": None, + }, + { + "sample_id": "keep", + "position": 1, + "modality": "image", + "content_type": "image/jpeg", + "text_content": None, + "binary_content": None, + "source_ref": None, + "materialize_error": None, + }, + { + "sample_id": "drop", + "position": 0, + "modality": "text", + "content_type": "text/plain", + "text_content": "one two three", + "binary_content": None, + "source_ref": None, + "materialize_error": None, + }, + { + "sample_id": "drop", + "position": 1, + "modality": "image", + "content_type": "image/jpeg", + "text_content": None, + "binary_content": None, + "source_ref": None, + "materialize_error": None, + }, + ] + task = interleaved_task(rows) + stage = InterleavedImageToTextRatioFilterStage(min_ratio=0.4, max_ratio=0.6) + out = stage.process(task) + out_frame = out.to_pandas() + assert len(out_frame) == 2 + assert set(out_frame["sample_id"].tolist()) == {"keep"} diff --git a/tests/stages/interleaved/filter/test_qrcode_filter.py b/tests/stages/interleaved/filter/test_qrcode_filter.py new file mode 100644 index 0000000000..bbd13ff7bb --- /dev/null +++ b/tests/stages/interleaved/filter/test_qrcode_filter.py @@ -0,0 +1,205 @@ +# 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 collections.abc import Iterator +from typing import Any +from unittest.mock import MagicMock, patch + +import cv2 +import numpy as np +import pandas as pd + +from nemo_curator.stages.interleaved.filter.qrcode_filter import ( + InterleavedQRCodeFilterStage, + _qr_code_ratio, +) + +from .conftest import interleaved_task, make_jpeg_bytes + + +def test_qr_code_ratio_no_qr_returns_zero() -> None: + rng = np.random.default_rng() + arr = rng.integers(0, 256, size=(50, 50, 3), dtype=np.uint8) + ratio = _qr_code_ratio(arr) + assert ratio == 0.0 + + +def test_qr_code_ratio_zero_image_area_returns_zero() -> None: + arr = np.zeros((0, 10, 3), dtype=np.uint8) + assert _qr_code_ratio(arr) == 0.0 + + +@patch("nemo_curator.stages.interleaved.filter.qrcode_filter.cv2.QRCodeDetector") +def test_qr_code_ratio_cv2_error_returns_zero(mock_detector_cls: MagicMock) -> None: + detector = MagicMock() + detector.detectAndDecodeMulti.side_effect = cv2.error("mock decode failure") + mock_detector_cls.return_value = detector + arr = np.ones((8, 8, 3), dtype=np.uint8) + assert _qr_code_ratio(arr) == 0.0 + + +def test_qrcode_filter_empty_task_unchanged() -> None: + task = interleaved_task([]) + stage = InterleavedQRCodeFilterStage(score_threshold=0.05) + out = stage.process(task) + assert out.num_items == 0 + + +def test_qrcode_filter_metadata_and_text_passthrough() -> None: + rows = [ + { + "sample_id": "s1", + "position": -1, + "modality": "metadata", + "content_type": "application/json", + "text_content": None, + "binary_content": None, + "source_ref": None, + "materialize_error": None, + }, + { + "sample_id": "s1", + "position": 0, + "modality": "text", + "content_type": "text/plain", + "text_content": "hello", + "binary_content": None, + "source_ref": None, + "materialize_error": None, + }, + ] + task = interleaved_task(rows) + stage = InterleavedQRCodeFilterStage(score_threshold=0.05) + out = stage.process(task) + out_frame = out.to_pandas() + assert len(out_frame) == 2 + + +def test_qrcode_filter_text_only_passthrough() -> None: + rows = [ + { + "sample_id": "s1", + "position": 0, + "modality": "text", + "content_type": "text/plain", + "text_content": "hello", + "binary_content": None, + "source_ref": None, + "materialize_error": None, + }, + ] + task = interleaved_task(rows) + stage = InterleavedQRCodeFilterStage(score_threshold=0.05) + out = stage.process(task) + out_frame = out.to_pandas() + assert len(out_frame) == 1 + + +@patch("nemo_curator.stages.interleaved.filter.qrcode_filter.image_bytes_to_array") +def test_qrcode_filter_image_decode_error_drops_row(mock_to_array: MagicMock) -> None: + mock_to_array.return_value = None + rows = [ + { + "sample_id": "s1", + "position": 0, + "modality": "image", + "content_type": "image/jpeg", + "text_content": None, + "binary_content": b"garbage", + "source_ref": None, + "materialize_error": None, + }, + ] + task = interleaved_task(rows) + stage = InterleavedQRCodeFilterStage(score_threshold=0.05) + out = stage.process(task) + out_frame = out.to_pandas() + assert len(out_frame) == 0 + + +def test_qrcode_filter_image_bytes_none_drops_row() -> None: + rows = [ + { + "sample_id": "s1", + "position": 0, + "modality": "image", + "content_type": "image/jpeg", + "text_content": None, + "binary_content": b"unused", + "source_ref": None, + "materialize_error": None, + }, + ] + task = interleaved_task(rows) + + def iter_materialized_bytes_none( + self: object, + task: object, + df: pd.DataFrame, + row_mask: pd.Series, + ) -> Iterator[tuple[Any, None]]: + del self, task + for idx in df[row_mask].index: + yield idx, None + + with patch.object(InterleavedQRCodeFilterStage, "iter_materialized_bytes", iter_materialized_bytes_none): + stage = InterleavedQRCodeFilterStage(score_threshold=0.05) + out = stage.process(task) + out_frame = out.to_pandas() + assert len(out_frame) == 0 + + +def test_qrcode_filter_image_below_threshold_kept() -> None: + jpeg = make_jpeg_bytes() + rows = [ + { + "sample_id": "s1", + "position": 0, + "modality": "image", + "content_type": "image/jpeg", + "text_content": None, + "binary_content": jpeg, + "source_ref": None, + "materialize_error": None, + }, + ] + task = interleaved_task(rows) + stage = InterleavedQRCodeFilterStage(score_threshold=1.0) + out = stage.process(task) + out_frame = out.to_pandas() + assert len(out_frame) == 1 + + +@patch("nemo_curator.stages.interleaved.filter.qrcode_filter._qr_code_ratio") +def test_qrcode_filter_image_above_threshold_dropped(mock_qr_ratio: MagicMock) -> None: + mock_qr_ratio.return_value = 0.5 + jpeg = make_jpeg_bytes() + rows = [ + { + "sample_id": "s1", + "position": 0, + "modality": "image", + "content_type": "image/jpeg", + "text_content": None, + "binary_content": jpeg, + "source_ref": None, + "materialize_error": None, + }, + ] + task = interleaved_task(rows) + stage = InterleavedQRCodeFilterStage(score_threshold=0.05) + out = stage.process(task) + out_frame = out.to_pandas() + assert len(out_frame) == 0 + mock_qr_ratio.assert_called() diff --git a/tests/stages/interleaved/utils/__init__.py b/tests/stages/interleaved/utils/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/stages/interleaved/utils/test_image_utils.py b/tests/stages/interleaved/utils/test_image_utils.py new file mode 100644 index 0000000000..cb5f0bbd91 --- /dev/null +++ b/tests/stages/interleaved/utils/test_image_utils.py @@ -0,0 +1,34 @@ +# 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 MagicMock, patch + +import cv2 +import numpy as np + +from nemo_curator.stages.interleaved.utils.image_utils import image_bytes_to_array + + +@patch("nemo_curator.stages.interleaved.utils.image_utils.cv2.imdecode") +def test_image_bytes_to_array_cv2_imdecode_error_returns_none(mock_imdecode: MagicMock) -> None: + mock_imdecode.side_effect = cv2.error("mock imdecode failure") + assert image_bytes_to_array(b"\x00\x01\x02", row_index=7) is None + + +@patch("nemo_curator.stages.interleaved.utils.image_utils.cv2.cvtColor") +@patch("nemo_curator.stages.interleaved.utils.image_utils.cv2.imdecode") +def test_image_bytes_to_array_cv2_cvtcolor_error_returns_none(mock_imdecode: MagicMock, mock_cvt: MagicMock) -> None: + mock_imdecode.return_value = np.zeros((2, 2, 3), dtype=np.uint8) + mock_cvt.side_effect = cv2.error("mock cvtColor failure") + assert image_bytes_to_array(b"\x00\x01\x02", row_index=None) is None