Nemotron OCR SDG Pipeline by suiyoubi · Pull Request #1899 · NVIDIA-NeMo/Curator

suiyoubi · 2026-04-30T15:49:47Z

Description

Adds the Nemotron OCR SDG pipeline — a multimodal synthetic data generation pipeline that converts images into structured OCR + QA conversation data for vision-language model training.

Pipeline stages

Stage	Model	Output
`NemotronOCRV2Stage`	NemotronOCR-v2	Dense word-level OCR with bounding boxes
`OCRScoringQAStage`	Gemini 3 Pro (NVIDIA Inference API)	Scoring, validation, and missing-region detection
`OCRConversationalizeStage`	—	11 output format variants → `ConversationSample`
`OCRDenseQAStage`	—	6 QA types (bbox↔text, point↔text, dense dump)

Key components

`nemo_curator/models/omni/` — `NVInferenceModel` base class for NVIDIA Inference API-backed VLMs; `Gemini3Pro` concrete implementation
`nemo_curator/models/client/nvinference_client.py` — thin streaming client helpers (`get_nvinference_api_key`, `create_openai_client`, `stream_chat_completion_text`)
`nemo_curator/stages/synthetic/omni/base.py` — `VLMProcessingStage` and `ModelProcessingStage` base classes with batched inference, per-prompt error isolation, and setup/teardown lifecycle
`nemo_curator/stages/synthetic/omni/io.py` — `HFDatasetImageReader`, `TarImageReader`, `ParquetReader`, `SkipProcessedStage`, `ResultWriterStage`
`nemo_curator/tasks/ocr.py` — `OCRDenseWord` and `OCRData` task data classes
`docker/Dockerfile` — installs `nemotron-ocr-v2` from source (no-build-isolation, CUDA arch list for A100/A10/RTX Ada/H100)
`tutorials/synthetic/omni/hf_ocr_pipeline.py` — end-to-end example using HuggingFace datasets

Tests

63 new unit tests across `tests/tasks/`, `tests/models/`, and `tests/stages/synthetic/omni/` — all CPU-only, no GPU required.

Checklist

I am familiar with the Contributing Guide.
New or Existing tests cover these changes.
The documentation is up to date with these changes.

Signed-off-by: Ao Tang <aot@nvidia.com>

- Updated the `transformers` dependency from `<=4.55.2` to `==4.57.0` in `pyproject.toml` and `uv.lock` to ensure compatibility with the Cosmos Embed imports. - Added a new `Gemini3Pro` model class in `gemini.py` utilizing the NVIDIA Inference API. - Introduced `DescriptionOutputStage` and `DescriptionValidatorStage` for processing and validating image descriptions, respectively. - Enhanced `VLMProcessingStage` to improve GPU resource handling and added a `num_workers` parameter to `DescriptionStage` for better scalability. This commit enhances the model's capabilities and ensures that dependencies are up-to-date for optimal performance. Signed-off-by: Ao Tang <aot@nvidia.com>

Signed-off-by: Ao Tang <aot@nvidia.com>

… aot/omni_sdg

Signed-off-by: Ao Tang <aot@nvidia.com>

Removes description-specific stages (description*.py, description pipeline tutorials) that belong on aot/omni_description. Adds OCR result inspection/review scripts and shared design docs. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

… pipeline tutorial Signed-off-by: Ao Tang <aot@nvidia.com>

Signed-off-by: Ao Tang <aot@nvidia.com>

…Data and OCRDenseWord classes Signed-off-by: Ao Tang <aot@nvidia.com>

--metrics-dir wires Ray metrics into the running Prometheus/Grafana instance. --run-name sets SLURM_JOB_NAME so Xenna labels the run on the ray_pipeline_input_tasks metric for human-readable identification in Grafana. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

… timing - OCRConversationData.to_dict(): call conversation.to_dict() explicitly instead of relying on dataclasses.asdict(), which bypasses the custom ConversationSample serialization and drops the "t" media-type field from image fragments. - RayClient: move Prometheus service-discovery registration to after Ray is started and responsive; add _wait_for_ray_service_discovery_file() so the SD file exists before Prometheus is told to watch it. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Remove internal-only I/O stages from io.py (InputFormat, ImageReaderStage, ImageFolderReaderStage, TarImageReaderStage, JsonlTarImageReaderStage, OcrJsonlReaderStage, JsonlPipelineOutputReaderStage, TarImageReader, ParquetImageReaderStage, ParquetImageReader). These classes are only used by the internal ocr_pipeline.py and will live on aot/omni_sdg_internal. Public io.py now exports: HFDatasetImageReaderStage, SkipProcessedStage, ResultWriterStage, merge_output_shards, ImageWriterStage, and the FileReader helpers (load_image_from_task, TarFileReader, etc.). Add tests/stages/test_hf_dataset_image_reader.py covering HFDatasetImageReaderStage. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

