diff --git a/nemo_retriever/src/nemo_retriever/common/modality/ocr/shared.py b/nemo_retriever/src/nemo_retriever/common/modality/ocr/shared.py index 15cdb59bf..2c1069226 100644 --- a/nemo_retriever/src/nemo_retriever/common/modality/ocr/shared.py +++ b/nemo_retriever/src/nemo_retriever/common/modality/ocr/shared.py @@ -781,7 +781,7 @@ def _append_ocr_prediction( row_result.chart_items.append(entry) elif label_name == "infographic": row_result.infographic_items.append(entry) - elif label_name in _TEXT_LABELS: + elif label_name in _TEXT_LABELS or label_name == "full_page": row_result.text_blocks.extend(blocks) @@ -820,6 +820,10 @@ def _run_remote_ocr( prepared.wanted_labels, as_b64=True, ) + meta = getattr(prepared.row, "metadata", None) or {} + needs_text = meta.get("needs_ocr_for_text", False) if isinstance(meta, dict) else False + if needs_text and not any(label in _TEXT_LABELS for label, _bbox, _crop_b64 in crops): + crops.append(("full_page", [0.0, 0.0, 1.0, 1.0], prepared.page_image_b64)) crop_b64s: List[str] = [crop_b64 for _label, _bbox, crop_b64 in crops] crop_metadata: List[Tuple[str, List[float]]] = [(label_name, bbox) for label_name, bbox, _crop_b64 in crops] if not crop_b64s: @@ -874,6 +878,17 @@ def _collect_local_crop_jobs( prepared.detections, prepared.wanted_labels, ) + meta = getattr(prepared.row, "metadata", None) or {} + needs_text = meta.get("needs_ocr_for_text", False) if isinstance(meta, dict) else False + if needs_text and not any(label in _TEXT_LABELS for label, _bbox, _crop_arr in crops): + try: + raw = base64.b64decode(prepared.page_image_b64) + with Image.open(io.BytesIO(raw)) as im0: + full_crop = np.asarray(im0.convert("RGB"), dtype=np.uint8).copy() + crops.append(("full_page", [0.0, 0.0, 1.0, 1.0], full_crop)) + except Exception: + if not crops: + crops = [] for label_name, bbox, crop_array in crops: merge_level = "word" if label_name == "table" else "paragraph" jobs_by_merge_level[merge_level].append( diff --git a/nemo_retriever/src/nemo_retriever/graph/ingestor_runtime.py b/nemo_retriever/src/nemo_retriever/graph/ingestor_runtime.py index 9c99d8537..b9cc64054 100644 --- a/nemo_retriever/src/nemo_retriever/graph/ingestor_runtime.py +++ b/nemo_retriever/src/nemo_retriever/graph/ingestor_runtime.py @@ -806,7 +806,7 @@ def build_graph( detect_kwargs["inference_batch_size"] = int(extract_params.inference_batch_size) ocr_kwargs: dict[str, Any] = {} - if extract_params.method in ("pdfium_hybrid", "ocr") and extract_params.extract_text: + if extract_params.extract_text: ocr_kwargs["extract_text"] = True if extract_params.extract_tables: ocr_kwargs["extract_tables"] = True diff --git a/nemo_retriever/src/nemo_retriever/operators/extract/pdf/extract.py b/nemo_retriever/src/nemo_retriever/operators/extract/pdf/extract.py index 8cc78961c..7bba2a641 100644 --- a/nemo_retriever/src/nemo_retriever/operators/extract/pdf/extract.py +++ b/nemo_retriever/src/nemo_retriever/operators/extract/pdf/extract.py @@ -326,7 +326,7 @@ def pdf_extraction( is_scanned_page = _is_scanned_page(page) ocr_extraction_needed_for_text = extract_text and ( - (text_extraction_method == "pdfium_hybrid" and is_scanned_page) + (text_extraction_method in ("pdfium", "pdfium_hybrid") and is_scanned_page) or text_extraction_method == "ocr" ) diff --git a/nemo_retriever/src/nemo_retriever/operators/graph_ops/multi_type_extract_operator.py b/nemo_retriever/src/nemo_retriever/operators/graph_ops/multi_type_extract_operator.py index 46592f02a..d18202037 100644 --- a/nemo_retriever/src/nemo_retriever/operators/graph_ops/multi_type_extract_operator.py +++ b/nemo_retriever/src/nemo_retriever/operators/graph_ops/multi_type_extract_operator.py @@ -103,7 +103,7 @@ def _parse_mode_enabled(extract_params: ExtractParams) -> bool: def _ocr_stage_needed(extract_params: ExtractParams) -> bool: - if extract_params.method in ("pdfium_hybrid", "ocr") and extract_params.extract_text: + if extract_params.extract_text: return True if extract_params.extract_tables: # OCR is always needed for table crops: either to produce pseudo-markdown @@ -481,7 +481,7 @@ def _run_detection_pipeline(self, batch_df: pd.DataFrame) -> pd.DataFrame: } if ocr_lang is not None: ocr_kwargs["ocr_lang"] = ocr_lang - if extract_params.method in ("pdfium_hybrid", "ocr") and extract_params.extract_text: + if extract_params.extract_text: ocr_kwargs["extract_text"] = True if extract_params.extract_tables: ocr_kwargs["extract_tables"] = True diff --git a/nemo_retriever/tests/test_ocr_version_selection.py b/nemo_retriever/tests/test_ocr_version_selection.py index 000421516..9c691dafd 100644 --- a/nemo_retriever/tests/test_ocr_version_selection.py +++ b/nemo_retriever/tests/test_ocr_version_selection.py @@ -333,3 +333,266 @@ def _fake_resolve(operator_class, resources, operator_kwargs=None): table_kwargs = next(kwargs for name, kwargs in captured_kwargs if name == "TableStructureActor") assert table_kwargs.get("ocr_version") == "v2" assert table_kwargs.get("ocr_lang") == "english" + + +# --------------------------------------------------------------------------- +# Regression tests for Issue #2443 — default pipeline skips OCR extract_text +# --------------------------------------------------------------------------- + + +def test_default_pdfium_method_includes_ocr_when_extract_text_true() -> None: + """Default method='pdfium' must still include OCRActor when extract_text=True. + + Regression test for https://github.com/NVIDIA/NeMo-Retriever/issues/2443. + Previously the OCR stage was only appended when method was 'pdfium_hybrid' + or 'ocr', causing the default pipeline to skip text extraction for scanned + images and PDFs. + """ + graph = build_graph( + extract_params=ExtractParams( + method="pdfium", + extract_text=True, + extract_tables=False, + extract_charts=False, + extract_infographics=False, + ), + embed_params=EmbedParams( + model_name="nvidia/llama-nemotron-embed-1b-v2", + embed_invoke_url="http://embed.example/v1", + ), + ) + + nodes = _linear_nodes(graph) + classes = [node.operator_class for node in nodes] + assert OCRActor in classes + ocr_node = next(node for node in nodes if node.operator_class is OCRActor) + assert ocr_node.operator_kwargs.get("extract_text") is True + + +def test_default_pdfium_method_excludes_ocr_when_extract_text_false() -> None: + """method='pdfium' with extract_text=False must not include OCRActor.""" + graph = build_graph( + extract_params=ExtractParams( + method="pdfium", + extract_text=False, + extract_tables=False, + extract_charts=False, + extract_infographics=False, + ), + embed_params=EmbedParams( + model_name="nvidia/llama-nemotron-embed-1b-v2", + embed_invoke_url="http://embed.example/v1", + ), + ) + + nodes = _linear_nodes(graph) + classes = [node.operator_class for node in nodes] + assert OCRActor not in classes + + +def test_ocr_stage_needed_true_for_default_method_extract_text() -> None: + """_ocr_stage_needed must return True when extract_text=True regardless of method.""" + from nemo_retriever.operators.graph_ops.multi_type_extract_operator import _ocr_stage_needed + + params = ExtractParams(method="pdfium", extract_text=True) + assert _ocr_stage_needed(params) is True + + +def test_ocr_stage_needed_false_when_extract_text_disabled() -> None: + """_ocr_stage_needed must return False when extract_text=False and no other OCR flags.""" + from nemo_retriever.operators.graph_ops.multi_type_extract_operator import _ocr_stage_needed + + params = ExtractParams( + method="pdfium", + extract_text=False, + extract_tables=False, + extract_charts=False, + extract_infographics=False, + ) + assert _ocr_stage_needed(params) is False + + +def test_detection_pipeline_includes_ocr_for_default_method(monkeypatch) -> None: + """In-process detection pipeline must include OCR when method='pdfium' and extract_text=True.""" + from nemo_retriever.operators.graph_ops.multi_type_extract_operator import MultiTypeExtractCPUActor + from nemo_retriever.common.ray_resource_hueristics import Resources + import pandas as pd + + calls = [] + + class _IdentityStage: + def __init__(self, **kwargs): + self.kwargs = kwargs + + def run(self, data): + return data + + def _fake_resolve(operator_class, resources, operator_kwargs=None): + calls.append(operator_class.__name__) + return _IdentityStage + + monkeypatch.setattr( + "nemo_retriever.operators.graph_ops.multi_type_extract_operator.resolve_operator_class", + _fake_resolve, + ) + monkeypatch.setattr( + "nemo_retriever.common.ray_resource_hueristics.gather_local_resources", + lambda: Resources(cpu_count=8, gpu_count=1), + ) + + op = MultiTypeExtractCPUActor( + extraction_mode="image", + extract_params=ExtractParams( + method="pdfium", + extract_text=True, + extract_tables=False, + extract_charts=False, + extract_infographics=False, + ), + ) + + op._run_detection_pipeline(pd.DataFrame({"page_image": ["x"]})) + assert "OCRActor" in calls + + +def test_pdf_extraction_sets_needs_ocr_for_scanned_default_method() -> None: + """PDF extraction must set needs_ocr_for_text=True for scanned pages under default method='pdfium'. + + Regression test for https://github.com/NVIDIA/NeMo-Retriever/issues/2443 + (Greptile review finding: PDF OCR eligibility remains disabled). + Previously ocr_extraction_needed_for_text was gated on + method in ('pdfium_hybrid', 'ocr'), so scanned PDF pages with + method='pdfium' never received OCR text extraction. + """ + import io + from unittest.mock import patch + + import pandas as pd + + pdfium = pytest.importorskip("pypdfium2") + from nemo_retriever.operators.extract.pdf.extract import pdf_extraction + + doc = pdfium.PdfDocument.new() + doc.new_page(612, 792) + buf = io.BytesIO() + doc.save(buf) + doc.close() + + def _fake_is_scanned(_page): + return True + + with patch( + "nemo_retriever.operators.extract.pdf.extract._is_scanned_page", + side_effect=_fake_is_scanned, + ): + result = pdf_extraction( + pd.DataFrame([{"bytes": buf.getvalue(), "path": "scanned.pdf", "page_number": 1}]), + extract_text=True, + text_extraction_method="pdfium", + ) + + row = result.iloc[0] + assert row["text"] == "" + assert row["metadata"]["needs_ocr_for_text"] is True + + +def test_full_page_text_fallback_when_no_detections() -> None: + """When no element detections exist but needs_ocr_for_text is True, OCR must still produce text. + + Regression test for https://github.com/NVIDIA/NeMo-Retriever/issues/2443 + (Greptile review: 'fast-text OCR still produces empty text'). + + fast-text profile sets use_page_elements=False, so PageElementDetectionActor + is absent. Without a full-page fallback, _crop_all_from_page returns [] + and scanned PDF text remains empty even though OCR is scheduled. + """ + import base64 + import io + + import numpy as np + from PIL import Image + + from nemo_retriever.common.modality.ocr.shared import ( + _PreparedOCRRow, + _collect_local_crop_jobs, + ) + + img = Image.fromarray(np.zeros((8, 8, 3), dtype=np.uint8), "RGB") + buf = io.BytesIO() + img.save(buf, format="PNG") + page_b64 = base64.b64encode(buf.getvalue()).decode("ascii") + + class _FakeRow: + def __init__(self, **kwargs): + for k, v in kwargs.items(): + setattr(self, k, v) + + prepared = _PreparedOCRRow( + row_index=0, + row=_FakeRow(metadata={"needs_ocr_for_text": True}), + page_image_b64=page_b64, + detections=[], + wanted_labels={"text", "title", "header_footer"}, + ) + + from nemo_retriever.common.modality.ocr.shared import _OCRRowResult + + row_results = [_OCRRowResult()] + jobs = _collect_local_crop_jobs([prepared], row_results) + + all_jobs = jobs["word"] + jobs["paragraph"] + assert len(all_jobs) == 1 + assert all_jobs[0].label_name == "full_page" + assert all_jobs[0].bbox == [0.0, 0.0, 1.0, 1.0] + + +def test_full_page_text_fallback_when_only_structured_detections() -> None: + """When detections contain only structured elements (no text regions) but + needs_ocr_for_text is True, a full-page text OCR crop must still be added. + + Regression test for https://github.com/NVIDIA/NeMo-Retriever/issues/2443 + (Greptile review: 'structured crops suppress full-page text fallback'). + + Scenario: page-elements detection finds a table but no text regions. + The table crop goes to table_items, not text_blocks. Without a full-page + fallback, text_blocks stays empty despite needs_ocr_for_text=True. + """ + import base64 + import io + + import numpy as np + from PIL import Image + + from nemo_retriever.common.modality.ocr.shared import ( + _PreparedOCRRow, + _collect_local_crop_jobs, + _OCRRowResult, + ) + + img = Image.fromarray(np.zeros((8, 8, 3), dtype=np.uint8), "RGB") + buf = io.BytesIO() + img.save(buf, format="PNG") + page_b64 = base64.b64encode(buf.getvalue()).decode("ascii") + + class _FakeRow: + def __init__(self, **kwargs): + for k, v in kwargs.items(): + setattr(self, k, v) + + # wanted_labels includes both structured and text labels (as _prepare_ocr_rows + # would set when extract_text=True + needs_ocr_for_text=True + extract_tables=True). + prepared = _PreparedOCRRow( + row_index=0, + row=_FakeRow(metadata={"needs_ocr_for_text": True}), + page_image_b64=page_b64, + detections=[{"label_name": "table", "bbox_xyxy_norm": [0.1, 0.1, 0.9, 0.9]}], + wanted_labels={"table", "text", "title", "header_footer"}, + ) + + row_results = [_OCRRowResult()] + jobs = _collect_local_crop_jobs([prepared], row_results) + + all_jobs = jobs["word"] + jobs["paragraph"] + label_names = [j.label_name for j in all_jobs] + assert "table" in label_names, "table crop must be present" + assert "full_page" in label_names, "full-page text crop must be added when needs_ocr_for_text=True"