(retriever) fix remote ocr and pe logic to match local behavior (#1810)

Author: edknv

nemo_retriever/src/nemo_retriever/nim/nim.py

@@ -161,6 +161,7 @@ def invoke_image_inference_batches(
     *,
     invoke_url: str,
     image_b64_list: Sequence[str],
+    merge_levels: Optional[Sequence[str]] = None,
     api_key: Optional[str] = None,
     timeout_s: float = 60.0,
     max_batch_size: int = 8,
@@ -174,6 +175,14 @@ def invoke_image_inference_batches(
     `invoke_url` may be a single URL or a comma-separated URL list.
     When multiple URLs are provided, batches are distributed round-robin.
 
+    Parameters
+    ----------
+    merge_levels
+        Optional per-image merge level (``"word"``, ``"sentence"``, or
+        ``"paragraph"``).  When provided, must have the same length as
+        *image_b64_list*.  Passed as ``merge_levels`` in the JSON payload
+        so the NIM can apply per-crop merging behaviour.
+
     Returns one response item per input image, in the same order.
     """
     invoke_urls = _parse_invoke_urls(invoke_url)
@@ -187,6 +196,9 @@ def invoke_image_inference_batches(
     if n == 0:
         return []
 
+    if merge_levels is not None and len(merge_levels) != n:
+        raise ValueError(f"merge_levels length ({len(merge_levels)}) must match image_b64_list length ({n})")
+
     ranges = _chunk_ranges(n, int(max_batch_size))
     flattened: List[Optional[Any]] = [None] * n
 
@@ -198,7 +210,9 @@ def _invoke_one_batch(start: int, end: int, endpoint_url: str) -> Tuple[int, int
             }
             for b64 in image_b64_list[start:end]
         ]
-        payload = {"input": inputs}
+        payload: Dict[str, Any] = {"input": inputs}
+        if merge_levels is not None:
+            payload["merge_levels"] = list(merge_levels[start:end])
         response_json = _post_with_retries(
             invoke_url=endpoint_url,
             payload=payload,

nemo_retriever/src/nemo_retriever/ocr/ocr.py

@@ -562,11 +562,16 @@ def ocr_page_elements(
                 crops = _crop_all_from_page(page_image_b64, dets, row_wanted, as_b64=True)
                 crop_b64s: List[str] = [b64 for _label, _bbox, b64 in crops]
                 crop_meta: List[Tuple[str, List[float]]] = [(label, bbox) for label, bbox, _b64 in crops]
+                # Tables need word-level merging; everything else uses paragraph.
+                crop_merge_levels: List[str] = [
+                    "word" if label == "table" else "paragraph" for label, _bbox, _b64 in crops
+                ]
 
                 if crop_b64s:
                     response_items = invoke_image_inference_batches(
                         invoke_url=invoke_url,
                         image_b64_list=crop_b64s,
+                        merge_levels=crop_merge_levels,
                         api_key=api_key,
                         timeout_s=float(request_timeout_s),
                         max_batch_size=int(kwargs.get("inference_batch_size", 8)),

nemo_retriever/src/nemo_retriever/page_elements/page_elements.py

@@ -417,7 +417,7 @@ def _remote_response_to_detections(
                 batch_size=1,
                 label_names=label_names,
             )[0]
-            return _apply_page_elements_v3_postprocess(dets)
+            return _apply_final_score_filter(_apply_page_elements_v3_postprocess(dets))
         except Exception:
             pass
 
@@ -430,7 +430,7 @@ def _remote_response_to_detections(
         if isinstance(bb, dict):
             try:
                 dets = _bounding_boxes_to_detections(bb)
-                return _apply_page_elements_v3_postprocess(dets)
+                return _apply_final_score_filter(_apply_page_elements_v3_postprocess(dets))
             except Exception:
                 pass
 
@@ -442,7 +442,7 @@ def _remote_response_to_detections(
         if all(isinstance(v, list) for v in cand.values()):
             try:
                 dets = _annotation_dict_to_detections(cand)  # type: ignore[arg-type]
-                return _apply_page_elements_v3_postprocess(dets)
+                return _apply_final_score_filter(_apply_page_elements_v3_postprocess(dets))
             except Exception:
                 pass