ImageWriterStage is not used by hf_ocr_pipeline.py and only makes sense with the internal JSONL+tar pipeline; moved to aot/omni_sdg_internal. TarFileReader, ParquetFileReader, _file_readers dispatcher, and the deprecated _parse_tar_slice_path wrapper are removed. In the public HF pipeline images are always regular JPEG files on disk, so load_image_from_task is simplified to a single RegularFileReader call. io.py: 870 → 631 lines. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

…e_url SUPPORTED_IMAGE_EXTENSIONS was only used by the internal reader classes removed in the previous commit. FileReader.read_image_url() was never called in this branch — drop it and its now-unused base64 import. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

…ng in nvinference_client.py Signed-off-by: Ao Tang <aot@nvidia.com>

Signed-off-by: Ao Tang <aot@nvidia.com>

… clean up OCRScoringQAStage by removing unused kwargs parameter. Signed-off-by: Ao Tang <aot@nvidia.com>

greptile-apps · 2026-05-20T19:21:31Z

+            data["image_path"] = self._get_image_path_str(task.data.image_path)
+            # Keep empty lists/strings/False (e.g. OCR may legitimately be []).
+            # Only drop fields that are explicitly None, and always omit is_valid.
+            self._file.write(
+                json.dumps({k: v for k, v in data.items() if v is not None and k != "is_valid"}, default=str) + "\n"
+            )
+        elif self.valid_only:
+            self._skipped_count += 1
+            return task
+        else:
+            data = task.data.to_dict()
+            data["image_path"] = self._get_image_path_str(task.data.image_path)
+            self._file.write(
+                json.dumps({k: v for k, v in data.items() if v is not None and k != "is_valid"}, default=str) + "\n"
+            )
+        self._file.flush()  # Flush after each write for safety
+        self._saved_count += 1
+        return task


Invalid records silently lose is_valid=False when written and re-read

ResultWriterStage.process always strips is_valid from the output via k != "is_valid". When valid_only=False (the CLI default — --valid-only is store_true defaulting to False), invalid records are written without an is_valid field. OCRData.from_dict defaults that field to True (data.get("is_valid", True)), so any code path that reloads the JSONL (e.g., a resume workflow or SkipProcessedStage) will treat every previously-invalid record as valid and re-attempt processing it.

The filter should preserve is_valid=False for invalid records while still omitting it for valid ones:

filtered = {k: v for k, v in data.items() if v is not None and k != "is_valid"} if not task.data.is_valid: filtered["is_valid"] = False self._file.write(json.dumps(filtered, default=str) + "\n")

Alternatively, include is_valid unconditionally and remove the k != "is_valid" filter.

praateekmahajan · 2026-05-26T21:07:18Z

+from nemo_curator.models.client import OpenAIClient
+
+
+class NVInferenceModel(ModelInterface):


why is this a ModelInterface..

Can this just be a ProcessingStage?

See https://github.com/NVIDIA-NeMo/Curator/blob/main/nemo_curator/stages/synthetic/qa_multilingual_synthetic.py#L33

To which we pass in a client which can be OpenAIClient or AsyncOpenAIClient

This is just the modeling class. Similar to how we have in video for example that we have QwenLM class and then we invoke the model in CaptionGenerationStage.

Here we just have the NVInferenceModel which is supposed to work for any API endpoint from build.nvidia.com . and we use this model in the ModelProcessingStage (which is a ProcessingStage)

praateekmahajan · 2026-05-26T21:23:07Z

+    logger.info(f"Tasks processed: {len(output_tasks)}")
+
+    merged = merge_output_shards(Path(args.output_path))
+    logger.info(f"Output: {merged}")


We should probably add this to benchmark.. PS for that you'll also need to copy dataset for which this will work (so that we don't re download from the internet each time)..

copy the data to eos and dgx-a100 and then update the benchmarking/script/....yaml with the right path..

Feel free to ask @rlratzel for more questions around that

I have the benchmark setup here in another branch: https://github.com/NVIDIA-NeMo/Curator/tree/aot/omni_sdg_benchmark
let's merge that later once we finalizing on the benchmark details so we know what to measure (i.e. GPU SKU)

…ncOpenAIClient. Introduced parameters for concurrent requests and updated related methods for async handling. Updated model processing stages to utilize ImageSampleTask instead of SingleDataTask for better image task management. Enhanced I/O stages to write JSONL results and read images from Hugging Face datasets. Improved OCR stages to handle image tasks more effectively. Updated tests to reflect changes in task handling and client initialization. Signed-off-by: Ao Tang <aot@nvidia.com>

Signed-off-by: Ao Tang <aot@nvidia.com>

suiyoubi · 2026-06-04T13:54:26Z

/ok to test 7d4d113

… ModelInterface wrapper The OCR scoring path wrapped the NVIDIA Inference endpoint in NVInferenceModel(ModelInterface). ModelInterface is for GPU weight-bearing models (download weights, model_id_names = HF id); an HTTP client doesn't fit, and ~80% of the class re-implemented the AsyncOpenAIClient it already held — including a bespoke asyncio semaphore — while bypassing query_model and so losing the 429/connection retry+backoff that AsyncLLMClient provides for free. Replace it with a thin NVInferenceClient(AsyncOpenAIClient) that overrides only _query_model_impl (stream + reassemble delta.content for reasoning models, drop reasoning_content) and setup (resolve the API key from the env on the worker). It inherits the concurrency semaphore + retry. ModelProcessingStage now takes a client + model_name and calls client.query_model — mirroring the other SDG stages (e.g. QAMultilingualSyntheticStage) — with image→messages assembly moved into the stage. Nothing in the old wrapper was NVIDIA-endpoint-specific. Also: - drop the dead sync path (use_async was test-only) and the duplicated stream method - make setup() idempotent and forward stop/extra_kwargs in the create call - simplify build_prompt: ocr_dense is always OCRDenseItem, so drop the dict fallback - update unit tests for the client-based API Signed-off-by: Ao Tang <aot@nvidia.com>

greptile-apps · 2026-06-11T15:48:35Z

+        self._file.write(
+            json.dumps({k: v for k, v in data.items() if v is not None and k != "is_valid"}, default=str) + "\n"
+        )


Invalid records written when valid_only=False silently lose their is_valid=False state. OCRData.from_dict defaults is_valid to True (line 125 of ocr.py), so any downstream code that reloads the JSONL — e.g. a resume workflow or SkipProcessedStage — will treat every previously-invalid record as valid and reprocess it.

Suggested change

self._file.write(

json.dumps({k: v for k, v in data.items() if v is not None and k != "is_valid"}, default=str) + "\n"

)

filtered = {k: v for k, v in data.items() if v is not None and k != "is_valid"}

if not task.data.is_valid:

filtered["is_valid"] = False

self._file.write(json.dumps(filtered, default=str) + "\n")

suiyoubi · 2026-06-11T16:41:22Z

/ok to test 025bc07

suiyoubi · 2026-06-11T20:14:04Z

/ok to test 85c82d5

Signed-off-by: Ao Tang <aot@nvidia.com>

suiyoubi · 2026-06-11T21:12:47Z

/ok to test fc82dc0

suiyoubi · 2026-06-12T15:41:44Z

/ok to test 146c38b

suiyoubi and others added 30 commits March 4, 2026 10:40

omni sdg

b6bf8c1

Signed-off-by: Ao Tang <aot@nvidia.com>

Refactor model initialization and enhance pipeline execution

2b9f157

Signed-off-by: Ao Tang <aot@nvidia.com>

Merge branch 'main' of github.com:NVIDIA-NeMo/Curator into aot/omni_sdg

21f417f

add metrics support

ec46518

Signed-off-by: Ao Tang <aot@nvidia.com>

Add support for per-stage pip specifications and virtual environments

5271c7c

Signed-off-by: Ao Tang <aot@nvidia.com>

Merge branch 'aot/runtime_env' of github.com:NVIDIA-NeMo/Curator into…

e4df9af

… aot/omni_sdg

Add test

fa9f9a5

Signed-off-by: Ao Tang <aot@nvidia.com>

Merge branch 'main' of github.com:NVIDIA-NeMo/Curator into aot/omni_sdg

2d34d90

Signed-off-by: Ao Tang <aot@nvidia.com>

Add ocr nemotron v2 pipeline

df58099

Signed-off-by: Ao Tang <aot@nvidia.com>

Enhance OCR pipeline with new scoring and verification stages

0cb4be7

Signed-off-by: Ao Tang <aot@nvidia.com>

Add OCR conversationalization and output formatting for dense OCR data

f41918c

Signed-off-by: Ao Tang <aot@nvidia.com>

Add HFDatasetImageReaderStage for HuggingFace datasets and create OCR…

cc19ce9

… pipeline tutorial Signed-off-by: Ao Tang <aot@nvidia.com>

Remove per-stage runtime environment design document and example script

94ec830

Signed-off-by: Ao Tang <aot@nvidia.com>

Refactor OCR pipeline and result inspection scripts

3b9a139

Signed-off-by: Ao Tang <aot@nvidia.com>

remove

effe3aa

Signed-off-by: Ao Tang <aot@nvidia.com>

remove rtx

965a87c

Signed-off-by: Ao Tang <aot@nvidia.com>

remove generate_stream

71d5c02

Signed-off-by: Ao Tang <aot@nvidia.com>

refactor

3330d71

Signed-off-by: Ao Tang <aot@nvidia.com>

Refactor ParquetImageReader

2a634d5

Signed-off-by: Ao Tang <aot@nvidia.com>

remove

37a710f

Signed-off-by: Ao Tang <aot@nvidia.com>

Refactor JSON serialization in ResultWriterStage and update ImageTask…

344157a

…Data and OCRDenseWord classes Signed-off-by: Ao Tang <aot@nvidia.com>

Refactor RayClient metrics service discovery and improve error handli…

f5a2326

…ng in nvinference_client.py Signed-off-by: Ao Tang <aot@nvidia.com>

init file

bb74e10

Signed-off-by: Ao Tang <aot@nvidia.com>

refactor

924c8b5

Signed-off-by: Ao Tang <aot@nvidia.com>

greptile-apps Bot reviewed May 14, 2026

View reviewed changes

Comment thread nemo_curator/stages/synthetic/omni/ocr_scoring_qa.py

remove nvinferencec_client

d9bf271

Signed-off-by: Ao Tang <aot@nvidia.com>

greptile-apps Bot reviewed May 14, 2026

View reviewed changes

Comment thread nemo_curator/stages/synthetic/omni/ocr_scoring_qa.py

secrect in test

cd1224a

Signed-off-by: Ao Tang <aot@nvidia.com>

greptile-apps Bot reviewed May 14, 2026

View reviewed changes

Comment thread docker/Dockerfile

suiyoubi added 2 commits May 20, 2026 12:12

Update Dockerfile to pin NEMOTRON_OCR_REV for reproducible builds and…

82aa19d

… clean up OCRScoringQAStage by removing unused kwargs parameter. Signed-off-by: Ao Tang <aot@nvidia.com>

Merge branch 'main' into aot/omni_sdg

1c728b2

greptile-apps Bot reviewed May 20, 2026

View reviewed changes

praateekmahajan requested changes May 26, 2026

View reviewed changes

suiyoubi added 3 commits June 3, 2026 08:31

change endpoint to public nv APIs

652f21c

Signed-off-by: Ao Tang <aot@nvidia.com>

Merge branch 'main' into aot/omni_sdg

7d4d113

greptile-apps Bot reviewed Jun 11, 2026

View reviewed changes

praateekmahajan approved these changes Jun 11, 2026

View reviewed changes

Merge branch 'main' into aot/omni_sdg

85c82d5

remove task_id

fc82dc0

Signed-off-by: Ao Tang <aot@nvidia.com>

Merge branch 'main' into aot/omni_sdg

146c38b

thomasdhc approved these changes Jun 12, 2026

View reviewed changes

remyx-ai Bot mentioned this pull request Jun 16, 2026

OmniVideo-100K: A Dataset for Audio-Visual Reasoning through Structured Scripts and Evidence Chains smellslikeml/NeMo-Curator#2

Closed

		from nemo_curator.models.client import OpenAIClient


		class NVInferenceModel(ModelInterface):

-        self._file.write(
-            json.dumps({k: v for k, v in data.items() if v is not None and k != "is_valid"}, default=str) + "\n"
-        )
+        filtered = {k: v for k, v in data.items() if v is not None and k != "is_valid"}
+        if not task.data.is_valid:
+            filtered["is_valid"] = False
+        self._file.write(json.dumps(filtered, default=str) + "\n")

Conversation

suiyoubi commented Apr 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Pipeline stages

Key components

Tests

Checklist

Uh oh!

Uh oh!

Uh oh!

Uh oh!

greptile-apps Bot May 20, 2026

Choose a reason for hiding this comment

Uh oh!

praateekmahajan May 26, 2026

Choose a reason for hiding this comment

Uh oh!

suiyoubi Jun 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

praateekmahajan May 26, 2026

Choose a reason for hiding this comment

Uh oh!

suiyoubi Jun 3, 2026

Choose a reason for hiding this comment

Uh oh!

suiyoubi commented Jun 4, 2026

Uh oh!

greptile-apps Bot Jun 11, 2026

Choose a reason for hiding this comment

Uh oh!

suiyoubi commented Jun 11, 2026

Uh oh!

suiyoubi commented Jun 11, 2026

Uh oh!

suiyoubi commented Jun 11, 2026

Uh oh!

suiyoubi commented Jun 12, 2026

Uh oh!

This branch was previously deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

suiyoubi commented Apr 30, 2026 •

edited

Loading

suiyoubi Jun 3, 2026 •

edited

Loading