From 9e42bd0a6d85d03de56da553f3b3cd6c05755809 Mon Sep 17 00:00:00 2001 From: Jeremy Dyer Date: Thu, 16 Apr 2026 19:01:01 -0400 Subject: [PATCH 1/7] First pass at async ray operators --- nemo_retriever/pyproject.toml | 1 + .../src/nemo_retriever/audio/asr_actor.py | 2 +- .../src/nemo_retriever/chart/cpu_actor.py | 20 +- .../src/nemo_retriever/chart/gpu_actor.py | 20 +- .../src/nemo_retriever/chart/shared.py | 228 ++++++++++ .../nemo_retriever/graph/abstract_operator.py | 24 +- .../nemo_retriever/graph/ingestor_runtime.py | 16 +- .../graph/operator_archetype.py | 11 +- .../tabular_fetch_embeddings_operator.py | 2 +- .../src/nemo_retriever/html/ray_data.py | 4 +- .../src/nemo_retriever/image/ray_data.py | 4 +- .../infographic/infographic_detection.py | 8 +- .../nemo_retriever/nim/chat_completions.py | 129 +++++- nemo_retriever/src/nemo_retriever/nim/nim.py | 207 ++++++++++ .../src/nemo_retriever/ocr/cpu_ocr.py | 15 +- .../src/nemo_retriever/ocr/cpu_parse.py | 18 +- .../src/nemo_retriever/ocr/gpu_ocr.py | 15 +- .../src/nemo_retriever/ocr/gpu_parse.py | 22 +- .../src/nemo_retriever/ocr/shared.py | 388 +++++++++++++++++- .../src/nemo_retriever/operators/__init__.py | 4 +- .../src/nemo_retriever/operators/embedding.py | 4 +- .../nemo_retriever/page_elements/cpu_actor.py | 14 +- .../nemo_retriever/page_elements/gpu_actor.py | 14 +- .../nemo_retriever/page_elements/shared.py | 129 +++++- .../nemo_retriever/parse/nemotron_parse.py | 8 +- .../src/nemo_retriever/pdf/extract.py | 4 +- .../src/nemo_retriever/pdf/split.py | 4 +- .../src/nemo_retriever/recall/core.py | 2 +- .../src/nemo_retriever/rerank/rerank.py | 8 +- .../src/nemo_retriever/retriever.py | 12 +- .../src/nemo_retriever/table/cpu_actor.py | 18 +- .../src/nemo_retriever/table/gpu_actor.py | 18 +- .../src/nemo_retriever/table/shared.py | 152 +++++++ .../nemo_retriever/text_embed/cpu_operator.py | 2 +- .../nemo_retriever/text_embed/gpu_operator.py | 6 +- .../nemo_retriever/text_embed/operators.py | 31 +- .../nemo_retriever/text_embed/text_embed.py | 4 +- .../src/nemo_retriever/txt/ray_data.py | 8 +- .../nemo_retriever/utils/convert/to_pdf.py | 4 +- nemo_retriever/tests/test_actor_operators.py | 92 +++-- nemo_retriever/tests/test_asr_actor.py | 14 +- .../tests/test_audio_chunk_actor.py | 10 +- .../tests/test_chart_graphic_elements.py | 8 +- nemo_retriever/tests/test_doc_to_pdf_actor.py | 8 +- nemo_retriever/tests/test_image_load.py | 15 +- nemo_retriever/tests/test_ingest_plans.py | 24 +- .../tests/test_nemotron_rerank_v2.py | 14 +- .../test_operator_flags_and_cpu_actors.py | 30 +- nemo_retriever/tests/test_table_structure.py | 8 +- nemo_retriever/uv.lock | 154 ++++++- 50 files changed, 1774 insertions(+), 183 deletions(-) diff --git a/nemo_retriever/pyproject.toml b/nemo_retriever/pyproject.toml index dbd0e2966..0cc09bd9b 100644 --- a/nemo_retriever/pyproject.toml +++ b/nemo_retriever/pyproject.toml @@ -40,6 +40,7 @@ dependencies = [ "uvicorn[standard]>=0.30.0", "python-multipart>=0.0.9", # HTTP clients + "aiohttp>=3.9.0", "httpx>=0.27.0", "requests>=2.32.5", "urllib3==2.6.3", diff --git a/nemo_retriever/src/nemo_retriever/audio/asr_actor.py b/nemo_retriever/src/nemo_retriever/audio/asr_actor.py index fe55bcd7c..283d8c6a3 100644 --- a/nemo_retriever/src/nemo_retriever/audio/asr_actor.py +++ b/nemo_retriever/src/nemo_retriever/audio/asr_actor.py @@ -393,4 +393,4 @@ def apply_asr_to_df( """ params = ASRParams(**(asr_params or {})) actor = ASRActor(params=params) - return actor(batch_df) + return actor.run(batch_df) diff --git a/nemo_retriever/src/nemo_retriever/chart/cpu_actor.py b/nemo_retriever/src/nemo_retriever/chart/cpu_actor.py index 021cf6fa8..64418564c 100644 --- a/nemo_retriever/src/nemo_retriever/chart/cpu_actor.py +++ b/nemo_retriever/src/nemo_retriever/chart/cpu_actor.py @@ -11,7 +11,7 @@ from nemo_retriever.graph.abstract_operator import AbstractOperator from nemo_retriever.graph.cpu_operator import CPUOperator from nemo_retriever.params import RemoteRetryParams -from nemo_retriever.chart.shared import graphic_elements_ocr_page_elements +from nemo_retriever.chart.shared import agraphic_elements_ocr_page_elements, graphic_elements_ocr_page_elements class GraphicElementsCPUActor(AbstractOperator, CPUOperator): @@ -73,9 +73,23 @@ def process(self, data: Any, **kwargs: Any) -> Any: def postprocess(self, data: Any, **kwargs: Any) -> Any: return data - def __call__(self, batch_df: Any, **override_kwargs: Any) -> Any: + async def aprocess(self, data: Any, **kwargs: Any) -> Any: + return await agraphic_elements_ocr_page_elements( + data, + graphic_elements_model=self._graphic_elements_model, + ocr_model=self._ocr_model, + graphic_elements_invoke_url=self._graphic_elements_invoke_url, + ocr_invoke_url=self._ocr_invoke_url, + api_key=self._api_key, + request_timeout_s=self._request_timeout_s, + remote_retry=self._remote_retry, + inference_batch_size=self._inference_batch_size, + **kwargs, + ) + + async def __call__(self, batch_df: Any, **override_kwargs: Any) -> Any: try: - return self.run(batch_df, **override_kwargs) + return await self.arun(batch_df, **override_kwargs) except BaseException as exc: if isinstance(batch_df, pd.DataFrame): out = batch_df.copy() diff --git a/nemo_retriever/src/nemo_retriever/chart/gpu_actor.py b/nemo_retriever/src/nemo_retriever/chart/gpu_actor.py index edf86de80..1c40530e9 100644 --- a/nemo_retriever/src/nemo_retriever/chart/gpu_actor.py +++ b/nemo_retriever/src/nemo_retriever/chart/gpu_actor.py @@ -11,7 +11,7 @@ from nemo_retriever.graph.abstract_operator import AbstractOperator from nemo_retriever.graph.gpu_operator import GPUOperator from nemo_retriever.params import RemoteRetryParams -from nemo_retriever.chart.shared import graphic_elements_ocr_page_elements +from nemo_retriever.chart.shared import agraphic_elements_ocr_page_elements, graphic_elements_ocr_page_elements class GraphicElementsActor(AbstractOperator, GPUOperator): @@ -79,9 +79,23 @@ def process(self, data: Any, **kwargs: Any) -> Any: def postprocess(self, data: Any, **kwargs: Any) -> Any: return data - def __call__(self, batch_df: Any, **override_kwargs: Any) -> Any: + async def aprocess(self, data: Any, **kwargs: Any) -> Any: + return await agraphic_elements_ocr_page_elements( + data, + graphic_elements_model=self._graphic_elements_model, + ocr_model=self._ocr_model, + graphic_elements_invoke_url=self._graphic_elements_invoke_url, + ocr_invoke_url=self._ocr_invoke_url, + api_key=self._api_key, + request_timeout_s=self._request_timeout_s, + remote_retry=self._remote_retry, + inference_batch_size=self._inference_batch_size, + **kwargs, + ) + + async def __call__(self, batch_df: Any, **override_kwargs: Any) -> Any: try: - return self.run(batch_df, **override_kwargs) + return await self.arun(batch_df, **override_kwargs) except BaseException as exc: if isinstance(batch_df, pd.DataFrame): out = batch_df.copy() diff --git a/nemo_retriever/src/nemo_retriever/chart/shared.py b/nemo_retriever/src/nemo_retriever/chart/shared.py index c34606234..8c0791b76 100644 --- a/nemo_retriever/src/nemo_retriever/chart/shared.py +++ b/nemo_retriever/src/nemo_retriever/chart/shared.py @@ -532,6 +532,234 @@ def graphic_elements_ocr_page_elements( return out +async def agraphic_elements_ocr_page_elements( + batch_df: Any, + *, + graphic_elements_model: Any = None, + ocr_model: Any = None, + graphic_elements_invoke_url: str = "", + ocr_invoke_url: str = "", + api_key: str = "", + request_timeout_s: float = 120.0, + remote_retry: RemoteRetryParams | None = None, + **kwargs: Any, +) -> Any: + """Async version of :func:`graphic_elements_ocr_page_elements`.""" + import asyncio + + from nemo_retriever.nim.nim import ainvoke_image_inference_batches + from nemo_retriever.ocr.ocr import ( + _blocks_to_text, + _crop_all_from_page, + _extract_remote_ocr_item, + _np_rgb_to_b64_png, + _parse_ocr_result, + ) + from nemo_retriever.utils.table_and_chart import join_graphic_elements_and_ocr_output + + retry = remote_retry or RemoteRetryParams( + remote_max_pool_workers=int(kwargs.get("remote_max_pool_workers", 16)), + remote_max_retries=int(kwargs.get("remote_max_retries", 10)), + remote_max_429_retries=int(kwargs.get("remote_max_429_retries", 5)), + ) + + if not isinstance(batch_df, pd.DataFrame): + raise NotImplementedError("agraphic_elements_ocr_page_elements currently only supports pandas.DataFrame input.") + + ge_url = (graphic_elements_invoke_url or kwargs.get("graphic_elements_invoke_url") or "").strip() + ocr_url = (ocr_invoke_url or kwargs.get("ocr_invoke_url") or "").strip() + use_remote_ge = bool(ge_url) + use_remote_ocr = bool(ocr_url) + + if not use_remote_ge and graphic_elements_model is None: + raise ValueError("A local `graphic_elements_model` is required when `graphic_elements_invoke_url` is not set.") + if not use_remote_ocr and ocr_model is None: + raise ValueError("A local `ocr_model` is required when `ocr_invoke_url` is not set.") + + label_names = _labels_from_model(graphic_elements_model) if graphic_elements_model is not None else [] + inference_batch_size = int(kwargs.get("inference_batch_size", 8)) + + all_chart: List[List[Dict[str, Any]]] = [] + all_meta: List[Dict[str, Any]] = [] + + t0_total = time.perf_counter() + + for row in batch_df.itertuples(index=False): + chart_items: List[Dict[str, Any]] = [] + row_error: Any = None + + try: + pe = getattr(row, "page_elements_v3", None) + dets: List[Dict[str, Any]] = [] + if isinstance(pe, dict): + dets = pe.get("detections") or [] + if not isinstance(dets, list): + dets = [] + + page_image = getattr(row, "page_image", None) or {} + page_image_b64 = page_image.get("image_b64") if isinstance(page_image, dict) else None + + if not isinstance(page_image_b64, str) or not page_image_b64: + all_chart.append(chart_items) + all_meta.append({"timing": None, "error": None}) + continue + + crops = _crop_all_from_page(page_image_b64, dets, {"chart"}) + + if not crops: + all_chart.append(chart_items) + all_meta.append({"timing": None, "error": None}) + continue + + crop_b64s = ( + [_np_rgb_to_b64_png(crop_array) for _, _, crop_array in crops] + if (use_remote_ge or use_remote_ocr) + else [] + ) + + ge_results: List[List[Dict[str, Any]]] = [] + ocr_results: List[Any] = [] + + if use_remote_ge and use_remote_ocr: + ge_task = ainvoke_image_inference_batches( + invoke_url=ge_url, + image_b64_list=crop_b64s, + api_key=api_key or None, + timeout_s=float(request_timeout_s), + max_batch_size=inference_batch_size, + max_concurrency=int(retry.remote_max_pool_workers), + max_retries=int(retry.remote_max_retries), + max_429_retries=int(retry.remote_max_429_retries), + ) + ocr_task = ainvoke_image_inference_batches( + invoke_url=ocr_url, + image_b64_list=crop_b64s, + api_key=api_key or None, + timeout_s=float(request_timeout_s), + max_batch_size=inference_batch_size, + max_concurrency=int(retry.remote_max_pool_workers), + max_retries=int(retry.remote_max_retries), + max_429_retries=int(retry.remote_max_429_retries), + ) + ge_items, ocr_items = await asyncio.gather(ge_task, ocr_task) + + if len(ge_items) != len(crops): + raise RuntimeError(f"Expected {len(crops)} GE responses, got {len(ge_items)}") + for resp in ge_items: + ge_results.append( + [ + d + for d in _remote_response_to_ge_detections(resp) + if (d.get("score") or 0.0) >= YOLOX_GRAPHIC_MIN_SCORE + ] + ) + if len(ocr_items) != len(crops): + raise RuntimeError(f"Expected {len(crops)} OCR responses, got {len(ocr_items)}") + for resp in ocr_items: + ocr_results.append(_extract_remote_ocr_item(resp)) + else: + if use_remote_ge: + ge_items = await ainvoke_image_inference_batches( + invoke_url=ge_url, + image_b64_list=crop_b64s, + api_key=api_key or None, + timeout_s=float(request_timeout_s), + max_batch_size=inference_batch_size, + max_concurrency=int(retry.remote_max_pool_workers), + max_retries=int(retry.remote_max_retries), + max_429_retries=int(retry.remote_max_429_retries), + ) + if len(ge_items) != len(crops): + raise RuntimeError(f"Expected {len(crops)} GE responses, got {len(ge_items)}") + for resp in ge_items: + ge_results.append( + [ + d + for d in _remote_response_to_ge_detections(resp) + if (d.get("score") or 0.0) >= YOLOX_GRAPHIC_MIN_SCORE + ] + ) + else: + + def _run_local_ge(): + results = [] + for _, _, crop_array in crops: + chw = torch.from_numpy(crop_array).permute(2, 0, 1).contiguous().to(dtype=torch.float32) + h, w = crop_array.shape[:2] + x = chw.unsqueeze(0) + try: + pre = graphic_elements_model.preprocess(x) + except Exception: + pre = x + if isinstance(pre, torch.Tensor) and pre.ndim == 3: + pre = pre.unsqueeze(0) + pred = graphic_elements_model.invoke(pre, (h, w)) + ge_dets = _prediction_to_detections(pred, label_names=label_names) + results.append([d for d in ge_dets if (d.get("score") or 0.0) >= YOLOX_GRAPHIC_MIN_SCORE]) + return results + + ge_results = await asyncio.to_thread(_run_local_ge) + + if use_remote_ocr: + ocr_items = await ainvoke_image_inference_batches( + invoke_url=ocr_url, + image_b64_list=crop_b64s, + api_key=api_key or None, + timeout_s=float(request_timeout_s), + max_batch_size=inference_batch_size, + max_concurrency=int(retry.remote_max_pool_workers), + max_retries=int(retry.remote_max_retries), + max_429_retries=int(retry.remote_max_429_retries), + ) + if len(ocr_items) != len(crops): + raise RuntimeError(f"Expected {len(crops)} OCR responses, got {len(ocr_items)}") + for resp in ocr_items: + ocr_results.append(_extract_remote_ocr_item(resp)) + else: + + def _run_local_ocr(): + results = [] + for _, _, crop_array in crops: + results.append(ocr_model.invoke(crop_array, merge_level="word")) + return results + + ocr_results = await asyncio.to_thread(_run_local_ocr) + + for crop_i, (label_name, bbox, crop_array) in enumerate(crops): + crop_hw = (int(crop_array.shape[0]), int(crop_array.shape[1])) + ge_dets = ge_results[crop_i] + ocr_preds = ocr_results[crop_i] + + text = join_graphic_elements_and_ocr_output(ge_dets, ocr_preds, crop_hw) + + if not text: + blocks = _parse_ocr_result(ocr_preds) + text = _blocks_to_text(blocks) + + chart_items.append({"bbox_xyxy_norm": bbox, "text": text}) + + except BaseException as e: + print(f"Warning: graphic-elements+OCR failed: {type(e).__name__}: {e}") + row_error = { + "stage": "graphic_elements_ocr_page_elements", + "type": e.__class__.__name__, + "message": str(e), + "traceback": "".join(traceback.format_exception(type(e), e, e.__traceback__)), + } + + all_chart.append(chart_items) + all_meta.append({"timing": None, "error": row_error}) + + elapsed = time.perf_counter() - t0_total + for meta in all_meta: + meta["timing"] = {"seconds": float(elapsed)} + + out = batch_df.copy() + out["chart"] = all_chart + out["graphic_elements_ocr_v1"] = all_meta + return out + + # --------------------------------------------------------------------------- # Combined graphic-elements + OCR Ray Actor # --------------------------------------------------------------------------- diff --git a/nemo_retriever/src/nemo_retriever/graph/abstract_operator.py b/nemo_retriever/src/nemo_retriever/graph/abstract_operator.py index b0f5ee508..6a97a3d6e 100644 --- a/nemo_retriever/src/nemo_retriever/graph/abstract_operator.py +++ b/nemo_retriever/src/nemo_retriever/graph/abstract_operator.py @@ -4,6 +4,7 @@ from __future__ import annotations +import asyncio from abc import ABC, abstractmethod import inspect from typing import Any, TYPE_CHECKING @@ -35,9 +36,26 @@ def run(self, data: Any, **kwargs: Any) -> Any: data = self.postprocess(data, **kwargs) return data - def __call__(self, data: Any, **kwargs: Any) -> Any: - """Make operators directly usable as Ray ``map_batches`` callables.""" - return self.run(data, **kwargs) + async def aprocess(self, data: Any, **kwargs: Any) -> Any: + """Async version of :meth:`process`. + + Default wraps the synchronous ``process`` in a thread so + compute-bound subclasses work without modification. I/O-bound + subclasses should override this with a proper ``await``-based + implementation. + """ + return await asyncio.to_thread(self.process, data, **kwargs) + + async def arun(self, data: Any, **kwargs: Any) -> Any: + """Async version of :meth:`run`.""" + data = self.preprocess(data, **kwargs) + data = await self.aprocess(data, **kwargs) + data = self.postprocess(data, **kwargs) + return data + + async def __call__(self, data: Any, **kwargs: Any) -> Any: + """Make operators directly usable as Ray ``map_batches`` async callables.""" + return await self.arun(data, **kwargs) def get_constructor_kwargs(self) -> dict[str, Any]: """Best-effort constructor kwargs for executor-side reconstruction.""" diff --git a/nemo_retriever/src/nemo_retriever/graph/ingestor_runtime.py b/nemo_retriever/src/nemo_retriever/graph/ingestor_runtime.py index 0963d96c4..02bd2270f 100644 --- a/nemo_retriever/src/nemo_retriever/graph/ingestor_runtime.py +++ b/nemo_retriever/src/nemo_retriever/graph/ingestor_runtime.py @@ -22,7 +22,7 @@ explode_content_to_rows, ) from nemo_retriever.graph.multi_type_extract_operator import MultiTypeExtractOperator -from nemo_retriever.text_embed.operators import _BatchEmbedActor +from nemo_retriever.text_embed.operators import BatchEmbedActor from nemo_retriever.ocr.ocr import OCRActor from nemo_retriever.parse.nemotron_parse import NemotronParseActor from nemo_retriever.page_elements.page_elements import PageElementDetectionActor @@ -107,9 +107,9 @@ def _force_cpu_only(node_name: str) -> None: embed_invoke_url = _positive(getattr(embed_params, "embed_invoke_url", None)) explicit_bs = getattr(embed_tuning, "embed_batch_size", None) if embed_tuning is not None else None embed_bs = _positive(explicit_bs) or (plan.embed_batch_size if plan else None) - _set(_BatchEmbedActor.__name__, "batch_size", embed_bs) + _set(BatchEmbedActor.__name__, "batch_size", embed_bs) if embed_bs: - overrides.setdefault(_BatchEmbedActor.__name__, {})["target_num_rows_per_block"] = embed_bs + overrides.setdefault(BatchEmbedActor.__name__, {})["target_num_rows_per_block"] = embed_bs embed_concurrency = ( _resolve( getattr(embed_tuning, "embed_workers", None) if embed_tuning is not None else None, @@ -117,19 +117,19 @@ def _force_cpu_only(node_name: str) -> None: ) or 0 ) - _set(_BatchEmbedActor.__name__, "concurrency", embed_concurrency or None) + _set(BatchEmbedActor.__name__, "concurrency", embed_concurrency or None) embed_cpus = ( _resolve( getattr(embed_tuning, "embed_cpus_per_actor", None) if embed_tuning is not None else None, ) or 1.0 ) - _set(_BatchEmbedActor.__name__, "num_cpus", embed_cpus if embed_cpus != 1.0 else None) + _set(BatchEmbedActor.__name__, "num_cpus", embed_cpus if embed_cpus != 1.0 else None) if effective_allow_no_gpu: - _force_cpu_only(_BatchEmbedActor.__name__) + _force_cpu_only(BatchEmbedActor.__name__) elif not embed_invoke_url: _set_gpu( - _BatchEmbedActor.__name__, + BatchEmbedActor.__name__, getattr(embed_tuning, "gpu_embed", None) if embed_tuning is not None else None, plan.embed_gpus_per_actor if plan else None, ) @@ -416,7 +416,7 @@ def _append_ordered_transform_stages( ), name="ExplodeContentToRows", ) - graph = graph >> _BatchEmbedActor(params=embed_params) + graph = graph >> BatchEmbedActor(params=embed_params) return graph diff --git a/nemo_retriever/src/nemo_retriever/graph/operator_archetype.py b/nemo_retriever/src/nemo_retriever/graph/operator_archetype.py index 3fb968263..c28e6e885 100644 --- a/nemo_retriever/src/nemo_retriever/graph/operator_archetype.py +++ b/nemo_retriever/src/nemo_retriever/graph/operator_archetype.py @@ -71,8 +71,15 @@ def postprocess(self, data: Any, **kwargs: Any) -> Any: def run(self, data: Any, **kwargs: Any) -> Any: return self._resolve_delegate().run(data, **kwargs) - def __call__(self, data: Any, **kwargs: Any) -> Any: - return self._resolve_delegate()(data, **kwargs) + async def aprocess(self, data: Any, **kwargs: Any) -> Any: + return await self._resolve_delegate().aprocess(data, **kwargs) + + async def arun(self, data: Any, **kwargs: Any) -> Any: + return await self._resolve_delegate().arun(data, **kwargs) + + async def __call__(self, data: Any, **kwargs: Any) -> Any: + delegate = self._resolve_delegate() + return await delegate(data, **kwargs) def _resolve_delegate(self, resources: ClusterResources | Resources | None = None) -> AbstractOperator: if not hasattr(self, "_resolved_delegate"): diff --git a/nemo_retriever/src/nemo_retriever/graph/tabular_fetch_embeddings_operator.py b/nemo_retriever/src/nemo_retriever/graph/tabular_fetch_embeddings_operator.py index a97cf22f1..0a45f76e9 100644 --- a/nemo_retriever/src/nemo_retriever/graph/tabular_fetch_embeddings_operator.py +++ b/nemo_retriever/src/nemo_retriever/graph/tabular_fetch_embeddings_operator.py @@ -22,7 +22,7 @@ class TabularFetchEmbeddingsOp(AbstractOperator, CPUOperator): ``text``, ``_embed_modality``, ``path``, ``page_number``, ``metadata``. The output schema matches the format produced by the unstructured pipeline, - so the standard :class:`~nemo_retriever.text_embed.operators._BatchEmbedActor` + so the standard :class:`~nemo_retriever.text_embed.operators.BatchEmbedActor` can be chained directly after this operator. """ diff --git a/nemo_retriever/src/nemo_retriever/html/ray_data.py b/nemo_retriever/src/nemo_retriever/html/ray_data.py index b872acfa5..3375c7d5f 100644 --- a/nemo_retriever/src/nemo_retriever/html/ray_data.py +++ b/nemo_retriever/src/nemo_retriever/html/ray_data.py @@ -63,8 +63,8 @@ def process(self, data: Any, **kwargs: Any) -> Any: def postprocess(self, data: Any, **kwargs: Any) -> Any: return data - def __call__(self, batch_df: pd.DataFrame) -> pd.DataFrame: - return self.run(batch_df) + async def __call__(self, batch_df: pd.DataFrame) -> pd.DataFrame: + return await self.arun(batch_df) class HtmlSplitActor(ArchetypeOperator): diff --git a/nemo_retriever/src/nemo_retriever/image/ray_data.py b/nemo_retriever/src/nemo_retriever/image/ray_data.py index 84e0e5f21..af9f488dc 100644 --- a/nemo_retriever/src/nemo_retriever/image/ray_data.py +++ b/nemo_retriever/src/nemo_retriever/image/ray_data.py @@ -75,8 +75,8 @@ def process(self, data: Any, **kwargs: Any) -> Any: def postprocess(self, data: Any, **kwargs: Any) -> Any: return data - def __call__(self, batch_df: pd.DataFrame) -> pd.DataFrame: - return self.run(batch_df) + async def __call__(self, batch_df: pd.DataFrame) -> pd.DataFrame: + return await self.arun(batch_df) class ImageLoadActor(ArchetypeOperator): diff --git a/nemo_retriever/src/nemo_retriever/infographic/infographic_detection.py b/nemo_retriever/src/nemo_retriever/infographic/infographic_detection.py index 0eba31ab2..6dafa73a4 100644 --- a/nemo_retriever/src/nemo_retriever/infographic/infographic_detection.py +++ b/nemo_retriever/src/nemo_retriever/infographic/infographic_detection.py @@ -775,9 +775,9 @@ def process(self, batch_df: Any, **override_kwargs: Any) -> Any: def postprocess(self, data: Any, **kwargs: Any) -> Any: return data - def __call__(self, batch_df: Any, **override_kwargs: Any) -> Any: + async def __call__(self, batch_df: Any, **override_kwargs: Any) -> Any: try: - return self.run(batch_df, **override_kwargs) + return await self.arun(batch_df, **override_kwargs) except BaseException as e: if isinstance(batch_df, pd.DataFrame): out = batch_df.copy() @@ -824,9 +824,9 @@ def process(self, batch_df: Any, **override_kwargs: Any) -> Any: def postprocess(self, data: Any, **kwargs: Any) -> Any: return data - def __call__(self, batch_df: Any, **override_kwargs: Any) -> Any: + async def __call__(self, batch_df: Any, **override_kwargs: Any) -> Any: try: - return self.run(batch_df, **override_kwargs) + return await self.arun(batch_df, **override_kwargs) except BaseException as e: if isinstance(batch_df, pd.DataFrame): out = batch_df.copy() diff --git a/nemo_retriever/src/nemo_retriever/nim/chat_completions.py b/nemo_retriever/src/nemo_retriever/nim/chat_completions.py index f6e9b6316..114a3205d 100644 --- a/nemo_retriever/src/nemo_retriever/nim/chat_completions.py +++ b/nemo_retriever/src/nemo_retriever/nim/chat_completions.py @@ -11,10 +11,18 @@ from __future__ import annotations +import asyncio from typing import Any, Dict, List, Optional, Sequence, Tuple from concurrent.futures import ThreadPoolExecutor, as_completed -from nemo_retriever.nim.nim import _parse_invoke_urls, _post_with_retries, _mime_from_b64 +import aiohttp + +from nemo_retriever.nim.nim import ( + _parse_invoke_urls, + _post_with_retries, + _apost_with_retries, + _mime_from_b64, +) def extract_chat_completion_text(response_json: Any) -> str: @@ -156,3 +164,122 @@ def invoke_chat_completions_images( max_retries=max_retries, max_429_retries=max_429_retries, ) + + +# --------------------------------------------------------------------------- +# Async variants (aiohttp) +# --------------------------------------------------------------------------- + + +async def ainvoke_chat_completions( + *, + invoke_url: str, + messages_list: Sequence[List[Dict[str, Any]]], + model: Optional[str] = None, + api_key: Optional[str] = None, + timeout_s: float = 120.0, + temperature: float = 0.0, + extra_body: Optional[Dict[str, Any]] = None, + max_concurrency: int = 16, + max_retries: int = 10, + max_429_retries: int = 5, +) -> List[str]: + """Async version of :func:`invoke_chat_completions` using aiohttp.""" + if not messages_list: + return [] + + token = (api_key or "").strip() + headers: Dict[str, str] = {"Accept": "application/json", "Content-Type": "application/json"} + if token: + headers["Authorization"] = f"Bearer {token}" + + invoke_urls = _parse_invoke_urls(invoke_url) + results: List[Optional[str]] = [None] * len(messages_list) + + async def _invoke_one( + session: aiohttp.ClientSession, + idx: int, + messages: List[Dict[str, Any]], + endpoint_url: str, + ) -> Tuple[int, str]: + payload: Dict[str, Any] = { + "messages": messages, + "temperature": temperature, + } + if model: + payload["model"] = model + if extra_body: + payload.update(extra_body) + response_json = await _apost_with_retries( + session=session, + invoke_url=endpoint_url, + payload=payload, + headers=headers, + timeout_s=float(timeout_s), + max_retries=int(max_retries), + max_429_retries=int(max_429_retries), + ) + return idx, extract_chat_completion_text(response_json) + + connector = aiohttp.TCPConnector(limit=max(1, int(max_concurrency))) + async with aiohttp.ClientSession(connector=connector) as session: + tasks = [ + _invoke_one(session, i, msgs, invoke_urls[i % len(invoke_urls)]) for i, msgs in enumerate(messages_list) + ] + completed = await asyncio.gather(*tasks) + + for i, text in completed: + results[i] = text + + return [r if r is not None else "" for r in results] + + +async def ainvoke_chat_completions_images( + *, + invoke_url: str, + image_b64_list: Sequence[str], + model: Optional[str] = None, + api_key: Optional[str] = None, + timeout_s: float = 120.0, + task_prompt: Optional[str] = None, + temperature: float = 0.0, + repetition_penalty: float = 1.1, + extra_body: Optional[Dict[str, Any]] = None, + max_concurrency: int = 16, + max_retries: int = 10, + max_429_retries: int = 5, +) -> List[str]: + """Async version of :func:`invoke_chat_completions_images`.""" + if not image_b64_list: + return [] + + messages_list: List[List[Dict[str, Any]]] = [] + for b64 in image_b64_list: + mime = _mime_from_b64(b64) + content: List[Dict[str, Any]] = [] + if task_prompt: + content.append({"type": "text", "text": task_prompt}) + content.append( + { + "type": "image_url", + "image_url": {"url": f"data:{mime};base64,{b64}"}, + } + ) + messages_list.append([{"role": "user", "content": content}]) + + merged_extra: Dict[str, Any] = {"repetition_penalty": repetition_penalty} + if extra_body: + merged_extra.update(extra_body) + + return await ainvoke_chat_completions( + invoke_url=invoke_url, + messages_list=messages_list, + model=model, + api_key=api_key, + timeout_s=timeout_s, + temperature=temperature, + extra_body=merged_extra, + max_concurrency=max_concurrency, + max_retries=max_retries, + max_429_retries=max_429_retries, + ) diff --git a/nemo_retriever/src/nemo_retriever/nim/nim.py b/nemo_retriever/src/nemo_retriever/nim/nim.py index 618a37aa1..37c824448 100644 --- a/nemo_retriever/src/nemo_retriever/nim/nim.py +++ b/nemo_retriever/src/nemo_retriever/nim/nim.py @@ -4,11 +4,13 @@ from __future__ import annotations +import asyncio import logging import time from concurrent.futures import ThreadPoolExecutor, as_completed from typing import Any, Dict, List, Optional, Sequence, Tuple +import aiohttp import requests logger = logging.getLogger(__name__) @@ -272,3 +274,208 @@ def invoke_page_elements_batches( max_retries=max_retries, max_429_retries=max_429_retries, ) + + +# --------------------------------------------------------------------------- +# Async variants (aiohttp) +# --------------------------------------------------------------------------- + + +async def _apost_with_retries( + *, + session: aiohttp.ClientSession, + invoke_url: str, + payload: Dict[str, Any], + headers: Dict[str, str], + timeout_s: float, + max_retries: int, + max_429_retries: int, +) -> Any: + base_delay = 2.0 + attempt = 0 + retries_429 = 0 + timeout = aiohttp.ClientTimeout(total=float(timeout_s)) + + while attempt < int(max_retries): + try: + async with session.post(invoke_url, json=payload, headers=headers, timeout=timeout) as response: + status_code = response.status + + if status_code == 429: + retries_429 += 1 + if retries_429 >= int(max_429_retries): + response.raise_for_status() + backoff_time = base_delay * (2**retries_429) + logger.warning( + "NIM endpoint %s returned 429 (rate limited). Retry %d/%d after %.1fs backoff.", + invoke_url, + retries_429, + max_429_retries, + backoff_time, + ) + await asyncio.sleep(backoff_time) + continue + + if status_code == 503 or (500 <= status_code < 600): + if attempt == int(max_retries) - 1: + response.raise_for_status() + backoff_time = base_delay * (2**attempt) + logger.warning( + "NIM endpoint %s returned %d. Retry %d/%d after %.1fs backoff.", + invoke_url, + status_code, + attempt + 1, + max_retries, + backoff_time, + ) + await asyncio.sleep(backoff_time) + attempt += 1 + continue + + if 400 <= status_code < 500: + body = await response.text() + raise aiohttp.ClientResponseError( + response.request_info, + response.history, + status=status_code, + message=f"HTTP {status_code} from {invoke_url}: {body}", + ) + + response.raise_for_status() + return await response.json() + + except asyncio.TimeoutError as exc: + if attempt == int(max_retries) - 1: + raise TimeoutError(f"Request timed out after {attempt + 1} attempts.") from exc + backoff_time = base_delay * (2**attempt) + logger.warning( + "NIM endpoint %s timed out (%.1fs). Retry %d/%d after %.1fs backoff.", + invoke_url, + timeout_s, + attempt + 1, + max_retries, + backoff_time, + ) + await asyncio.sleep(backoff_time) + attempt += 1 + except aiohttp.ClientError as exc: + resp_status = getattr(exc, "status", None) + if resp_status is not None and 400 <= resp_status < 500: + raise + if attempt == int(max_retries) - 1: + raise + backoff_time = base_delay * (2**attempt) + logger.warning( + "NIM endpoint %s request failed: %s. Retry %d/%d after %.1fs backoff.", + invoke_url, + exc, + attempt + 1, + max_retries, + backoff_time, + ) + await asyncio.sleep(backoff_time) + attempt += 1 + + raise RuntimeError(f"Failed to get a successful response after {max_retries} retries.") + + +async def ainvoke_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, + max_concurrency: int = 8, + max_retries: int = 5, + max_429_retries: int = 3, +) -> List[Any]: + """Async version of :func:`invoke_image_inference_batches` using aiohttp.""" + invoke_urls = _parse_invoke_urls(invoke_url) + + token = (api_key or "").strip() + headers: Dict[str, str] = {"Accept": "application/json", "Content-Type": "application/json"} + if token: + headers["Authorization"] = f"Bearer {token}" + + n = len(image_b64_list) + 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 + + async def _invoke_one_batch( + session: aiohttp.ClientSession, + start: int, + end: int, + endpoint_url: str, + ) -> Tuple[int, int, List[Any]]: + inputs = [ + { + "type": "image_url", + "url": f"data:{_mime_from_b64(b64)};base64,{b64}", + } + for b64 in image_b64_list[start:end] + ] + payload: Dict[str, Any] = {"input": inputs} + if merge_levels is not None: + payload["merge_levels"] = list(merge_levels[start:end]) + response_json = await _apost_with_retries( + session=session, + invoke_url=endpoint_url, + payload=payload, + headers=headers, + timeout_s=float(timeout_s), + max_retries=int(max_retries), + max_429_retries=int(max_429_retries), + ) + per_image = _normalize_batch_response(response_json, end - start) + return start, end, per_image + + connector = aiohttp.TCPConnector(limit=max(1, int(max_concurrency))) + async with aiohttp.ClientSession(connector=connector) as session: + tasks = [ + _invoke_one_batch(session, start, end, invoke_urls[idx % len(invoke_urls)]) + for idx, (start, end) in enumerate(ranges) + ] + results = await asyncio.gather(*tasks) + + for start, end, per_image in results: + for i, item in enumerate(per_image): + flattened[start + i] = item + + out: List[Any] = [] + for idx, item in enumerate(flattened): + if item is None: + raise RuntimeError(f"Missing response for item index {idx}") + out.append(item) + return out + + +async def ainvoke_page_elements_batches( + *, + invoke_url: str, + image_b64_list: Sequence[str], + api_key: Optional[str] = None, + timeout_s: float = 60.0, + max_batch_size: int = 8, + max_concurrency: int = 8, + max_retries: int = 5, + max_429_retries: int = 3, +) -> List[Any]: + """Async backward-compatible alias for page-elements callers.""" + return await ainvoke_image_inference_batches( + invoke_url=invoke_url, + image_b64_list=image_b64_list, + api_key=api_key, + timeout_s=timeout_s, + max_batch_size=max_batch_size, + max_concurrency=max_concurrency, + max_retries=max_retries, + max_429_retries=max_429_retries, + ) diff --git a/nemo_retriever/src/nemo_retriever/ocr/cpu_ocr.py b/nemo_retriever/src/nemo_retriever/ocr/cpu_ocr.py index 8d3f8a708..e367c6f35 100644 --- a/nemo_retriever/src/nemo_retriever/ocr/cpu_ocr.py +++ b/nemo_retriever/src/nemo_retriever/ocr/cpu_ocr.py @@ -12,7 +12,7 @@ from nemo_retriever.graph.cpu_operator import CPUOperator from nemo_retriever.params import RemoteRetryParams from nemo_retriever.ocr.shared import _error_payload -from nemo_retriever.ocr.shared import ocr_page_elements +from nemo_retriever.ocr.shared import aocr_page_elements, ocr_page_elements class OCRCPUActor(AbstractOperator, CPUOperator): @@ -59,9 +59,18 @@ def process(self, data: Any, **kwargs: Any) -> Any: def postprocess(self, data: Any, **kwargs: Any) -> Any: return data - def __call__(self, batch_df: Any, **override_kwargs: Any) -> Any: + async def aprocess(self, data: Any, **kwargs: Any) -> Any: + return await aocr_page_elements( + data, + model=self._model, + remote_retry=self._remote_retry, + **self.ocr_kwargs, + **kwargs, + ) + + async def __call__(self, batch_df: Any, **override_kwargs: Any) -> Any: try: - return self.run(batch_df, **override_kwargs) + return await self.arun(batch_df, **override_kwargs) except BaseException as exc: if isinstance(batch_df, pd.DataFrame): out = batch_df.copy() diff --git a/nemo_retriever/src/nemo_retriever/ocr/cpu_parse.py b/nemo_retriever/src/nemo_retriever/ocr/cpu_parse.py index ce31bb9d6..30d7b277c 100644 --- a/nemo_retriever/src/nemo_retriever/ocr/cpu_parse.py +++ b/nemo_retriever/src/nemo_retriever/ocr/cpu_parse.py @@ -9,7 +9,7 @@ from nemo_retriever.graph.abstract_operator import AbstractOperator from nemo_retriever.graph.cpu_operator import CPUOperator from nemo_retriever.params import RemoteRetryParams -from nemo_retriever.ocr.shared import nemotron_parse_page_elements +from nemo_retriever.ocr.shared import anemotron_parse_page_elements, nemotron_parse_page_elements class NemotronParseCPUActor(AbstractOperator, CPUOperator): @@ -70,3 +70,19 @@ def process(self, data: Any, **kwargs: Any) -> Any: def postprocess(self, data: Any, **kwargs: Any) -> Any: return data + + async def aprocess(self, data: Any, **kwargs: Any) -> Any: + return await anemotron_parse_page_elements( + data, + model=self._model, + invoke_url=self._invoke_url, + api_key=self._api_key, + request_timeout_s=self._request_timeout_s, + task_prompt=self._task_prompt, + extract_text=self._extract_text, + extract_tables=self._extract_tables, + extract_charts=self._extract_charts, + extract_infographics=self._extract_infographics, + remote_retry=self._remote_retry, + **kwargs, + ) diff --git a/nemo_retriever/src/nemo_retriever/ocr/gpu_ocr.py b/nemo_retriever/src/nemo_retriever/ocr/gpu_ocr.py index e26203992..3f7c425d2 100644 --- a/nemo_retriever/src/nemo_retriever/ocr/gpu_ocr.py +++ b/nemo_retriever/src/nemo_retriever/ocr/gpu_ocr.py @@ -11,7 +11,7 @@ from nemo_retriever.graph.abstract_operator import AbstractOperator from nemo_retriever.graph.gpu_operator import GPUOperator from nemo_retriever.params import RemoteRetryParams -from nemo_retriever.ocr.shared import Image, _error_payload, ocr_page_elements +from nemo_retriever.ocr.shared import Image, _error_payload, aocr_page_elements, ocr_page_elements class OCRActor(AbstractOperator, GPUOperator): @@ -64,9 +64,18 @@ def process(self, data: Any, **kwargs: Any) -> Any: def postprocess(self, data: Any, **kwargs: Any) -> Any: return data - def __call__(self, batch_df: Any, **override_kwargs: Any) -> Any: + async def aprocess(self, data: Any, **kwargs: Any) -> Any: + return await aocr_page_elements( + data, + model=self._model, + remote_retry=self._remote_retry, + **self.ocr_kwargs, + **kwargs, + ) + + async def __call__(self, batch_df: Any, **override_kwargs: Any) -> Any: try: - return self.run(batch_df, **override_kwargs) + return await self.arun(batch_df, **override_kwargs) except BaseException as exc: if isinstance(batch_df, pd.DataFrame): out = batch_df.copy() diff --git a/nemo_retriever/src/nemo_retriever/ocr/gpu_parse.py b/nemo_retriever/src/nemo_retriever/ocr/gpu_parse.py index eee10b306..3b731b7d7 100644 --- a/nemo_retriever/src/nemo_retriever/ocr/gpu_parse.py +++ b/nemo_retriever/src/nemo_retriever/ocr/gpu_parse.py @@ -11,7 +11,7 @@ from nemo_retriever.graph.abstract_operator import AbstractOperator from nemo_retriever.graph.gpu_operator import GPUOperator from nemo_retriever.params import RemoteRetryParams -from nemo_retriever.ocr.shared import _error_payload, nemotron_parse_page_elements +from nemo_retriever.ocr.shared import _error_payload, anemotron_parse_page_elements, nemotron_parse_page_elements class NemotronParseActor(AbstractOperator, GPUOperator): @@ -76,9 +76,25 @@ def process(self, data: Any, **kwargs: Any) -> Any: def postprocess(self, data: Any, **kwargs: Any) -> Any: return data - def __call__(self, batch_df: Any, **override_kwargs: Any) -> Any: + async def aprocess(self, data: Any, **kwargs: Any) -> Any: + return await anemotron_parse_page_elements( + data, + model=self._model, + invoke_url=self._invoke_url, + api_key=self._api_key, + request_timeout_s=self._request_timeout_s, + task_prompt=self._task_prompt, + extract_text=self._extract_text, + extract_tables=self._extract_tables, + extract_charts=self._extract_charts, + extract_infographics=self._extract_infographics, + remote_retry=self._remote_retry, + **kwargs, + ) + + async def __call__(self, batch_df: Any, **override_kwargs: Any) -> Any: try: - return self.run(batch_df, **override_kwargs) + return await self.arun(batch_df, **override_kwargs) except BaseException as exc: if isinstance(batch_df, pd.DataFrame): out = batch_df.copy() diff --git a/nemo_retriever/src/nemo_retriever/ocr/shared.py b/nemo_retriever/src/nemo_retriever/ocr/shared.py index 73a59ae26..9d8c190d3 100644 --- a/nemo_retriever/src/nemo_retriever/ocr/shared.py +++ b/nemo_retriever/src/nemo_retriever/ocr/shared.py @@ -22,7 +22,7 @@ import numpy as np import pandas as pd from nemo_retriever.params import RemoteRetryParams -from nemo_retriever.nim.nim import invoke_image_inference_batches +from nemo_retriever.nim.nim import ainvoke_image_inference_batches, invoke_image_inference_batches from nemo_retriever.utils.table_and_chart import join_graphic_elements_and_ocr_output try: @@ -975,3 +975,389 @@ def nemotron_parse_page_elements( out["infographic_parse"] = all_infographic out["nemotron_parse_v1_2"] = all_meta return out + + +async def aocr_page_elements( + batch_df: Any, + *, + model: Any = None, + invoke_url: Optional[str] = None, + api_key: Optional[str] = None, + request_timeout_s: float = 120.0, + extract_text: bool = False, + extract_tables: bool = False, + extract_charts: bool = False, + extract_infographics: bool = False, + use_graphic_elements: bool = False, + inference_batch_size: int = 8, + remote_retry: RemoteRetryParams | None = None, + **kwargs: Any, +) -> Any: + """Async version of :func:`ocr_page_elements`. + + Remote NIM calls use ``ainvoke_image_inference_batches``; local GPU + inference is delegated to ``asyncio.to_thread`` wrapping the sync version. + """ + import asyncio + + invoke_url_resolved = (invoke_url or kwargs.get("ocr_invoke_url") or "").strip() + use_remote = bool(invoke_url_resolved) + + if not use_remote: + return await asyncio.to_thread( + ocr_page_elements, + batch_df, + model=model, + invoke_url=invoke_url, + api_key=api_key, + request_timeout_s=request_timeout_s, + extract_text=extract_text, + extract_tables=extract_tables, + extract_charts=extract_charts, + extract_infographics=extract_infographics, + use_graphic_elements=use_graphic_elements, + inference_batch_size=inference_batch_size, + remote_retry=remote_retry, + **kwargs, + ) + + retry = remote_retry or RemoteRetryParams( + remote_max_pool_workers=int(kwargs.get("remote_max_pool_workers", 16)), + remote_max_retries=int(kwargs.get("remote_max_retries", 10)), + remote_max_429_retries=int(kwargs.get("remote_max_429_retries", 5)), + ) + + if not isinstance(batch_df, pd.DataFrame): + raise NotImplementedError("aocr_page_elements currently only supports pandas.DataFrame input.") + + wanted_labels: set[str] = set() + if extract_tables: + wanted_labels.add("table") + if extract_charts: + wanted_labels.add("chart") + if extract_infographics: + wanted_labels.add("infographic") + + all_table: List[List[Dict[str, Any]]] = [] + all_chart: List[List[Dict[str, Any]]] = [] + all_infographic: List[List[Dict[str, Any]]] = [] + all_text: List[str] = [] + all_ocr_meta: List[Dict[str, Any]] = [] + + t0_total = time.perf_counter() + + for row in batch_df.itertuples(index=False): + table_items: List[Dict[str, Any]] = [] + chart_items: List[Dict[str, Any]] = [] + infographic_items: List[Dict[str, Any]] = [] + row_ocr_text_blocks: List[Dict[str, Any]] = [] + row_error: Any = None + + try: + pe = getattr(row, "page_elements_v3", None) + dets: List[Dict[str, Any]] = [] + if isinstance(pe, dict): + dets = pe.get("detections") or [] + if not isinstance(dets, list): + dets = [] + + page_image = getattr(row, "page_image", None) or {} + page_image_b64 = page_image.get("image_b64") if isinstance(page_image, dict) else None + + if not isinstance(page_image_b64, str) or not page_image_b64: + all_table.append(table_items) + all_chart.append(chart_items) + all_infographic.append(infographic_items) + all_text.append(None) + all_ocr_meta.append({"timing": None, "error": None}) + continue + + row_wanted = wanted_labels + if extract_text: + meta = getattr(row, "metadata", None) or {} + needs_ocr = meta.get("needs_ocr_for_text", False) if isinstance(meta, dict) else False + if needs_ocr: + row_wanted = wanted_labels | _TEXT_LABELS + + 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] + + if crop_b64s: + response_items = await ainvoke_image_inference_batches( + invoke_url=invoke_url_resolved, + image_b64_list=crop_b64s, + api_key=api_key, + timeout_s=float(request_timeout_s), + max_batch_size=int(kwargs.get("inference_batch_size", 8)), + max_concurrency=int(retry.remote_max_pool_workers), + max_retries=int(retry.remote_max_retries), + max_429_retries=int(retry.remote_max_429_retries), + ) + if len(response_items) != len(crop_meta): + raise RuntimeError(f"Expected {len(crop_meta)} OCR responses, got {len(response_items)}") + + for i, (label_name, bbox) in enumerate(crop_meta): + preds = _extract_remote_ocr_item(response_items[i]) + + if label_name == "chart" and use_graphic_elements: + ge_dets = _find_ge_detections_for_bbox(row, bbox) + if ge_dets: + crop_hw = (0, 0) + try: + _raw = base64.b64decode(crop_b64s[i]) + with Image.open(io.BytesIO(_raw)) as _cim: + _cw, _ch = _cim.size + crop_hw = (_ch, _cw) + except Exception: + pass + text = join_graphic_elements_and_ocr_output(ge_dets, preds, crop_hw) + if text: + chart_items.append({"bbox_xyxy_norm": bbox, "text": text}) + continue + + blocks = _parse_ocr_result(preds) + if label_name == "table": + crop_hw_table: Tuple[int, int] = (0, 0) + try: + _raw = base64.b64decode(crop_b64s[i]) + with Image.open(io.BytesIO(_raw)) as _cim: + _cw, _ch = _cim.size + crop_hw_table = (_ch, _cw) + except Exception: + pass + text = _blocks_to_pseudo_markdown(blocks, crop_hw=crop_hw_table) or _blocks_to_text(blocks) + else: + text = _blocks_to_text(blocks) + entry = {"bbox_xyxy_norm": bbox, "text": text} + if label_name == "table": + table_items.append(entry) + elif label_name == "chart": + chart_items.append(entry) + elif label_name == "infographic": + infographic_items.append(entry) + elif label_name in _TEXT_LABELS: + row_ocr_text_blocks.extend(blocks) + + except BaseException as e: + print(f"Warning: OCR failed: {type(e).__name__}: {e}") + row_error = { + "stage": "ocr_page_elements", + "type": e.__class__.__name__, + "message": str(e), + "traceback": "".join(traceback.format_exception(type(e), e, e.__traceback__)), + } + + if extract_text and row_ocr_text_blocks: + all_text.append(_blocks_to_text(row_ocr_text_blocks)) + else: + all_text.append(None) + + all_table.append(table_items) + all_chart.append(chart_items) + all_infographic.append(infographic_items) + all_ocr_meta.append({"timing": None, "error": row_error}) + + elapsed = time.perf_counter() - t0_total + for meta in all_ocr_meta: + meta["timing"] = {"seconds": float(elapsed)} + + out = batch_df.copy() + if extract_tables or "table" not in out.columns: + out["table"] = all_table + if extract_charts or "chart" not in out.columns: + out["chart"] = all_chart + if extract_infographics or "infographic" not in out.columns: + out["infographic"] = all_infographic + if extract_text and "text" in out.columns: + for i, ocr_text in enumerate(all_text): + if ocr_text is not None: + out.iat[i, out.columns.get_loc("text")] = ocr_text + elif extract_text: + out["text"] = [t if t is not None else "" for t in all_text] + out["ocr_v1"] = all_ocr_meta + return out + + +async def anemotron_parse_page_elements( + batch_df: Any, + *, + model: Any = None, + invoke_url: Optional[str] = None, + api_key: Optional[str] = None, + request_timeout_s: float = 120.0, + extract_text: bool = False, + extract_tables: bool = False, + extract_charts: bool = False, + extract_infographics: bool = False, + task_prompt: str = "", + remote_retry: RemoteRetryParams | None = None, + **kwargs: Any, +) -> Any: + """Async version of :func:`nemotron_parse_page_elements`. + + Remote NIM calls use ``ainvoke_image_inference_batches``; local inference + is delegated to ``asyncio.to_thread`` wrapping the sync version. + """ + import asyncio + + invoke_url_resolved = (invoke_url or kwargs.get("nemotron_parse_invoke_url") or "").strip() + use_remote = bool(invoke_url_resolved) + + if not use_remote: + return await asyncio.to_thread( + nemotron_parse_page_elements, + batch_df, + model=model, + invoke_url=invoke_url, + api_key=api_key, + request_timeout_s=request_timeout_s, + extract_text=extract_text, + extract_tables=extract_tables, + extract_charts=extract_charts, + extract_infographics=extract_infographics, + task_prompt=task_prompt, + remote_retry=remote_retry, + **kwargs, + ) + + retry = remote_retry or RemoteRetryParams( + remote_max_pool_workers=int(kwargs.get("remote_max_pool_workers", 16)), + remote_max_retries=int(kwargs.get("remote_max_retries", 10)), + remote_max_429_retries=int(kwargs.get("remote_max_429_retries", 5)), + ) + + if not isinstance(batch_df, pd.DataFrame): + raise NotImplementedError("anemotron_parse_page_elements currently only supports pandas.DataFrame input.") + + wanted_labels: set[str] = set() + if extract_tables: + wanted_labels.add("table") + if extract_charts: + wanted_labels.add("chart") + if extract_infographics: + wanted_labels.add("infographic") + + all_table: List[List[Dict[str, Any]]] = [] + all_chart: List[List[Dict[str, Any]]] = [] + all_infographic: List[List[Dict[str, Any]]] = [] + all_text: List[str] = [] + all_meta: List[Dict[str, Any]] = [] + + t0_total = time.perf_counter() + + for row in batch_df.itertuples(index=False): + table_items: List[Dict[str, Any]] = [] + chart_items: List[Dict[str, Any]] = [] + infographic_items: List[Dict[str, Any]] = [] + row_text: Optional[str] = None + row_error: Any = None + + try: + pe = getattr(row, "page_elements_v3", None) + dets_list: List[Dict[str, Any]] = [] + if isinstance(pe, dict): + dets_list = pe.get("detections") or [] + if not isinstance(dets_list, list): + dets_list = [] + + page_image = getattr(row, "page_image", None) or {} + page_image_b64 = page_image.get("image_b64") if isinstance(page_image, dict) else None + if not isinstance(page_image_b64, str) or not page_image_b64: + all_table.append(table_items) + all_chart.append(chart_items) + all_infographic.append(infographic_items) + all_text.append(None) + all_meta.append({"timing": None, "error": None}) + continue + + crops = _crop_all_from_page(page_image_b64, dets_list, wanted_labels, as_b64=True) + if not crops and wanted_labels: + crops = [("full_page", [0.0, 0.0, 1.0, 1.0], page_image_b64)] + + crop_b64s: List[str] = [b64 for _label, _bbox, b64 in crops] + crop_meta_items: List[Tuple[str, List[float]]] = [(label, bbox) for label, bbox, _b64 in crops] + + if crop_b64s: + response_items = await ainvoke_image_inference_batches( + invoke_url=invoke_url_resolved, + image_b64_list=crop_b64s, + api_key=api_key, + timeout_s=float(request_timeout_s), + max_batch_size=int(kwargs.get("inference_batch_size", 8)), + max_concurrency=int(retry.remote_max_pool_workers), + max_retries=int(retry.remote_max_retries), + max_429_retries=int(retry.remote_max_429_retries), + ) + if len(response_items) != len(crop_meta_items): + raise RuntimeError(f"Expected {len(crop_meta_items)} Parse responses, got {len(response_items)}") + + for i, (label_name, bbox) in enumerate(crop_meta_items): + text = _extract_parse_text(response_items[i]) + entry = {"bbox_xyxy_norm": bbox, "text": text} + if label_name == "table": + table_items.append(entry) + elif label_name == "chart": + chart_items.append(entry) + elif label_name == "infographic": + infographic_items.append(entry) + elif label_name == "full_page": + if extract_tables: + table_items.append(dict(entry)) + if extract_charts: + chart_items.append(dict(entry)) + if extract_infographics: + infographic_items.append(dict(entry)) + + meta_row = getattr(row, "metadata", None) or {} + needs_ocr = meta_row.get("needs_ocr_for_text", False) if isinstance(meta_row, dict) else False + if extract_text and needs_ocr: + try: + resp = await ainvoke_image_inference_batches( + invoke_url=invoke_url_resolved, + image_b64_list=[page_image_b64], + api_key=api_key, + timeout_s=float(request_timeout_s), + max_batch_size=1, + max_concurrency=int(retry.remote_max_pool_workers), + max_retries=int(retry.remote_max_retries), + max_429_retries=int(retry.remote_max_429_retries), + ) + row_text = _extract_parse_text(resp[0]) if resp else "" + except Exception: + row_text = "" + + except BaseException as e: + print(f"Warning: Nemotron Parse failed: {type(e).__name__}: {e}") + row_error = { + "stage": "nemotron_parse_page_elements", + "type": e.__class__.__name__, + "message": str(e), + "traceback": "".join(traceback.format_exception(type(e), e, e.__traceback__)), + } + + all_text.append(row_text) + all_table.append(table_items) + all_chart.append(chart_items) + all_infographic.append(infographic_items) + all_meta.append({"timing": None, "error": row_error}) + + elapsed = time.perf_counter() - t0_total + for meta in all_meta: + meta["timing"] = {"seconds": float(elapsed)} + + out = batch_df.copy() + if extract_text and "text" in out.columns: + for i, parse_text in enumerate(all_text): + if parse_text is not None: + out.iat[i, out.columns.get_loc("text")] = parse_text + elif extract_text: + out["text"] = [t if t is not None else "" for t in all_text] + out["table"] = all_table + out["chart"] = all_chart + out["infographic"] = all_infographic + out["table_parse"] = all_table + out["chart_parse"] = all_chart + out["infographic_parse"] = all_infographic + out["nemotron_parse_v1_2"] = all_meta + return out diff --git a/nemo_retriever/src/nemo_retriever/operators/__init__.py b/nemo_retriever/src/nemo_retriever/operators/__init__.py index a0c692ae7..4a67f7b52 100644 --- a/nemo_retriever/src/nemo_retriever/operators/__init__.py +++ b/nemo_retriever/src/nemo_retriever/operators/__init__.py @@ -6,6 +6,6 @@ from nemo_retriever.operators.base import AbstractOperator, CPUOperator, GPUOperator from nemo_retriever.operators.content import ExplodeContentActor -from nemo_retriever.operators.embedding import _BatchEmbedActor +from nemo_retriever.operators.embedding import BatchEmbedActor -__all__ = ["AbstractOperator", "CPUOperator", "GPUOperator", "ExplodeContentActor", "_BatchEmbedActor"] +__all__ = ["AbstractOperator", "CPUOperator", "GPUOperator", "ExplodeContentActor", "BatchEmbedActor"] diff --git a/nemo_retriever/src/nemo_retriever/operators/embedding.py b/nemo_retriever/src/nemo_retriever/operators/embedding.py index 46543675b..09ba48e79 100644 --- a/nemo_retriever/src/nemo_retriever/operators/embedding.py +++ b/nemo_retriever/src/nemo_retriever/operators/embedding.py @@ -6,6 +6,6 @@ from __future__ import annotations -from nemo_retriever.text_embed.operators import _BatchEmbedActor, _BatchEmbedCPUActor, _BatchEmbedGPUActor +from nemo_retriever.text_embed.operators import BatchEmbedActor, BatchEmbedCPUActor, BatchEmbedGPUActor -__all__ = ["_BatchEmbedActor", "_BatchEmbedCPUActor", "_BatchEmbedGPUActor"] +__all__ = ["BatchEmbedActor", "BatchEmbedCPUActor", "BatchEmbedGPUActor"] diff --git a/nemo_retriever/src/nemo_retriever/page_elements/cpu_actor.py b/nemo_retriever/src/nemo_retriever/page_elements/cpu_actor.py index 74c6fec20..a5e16058c 100644 --- a/nemo_retriever/src/nemo_retriever/page_elements/cpu_actor.py +++ b/nemo_retriever/src/nemo_retriever/page_elements/cpu_actor.py @@ -10,7 +10,7 @@ from nemo_retriever.graph.abstract_operator import AbstractOperator from nemo_retriever.graph.cpu_operator import CPUOperator -from nemo_retriever.page_elements.shared import _error_payload, detect_page_elements_v3 +from nemo_retriever.page_elements.shared import _error_payload, adetect_page_elements_v3, detect_page_elements_v3 class PageElementDetectionCPUActor(AbstractOperator, CPUOperator): @@ -48,9 +48,17 @@ def process(self, data: Any, **kwargs: Any) -> Any: def postprocess(self, data: Any, **kwargs: Any) -> Any: return data - def __call__(self, pages_df: Any, **override_kwargs: Any) -> Any: + async def aprocess(self, data: Any, **kwargs: Any) -> Any: + return await adetect_page_elements_v3( + data, + model=self._model, + **self.detect_kwargs, + **kwargs, + ) + + async def __call__(self, pages_df: Any, **override_kwargs: Any) -> Any: try: - return self.run(pages_df, **override_kwargs) + return await self.arun(pages_df, **override_kwargs) except Exception as exc: if isinstance(pages_df, pd.DataFrame): out = pages_df.copy() diff --git a/nemo_retriever/src/nemo_retriever/page_elements/gpu_actor.py b/nemo_retriever/src/nemo_retriever/page_elements/gpu_actor.py index e730c193e..0c7e23135 100644 --- a/nemo_retriever/src/nemo_retriever/page_elements/gpu_actor.py +++ b/nemo_retriever/src/nemo_retriever/page_elements/gpu_actor.py @@ -10,7 +10,7 @@ from nemo_retriever.graph.abstract_operator import AbstractOperator from nemo_retriever.graph.gpu_operator import GPUOperator -from nemo_retriever.page_elements.shared import _error_payload, detect_page_elements_v3 +from nemo_retriever.page_elements.shared import _error_payload, adetect_page_elements_v3, detect_page_elements_v3 class PageElementDetectionActor(AbstractOperator, GPUOperator): @@ -50,9 +50,17 @@ def process(self, data: Any, **kwargs: Any) -> Any: def postprocess(self, data: Any, **kwargs: Any) -> Any: return data - def __call__(self, pages_df: Any, **override_kwargs: Any) -> Any: + async def aprocess(self, data: Any, **kwargs: Any) -> Any: + return await adetect_page_elements_v3( + data, + model=self._model, + **self.detect_kwargs, + **kwargs, + ) + + async def __call__(self, pages_df: Any, **override_kwargs: Any) -> Any: try: - return self.run(pages_df, **override_kwargs) + return await self.arun(pages_df, **override_kwargs) except Exception as exc: if isinstance(pages_df, pd.DataFrame): out = pages_df.copy() diff --git a/nemo_retriever/src/nemo_retriever/page_elements/shared.py b/nemo_retriever/src/nemo_retriever/page_elements/shared.py index e02dc163e..9502ce5a7 100644 --- a/nemo_retriever/src/nemo_retriever/page_elements/shared.py +++ b/nemo_retriever/src/nemo_retriever/page_elements/shared.py @@ -12,7 +12,7 @@ import traceback import pandas as pd -from nemo_retriever.nim.nim import invoke_page_elements_batches +from nemo_retriever.nim.nim import ainvoke_page_elements_batches, invoke_page_elements_batches from nemo_retriever.params import RemoteRetryParams from nemo_retriever.page_elements.local import ( YOLOX_PAGE_V3_CLASS_LABELS, @@ -727,3 +727,130 @@ def detect_page_elements_v3( _counts_by_label(p.get("detections") or []) if isinstance(p, dict) else {} for p in row_payloads ] return out + + +async def adetect_page_elements_v3( + pages_df: Any, + *, + model: Any = None, + invoke_url: Optional[str] = None, + api_key: Optional[str] = None, + request_timeout_s: float = 120.0, + inference_batch_size: int = 8, + output_column: str = "page_elements_v3", + num_detections_column: str = "page_elements_v3_num_detections", + counts_by_label_column: str = "page_elements_v3_counts_by_label", + remote_retry: RemoteRetryParams | None = None, + **kwargs: Any, +) -> Any: + """Async version of :func:`detect_page_elements_v3`. + + Uses ``ainvoke_page_elements_batches`` for non-blocking remote NIM calls. + Delegates local GPU inference to ``asyncio.to_thread`` wrapping the sync + version. + """ + import asyncio + + invoke_url_resolved = (invoke_url or kwargs.get("page_elements_invoke_url") or "").strip() + use_remote = bool(invoke_url_resolved) + + if not use_remote: + return await asyncio.to_thread( + detect_page_elements_v3, + pages_df, + model=model, + invoke_url=invoke_url, + api_key=api_key, + request_timeout_s=request_timeout_s, + inference_batch_size=inference_batch_size, + output_column=output_column, + num_detections_column=num_detections_column, + counts_by_label_column=counts_by_label_column, + remote_retry=remote_retry, + **kwargs, + ) + + retry = remote_retry or RemoteRetryParams( + remote_max_pool_workers=int(kwargs.get("remote_max_pool_workers", 16)), + remote_max_retries=int(kwargs.get("remote_max_retries", 10)), + remote_max_429_retries=int(kwargs.get("remote_max_429_retries", 5)), + ) + + if not isinstance(pages_df, pd.DataFrame): + raise NotImplementedError("adetect_page_elements_v3 currently only supports pandas.DataFrame input.") + + label_names = list(_RETRIEVER_LABEL_NAMES) + thresholds_per_class = [ + YOLOX_PAGE_V3_FINAL_SCORE.get(_RETRIEVER_TO_API.get(name, name), 0.0) for name in label_names + ] + + row_b64: List[Optional[str]] = [] + row_payloads: List[Dict[str, Any]] = [] + + for _, row in pages_df.iterrows(): + try: + b64 = row.get("page_image")["image_b64"] + if not b64: + raise ValueError("No usable image_b64 found in row.") + row_b64.append(b64) + row_payloads.append({"detections": []}) + except BaseException as e: + row_b64.append(None) + row_payloads.append(_error_payload(stage="decode_image", exc=e)) + + valid_indices = [i for i, b64 in enumerate(row_b64) if b64] + + if valid_indices: + valid_b64: List[str] = [] + for row_i in valid_indices: + b64 = row_b64[row_i] + if b64: + valid_b64.append(b64) + + t0 = time.perf_counter() + try: + response_items = await ainvoke_page_elements_batches( + invoke_url=invoke_url_resolved, + image_b64_list=valid_b64, + api_key=api_key, + timeout_s=float(request_timeout_s), + max_batch_size=int(inference_batch_size), + max_concurrency=int(retry.remote_max_pool_workers), + max_retries=int(retry.remote_max_retries), + max_429_retries=int(retry.remote_max_429_retries), + ) + elapsed = time.perf_counter() - t0 + + if len(response_items) != len(valid_indices): + raise RuntimeError( + "Remote response count mismatch: " f"expected {len(valid_indices)}, got {len(response_items)}" + ) + + for local_i, row_i in enumerate(valid_indices): + dets = _remote_response_to_detections( + response_json=response_items[local_i], + label_names=label_names, + thresholds_per_class=thresholds_per_class, + ) + row_payloads[row_i] = { + "detections": dets, + "timing": {"seconds": float(elapsed)}, + "error": None, + } + except BaseException as e: + elapsed = time.perf_counter() - t0 + print(f"Warning: page_elements remote inference failed: {type(e).__name__}: {e}") + for row_i in valid_indices: + row_payloads[row_i] = _error_payload(stage="remote_inference", exc=e) | { + "timing": {"seconds": float(elapsed)} + } + + out = pages_df.copy() + out[output_column] = row_payloads + out[num_detections_column] = [ + int(len(p.get("detections") or [])) if isinstance(p, dict) else 0 for p in row_payloads + ] + out[counts_by_label_column] = [ + _counts_by_label(p.get("detections") or []) if isinstance(p, dict) else {} for p in row_payloads + ] + return out diff --git a/nemo_retriever/src/nemo_retriever/parse/nemotron_parse.py b/nemo_retriever/src/nemo_retriever/parse/nemotron_parse.py index 4c475e39a..824efb67a 100644 --- a/nemo_retriever/src/nemo_retriever/parse/nemotron_parse.py +++ b/nemo_retriever/src/nemo_retriever/parse/nemotron_parse.py @@ -483,9 +483,9 @@ def process(self, data: Any, **kwargs: Any) -> Any: def postprocess(self, data: Any, **kwargs: Any) -> Any: return data - def __call__(self, batch_df: Any, **override_kwargs: Any) -> Any: + async def __call__(self, batch_df: Any, **override_kwargs: Any) -> Any: try: - return self.run(batch_df, **override_kwargs) + return await self.arun(batch_df, **override_kwargs) except BaseException as e: if isinstance(batch_df, pd.DataFrame): out = batch_df.copy() @@ -569,9 +569,9 @@ def process(self, data: Any, **kwargs: Any) -> Any: def postprocess(self, data: Any, **kwargs: Any) -> Any: return data - def __call__(self, batch_df: Any, **override_kwargs: Any) -> Any: + async def __call__(self, batch_df: Any, **override_kwargs: Any) -> Any: try: - return self.run(batch_df, **override_kwargs) + return await self.arun(batch_df, **override_kwargs) except BaseException as e: if isinstance(batch_df, pd.DataFrame): out = batch_df.copy() diff --git a/nemo_retriever/src/nemo_retriever/pdf/extract.py b/nemo_retriever/src/nemo_retriever/pdf/extract.py index 2b41183dc..2cee473fd 100644 --- a/nemo_retriever/src/nemo_retriever/pdf/extract.py +++ b/nemo_retriever/src/nemo_retriever/pdf/extract.py @@ -413,9 +413,9 @@ def process(self, data: Any, **kwargs: Any) -> Any: def postprocess(self, data: Any, **kwargs: Any) -> Any: return data - def __call__(self, pdf: Any, **override_kwargs: Any) -> Optional[Any]: + async def __call__(self, pdf: Any, **override_kwargs: Any) -> Optional[Any]: try: - return self.run(pdf, **override_kwargs) + return await self.arun(pdf, **override_kwargs) except BaseException as e: # As a last line of defense, never let the Ray UDF raise. source_path = None diff --git a/nemo_retriever/src/nemo_retriever/pdf/split.py b/nemo_retriever/src/nemo_retriever/pdf/split.py index 1c92cca1b..2858684d6 100644 --- a/nemo_retriever/src/nemo_retriever/pdf/split.py +++ b/nemo_retriever/src/nemo_retriever/pdf/split.py @@ -186,8 +186,8 @@ def process(self, data: Any, **kwargs: Any) -> Any: def postprocess(self, data: Any, **kwargs: Any) -> Any: return data - def __call__(self, pdf_batch: Any) -> Any: - return self.run(pdf_batch) + async def __call__(self, pdf_batch: Any) -> Any: + return await self.arun(pdf_batch) class PDFSplitActor(ArchetypeOperator): diff --git a/nemo_retriever/src/nemo_retriever/recall/core.py b/nemo_retriever/src/nemo_retriever/recall/core.py index e5be26dea..f70255024 100644 --- a/nemo_retriever/src/nemo_retriever/recall/core.py +++ b/nemo_retriever/src/nemo_retriever/recall/core.py @@ -522,7 +522,7 @@ def retrieve_and_score( lancedb_uri=cfg.lancedb_uri, lancedb_table=cfg.lancedb_table, embedder=cfg.embedding_model or VL_EMBED_MODEL, - embedding_http_endpoint=cfg.embedding_http_endpoint, + embedding_http_endpoint=endpoint, embedding_api_key=cfg.embedding_api_key, top_k=cfg.top_k, nprobes=cfg.nprobes, diff --git a/nemo_retriever/src/nemo_retriever/rerank/rerank.py b/nemo_retriever/src/nemo_retriever/rerank/rerank.py index abb33bf41..d1920e08c 100644 --- a/nemo_retriever/src/nemo_retriever/rerank/rerank.py +++ b/nemo_retriever/src/nemo_retriever/rerank/rerank.py @@ -390,9 +390,9 @@ def process(self, batch_df: Any, **override_kwargs: Any) -> Any: def postprocess(self, data: Any, **kwargs: Any) -> Any: return data - def __call__(self, batch_df: Any, **override_kwargs: Any) -> Any: + async def __call__(self, batch_df: Any, **override_kwargs: Any) -> Any: try: - return self.run(batch_df, **override_kwargs) + return await self.arun(batch_df, **override_kwargs) except BaseException as exc: if isinstance(batch_df, pd.DataFrame): out = batch_df.copy() @@ -424,9 +424,9 @@ def process(self, batch_df: Any, **override_kwargs: Any) -> Any: def postprocess(self, data: Any, **kwargs: Any) -> Any: return data - def __call__(self, batch_df: Any, **override_kwargs: Any) -> Any: + async def __call__(self, batch_df: Any, **override_kwargs: Any) -> Any: try: - return self.run(batch_df, **override_kwargs) + return await self.arun(batch_df, **override_kwargs) except BaseException as exc: if isinstance(batch_df, pd.DataFrame): out = batch_df.copy() diff --git a/nemo_retriever/src/nemo_retriever/retriever.py b/nemo_retriever/src/nemo_retriever/retriever.py index 59d3e585a..5ed2b0366 100644 --- a/nemo_retriever/src/nemo_retriever/retriever.py +++ b/nemo_retriever/src/nemo_retriever/retriever.py @@ -4,7 +4,7 @@ from __future__ import annotations -from dataclasses import dataclass, field +from dataclasses import dataclass, field, fields from pathlib import Path from typing import Any, Optional, Sequence from tqdm import tqdm @@ -94,6 +94,16 @@ class Retriever: # Internal cache for local HF embedders, keyed by model name. _embedder_cache: dict = field(default_factory=dict, init=False, repr=False, compare=False) + def __str__(self) -> str: + lines = [f"{self.__class__.__name__}("] + for f in fields(self): + if f.name.startswith("_"): + continue + val = getattr(self, f.name) + lines.append(f" {f.name}={val!r},") + lines.append(")") + return "\n".join(lines) + def _resolve_embedding_endpoint(self) -> Optional[str]: http_ep = self.embedding_http_endpoint.strip() if isinstance(self.embedding_http_endpoint, str) else None single = self.embedding_endpoint.strip() if isinstance(self.embedding_endpoint, str) else None diff --git a/nemo_retriever/src/nemo_retriever/table/cpu_actor.py b/nemo_retriever/src/nemo_retriever/table/cpu_actor.py index d38fc765e..f89b8f0f7 100644 --- a/nemo_retriever/src/nemo_retriever/table/cpu_actor.py +++ b/nemo_retriever/src/nemo_retriever/table/cpu_actor.py @@ -11,7 +11,7 @@ from nemo_retriever.graph.abstract_operator import AbstractOperator from nemo_retriever.graph.cpu_operator import CPUOperator from nemo_retriever.params import RemoteRetryParams -from nemo_retriever.table.shared import table_structure_ocr_page_elements +from nemo_retriever.table.shared import atable_structure_ocr_page_elements, table_structure_ocr_page_elements class TableStructureCPUActor(AbstractOperator, CPUOperator): @@ -68,9 +68,21 @@ def process(self, data: Any, **kwargs: Any) -> Any: def postprocess(self, data: Any, **kwargs: Any) -> Any: return data - def __call__(self, batch_df: Any, **override_kwargs: Any) -> Any: + async def aprocess(self, data: Any, **kwargs: Any) -> Any: + return await atable_structure_ocr_page_elements( + data, + table_structure_model=self._table_structure_model, + table_structure_invoke_url=self._table_structure_invoke_url, + api_key=self._api_key, + table_output_format=self._table_output_format, + request_timeout_s=self._request_timeout_s, + remote_retry=self._remote_retry, + **kwargs, + ) + + async def __call__(self, batch_df: Any, **override_kwargs: Any) -> Any: try: - return self.run(batch_df, **override_kwargs) + return await self.arun(batch_df, **override_kwargs) except BaseException as exc: if isinstance(batch_df, pd.DataFrame): out = batch_df.copy() diff --git a/nemo_retriever/src/nemo_retriever/table/gpu_actor.py b/nemo_retriever/src/nemo_retriever/table/gpu_actor.py index 31c7919a2..a2cca2066 100644 --- a/nemo_retriever/src/nemo_retriever/table/gpu_actor.py +++ b/nemo_retriever/src/nemo_retriever/table/gpu_actor.py @@ -11,7 +11,7 @@ from nemo_retriever.graph.abstract_operator import AbstractOperator from nemo_retriever.graph.gpu_operator import GPUOperator from nemo_retriever.params import RemoteRetryParams -from nemo_retriever.table.shared import table_structure_ocr_page_elements +from nemo_retriever.table.shared import atable_structure_ocr_page_elements, table_structure_ocr_page_elements class TableStructureActor(AbstractOperator, GPUOperator): @@ -69,9 +69,21 @@ def process(self, data: Any, **kwargs: Any) -> Any: def postprocess(self, data: Any, **kwargs: Any) -> Any: return data - def __call__(self, batch_df: Any, **override_kwargs: Any) -> Any: + async def aprocess(self, data: Any, **kwargs: Any) -> Any: + return await atable_structure_ocr_page_elements( + data, + table_structure_model=self._table_structure_model, + table_structure_invoke_url=self._table_structure_invoke_url, + api_key=self._api_key, + table_output_format=self._table_output_format, + request_timeout_s=self._request_timeout_s, + remote_retry=self._remote_retry, + **kwargs, + ) + + async def __call__(self, batch_df: Any, **override_kwargs: Any) -> Any: try: - return self.run(batch_df, **override_kwargs) + return await self.arun(batch_df, **override_kwargs) except BaseException as exc: if isinstance(batch_df, pd.DataFrame): out = batch_df.copy() diff --git a/nemo_retriever/src/nemo_retriever/table/shared.py b/nemo_retriever/src/nemo_retriever/table/shared.py index 59579cf3b..219910892 100644 --- a/nemo_retriever/src/nemo_retriever/table/shared.py +++ b/nemo_retriever/src/nemo_retriever/table/shared.py @@ -422,6 +422,158 @@ def table_structure_ocr_page_elements( return out +async def atable_structure_ocr_page_elements( + batch_df: Any, + *, + table_structure_model: Any = None, + table_structure_invoke_url: str = "", + api_key: str = "", + request_timeout_s: float = 120.0, + remote_retry: RemoteRetryParams | None = None, + **kwargs: Any, +) -> Any: + """Async version of :func:`table_structure_ocr_page_elements`. + + Uses ``ainvoke_image_inference_batches`` for non-blocking remote NIM calls. + Falls back to ``asyncio.to_thread`` for local GPU inference. + """ + import asyncio + + from nemo_retriever.nim.nim import ainvoke_image_inference_batches + from nemo_retriever.ocr.ocr import _crop_all_from_page, _np_rgb_to_b64_png + + retry = remote_retry or RemoteRetryParams( + remote_max_pool_workers=int(kwargs.get("remote_max_pool_workers", 16)), + remote_max_retries=int(kwargs.get("remote_max_retries", 10)), + remote_max_429_retries=int(kwargs.get("remote_max_429_retries", 5)), + ) + + if not isinstance(batch_df, pd.DataFrame): + raise NotImplementedError("atable_structure_ocr_page_elements currently only supports pandas.DataFrame input.") + + ts_url = (table_structure_invoke_url or kwargs.get("table_structure_invoke_url") or "").strip() + use_remote_ts = bool(ts_url) + table_output_format = kwargs.get("table_output_format") + + if not use_remote_ts and table_structure_model is None: + raise ValueError("A local `table_structure_model` is required when `table_structure_invoke_url` is not set.") + + label_names = _labels_from_model(table_structure_model) if table_structure_model is not None else [] + if not label_names: + label_names = _DEFAULT_TABLE_STRUCTURE_LABELS + inference_batch_size = int(kwargs.get("inference_batch_size", 8)) + + all_table: List[List[Dict[str, Any]]] = [] + all_meta: List[Dict[str, Any]] = [] + + t0_total = time.perf_counter() + + for row in batch_df.itertuples(index=False): + table_items: List[Dict[str, Any]] = [] + row_error: Any = None + + try: + pe = getattr(row, "page_elements_v3", None) + dets: List[Dict[str, Any]] = [] + if isinstance(pe, dict): + dets = pe.get("detections") or [] + if not isinstance(dets, list): + dets = [] + + page_image = getattr(row, "page_image", None) or {} + page_image_b64 = page_image.get("image_b64") if isinstance(page_image, dict) else None + + if not isinstance(page_image_b64, str) or not page_image_b64: + all_table.append(table_items) + all_meta.append({"timing": None, "error": None}) + continue + + crops = _crop_all_from_page(page_image_b64, dets, {"table"}) + + if not crops: + all_table.append(table_items) + all_meta.append({"timing": None, "error": None}) + continue + + crop_b64s = [_np_rgb_to_b64_png(crop_array) for _, _, crop_array in crops] if use_remote_ts else [] + + structure_results: List[List[Dict[str, Any]]] = [] + if use_remote_ts: + response_items = await ainvoke_image_inference_batches( + invoke_url=ts_url, + image_b64_list=crop_b64s, + api_key=api_key or None, + timeout_s=float(request_timeout_s), + max_batch_size=inference_batch_size, + max_concurrency=int(retry.remote_max_pool_workers), + max_retries=int(retry.remote_max_retries), + max_429_retries=int(retry.remote_max_429_retries), + ) + if len(response_items) != len(crops): + raise RuntimeError(f"Expected {len(crops)} table-structure responses, got {len(response_items)}") + for resp in response_items: + parsed = _parse_nim_bounding_boxes(resp) + if not parsed: + pred_item = _extract_remote_pred_item(resp) + parsed = _prediction_to_detections(pred_item, label_names=label_names) + structure_results.append([d for d in parsed if (d.get("score") or 0.0) >= YOLOX_TABLE_MIN_SCORE]) + else: + + def _run_local_ts(): + results = [] + for _, _, crop_array in crops: + chw = torch.from_numpy(crop_array).permute(2, 0, 1).contiguous().to(dtype=torch.float32) + h, w = crop_array.shape[:2] + x = chw.unsqueeze(0) + try: + pre = table_structure_model.preprocess(x, (h, w)) + except TypeError: + pre = table_structure_model.preprocess(x) + if isinstance(pre, torch.Tensor) and pre.ndim == 3: + pre = pre.unsqueeze(0) + pred = table_structure_model.invoke(pre, (h, w)) + dets_local = _prediction_to_detections(pred, label_names=label_names) + results.append([d for d in dets_local if (d.get("score") or 0.0) >= YOLOX_TABLE_MIN_SCORE]) + return results + + structure_results = await asyncio.to_thread(_run_local_ts) + + for crop_i, (_, bbox, _) in enumerate(crops): + structure_dets = structure_results[crop_i] + table_items.append( + { + "bbox_xyxy_norm": bbox, + "text": _render_structure_only_text( + structure_dets, + table_output_format=table_output_format, + ), + "structure_detections": structure_dets, + "structure_counts": _count_structure_labels(structure_dets), + } + ) + + except BaseException as e: + print(f"Warning: table-structure failed: {type(e).__name__}: {e}") + row_error = { + "stage": "table_structure_ocr_page_elements", + "type": e.__class__.__name__, + "message": str(e), + "traceback": "".join(traceback.format_exception(type(e), e, e.__traceback__)), + } + + all_table.append(table_items) + all_meta.append({"timing": None, "error": row_error}) + + elapsed = time.perf_counter() - t0_total + for meta in all_meta: + meta["timing"] = {"seconds": float(elapsed)} + + out = batch_df.copy() + out["table"] = all_table + out["table_structure_ocr_v1"] = all_meta + return out + + # --------------------------------------------------------------------------- # Combined table-structure + OCR Ray Actor # --------------------------------------------------------------------------- diff --git a/nemo_retriever/src/nemo_retriever/text_embed/cpu_operator.py b/nemo_retriever/src/nemo_retriever/text_embed/cpu_operator.py index c9517aa9f..2216f87e7 100644 --- a/nemo_retriever/src/nemo_retriever/text_embed/cpu_operator.py +++ b/nemo_retriever/src/nemo_retriever/text_embed/cpu_operator.py @@ -15,7 +15,7 @@ from nemo_retriever.text_embed.shared import build_embed_kwargs -class _BatchEmbedCPUActor(AbstractOperator, CPUOperator): +class BatchEmbedCPUActor(AbstractOperator, CPUOperator): """CPU-only embedding actor that always targets a remote endpoint.""" DEFAULT_EMBED_INVOKE_URL = "https://integrate.api.nvidia.com/v1/embeddings" diff --git a/nemo_retriever/src/nemo_retriever/text_embed/gpu_operator.py b/nemo_retriever/src/nemo_retriever/text_embed/gpu_operator.py index 0eb60bddc..090f490c0 100644 --- a/nemo_retriever/src/nemo_retriever/text_embed/gpu_operator.py +++ b/nemo_retriever/src/nemo_retriever/text_embed/gpu_operator.py @@ -15,7 +15,7 @@ from nemo_retriever.text_embed.shared import build_embed_kwargs -class _BatchEmbedActor(AbstractOperator, GPUOperator): +class BatchEmbedGPUActor(AbstractOperator, GPUOperator): """Graph embedding actor that loads a local embedder or calls a remote endpoint.""" def __init__(self, params: EmbedParams) -> None: @@ -55,5 +55,5 @@ def process(self, data: Any, **kwargs: Any) -> Any: def postprocess(self, data: Any, **kwargs: Any) -> Any: return data - def __call__(self, batch_df: Any) -> Any: - return self.run(batch_df) + async def __call__(self, batch_df: Any) -> Any: + return await self.arun(batch_df) diff --git a/nemo_retriever/src/nemo_retriever/text_embed/operators.py b/nemo_retriever/src/nemo_retriever/text_embed/operators.py index c30774454..aa7eb8bc7 100644 --- a/nemo_retriever/src/nemo_retriever/text_embed/operators.py +++ b/nemo_retriever/src/nemo_retriever/text_embed/operators.py @@ -11,10 +11,10 @@ from nemo_retriever.graph.operator_archetype import ArchetypeOperator from nemo_retriever.text_embed.runtime import embed_text_main_text_embed -__all__ = ["_BatchEmbedActor", "embed_text_main_text_embed"] +__all__ = ["BatchEmbedActor", "embed_text_main_text_embed"] -class _BatchEmbedActor(ArchetypeOperator): +class BatchEmbedActor(ArchetypeOperator): """Graph-facing batch embedding archetype.""" @classmethod @@ -25,27 +25,38 @@ def prefers_cpu_variant(cls, operator_kwargs: dict[str, Any] | None = None) -> b @classmethod def cpu_variant_class(cls): - from nemo_retriever.text_embed.cpu_operator import _BatchEmbedCPUActor + from nemo_retriever.text_embed.cpu_operator import BatchEmbedCPUActor - return _BatchEmbedCPUActor + return BatchEmbedCPUActor @classmethod def gpu_variant_class(cls): - from nemo_retriever.text_embed.gpu_operator import _BatchEmbedActor as _BatchEmbedGPUActor + from nemo_retriever.text_embed.gpu_operator import BatchEmbedGPUActor - return _BatchEmbedGPUActor + return BatchEmbedGPUActor def __init__(self, params: Any) -> None: super().__init__(params=params) def __getattr__(name: str): + if name == "BatchEmbedCPUActor": + from nemo_retriever.text_embed.cpu_operator import BatchEmbedCPUActor + + return BatchEmbedCPUActor + if name == "BatchEmbedGPUActor": + from nemo_retriever.text_embed.gpu_operator import BatchEmbedGPUActor + + return BatchEmbedGPUActor + # Backward compatibility aliases + if name == "_BatchEmbedActor": + return BatchEmbedActor if name == "_BatchEmbedCPUActor": - from nemo_retriever.text_embed.cpu_operator import _BatchEmbedCPUActor + from nemo_retriever.text_embed.cpu_operator import BatchEmbedCPUActor - return _BatchEmbedCPUActor + return BatchEmbedCPUActor if name == "_BatchEmbedGPUActor": - from nemo_retriever.text_embed.gpu_operator import _BatchEmbedActor as _BatchEmbedGPUActor + from nemo_retriever.text_embed.gpu_operator import BatchEmbedGPUActor - return _BatchEmbedGPUActor + return BatchEmbedGPUActor raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/nemo_retriever/src/nemo_retriever/text_embed/text_embed.py b/nemo_retriever/src/nemo_retriever/text_embed/text_embed.py index 6c5fb8086..a80029829 100644 --- a/nemo_retriever/src/nemo_retriever/text_embed/text_embed.py +++ b/nemo_retriever/src/nemo_retriever/text_embed/text_embed.py @@ -215,9 +215,9 @@ def process(self, batch_df: Any, **override_kwargs: Any) -> Any: def postprocess(self, data: Any, **kwargs: Any) -> Any: return data - def __call__(self, batch_df: Any, **override_kwargs: Any) -> Any: + async def __call__(self, batch_df: Any, **override_kwargs: Any) -> Any: try: - return self.run(batch_df, **override_kwargs) + return await self.arun(batch_df, **override_kwargs) except BaseException as e: if isinstance(batch_df, pd.DataFrame): out = batch_df.copy() diff --git a/nemo_retriever/src/nemo_retriever/txt/ray_data.py b/nemo_retriever/src/nemo_retriever/txt/ray_data.py index 2a6f8b64a..f73c5a704 100644 --- a/nemo_retriever/src/nemo_retriever/txt/ray_data.py +++ b/nemo_retriever/src/nemo_retriever/txt/ray_data.py @@ -49,8 +49,8 @@ def process(self, data: Any, **kwargs: Any) -> Any: def postprocess(self, data: Any, **kwargs: Any) -> Any: return data - def __call__(self, batch_df: pd.DataFrame) -> pd.DataFrame: - return self.run(batch_df) + async def __call__(self, batch_df: pd.DataFrame) -> pd.DataFrame: + return await self.arun(batch_df) class TxtSplitCPUActor(AbstractOperator, CPUOperator): @@ -96,8 +96,8 @@ def process(self, data: Any, **kwargs: Any) -> Any: def postprocess(self, data: Any, **kwargs: Any) -> Any: return data - def __call__(self, batch_df: pd.DataFrame) -> pd.DataFrame: - return self.run(batch_df) + async def __call__(self, batch_df: pd.DataFrame) -> pd.DataFrame: + return await self.arun(batch_df) class TextChunkActor(ArchetypeOperator): diff --git a/nemo_retriever/src/nemo_retriever/utils/convert/to_pdf.py b/nemo_retriever/src/nemo_retriever/utils/convert/to_pdf.py index 4a8f5ed22..bff476218 100644 --- a/nemo_retriever/src/nemo_retriever/utils/convert/to_pdf.py +++ b/nemo_retriever/src/nemo_retriever/utils/convert/to_pdf.py @@ -171,8 +171,8 @@ def process(self, data: Any, **kwargs: Any) -> Any: def postprocess(self, data: Any, **kwargs: Any) -> Any: return data - def __call__(self, batch_df: Any) -> Any: - return self.run(batch_df) + async def __call__(self, batch_df: Any) -> Any: + return await self.arun(batch_df) class DocToPdfConversionActor(ArchetypeOperator): diff --git a/nemo_retriever/tests/test_actor_operators.py b/nemo_retriever/tests/test_actor_operators.py index 1f857e63a..26d907249 100644 --- a/nemo_retriever/tests/test_actor_operators.py +++ b/nemo_retriever/tests/test_actor_operators.py @@ -4,8 +4,9 @@ """Unit tests verifying all pipeline actors inherit from AbstractOperator.""" +import asyncio from pathlib import Path -from unittest.mock import patch +from unittest.mock import AsyncMock, patch import pandas as pd import pytest @@ -13,6 +14,15 @@ from nemo_retriever.graph.abstract_operator import AbstractOperator +def _run(coro): + """Run a coroutine synchronously in tests.""" + loop = asyncio.new_event_loop() + try: + return loop.run_until_complete(coro) + finally: + loop.close() + + # --------------------------------------------------------------------------- # 1. PDFSplitActor # --------------------------------------------------------------------------- @@ -52,7 +62,7 @@ def test_call_delegates_to_run(self, mock_fn): expected = pd.DataFrame({"page": [1]}) mock_fn.return_value = expected actor = self._make() - result = actor(pd.DataFrame({"bytes": [b"x"]})) + result = _run(actor(pd.DataFrame({"bytes": [b"x"]}))) pd.testing.assert_frame_equal(result, expected) @@ -89,14 +99,14 @@ def test_call_delegates_to_run(self, mock_fn): expected = pd.DataFrame({"text": ["hello"]}) mock_fn.return_value = expected actor = self._make() - result = actor(pd.DataFrame({"bytes": [b"x"]})) + result = _run(actor(pd.DataFrame({"bytes": [b"x"]}))) pd.testing.assert_frame_equal(result, expected) @patch("nemo_retriever.pdf.extract.pdf_extraction", side_effect=RuntimeError("boom")) def test_call_error_handling(self, mock_fn): actor = self._make() df = pd.DataFrame({"bytes": [b"x"], "path": ["/tmp/a.pdf"]}) - result = actor(df) + result = _run(actor(df)) assert isinstance(result, list) record = result[0] assert record["metadata"]["error"]["type"] == "RuntimeError" @@ -110,15 +120,17 @@ def test_pdfium_output_can_have_empty_text_without_ocr_flag(self): pytest.skip(f"External regression fixture not available: {pdf_path}") source_df = pd.DataFrame({"path": [str(pdf_path)], "bytes": [pdf_path.read_bytes()]}) - split_df = PDFSplitActor()(source_df) - - result = PDFExtractionActor( - method="pdfium", - extract_text=True, - extract_tables=True, - extract_charts=True, - extract_infographics=True, - )(split_df.head(5)) + split_df = _run(PDFSplitActor()(source_df)) + + result = _run( + PDFExtractionActor( + method="pdfium", + extract_text=True, + extract_tables=True, + extract_charts=True, + extract_infographics=True, + )(split_df.head(5)) + ) first_page = result[result["page_number"] == 1].iloc[0] metadata = first_page["metadata"] @@ -157,19 +169,23 @@ def test_process(self, mock_fn): mock_fn.assert_called_once() pd.testing.assert_frame_equal(result, expected) - @patch("nemo_retriever.page_elements.cpu_actor.detect_page_elements_v3") + @patch("nemo_retriever.page_elements.cpu_actor.adetect_page_elements_v3", new_callable=AsyncMock) def test_call_delegates(self, mock_fn): expected = pd.DataFrame({"page_elements_v3": ["det"]}) mock_fn.return_value = expected actor = self._make() - result = actor(pd.DataFrame({"page_image": ["x"]})) + result = _run(actor(pd.DataFrame({"page_image": ["x"]}))) pd.testing.assert_frame_equal(result, expected) - @patch("nemo_retriever.page_elements.cpu_actor.detect_page_elements_v3", side_effect=RuntimeError("boom")) + @patch( + "nemo_retriever.page_elements.cpu_actor.adetect_page_elements_v3", + new_callable=AsyncMock, + side_effect=RuntimeError("boom"), + ) def test_call_error_handling(self, mock_fn): actor = self._make() df = pd.DataFrame({"page_image": ["x"]}) - result = actor(df) + result = _run(actor(df)) assert isinstance(result, pd.DataFrame) assert "page_elements_v3" in result.columns @@ -205,11 +221,15 @@ def test_process(self, mock_fn): mock_fn.assert_called_once() pd.testing.assert_frame_equal(result, expected) - @patch("nemo_retriever.chart.cpu_actor.graphic_elements_ocr_page_elements", side_effect=RuntimeError("boom")) + @patch( + "nemo_retriever.chart.cpu_actor.agraphic_elements_ocr_page_elements", + new_callable=AsyncMock, + side_effect=RuntimeError("boom"), + ) def test_call_error_handling(self, mock_fn): actor = self._make() df = pd.DataFrame({"page_image": ["x"]}) - result = actor(df) + result = _run(actor(df)) assert isinstance(result, pd.DataFrame) assert "graphic_elements_ocr_v1" in result.columns @@ -242,11 +262,15 @@ def test_process(self, mock_fn): mock_fn.assert_called_once() pd.testing.assert_frame_equal(result, expected) - @patch("nemo_retriever.table.cpu_actor.table_structure_ocr_page_elements", side_effect=RuntimeError("boom")) + @patch( + "nemo_retriever.table.cpu_actor.atable_structure_ocr_page_elements", + new_callable=AsyncMock, + side_effect=RuntimeError("boom"), + ) def test_call_error_handling(self, mock_fn): actor = self._make() df = pd.DataFrame({"page_image": ["x"]}) - result = actor(df) + result = _run(actor(df)) assert isinstance(result, pd.DataFrame) assert "table_structure_ocr_v1" in result.columns @@ -279,11 +303,11 @@ def test_process(self, mock_fn): mock_fn.assert_called_once() pd.testing.assert_frame_equal(result, expected) - @patch("nemo_retriever.ocr.cpu_ocr.ocr_page_elements", side_effect=RuntimeError("boom")) + @patch("nemo_retriever.ocr.cpu_ocr.aocr_page_elements", new_callable=AsyncMock, side_effect=RuntimeError("boom")) def test_call_error_handling(self, mock_fn): actor = self._make() df = pd.DataFrame({"page_image": ["x"]}) - result = actor(df) + result = _run(actor(df)) assert isinstance(result, pd.DataFrame) assert "ocr_v1" in result.columns @@ -320,7 +344,7 @@ def test_process(self, mock_fn): def test_call_error_handling(self, mock_fn): actor = self._make() df = pd.DataFrame({"page_image": ["x"]}) - result = actor(df) + result = _run(actor(df)) assert isinstance(result, pd.DataFrame) assert "nemotron_parse_v1_2" in result.columns @@ -360,7 +384,7 @@ def test_call_delegates(self, mock_fn): expected = pd.DataFrame({"text": ["chunk1"]}) mock_fn.return_value = expected actor = self._make() - result = actor(pd.DataFrame({"text": ["hello world"]})) + result = _run(actor(pd.DataFrame({"text": ["hello world"]}))) pd.testing.assert_frame_equal(result, expected) @@ -403,7 +427,7 @@ def test_call_delegates(self, mock_fn): expected = pd.DataFrame({"path": ["/tmp/a.png"], "page_number": [0]}) mock_fn.return_value = expected actor = self._make() - result = actor(pd.DataFrame({"bytes": [b"img"], "path": ["/tmp/a.png"]})) + result = _run(actor(pd.DataFrame({"bytes": [b"img"], "path": ["/tmp/a.png"]}))) pd.testing.assert_frame_equal(result, expected) @@ -446,7 +470,7 @@ def test_call_delegates(self, mock_fn): expected = pd.DataFrame({"text": ["chunk"], "path": ["/a.txt"], "page_number": [0], "metadata": [{}]}) mock_fn.return_value = expected actor = self._make() - result = actor(pd.DataFrame({"bytes": [b"hello"], "path": ["/a.txt"]})) + result = _run(actor(pd.DataFrame({"bytes": [b"hello"], "path": ["/a.txt"]}))) pd.testing.assert_frame_equal(result, expected) @@ -484,25 +508,25 @@ def test_call_delegates(self, mock_fn): expected = pd.DataFrame({"text": ["chunk"], "path": ["/a.html"], "page_number": [0], "metadata": [{}]}) mock_fn.return_value = expected actor = self._make() - result = actor(pd.DataFrame({"bytes": [b"

hi

"], "path": ["/a.html"]})) + result = _run(actor(pd.DataFrame({"bytes": [b"

hi

"], "path": ["/a.html"]}))) pd.testing.assert_frame_equal(result, expected) # --------------------------------------------------------------------------- -# 12. _BatchEmbedActor +# 12. BatchEmbedActor # --------------------------------------------------------------------------- class TestBatchEmbedActor: def _make(self): from nemo_retriever.params import EmbedParams - from nemo_retriever.text_embed.operators import _BatchEmbedActor + from nemo_retriever.text_embed.operators import BatchEmbedActor params = EmbedParams(model_name="test-model", embed_invoke_url="http://fake") - return _BatchEmbedActor(params=params) + return BatchEmbedActor(params=params) def test_inherits(self): - from nemo_retriever.text_embed.operators import _BatchEmbedActor + from nemo_retriever.text_embed.operators import BatchEmbedActor - assert issubclass(_BatchEmbedActor, AbstractOperator) + assert issubclass(BatchEmbedActor, AbstractOperator) def test_preprocess_passthrough(self): actor = self._make() @@ -528,5 +552,5 @@ def test_call_delegates(self, mock_fn): expected = pd.DataFrame({"text": ["hello"], "embedding": [[0.1, 0.2]]}) mock_fn.return_value = expected actor = self._make() - result = actor(pd.DataFrame({"text": ["hello"]})) + result = _run(actor(pd.DataFrame({"text": ["hello"]}))) pd.testing.assert_frame_equal(result, expected) diff --git a/nemo_retriever/tests/test_asr_actor.py b/nemo_retriever/tests/test_asr_actor.py index c7297e391..72baa6b75 100644 --- a/nemo_retriever/tests/test_asr_actor.py +++ b/nemo_retriever/tests/test_asr_actor.py @@ -10,6 +10,7 @@ into sys.modules so the real module is never loaded. """ +import asyncio import base64 import sys from unittest.mock import MagicMock @@ -22,6 +23,11 @@ from nemo_retriever.params import ASRParams +def _run(coro): + """Run a coroutine synchronously in tests.""" + return asyncio.get_event_loop().run_until_complete(coro) + + def test_strip_pad_from_transcript(): """Transformers backend post-process removes and normalizes spaces.""" # Some tests monkeypatch nemo_retriever.model.local with a mock module object. @@ -49,7 +55,7 @@ def test_asr_actor_empty_batch(): params = ASRParams(audio_endpoints=("localhost:50051", None)) actor = ASRActor(params=params) empty = pd.DataFrame(columns=["path", "bytes"]) - out = actor(empty) + out = _run(actor(empty)) assert isinstance(out, pd.DataFrame) assert "text" in out.columns @@ -79,7 +85,7 @@ def test_asr_actor_mock_transcribe(): } ] ) - out = actor(batch) + out = _run(actor(batch)) assert len(out) == 1 assert out["text"].iloc[0] == "hello world transcript" @@ -142,7 +148,7 @@ def test_asr_actor_remote_segment_audio(): } ] ) - out = actor(batch) + out = _run(actor(batch)) assert len(out) == 2 assert out["text"].tolist() == ["Hello world.", "How are you?"] @@ -223,7 +229,7 @@ def test_local_asr_does_not_call_get_client(): } ] ) - out = actor(batch) + out = _run(actor(batch)) assert len(out) == 1 assert out["text"].iloc[0] == "mocked local transcript" diff --git a/nemo_retriever/tests/test_audio_chunk_actor.py b/nemo_retriever/tests/test_audio_chunk_actor.py index 293825488..edd0b5502 100644 --- a/nemo_retriever/tests/test_audio_chunk_actor.py +++ b/nemo_retriever/tests/test_audio_chunk_actor.py @@ -6,6 +6,7 @@ Unit tests for nemo_retriever.audio: MediaChunkActor and audio_path_to_chunks_df. """ +import asyncio import wave from pathlib import Path @@ -19,6 +20,11 @@ from nemo_retriever.params import AudioChunkParams +def _run(coro): + """Run a coroutine synchronously in tests.""" + return asyncio.get_event_loop().run_until_complete(coro) + + def _make_small_wav(path: Path, duration_sec: float = 0.5, sample_rate: int = 8000) -> None: """Write a minimal WAV file (e.g. 0.5s mono 8kHz) for tests.""" n_frames = int(sample_rate * duration_sec) @@ -36,7 +42,7 @@ def test_media_chunk_actor_empty_batch(): params = AudioChunkParams(split_type="size", split_interval=1000) actor = MediaChunkActor(params=params) empty = pd.DataFrame(columns=["path", "bytes"]) - out = actor(empty) + out = _run(actor(empty)) assert isinstance(out, pd.DataFrame) assert list(out.columns) == CHUNK_COLUMNS assert len(out) == 0 @@ -54,7 +60,7 @@ def test_media_chunk_actor_single_small_file(tmp_path: Path): params = AudioChunkParams(split_type="size", split_interval=1_000_000) actor = MediaChunkActor(params=params) batch = pd.DataFrame([{"path": str(wav.resolve()), "bytes": body}]) - out = actor(batch) + out = _run(actor(batch)) assert isinstance(out, pd.DataFrame) for col in ["path", "source_path", "duration", "chunk_index", "metadata", "page_number", "bytes"]: diff --git a/nemo_retriever/tests/test_chart_graphic_elements.py b/nemo_retriever/tests/test_chart_graphic_elements.py index 88e4433a4..2d73eb255 100644 --- a/nemo_retriever/tests/test_chart_graphic_elements.py +++ b/nemo_retriever/tests/test_chart_graphic_elements.py @@ -6,6 +6,7 @@ from __future__ import annotations +import asyncio import base64 import importlib import io @@ -17,6 +18,11 @@ from nemo_retriever.utils.table_and_chart import join_graphic_elements_and_ocr_output +def _run(coro): + """Run a coroutine synchronously in tests.""" + return asyncio.get_event_loop().run_until_complete(coro) + + def _can_import(mod: str) -> bool: return importlib.util.find_spec(mod) is not None @@ -291,7 +297,7 @@ def test_actor_error_returns_dataframe_with_error(self) -> None: df = _make_chart_page_df() # This will fail because both models are None and no URLs set. - result = actor(df) + result = _run(actor(df)) assert "chart" in result.columns assert "graphic_elements_ocr_v1" in result.columns meta = result.iloc[0]["graphic_elements_ocr_v1"] diff --git a/nemo_retriever/tests/test_doc_to_pdf_actor.py b/nemo_retriever/tests/test_doc_to_pdf_actor.py index c282aa9f2..32c159de7 100644 --- a/nemo_retriever/tests/test_doc_to_pdf_actor.py +++ b/nemo_retriever/tests/test_doc_to_pdf_actor.py @@ -2,6 +2,7 @@ # All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import asyncio from unittest.mock import patch import pandas as pd @@ -11,6 +12,11 @@ from nemo_retriever.utils.convert.to_pdf import DocToPdfConversionActor, convert_to_pdf_bytes, convert_batch_to_pdf +def _run(coro): + """Run a coroutine synchronously in tests.""" + return asyncio.get_event_loop().run_until_complete(coro) + + class TestConvertToPdfBytes: @patch("nemo_retriever.utils.convert.to_pdf.shutil.which", return_value=None) def test_raises_when_libreoffice_missing(self, _mock_which): @@ -79,6 +85,6 @@ def test_call_delegates_to_run(self, mock_convert): mock_convert.return_value = expected actor = DocToPdfConversionActor() df = pd.DataFrame({"bytes": [b"docx"], "path": ["/tmp/test.docx"]}) - result = actor(df) + result = _run(actor(df)) mock_convert.assert_called_once_with(df) pd.testing.assert_frame_equal(result, expected) diff --git a/nemo_retriever/tests/test_image_load.py b/nemo_retriever/tests/test_image_load.py index 623056059..a489c22be 100644 --- a/nemo_retriever/tests/test_image_load.py +++ b/nemo_retriever/tests/test_image_load.py @@ -6,6 +6,7 @@ from __future__ import annotations +import asyncio import base64 from io import BytesIO from unittest.mock import patch @@ -23,6 +24,12 @@ ) from nemo_retriever.image.ray_data import ImageLoadActor # noqa: E402 + +def _run(coro): + """Run a coroutine synchronously in tests.""" + return asyncio.get_event_loop().run_until_complete(coro) + + # -- Helpers ------------------------------------------------------------------ _PAGE_SCHEMA_COLUMNS = { @@ -175,7 +182,7 @@ def test_batch_processing(self) -> None: {"bytes": img2, "path": "/b/img2.png"}, ] ) - result = actor(batch) + result = _run(actor(batch)) assert isinstance(result, pd.DataFrame) assert len(result) == 2 assert list(result["path"]) == ["/a/img1.png", "/b/img2.png"] @@ -184,14 +191,14 @@ def test_batch_processing(self) -> None: def test_empty_batch(self) -> None: actor = ImageLoadActor() - result = actor(pd.DataFrame()) + result = _run(actor(pd.DataFrame())) assert isinstance(result, pd.DataFrame) assert len(result) == 0 def test_missing_columns_skipped(self) -> None: actor = ImageLoadActor() batch = pd.DataFrame([{"bytes": b"data"}]) # no 'path' column - result = actor(batch) + result = _run(actor(batch)) assert len(result) == 0 def test_corrupt_row_skipped(self) -> None: @@ -203,7 +210,7 @@ def test_corrupt_row_skipped(self) -> None: {"bytes": good, "path": "/good.png"}, ] ) - result = actor(batch) + result = _run(actor(batch)) # Corrupt row produces an error record, good row succeeds. assert len(result) == 2 diff --git a/nemo_retriever/tests/test_ingest_plans.py b/nemo_retriever/tests/test_ingest_plans.py index ae9e991bd..c130bc02c 100644 --- a/nemo_retriever/tests/test_ingest_plans.py +++ b/nemo_retriever/tests/test_ingest_plans.py @@ -7,7 +7,7 @@ from nemo_retriever.graph.pipeline_graph import Graph from nemo_retriever.ocr.ocr import OCRActor from nemo_retriever.page_elements.page_elements import PageElementDetectionActor -from nemo_retriever.text_embed.operators import _BatchEmbedActor +from nemo_retriever.text_embed.operators import BatchEmbedActor from nemo_retriever.graph.operator_archetype import ArchetypeOperator from nemo_retriever.graph.cpu_operator import CPUOperator from nemo_retriever.graph.gpu_operator import GPUOperator @@ -101,7 +101,7 @@ def test_build_graph_accepts_execution_plan() -> None: break node = node.children[0] - assert names == ["MultiTypeExtractOperator", "TextChunkActor", "_BatchEmbedActor"] + assert names == ["MultiTypeExtractOperator", "TextChunkActor", "BatchEmbedActor"] def test_build_graph_keeps_archetype_operator_classes() -> None: @@ -127,11 +127,11 @@ def test_build_graph_keeps_archetype_operator_classes() -> None: "PageElementDetectionActor", "OCRActor", "UDFOperator", - "_BatchEmbedActor", + "BatchEmbedActor", ] assert nodes[3].operator_class is PageElementDetectionActor assert nodes[4].operator_class is OCRActor - assert nodes[-1].operator_class is _BatchEmbedActor + assert nodes[-1].operator_class is BatchEmbedActor assert issubclass(nodes[3].operator_class, ArchetypeOperator) assert issubclass(nodes[4].operator_class, ArchetypeOperator) assert issubclass(nodes[-1].operator_class, ArchetypeOperator) @@ -162,10 +162,10 @@ def test_build_graph_resolves_endpoint_configured_nodes_to_cpu_variants() -> Non assert classes["TableStructureActor"].__name__ == "TableStructureCPUActor" assert classes["GraphicElementsActor"].__name__ == "GraphicElementsCPUActor" assert classes["OCRActor"].__name__ == "OCRCPUActor" - assert classes["_BatchEmbedActor"].__name__ == "_BatchEmbedCPUActor" + assert classes["BatchEmbedActor"].__name__ == "BatchEmbedCPUActor" assert issubclass(classes["PageElementDetectionActor"], CPUOperator) assert issubclass(classes["OCRActor"], CPUOperator) - assert issubclass(classes["_BatchEmbedActor"], CPUOperator) + assert issubclass(classes["BatchEmbedActor"], CPUOperator) def test_build_graph_resolves_local_nodes_to_gpu_variants_when_gpus_available() -> None: @@ -185,10 +185,10 @@ def test_build_graph_resolves_local_nodes_to_gpu_variants_when_gpus_available() assert classes["PageElementDetectionActor"] is not PageElementDetectionActor assert classes["OCRActor"] is not OCRActor - assert classes["_BatchEmbedActor"] is not _BatchEmbedActor + assert classes["BatchEmbedActor"] is not BatchEmbedActor assert issubclass(classes["PageElementDetectionActor"], GPUOperator) assert issubclass(classes["OCRActor"], GPUOperator) - assert issubclass(classes["_BatchEmbedActor"], GPUOperator) + assert issubclass(classes["BatchEmbedActor"], GPUOperator) def test_batch_tuning_to_node_overrides_auto_cpu_only_when_no_gpus() -> None: @@ -221,11 +221,11 @@ def test_batch_tuning_to_node_overrides_auto_cpu_only_when_no_gpus() -> None: cluster_resources=cluster, ) - assert overrides["_BatchEmbedActor"]["num_gpus"] == 0.0 + assert overrides["BatchEmbedActor"]["num_gpus"] == 0.0 assert overrides["OCRActor"]["num_gpus"] == 0.0 assert overrides["PageElementDetectionActor"]["num_gpus"] == 0.0 assert overrides["NemotronParseActor"]["num_gpus"] == 0.0 - assert overrides["_BatchEmbedActor"]["concurrency"] == 5 + assert overrides["BatchEmbedActor"]["concurrency"] == 5 assert overrides["OCRActor"]["concurrency"] == 4 assert overrides["PageElementDetectionActor"]["concurrency"] == 3 assert overrides["NemotronParseActor"]["concurrency"] == 2 @@ -310,7 +310,7 @@ def test_build_inprocess_graph_accepts_execution_plan() -> None: "TextChunkActor", "CaptionActor", "UDFOperator", - "_BatchEmbedActor", + "BatchEmbedActor", ] @@ -335,7 +335,7 @@ def test_build_inprocess_graph_supports_text_execution_plan() -> None: break node = node.children[0] - assert names == ["MultiTypeExtractOperator", "TextChunkActor", "_BatchEmbedActor"] + assert names == ["MultiTypeExtractOperator", "TextChunkActor", "BatchEmbedActor"] @pytest.mark.skipif(not is_media_available(), reason="ffmpeg not available") diff --git a/nemo_retriever/tests/test_nemotron_rerank_v2.py b/nemo_retriever/tests/test_nemotron_rerank_v2.py index 5a412bb44..5ac665e1c 100644 --- a/nemo_retriever/tests/test_nemotron_rerank_v2.py +++ b/nemo_retriever/tests/test_nemotron_rerank_v2.py @@ -11,6 +11,7 @@ from __future__ import annotations +import asyncio import sys from types import ModuleType from unittest.mock import MagicMock, patch @@ -18,6 +19,11 @@ import pytest +def _run(coro): + """Run a coroutine synchronously in tests.""" + return asyncio.get_event_loop().run_until_complete(coro) + + # --------------------------------------------------------------------------- # Helpers to build lightweight torch / transformers stubs # --------------------------------------------------------------------------- @@ -539,7 +545,7 @@ def test_actor_call_scores_dataframe(self): ] with patch("requests.post", return_value=mock_resp): - out = actor(df) + out = _run(actor(df)) assert "rerank_score" in out.columns assert len(out) == 2 @@ -559,7 +565,7 @@ def test_actor_call_sorts_descending_by_default(self): ] with patch("requests.post", return_value=mock_resp): - out = actor(df) + out = _run(actor(df)) scores = out["rerank_score"].tolist() assert scores == sorted(scores, reverse=True) @@ -572,7 +578,7 @@ def test_actor_call_returns_error_payload_on_exception(self): df = pd.DataFrame({"query": ["q"], "text": ["doc"]}) with patch("requests.post", side_effect=RuntimeError("connection failed")): - out = actor(df) + out = _run(actor(df)) # Should not raise; should return a DataFrame with error payload assert isinstance(out, pd.DataFrame) @@ -592,6 +598,6 @@ def test_actor_custom_score_column_name(self): mock_resp.json.return_value = {"rankings": [{"index": 0, "logit": 0.7}]} with patch("requests.post", return_value=mock_resp): - out = actor(df) + out = _run(actor(df)) assert "my_score" in out.columns diff --git a/nemo_retriever/tests/test_operator_flags_and_cpu_actors.py b/nemo_retriever/tests/test_operator_flags_and_cpu_actors.py index 7baf23ff0..087f7207a 100644 --- a/nemo_retriever/tests/test_operator_flags_and_cpu_actors.py +++ b/nemo_retriever/tests/test_operator_flags_and_cpu_actors.py @@ -27,7 +27,7 @@ def test_gpu_operators_have_flag(self): from nemo_retriever.table.table_detection import TableStructureGPUActor from nemo_retriever.ocr.ocr import OCRGPUActor from nemo_retriever.parse.nemotron_parse import NemotronParseGPUActor - from nemo_retriever.text_embed.operators import _BatchEmbedGPUActor + from nemo_retriever.text_embed.operators import BatchEmbedGPUActor from nemo_retriever.caption.caption import CaptionGPUActor from nemo_retriever.infographic.infographic_detection import InfographicDetectionGPUActor from nemo_retriever.rerank.rerank import NemotronRerankGPUActor @@ -38,7 +38,7 @@ def test_gpu_operators_have_flag(self): assert issubclass(TableStructureGPUActor, GPUOperator) assert issubclass(OCRGPUActor, GPUOperator) assert issubclass(NemotronParseGPUActor, GPUOperator) - assert issubclass(_BatchEmbedGPUActor, GPUOperator) + assert issubclass(BatchEmbedGPUActor, GPUOperator) assert issubclass(CaptionGPUActor, GPUOperator) assert issubclass(InfographicDetectionGPUActor, GPUOperator) assert issubclass(NemotronRerankGPUActor, GPUOperator) @@ -311,47 +311,47 @@ def _make_params(self): return EmbedParams(model_name="test-model", embed_invoke_url="http://fake") def test_inherits_cpu_operator(self): - from nemo_retriever.text_embed.cpu_operator import _BatchEmbedCPUActor + from nemo_retriever.text_embed.cpu_operator import BatchEmbedCPUActor - assert issubclass(_BatchEmbedCPUActor, CPUOperator) - assert not issubclass(_BatchEmbedCPUActor, GPUOperator) + assert issubclass(BatchEmbedCPUActor, CPUOperator) + assert not issubclass(BatchEmbedCPUActor, GPUOperator) def test_uses_default_invoke_url(self): - from nemo_retriever.text_embed.cpu_operator import _BatchEmbedCPUActor + from nemo_retriever.text_embed.cpu_operator import BatchEmbedCPUActor from nemo_retriever.params import EmbedParams - actor = _BatchEmbedCPUActor(params=EmbedParams(model_name="test-model")) + actor = BatchEmbedCPUActor(params=EmbedParams(model_name="test-model")) assert actor._model is None assert "integrate.api.nvidia.com" in actor._kwargs["embedding_endpoint"] def test_creates_with_custom_invoke_url(self): - from nemo_retriever.text_embed.cpu_operator import _BatchEmbedCPUActor + from nemo_retriever.text_embed.cpu_operator import BatchEmbedCPUActor - actor = _BatchEmbedCPUActor(params=self._make_params()) + actor = BatchEmbedCPUActor(params=self._make_params()) assert actor._model is None assert actor._kwargs["embedding_endpoint"] == "http://fake" @patch("nemo_retriever.text_embed.cpu_operator.embed_text_main_text_embed") def test_process(self, mock_fn): - from nemo_retriever.text_embed.cpu_operator import _BatchEmbedCPUActor + from nemo_retriever.text_embed.cpu_operator import BatchEmbedCPUActor expected = pd.DataFrame({"text": ["hello"], "embedding": [[0.1, 0.2]]}) mock_fn.return_value = expected - actor = _BatchEmbedCPUActor(params=self._make_params()) + actor = BatchEmbedCPUActor(params=self._make_params()) result = actor.process(pd.DataFrame({"text": ["hello"]})) mock_fn.assert_called_once() pd.testing.assert_frame_equal(result, expected) def test_preprocess_passthrough(self): - from nemo_retriever.text_embed.cpu_operator import _BatchEmbedCPUActor + from nemo_retriever.text_embed.cpu_operator import BatchEmbedCPUActor - actor = _BatchEmbedCPUActor(params=self._make_params()) + actor = BatchEmbedCPUActor(params=self._make_params()) df = pd.DataFrame({"text": ["hello"]}) pd.testing.assert_frame_equal(actor.preprocess(df), df) def test_postprocess_passthrough(self): - from nemo_retriever.text_embed.cpu_operator import _BatchEmbedCPUActor + from nemo_retriever.text_embed.cpu_operator import BatchEmbedCPUActor - actor = _BatchEmbedCPUActor(params=self._make_params()) + actor = BatchEmbedCPUActor(params=self._make_params()) df = pd.DataFrame({"text": ["hello"]}) pd.testing.assert_frame_equal(actor.postprocess(df), df) diff --git a/nemo_retriever/tests/test_table_structure.py b/nemo_retriever/tests/test_table_structure.py index 87e41429f..713c6d944 100644 --- a/nemo_retriever/tests/test_table_structure.py +++ b/nemo_retriever/tests/test_table_structure.py @@ -6,6 +6,7 @@ from __future__ import annotations +import asyncio import base64 import importlib import io @@ -17,6 +18,11 @@ from nemo_retriever.utils.table_and_chart import join_table_structure_and_ocr_output +def _run(coro): + """Run a coroutine synchronously in tests.""" + return asyncio.get_event_loop().run_until_complete(coro) + + def _can_import(mod: str) -> bool: return importlib.util.find_spec(mod) is not None @@ -313,7 +319,7 @@ def test_actor_error_returns_dataframe_with_error(self) -> None: df = _make_page_df() # This will fail because both models are None and no URLs set. - result = actor(df) + result = _run(actor(df)) assert "table" in result.columns assert "table_structure_ocr_v1" in result.columns meta = result.iloc[0]["table_structure_ocr_v1"] diff --git a/nemo_retriever/uv.lock b/nemo_retriever/uv.lock index de09d559a..1379d7294 100644 --- a/nemo_retriever/uv.lock +++ b/nemo_retriever/uv.lock @@ -1198,6 +1198,47 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ef/b4/f0b121a2300b629d09766aa3ffc2e755d8d72f31fe2bcf0b1055dbda1cbd/fastar-0.10.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:86a1805316324eeb98b05f6b1db921bc3a9d9c9c6f535b2204b2e039a29048c4", size = 1025819, upload-time = "2026-04-08T01:01:56.008Z" }, ] +[[package]] +name = "fastuuid" +version = "0.14.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/c3/7d/d9daedf0f2ebcacd20d599928f8913e9d2aea1d56d2d355a93bfa2b611d7/fastuuid-0.14.0.tar.gz", hash = "sha256:178947fc2f995b38497a74172adee64fdeb8b7ec18f2a5934d037641ba265d26", size = 18232, upload-time = "2025-10-19T22:19:22.402Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/02/a2/e78fcc5df65467f0d207661b7ef86c5b7ac62eea337c0c0fcedbeee6fb13/fastuuid-0.14.0-cp312-cp312-macosx_10_12_x86_64.macosx_11_0_arm64.macosx_10_12_universal2.whl", hash = "sha256:77e94728324b63660ebf8adb27055e92d2e4611645bf12ed9d88d30486471d0a", size = 510164, upload-time = "2025-10-19T22:31:45.635Z" }, + { url = "https://files.pythonhosted.org/packages/2b/b3/c846f933f22f581f558ee63f81f29fa924acd971ce903dab1a9b6701816e/fastuuid-0.14.0-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:caa1f14d2102cb8d353096bc6ef6c13b2c81f347e6ab9d6fbd48b9dea41c153d", size = 261837, upload-time = "2025-10-19T22:38:38.53Z" }, + { url = "https://files.pythonhosted.org/packages/54/ea/682551030f8c4fa9a769d9825570ad28c0c71e30cf34020b85c1f7ee7382/fastuuid-0.14.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:d23ef06f9e67163be38cece704170486715b177f6baae338110983f99a72c070", size = 251370, upload-time = "2025-10-19T22:40:26.07Z" }, + { url = "https://files.pythonhosted.org/packages/14/dd/5927f0a523d8e6a76b70968e6004966ee7df30322f5fc9b6cdfb0276646a/fastuuid-0.14.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:0c9ec605ace243b6dbe3bd27ebdd5d33b00d8d1d3f580b39fdd15cd96fd71796", size = 277766, upload-time = "2025-10-19T22:37:23.779Z" }, + { url = "https://files.pythonhosted.org/packages/16/6e/c0fb547eef61293153348f12e0f75a06abb322664b34a1573a7760501336/fastuuid-0.14.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:808527f2407f58a76c916d6aa15d58692a4a019fdf8d4c32ac7ff303b7d7af09", size = 278105, upload-time = "2025-10-19T22:26:56.821Z" }, + { url = "https://files.pythonhosted.org/packages/2d/b1/b9c75e03b768f61cf2e84ee193dc18601aeaf89a4684b20f2f0e9f52b62c/fastuuid-0.14.0-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:2fb3c0d7fef6674bbeacdd6dbd386924a7b60b26de849266d1ff6602937675c8", size = 301564, upload-time = "2025-10-19T22:30:31.604Z" }, + { url = "https://files.pythonhosted.org/packages/fc/fa/f7395fdac07c7a54f18f801744573707321ca0cee082e638e36452355a9d/fastuuid-0.14.0-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:ab3f5d36e4393e628a4df337c2c039069344db5f4b9d2a3c9cea48284f1dd741", size = 459659, upload-time = "2025-10-19T22:31:32.341Z" }, + { url = "https://files.pythonhosted.org/packages/66/49/c9fd06a4a0b1f0f048aacb6599e7d96e5d6bc6fa680ed0d46bf111929d1b/fastuuid-0.14.0-cp312-cp312-musllinux_1_1_i686.whl", hash = "sha256:b9a0ca4f03b7e0b01425281ffd44e99d360e15c895f1907ca105854ed85e2057", size = 478430, upload-time = "2025-10-19T22:26:22.962Z" }, + { url = "https://files.pythonhosted.org/packages/be/9c/909e8c95b494e8e140e8be6165d5fc3f61fdc46198c1554df7b3e1764471/fastuuid-0.14.0-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:3acdf655684cc09e60fb7e4cf524e8f42ea760031945aa8086c7eae2eeeabeb8", size = 450894, upload-time = "2025-10-19T22:27:01.647Z" }, + { url = "https://files.pythonhosted.org/packages/90/eb/d29d17521976e673c55ef7f210d4cdd72091a9ec6755d0fd4710d9b3c871/fastuuid-0.14.0-cp312-cp312-win32.whl", hash = "sha256:9579618be6280700ae36ac42c3efd157049fe4dd40ca49b021280481c78c3176", size = 154374, upload-time = "2025-10-19T22:29:19.879Z" }, + { url = "https://files.pythonhosted.org/packages/cc/fc/f5c799a6ea6d877faec0472d0b27c079b47c86b1cdc577720a5386483b36/fastuuid-0.14.0-cp312-cp312-win_amd64.whl", hash = "sha256:d9e4332dc4ba054434a9594cbfaf7823b57993d7d8e7267831c3e059857cf397", size = 156550, upload-time = "2025-10-19T22:27:49.658Z" }, + { url = "https://files.pythonhosted.org/packages/a5/83/ae12dd39b9a39b55d7f90abb8971f1a5f3c321fd72d5aa83f90dc67fe9ed/fastuuid-0.14.0-cp313-cp313-macosx_10_12_x86_64.macosx_11_0_arm64.macosx_10_12_universal2.whl", hash = "sha256:77a09cb7427e7af74c594e409f7731a0cf887221de2f698e1ca0ebf0f3139021", size = 510720, upload-time = "2025-10-19T22:42:34.633Z" }, + { url = "https://files.pythonhosted.org/packages/53/b0/a4b03ff5d00f563cc7546b933c28cb3f2a07344b2aec5834e874f7d44143/fastuuid-0.14.0-cp313-cp313-macosx_10_12_x86_64.whl", hash = "sha256:9bd57289daf7b153bfa3e8013446aa144ce5e8c825e9e366d455155ede5ea2dc", size = 262024, upload-time = "2025-10-19T22:30:25.482Z" }, + { url = "https://files.pythonhosted.org/packages/9c/6d/64aee0a0f6a58eeabadd582e55d0d7d70258ffdd01d093b30c53d668303b/fastuuid-0.14.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:ac60fc860cdf3c3f327374db87ab8e064c86566ca8c49d2e30df15eda1b0c2d5", size = 251679, upload-time = "2025-10-19T22:36:14.096Z" }, + { url = "https://files.pythonhosted.org/packages/60/f5/a7e9cda8369e4f7919d36552db9b2ae21db7915083bc6336f1b0082c8b2e/fastuuid-0.14.0-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ab32f74bd56565b186f036e33129da77db8be09178cd2f5206a5d4035fb2a23f", size = 277862, upload-time = "2025-10-19T22:36:23.302Z" }, + { url = "https://files.pythonhosted.org/packages/f0/d3/8ce11827c783affffd5bd4d6378b28eb6cc6d2ddf41474006b8d62e7448e/fastuuid-0.14.0-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:33e678459cf4addaedd9936bbb038e35b3f6b2061330fd8f2f6a1d80414c0f87", size = 278278, upload-time = "2025-10-19T22:29:43.809Z" }, + { url = "https://files.pythonhosted.org/packages/a2/51/680fb6352d0bbade04036da46264a8001f74b7484e2fd1f4da9e3db1c666/fastuuid-0.14.0-cp313-cp313-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:1e3cc56742f76cd25ecb98e4b82a25f978ccffba02e4bdce8aba857b6d85d87b", size = 301788, upload-time = "2025-10-19T22:36:06.825Z" }, + { url = "https://files.pythonhosted.org/packages/fa/7c/2014b5785bd8ebdab04ec857635ebd84d5ee4950186a577db9eff0fb8ff6/fastuuid-0.14.0-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:cb9a030f609194b679e1660f7e32733b7a0f332d519c5d5a6a0a580991290022", size = 459819, upload-time = "2025-10-19T22:35:31.623Z" }, + { url = "https://files.pythonhosted.org/packages/01/d2/524d4ceeba9160e7a9bc2ea3e8f4ccf1ad78f3bde34090ca0c51f09a5e91/fastuuid-0.14.0-cp313-cp313-musllinux_1_1_i686.whl", hash = "sha256:09098762aad4f8da3a888eb9ae01c84430c907a297b97166b8abc07b640f2995", size = 478546, upload-time = "2025-10-19T22:26:03.023Z" }, + { url = "https://files.pythonhosted.org/packages/bc/17/354d04951ce114bf4afc78e27a18cfbd6ee319ab1829c2d5fb5e94063ac6/fastuuid-0.14.0-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:1383fff584fa249b16329a059c68ad45d030d5a4b70fb7c73a08d98fd53bcdab", size = 450921, upload-time = "2025-10-19T22:31:02.151Z" }, + { url = "https://files.pythonhosted.org/packages/fb/be/d7be8670151d16d88f15bb121c5b66cdb5ea6a0c2a362d0dcf30276ade53/fastuuid-0.14.0-cp313-cp313-win32.whl", hash = "sha256:a0809f8cc5731c066c909047f9a314d5f536c871a7a22e815cc4967c110ac9ad", size = 154559, upload-time = "2025-10-19T22:36:36.011Z" }, + { url = "https://files.pythonhosted.org/packages/22/1d/5573ef3624ceb7abf4a46073d3554e37191c868abc3aecd5289a72f9810a/fastuuid-0.14.0-cp313-cp313-win_amd64.whl", hash = "sha256:0df14e92e7ad3276327631c9e7cec09e32572ce82089c55cb1bb8df71cf394ed", size = 156539, upload-time = "2025-10-19T22:33:35.898Z" }, + { url = "https://files.pythonhosted.org/packages/16/c9/8c7660d1fe3862e3f8acabd9be7fc9ad71eb270f1c65cce9a2b7a31329ab/fastuuid-0.14.0-cp314-cp314-macosx_10_12_x86_64.macosx_11_0_arm64.macosx_10_12_universal2.whl", hash = "sha256:b852a870a61cfc26c884af205d502881a2e59cc07076b60ab4a951cc0c94d1ad", size = 510600, upload-time = "2025-10-19T22:43:44.17Z" }, + { url = "https://files.pythonhosted.org/packages/4c/f4/a989c82f9a90d0ad995aa957b3e572ebef163c5299823b4027986f133dfb/fastuuid-0.14.0-cp314-cp314-macosx_10_12_x86_64.whl", hash = "sha256:c7502d6f54cd08024c3ea9b3514e2d6f190feb2f46e6dbcd3747882264bb5f7b", size = 262069, upload-time = "2025-10-19T22:43:38.38Z" }, + { url = "https://files.pythonhosted.org/packages/da/6c/a1a24f73574ac995482b1326cf7ab41301af0fabaa3e37eeb6b3df00e6e2/fastuuid-0.14.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:1ca61b592120cf314cfd66e662a5b54a578c5a15b26305e1b8b618a6f22df714", size = 251543, upload-time = "2025-10-19T22:32:22.537Z" }, + { url = "https://files.pythonhosted.org/packages/1a/20/2a9b59185ba7a6c7b37808431477c2d739fcbdabbf63e00243e37bd6bf49/fastuuid-0.14.0-cp314-cp314-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:aa75b6657ec129d0abded3bec745e6f7ab642e6dba3a5272a68247e85f5f316f", size = 277798, upload-time = "2025-10-19T22:33:53.821Z" }, + { url = "https://files.pythonhosted.org/packages/ef/33/4105ca574f6ded0af6a797d39add041bcfb468a1255fbbe82fcb6f592da2/fastuuid-0.14.0-cp314-cp314-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a8a0dfea3972200f72d4c7df02c8ac70bad1bb4c58d7e0ec1e6f341679073a7f", size = 278283, upload-time = "2025-10-19T22:29:02.812Z" }, + { url = "https://files.pythonhosted.org/packages/fe/8c/fca59f8e21c4deb013f574eae05723737ddb1d2937ce87cb2a5d20992dc3/fastuuid-0.14.0-cp314-cp314-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:1bf539a7a95f35b419f9ad105d5a8a35036df35fdafae48fb2fd2e5f318f0d75", size = 301627, upload-time = "2025-10-19T22:35:54.985Z" }, + { url = "https://files.pythonhosted.org/packages/cb/e2/f78c271b909c034d429218f2798ca4e89eeda7983f4257d7865976ddbb6c/fastuuid-0.14.0-cp314-cp314-musllinux_1_1_aarch64.whl", hash = "sha256:9a133bf9cc78fdbd1179cb58a59ad0100aa32d8675508150f3658814aeefeaa4", size = 459778, upload-time = "2025-10-19T22:28:00.999Z" }, + { url = "https://files.pythonhosted.org/packages/1e/f0/5ff209d865897667a2ff3e7a572267a9ced8f7313919f6d6043aed8b1caa/fastuuid-0.14.0-cp314-cp314-musllinux_1_1_i686.whl", hash = "sha256:f54d5b36c56a2d5e1a31e73b950b28a0d83eb0c37b91d10408875a5a29494bad", size = 478605, upload-time = "2025-10-19T22:36:21.764Z" }, + { url = "https://files.pythonhosted.org/packages/e0/c8/2ce1c78f983a2c4987ea865d9516dbdfb141a120fd3abb977ae6f02ba7ca/fastuuid-0.14.0-cp314-cp314-musllinux_1_1_x86_64.whl", hash = "sha256:ec27778c6ca3393ef662e2762dba8af13f4ec1aaa32d08d77f71f2a70ae9feb8", size = 450837, upload-time = "2025-10-19T22:34:37.178Z" }, + { url = "https://files.pythonhosted.org/packages/df/60/dad662ec9a33b4a5fe44f60699258da64172c39bd041da2994422cdc40fe/fastuuid-0.14.0-cp314-cp314-win32.whl", hash = "sha256:e23fc6a83f112de4be0cc1990e5b127c27663ae43f866353166f87df58e73d06", size = 154532, upload-time = "2025-10-19T22:35:18.217Z" }, + { url = "https://files.pythonhosted.org/packages/1f/f6/da4db31001e854025ffd26bc9ba0740a9cbba2c3259695f7c5834908b336/fastuuid-0.14.0-cp314-cp314-win_amd64.whl", hash = "sha256:df61342889d0f5e7a32f7284e55ef95103f2110fee433c2ae7c2c0956d76ac8a", size = 156457, upload-time = "2025-10-19T22:33:44.579Z" }, +] + [[package]] name = "ffmpeg-python" version = "0.2.0" @@ -1505,7 +1546,9 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/65/8b/3669ad3b3f247a791b2b4aceb3aa5a31f5f6817bf547e4e1ff712338145a/greenlet-3.4.0-cp312-cp312-macosx_11_0_universal2.whl", hash = "sha256:1a54a921561dd9518d31d2d3db4d7f80e589083063ab4d3e2e950756ef809e1a", size = 286902, upload-time = "2026-04-08T15:52:12.138Z" }, { url = "https://files.pythonhosted.org/packages/38/3e/3c0e19b82900873e2d8469b590a6c4b3dfd2b316d0591f1c26b38a4879a5/greenlet-3.4.0-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:16dec271460a9a2b154e3b1c2fa1050ce6280878430320e85e08c166772e3f97", size = 606099, upload-time = "2026-04-08T16:24:38.408Z" }, { url = "https://files.pythonhosted.org/packages/b5/33/99fef65e7754fc76a4ed14794074c38c9ed3394a5bd129d7f61b705f3168/greenlet-3.4.0-cp312-cp312-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:90036ce224ed6fe75508c1907a77e4540176dcf0744473627785dd519c6f9996", size = 618837, upload-time = "2026-04-08T16:30:58.298Z" }, + { url = "https://files.pythonhosted.org/packages/44/57/eae2cac10421feae6c0987e3dc106c6d86262b1cb379e171b017aba893a6/greenlet-3.4.0-cp312-cp312-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:6f0def07ec9a71d72315cf26c061aceee53b306c36ed38c35caba952ea1b319d", size = 624901, upload-time = "2026-04-08T16:40:38.981Z" }, { url = "https://files.pythonhosted.org/packages/36/f7/229f3aed6948faa20e0616a0b8568da22e365ede6a54d7d369058b128afd/greenlet-3.4.0-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a1c4f6b453006efb8310affb2d132832e9bbb4fc01ce6df6b70d810d38f1f6dc", size = 615062, upload-time = "2026-04-08T15:56:33.766Z" }, + { url = "https://files.pythonhosted.org/packages/6a/8a/0e73c9b94f31d1cc257fe79a0eff621674141cdae7d6d00f40de378a1e42/greenlet-3.4.0-cp312-cp312-manylinux_2_39_riscv64.whl", hash = "sha256:0e1254cf0cbaa17b04320c3a78575f29f3c161ef38f59c977108f19ffddaf077", size = 423927, upload-time = "2026-04-08T16:43:05.293Z" }, { url = "https://files.pythonhosted.org/packages/08/97/d988180011aa40135c46cd0d0cf01dd97f7162bae14139b4a3ef54889ba5/greenlet-3.4.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:9b2d9a138ffa0e306d0e2b72976d2fb10b97e690d40ab36a472acaab0838e2de", size = 1573511, upload-time = "2026-04-08T16:26:20.058Z" }, { url = "https://files.pythonhosted.org/packages/d4/0f/a5a26fe152fb3d12e6a474181f6e9848283504d0afd095f353d85726374b/greenlet-3.4.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:8424683caf46eb0eb6f626cb95e008e8cc30d0cb675bdfa48200925c79b38a08", size = 1640396, upload-time = "2026-04-08T15:57:30.88Z" }, { url = "https://files.pythonhosted.org/packages/42/cf/bb2c32d9a100e36ee9f6e38fad6b1e082b8184010cb06259b49e1266ca01/greenlet-3.4.0-cp312-cp312-win_amd64.whl", hash = "sha256:a0a53fb071531d003b075c444014ff8f8b1a9898d36bb88abd9ac7b3524648a2", size = 238892, upload-time = "2026-04-08T17:03:10.094Z" }, @@ -1513,7 +1556,9 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/7a/75/7e9cd1126a1e1f0cd67b0eda02e5221b28488d352684704a78ed505bd719/greenlet-3.4.0-cp313-cp313-macosx_11_0_universal2.whl", hash = "sha256:43748988b097f9c6f09364f260741aa73c80747f63389824435c7a50bfdfd5c1", size = 285856, upload-time = "2026-04-08T15:52:45.82Z" }, { url = "https://files.pythonhosted.org/packages/9d/c4/3e2df392e5cb199527c4d9dbcaa75c14edcc394b45040f0189f649631e3c/greenlet-3.4.0-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5566e4e2cd7a880e8c27618e3eab20f3494452d12fd5129edef7b2f7aa9a36d1", size = 610208, upload-time = "2026-04-08T16:24:39.674Z" }, { url = "https://files.pythonhosted.org/packages/da/af/750cdfda1d1bd30a6c28080245be8d0346e669a98fdbae7f4102aa95fff3/greenlet-3.4.0-cp313-cp313-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:1054c5a3c78e2ab599d452f23f7adafef55062a783a8e241d24f3b633ba6ff82", size = 621269, upload-time = "2026-04-08T16:30:59.767Z" }, + { url = "https://files.pythonhosted.org/packages/e0/93/c8c508d68ba93232784bbc1b5474d92371f2897dfc6bc281b419f2e0d492/greenlet-3.4.0-cp313-cp313-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:98eedd1803353daf1cd9ef23eef23eda5a4d22f99b1f998d273a8b78b70dd47f", size = 628455, upload-time = "2026-04-08T16:40:40.698Z" }, { url = "https://files.pythonhosted.org/packages/54/78/0cbc693622cd54ebe25207efbb3a0eb07c2639cb8594f6e3aaaa0bb077a8/greenlet-3.4.0-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f82cb6cddc27dd81c96b1506f4aa7def15070c3b2a67d4e46fd19016aacce6cf", size = 617549, upload-time = "2026-04-08T15:56:34.893Z" }, + { url = "https://files.pythonhosted.org/packages/7f/46/cfaaa0ade435a60550fd83d07dfd5c41f873a01da17ede5c4cade0b9bab8/greenlet-3.4.0-cp313-cp313-manylinux_2_39_riscv64.whl", hash = "sha256:b7857e2202aae67bc5725e0c1f6403c20a8ff46094ece015e7d474f5f7020b55", size = 426238, upload-time = "2026-04-08T16:43:06.865Z" }, { url = "https://files.pythonhosted.org/packages/ba/c0/8966767de01343c1ff47e8b855dc78e7d1a8ed2b7b9c83576a57e289f81d/greenlet-3.4.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:227a46251ecba4ff46ae742bc5ce95c91d5aceb4b02f885487aff269c127a729", size = 1575310, upload-time = "2026-04-08T16:26:21.671Z" }, { url = "https://files.pythonhosted.org/packages/b8/38/bcdc71ba05e9a5fda87f63ffc2abcd1f15693b659346df994a48c968003d/greenlet-3.4.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:5b99e87be7eba788dd5b75ba1cde5639edffdec5f91fe0d734a249535ec3408c", size = 1640435, upload-time = "2026-04-08T15:57:32.572Z" }, { url = "https://files.pythonhosted.org/packages/a1/c2/19b664b7173b9e4ef5f77e8cef9f14c20ec7fce7920dc1ccd7afd955d093/greenlet-3.4.0-cp313-cp313-win_amd64.whl", hash = "sha256:849f8bc17acd6295fcb5de8e46d55cc0e52381c56eaf50a2afd258e97bc65940", size = 238760, upload-time = "2026-04-08T17:04:03.878Z" }, @@ -1521,7 +1566,9 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/78/02/bde66806e8f169cf90b14d02c500c44cdbe02c8e224c9c67bafd1b8cadd1/greenlet-3.4.0-cp314-cp314-macosx_11_0_universal2.whl", hash = "sha256:10a07aca6babdd18c16a3f4f8880acfffc2b88dfe431ad6aa5f5740759d7d75e", size = 286291, upload-time = "2026-04-08T17:09:34.307Z" }, { url = "https://files.pythonhosted.org/packages/05/1f/39da1c336a87d47c58352fb8a78541ce63d63ae57c5b9dae1fe02801bbc2/greenlet-3.4.0-cp314-cp314-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:076e21040b3a917d3ce4ad68fb5c3c6b32f1405616c4a57aa83120979649bd3d", size = 656749, upload-time = "2026-04-08T16:24:41.721Z" }, { url = "https://files.pythonhosted.org/packages/d3/6c/90ee29a4ee27af7aa2e2ec408799eeb69ee3fcc5abcecac6ddd07a5cd0f2/greenlet-3.4.0-cp314-cp314-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:e82689eea4a237e530bb5cb41b180ef81fa2160e1f89422a67be7d90da67f615", size = 669084, upload-time = "2026-04-08T16:31:01.372Z" }, + { url = "https://files.pythonhosted.org/packages/d2/4a/74078d3936712cff6d3c91a930016f476ce4198d84e224fe6d81d3e02880/greenlet-3.4.0-cp314-cp314-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:06c2d3b89e0c62ba50bd7adf491b14f39da9e7e701647cb7b9ff4c99bee04b19", size = 673405, upload-time = "2026-04-08T16:40:42.527Z" }, { url = "https://files.pythonhosted.org/packages/07/49/d4cad6e5381a50947bb973d2f6cf6592621451b09368b8c20d9b8af49c5b/greenlet-3.4.0-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4df3b0b2289ec686d3c821a5fee44259c05cfe824dd5e6e12c8e5f5df23085cf", size = 665621, upload-time = "2026-04-08T15:56:35.995Z" }, + { url = "https://files.pythonhosted.org/packages/79/3e/df8a83ab894751bc31e1106fdfaa80ca9753222f106b04de93faaa55feb7/greenlet-3.4.0-cp314-cp314-manylinux_2_39_riscv64.whl", hash = "sha256:070b8bac2ff3b4d9e0ff36a0d19e42103331d9737e8504747cd1e659f76297bd", size = 471670, upload-time = "2026-04-08T16:43:08.512Z" }, { url = "https://files.pythonhosted.org/packages/37/31/d1edd54f424761b5d47718822f506b435b6aab2f3f93b465441143ea5119/greenlet-3.4.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:8bff29d586ea415688f4cec96a591fcc3bf762d046a796cdadc1fdb6e7f2d5bf", size = 1622259, upload-time = "2026-04-08T16:26:23.201Z" }, { url = "https://files.pythonhosted.org/packages/b0/c6/6d3f9cdcb21c4e12a79cb332579f1c6aa1af78eb68059c5a957c7812d95e/greenlet-3.4.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:8a569c2fb840c53c13a2b8967c63621fafbd1a0e015b9c82f408c33d626a2fda", size = 1686916, upload-time = "2026-04-08T15:57:34.282Z" }, { url = "https://files.pythonhosted.org/packages/63/45/c1ca4a1ad975de4727e52d3ffe641ae23e1d7a8ffaa8ff7a0477e1827b92/greenlet-3.4.0-cp314-cp314-win_amd64.whl", hash = "sha256:207ba5b97ea8b0b60eb43ffcacf26969dd83726095161d676aac03ff913ee50d", size = 239821, upload-time = "2026-04-08T17:03:48.423Z" }, @@ -1529,7 +1576,9 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d4/8f/18d72b629783f5e8d045a76f5325c1e938e659a9e4da79c7dcd10169a48d/greenlet-3.4.0-cp314-cp314t-macosx_11_0_universal2.whl", hash = "sha256:d70012e51df2dbbccfaf63a40aaf9b40c8bed37c3e3a38751c926301ce538ece", size = 294681, upload-time = "2026-04-08T15:52:35.778Z" }, { url = "https://files.pythonhosted.org/packages/9e/ad/5fa86ec46769c4153820d58a04062285b3b9e10ba3d461ee257b68dcbf53/greenlet-3.4.0-cp314-cp314t-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a58bec0751f43068cd40cff31bb3ca02ad6000b3a51ca81367af4eb5abc480c8", size = 658899, upload-time = "2026-04-08T16:24:43.32Z" }, { url = "https://files.pythonhosted.org/packages/43/f0/4e8174ca0e87ae748c409f055a1ba161038c43cc0a5a6f1433a26ac2e5bf/greenlet-3.4.0-cp314-cp314t-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:05fa0803561028f4b2e3b490ee41216a842eaee11aed004cc343a996d9523aa2", size = 665284, upload-time = "2026-04-08T16:31:02.833Z" }, + { url = "https://files.pythonhosted.org/packages/ef/92/466b0d9afd44b8af623139a3599d651c7564fa4152f25f117e1ee5949ffb/greenlet-3.4.0-cp314-cp314t-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:c4cd56a9eb7a6444edbc19062f7b6fbc8f287c663b946e3171d899693b1c19fa", size = 665872, upload-time = "2026-04-08T16:40:43.912Z" }, { url = "https://files.pythonhosted.org/packages/19/da/991cf7cd33662e2df92a1274b7eb4d61769294d38a1bba8a45f31364845e/greenlet-3.4.0-cp314-cp314t-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e60d38719cb80b3ab5e85f9f1aed4960acfde09868af6762ccb27b260d68f4ed", size = 661861, upload-time = "2026-04-08T15:56:37.269Z" }, + { url = "https://files.pythonhosted.org/packages/0d/14/3395a7ef3e260de0325152ddfe19dffb3e49fe10873b94654352b53ad48e/greenlet-3.4.0-cp314-cp314t-manylinux_2_39_riscv64.whl", hash = "sha256:1f85f204c4d54134ae850d401fa435c89cd667d5ce9dc567571776b45941af72", size = 489237, upload-time = "2026-04-08T16:43:09.993Z" }, { url = "https://files.pythonhosted.org/packages/36/c5/6c2c708e14db3d9caea4b459d8464f58c32047451142fe2cfd90e7458f41/greenlet-3.4.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:7f50c804733b43eded05ae694691c9aa68bca7d0a867d67d4a3f514742a2d53f", size = 1622182, upload-time = "2026-04-08T16:26:24.777Z" }, { url = "https://files.pythonhosted.org/packages/7a/4c/50c5fed19378e11a29fabab1f6be39ea95358f4a0a07e115a51ca93385d8/greenlet-3.4.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:2d4f0635dc4aa638cda4b2f5a07ae9a2cff9280327b581a3fcb6f317b4fbc38a", size = 1685050, upload-time = "2026-04-08T15:57:36.453Z" }, { url = "https://files.pythonhosted.org/packages/db/72/85ae954d734703ab48e622c59d4ce35d77ce840c265814af9c078cacc7aa/greenlet-3.4.0-cp314-cp314t-win_amd64.whl", hash = "sha256:1a4a48f24681300c640f143ba7c404270e1ebbbcf34331d7104a4ff40f8ea705", size = 245554, upload-time = "2026-04-08T17:03:50.044Z" }, @@ -1807,6 +1856,8 @@ version = "0.13.0" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/0d/5e/4ec91646aee381d01cdb9974e30882c9cd3b8c5d1079d6b5ff4af522439a/jiter-0.13.0.tar.gz", hash = "sha256:f2839f9c2c7e2dffc1bc5929a510e14ce0a946be9365fd1219e7ef342dae14f4", size = 164847, upload-time = "2026-02-02T12:37:56.441Z" } wheels = [ + { url = "https://files.pythonhosted.org/packages/2e/30/7687e4f87086829955013ca12a9233523349767f69653ebc27036313def9/jiter-0.13.0-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:0a2bd69fc1d902e89925fc34d1da51b2128019423d7b339a45d9e99c894e0663", size = 307958, upload-time = "2026-02-02T12:35:57.165Z" }, + { url = "https://files.pythonhosted.org/packages/c3/27/e57f9a783246ed95481e6749cc5002a8a767a73177a83c63ea71f0528b90/jiter-0.13.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:f917a04240ef31898182f76a332f508f2cc4b57d2b4d7ad2dbfebbfe167eb505", size = 318597, upload-time = "2026-02-02T12:35:58.591Z" }, { url = "https://files.pythonhosted.org/packages/cf/52/e5719a60ac5d4d7c5995461a94ad5ef962a37c8bf5b088390e6fad59b2ff/jiter-0.13.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c1e2b199f446d3e82246b4fd9236d7cb502dc2222b18698ba0d986d2fecc6152", size = 348821, upload-time = "2026-02-02T12:36:00.093Z" }, { url = "https://files.pythonhosted.org/packages/61/db/c1efc32b8ba4c740ab3fc2d037d8753f67685f475e26b9d6536a4322bcdd/jiter-0.13.0-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:04670992b576fa65bd056dbac0c39fe8bd67681c380cb2b48efa885711d9d726", size = 364163, upload-time = "2026-02-02T12:36:01.937Z" }, { url = "https://files.pythonhosted.org/packages/55/8a/fb75556236047c8806995671a18e4a0ad646ed255276f51a20f32dceaeec/jiter-0.13.0-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:5a1aff1fbdb803a376d4d22a8f63f8e7ccbce0b4890c26cc7af9e501ab339ef0", size = 483709, upload-time = "2026-02-02T12:36:03.41Z" }, @@ -1815,6 +1866,11 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/1a/1b/46c5e349019874ec5dfa508c14c37e29864ea108d376ae26d90bee238cd7/jiter-0.13.0-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:7b88d649135aca526da172e48083da915ec086b54e8e73a425ba50999468cc08", size = 391814, upload-time = "2026-02-02T12:36:08.368Z" }, { url = "https://files.pythonhosted.org/packages/15/9e/26184760e85baee7162ad37b7912797d2077718476bf91517641c92b3639/jiter-0.13.0-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:e404ea551d35438013c64b4f357b0474c7abf9f781c06d44fcaf7a14c69ff9e2", size = 513990, upload-time = "2026-02-02T12:36:09.993Z" }, { url = "https://files.pythonhosted.org/packages/e9/34/2c9355247d6debad57a0a15e76ab1566ab799388042743656e566b3b7de1/jiter-0.13.0-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:1f4748aad1b4a93c8bdd70f604d0f748cdc0e8744c5547798acfa52f10e79228", size = 548021, upload-time = "2026-02-02T12:36:11.376Z" }, + { url = "https://files.pythonhosted.org/packages/ac/4a/9f2c23255d04a834398b9c2e0e665382116911dc4d06b795710503cdad25/jiter-0.13.0-cp312-cp312-win32.whl", hash = "sha256:0bf670e3b1445fc4d31612199f1744f67f889ee1bbae703c4b54dc097e5dd394", size = 203024, upload-time = "2026-02-02T12:36:12.682Z" }, + { url = "https://files.pythonhosted.org/packages/09/ee/f0ae675a957ae5a8f160be3e87acea6b11dc7b89f6b7ab057e77b2d2b13a/jiter-0.13.0-cp312-cp312-win_amd64.whl", hash = "sha256:15db60e121e11fe186c0b15236bd5d18381b9ddacdcf4e659feb96fc6c969c92", size = 205424, upload-time = "2026-02-02T12:36:13.93Z" }, + { url = "https://files.pythonhosted.org/packages/1b/02/ae611edf913d3cbf02c97cdb90374af2082c48d7190d74c1111dde08bcdd/jiter-0.13.0-cp312-cp312-win_arm64.whl", hash = "sha256:41f92313d17989102f3cb5dd533a02787cdb99454d494344b0361355da52fcb9", size = 186818, upload-time = "2026-02-02T12:36:15.308Z" }, + { url = "https://files.pythonhosted.org/packages/91/9c/7ee5a6ff4b9991e1a45263bfc46731634c4a2bde27dfda6c8251df2d958c/jiter-0.13.0-cp313-cp313-macosx_10_12_x86_64.whl", hash = "sha256:1f8a55b848cbabf97d861495cd65f1e5c590246fabca8b48e1747c4dfc8f85bf", size = 306897, upload-time = "2026-02-02T12:36:16.748Z" }, + { url = "https://files.pythonhosted.org/packages/7c/02/be5b870d1d2be5dd6a91bdfb90f248fbb7dcbd21338f092c6b89817c3dbf/jiter-0.13.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:f556aa591c00f2c45eb1b89f68f52441a016034d18b65da60e2d2875bbbf344a", size = 317507, upload-time = "2026-02-02T12:36:18.351Z" }, { url = "https://files.pythonhosted.org/packages/da/92/b25d2ec333615f5f284f3a4024f7ce68cfa0604c322c6808b2344c7f5d2b/jiter-0.13.0-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f7e1d61da332ec412350463891923f960c3073cf1aae93b538f0bb4c8cd46efb", size = 350560, upload-time = "2026-02-02T12:36:19.746Z" }, { url = "https://files.pythonhosted.org/packages/be/ec/74dcb99fef0aca9fbe56b303bf79f6bd839010cb18ad41000bf6cc71eec0/jiter-0.13.0-cp313-cp313-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:3097d665a27bc96fd9bbf7f86178037db139f319f785e4757ce7ccbf390db6c2", size = 363232, upload-time = "2026-02-02T12:36:21.243Z" }, { url = "https://files.pythonhosted.org/packages/1b/37/f17375e0bb2f6a812d4dd92d7616e41917f740f3e71343627da9db2824ce/jiter-0.13.0-cp313-cp313-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:9d01ecc3a8cbdb6f25a37bd500510550b64ddf9f7d64a107d92f3ccb25035d0f", size = 483727, upload-time = "2026-02-02T12:36:22.688Z" }, @@ -1823,8 +1879,16 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/16/be/7ffd08203277a813f732ba897352797fa9493faf8dc7995b31f3d9cb9488/jiter-0.13.0-cp313-cp313-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:5467696f6b827f1116556cb0db620440380434591e93ecee7fd14d1a491b6daa", size = 390664, upload-time = "2026-02-02T12:36:26.866Z" }, { url = "https://files.pythonhosted.org/packages/d1/84/e0787856196d6d346264d6dcccb01f741e5f0bd014c1d9a2ebe149caf4f3/jiter-0.13.0-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:2d08c9475d48b92892583df9da592a0e2ac49bcd41fae1fec4f39ba6cf107820", size = 513543, upload-time = "2026-02-02T12:36:28.217Z" }, { url = "https://files.pythonhosted.org/packages/65/50/ecbd258181c4313cf79bca6c88fb63207d04d5bf5e4f65174114d072aa55/jiter-0.13.0-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:aed40e099404721d7fcaf5b89bd3b4568a4666358bcac7b6b15c09fb6252ab68", size = 547262, upload-time = "2026-02-02T12:36:29.678Z" }, + { url = "https://files.pythonhosted.org/packages/27/da/68f38d12e7111d2016cd198161b36e1f042bd115c169255bcb7ec823a3bf/jiter-0.13.0-cp313-cp313-win32.whl", hash = "sha256:36ebfbcffafb146d0e6ffb3e74d51e03d9c35ce7c625c8066cdbfc7b953bdc72", size = 200630, upload-time = "2026-02-02T12:36:31.808Z" }, + { url = "https://files.pythonhosted.org/packages/25/65/3bd1a972c9a08ecd22eb3b08a95d1941ebe6938aea620c246cf426ae09c2/jiter-0.13.0-cp313-cp313-win_amd64.whl", hash = "sha256:8d76029f077379374cf0dbc78dbe45b38dec4a2eb78b08b5194ce836b2517afc", size = 202602, upload-time = "2026-02-02T12:36:33.679Z" }, + { url = "https://files.pythonhosted.org/packages/15/fe/13bd3678a311aa67686bb303654792c48206a112068f8b0b21426eb6851e/jiter-0.13.0-cp313-cp313-win_arm64.whl", hash = "sha256:bb7613e1a427cfcb6ea4544f9ac566b93d5bf67e0d48c787eca673ff9c9dff2b", size = 185939, upload-time = "2026-02-02T12:36:35.065Z" }, + { url = "https://files.pythonhosted.org/packages/49/19/a929ec002ad3228bc97ca01dbb14f7632fffdc84a95ec92ceaf4145688ae/jiter-0.13.0-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:fa476ab5dd49f3bf3a168e05f89358c75a17608dbabb080ef65f96b27c19ab10", size = 316616, upload-time = "2026-02-02T12:36:36.579Z" }, { url = "https://files.pythonhosted.org/packages/52/56/d19a9a194afa37c1728831e5fb81b7722c3de18a3109e8f282bfc23e587a/jiter-0.13.0-cp313-cp313t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ade8cb6ff5632a62b7dbd4757d8c5573f7a2e9ae285d6b5b841707d8363205ef", size = 346850, upload-time = "2026-02-02T12:36:38.058Z" }, { url = "https://files.pythonhosted.org/packages/36/4a/94e831c6bf287754a8a019cb966ed39ff8be6ab78cadecf08df3bb02d505/jiter-0.13.0-cp313-cp313t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9950290340acc1adaded363edd94baebcee7dabdfa8bee4790794cd5cfad2af6", size = 358551, upload-time = "2026-02-02T12:36:39.417Z" }, + { url = "https://files.pythonhosted.org/packages/a2/ec/a4c72c822695fa80e55d2b4142b73f0012035d9fcf90eccc56bc060db37c/jiter-0.13.0-cp313-cp313t-win_amd64.whl", hash = "sha256:2b4972c6df33731aac0742b64fd0d18e0a69bc7d6e03108ce7d40c85fd9e3e6d", size = 201950, upload-time = "2026-02-02T12:36:40.791Z" }, + { url = "https://files.pythonhosted.org/packages/b6/00/393553ec27b824fbc29047e9c7cd4a3951d7fbe4a76743f17e44034fa4e4/jiter-0.13.0-cp313-cp313t-win_arm64.whl", hash = "sha256:701a1e77d1e593c1b435315ff625fd071f0998c5f02792038a5ca98899261b7d", size = 185852, upload-time = "2026-02-02T12:36:42.077Z" }, + { url = "https://files.pythonhosted.org/packages/6e/f5/f1997e987211f6f9bd71b8083047b316208b4aca0b529bb5f8c96c89ef3e/jiter-0.13.0-cp314-cp314-macosx_10_12_x86_64.whl", hash = "sha256:cc5223ab19fe25e2f0bf2643204ad7318896fe3729bf12fde41b77bfc4fafff0", size = 308804, upload-time = "2026-02-02T12:36:43.496Z" }, + { url = "https://files.pythonhosted.org/packages/cd/8f/5482a7677731fd44881f0204981ce2d7175db271f82cba2085dd2212e095/jiter-0.13.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:9776ebe51713acf438fd9b4405fcd86893ae5d03487546dae7f34993217f8a91", size = 318787, upload-time = "2026-02-02T12:36:45.071Z" }, { url = "https://files.pythonhosted.org/packages/f3/b9/7257ac59778f1cd025b26a23c5520a36a424f7f1b068f2442a5b499b7464/jiter-0.13.0-cp314-cp314-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:879e768938e7b49b5e90b7e3fecc0dbec01b8cb89595861fb39a8967c5220d09", size = 353880, upload-time = "2026-02-02T12:36:47.365Z" }, { url = "https://files.pythonhosted.org/packages/c3/87/719eec4a3f0841dad99e3d3604ee4cba36af4419a76f3cb0b8e2e691ad67/jiter-0.13.0-cp314-cp314-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:682161a67adea11e3aae9038c06c8b4a9a71023228767477d683f69903ebc607", size = 366702, upload-time = "2026-02-02T12:36:48.871Z" }, { url = "https://files.pythonhosted.org/packages/d2/65/415f0a75cf6921e43365a1bc227c565cb949caca8b7532776e430cbaa530/jiter-0.13.0-cp314-cp314-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:a13b68cd1cd8cc9de8f244ebae18ccb3e4067ad205220ef324c39181e23bbf66", size = 486319, upload-time = "2026-02-02T12:36:53.006Z" }, @@ -1833,6 +1897,10 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/17/08/91b9ea976c1c758240614bd88442681a87672eebc3d9a6dde476874e706b/jiter-0.13.0-cp314-cp314-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:1317fdffd16f5873e46ce27d0e0f7f4f90f0cdf1d86bf6abeaea9f63ca2c401d", size = 389634, upload-time = "2026-02-02T12:36:57.495Z" }, { url = "https://files.pythonhosted.org/packages/18/23/58325ef99390d6d40427ed6005bf1ad54f2577866594bcf13ce55675f87d/jiter-0.13.0-cp314-cp314-musllinux_1_1_aarch64.whl", hash = "sha256:c05b450d37ba0c9e21c77fef1f205f56bcee2330bddca68d344baebfc55ae0df", size = 514933, upload-time = "2026-02-02T12:36:58.909Z" }, { url = "https://files.pythonhosted.org/packages/5b/25/69f1120c7c395fd276c3996bb8adefa9c6b84c12bb7111e5c6ccdcd8526d/jiter-0.13.0-cp314-cp314-musllinux_1_1_x86_64.whl", hash = "sha256:775e10de3849d0631a97c603f996f518159272db00fdda0a780f81752255ee9d", size = 548842, upload-time = "2026-02-02T12:37:00.433Z" }, + { url = "https://files.pythonhosted.org/packages/18/05/981c9669d86850c5fbb0d9e62bba144787f9fba84546ba43d624ee27ef29/jiter-0.13.0-cp314-cp314-win32.whl", hash = "sha256:632bf7c1d28421c00dd8bbb8a3bac5663e1f57d5cd5ed962bce3c73bf62608e6", size = 202108, upload-time = "2026-02-02T12:37:01.718Z" }, + { url = "https://files.pythonhosted.org/packages/8d/96/cdcf54dd0b0341db7d25413229888a346c7130bd20820530905fdb65727b/jiter-0.13.0-cp314-cp314-win_amd64.whl", hash = "sha256:f22ef501c3f87ede88f23f9b11e608581c14f04db59b6a801f354397ae13739f", size = 204027, upload-time = "2026-02-02T12:37:03.075Z" }, + { url = "https://files.pythonhosted.org/packages/fb/f9/724bcaaab7a3cd727031fe4f6995cb86c4bd344909177c186699c8dec51a/jiter-0.13.0-cp314-cp314-win_arm64.whl", hash = "sha256:07b75fe09a4ee8e0c606200622e571e44943f47254f95e2436c8bdcaceb36d7d", size = 187199, upload-time = "2026-02-02T12:37:04.414Z" }, + { url = "https://files.pythonhosted.org/packages/62/92/1661d8b9fd6a3d7a2d89831db26fe3c1509a287d83ad7838831c7b7a5c7e/jiter-0.13.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:964538479359059a35fb400e769295d4b315ae61e4105396d355a12f7fef09f0", size = 318423, upload-time = "2026-02-02T12:37:05.806Z" }, { url = "https://files.pythonhosted.org/packages/4f/3b/f77d342a54d4ebcd128e520fc58ec2f5b30a423b0fd26acdfc0c6fef8e26/jiter-0.13.0-cp314-cp314t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e104da1db1c0991b3eaed391ccd650ae8d947eab1480c733e5a3fb28d4313e40", size = 351438, upload-time = "2026-02-02T12:37:07.189Z" }, { url = "https://files.pythonhosted.org/packages/76/b3/ba9a69f0e4209bd3331470c723c2f5509e6f0482e416b612431a5061ed71/jiter-0.13.0-cp314-cp314t-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:0e3a5f0cde8ff433b8e88e41aa40131455420fb3649a3c7abdda6145f8cb7202", size = 364774, upload-time = "2026-02-02T12:37:08.579Z" }, { url = "https://files.pythonhosted.org/packages/b3/16/6cdb31fa342932602458dbb631bfbd47f601e03d2e4950740e0b2100b570/jiter-0.13.0-cp314-cp314t-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:57aab48f40be1db920a582b30b116fe2435d184f77f0e4226f546794cedd9cf0", size = 487238, upload-time = "2026-02-02T12:37:10.066Z" }, @@ -1841,6 +1909,11 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/4b/d7/eabe3cf46715854ccc80be2cd78dd4c36aedeb30751dbf85a1d08c14373c/jiter-0.13.0-cp314-cp314t-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:7beae3a3d3b5212d3a55d2961db3c292e02e302feb43fce6a3f7a31b90ea6dfe", size = 389607, upload-time = "2026-02-02T12:37:14.881Z" }, { url = "https://files.pythonhosted.org/packages/df/2d/03963fc0804e6109b82decfb9974eb92df3797fe7222428cae12f8ccaa0c/jiter-0.13.0-cp314-cp314t-musllinux_1_1_aarch64.whl", hash = "sha256:e5562a0f0e90a6223b704163ea28e831bd3a9faa3512a711f031611e6b06c939", size = 514986, upload-time = "2026-02-02T12:37:16.326Z" }, { url = "https://files.pythonhosted.org/packages/f6/6c/8c83b45eb3eb1c1e18d841fe30b4b5bc5619d781267ca9bc03e005d8fd0a/jiter-0.13.0-cp314-cp314t-musllinux_1_1_x86_64.whl", hash = "sha256:6c26a424569a59140fb51160a56df13f438a2b0967365e987889186d5fc2f6f9", size = 548756, upload-time = "2026-02-02T12:37:17.736Z" }, + { url = "https://files.pythonhosted.org/packages/47/66/eea81dfff765ed66c68fd2ed8c96245109e13c896c2a5015c7839c92367e/jiter-0.13.0-cp314-cp314t-win32.whl", hash = "sha256:24dc96eca9f84da4131cdf87a95e6ce36765c3b156fc9ae33280873b1c32d5f6", size = 201196, upload-time = "2026-02-02T12:37:19.101Z" }, + { url = "https://files.pythonhosted.org/packages/ff/32/4ac9c7a76402f8f00d00842a7f6b83b284d0cf7c1e9d4227bc95aa6d17fa/jiter-0.13.0-cp314-cp314t-win_amd64.whl", hash = "sha256:0a8d76c7524087272c8ae913f5d9d608bd839154b62c4322ef65723d2e5bb0b8", size = 204215, upload-time = "2026-02-02T12:37:20.495Z" }, + { url = "https://files.pythonhosted.org/packages/f9/8e/7def204fea9f9be8b3c21a6f2dd6c020cf56c7d5ff753e0e23ed7f9ea57e/jiter-0.13.0-cp314-cp314t-win_arm64.whl", hash = "sha256:2c26cf47e2cad140fa23b6d58d435a7c0161f5c514284802f25e87fddfe11024", size = 187152, upload-time = "2026-02-02T12:37:22.124Z" }, + { url = "https://files.pythonhosted.org/packages/80/60/e50fa45dd7e2eae049f0ce964663849e897300433921198aef94b6ffa23a/jiter-0.13.0-graalpy312-graalpy250_312_native-macosx_10_12_x86_64.whl", hash = "sha256:3d744a6061afba08dd7ae375dcde870cffb14429b7477e10f67e9e6d68772a0a", size = 305169, upload-time = "2026-02-02T12:37:50.376Z" }, + { url = "https://files.pythonhosted.org/packages/d2/73/a009f41c5eed71c49bec53036c4b33555afcdee70682a18c6f66e396c039/jiter-0.13.0-graalpy312-graalpy250_312_native-macosx_11_0_arm64.whl", hash = "sha256:ff732bd0a0e778f43d5009840f20b935e79087b4dc65bd36f1cd0f9b04b8ff7f", size = 303808, upload-time = "2026-02-02T12:37:52.092Z" }, { url = "https://files.pythonhosted.org/packages/c4/10/528b439290763bff3d939268085d03382471b442f212dca4ff5f12802d43/jiter-0.13.0-graalpy312-graalpy250_312_native-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ab44b178f7981fcaea7e0a5df20e773c663d06ffda0198f1a524e91b2fde7e59", size = 337384, upload-time = "2026-02-02T12:37:53.582Z" }, { url = "https://files.pythonhosted.org/packages/67/8a/a342b2f0251f3dac4ca17618265d93bf244a2a4d089126e81e4c1056ac50/jiter-0.13.0-graalpy312-graalpy250_312_native-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7bb00b6d26db67a05fe3e12c76edc75f32077fb51deed13822dc648fa373bc19", size = 343768, upload-time = "2026-02-02T12:37:55.055Z" }, ] @@ -2096,6 +2169,29 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/2d/00/d90b10b962b4277f5e64a78b6609968859ff86889f5b898c1a778c06ec00/lark-1.2.2-py3-none-any.whl", hash = "sha256:c2276486b02f0f1b90be155f2c8ba4a8e194d42775786db622faccd652d8e80c", size = 111036, upload-time = "2024-08-13T19:48:58.603Z" }, ] +[[package]] +name = "litellm" +version = "1.83.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "aiohttp" }, + { name = "click" }, + { name = "fastuuid" }, + { name = "httpx" }, + { name = "importlib-metadata" }, + { name = "jinja2" }, + { name = "jsonschema" }, + { name = "openai" }, + { name = "pydantic" }, + { name = "python-dotenv" }, + { name = "tiktoken" }, + { name = "tokenizers" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/22/92/6ce9737554994ca8e536e5f4f6a87cc7c4774b656c9eb9add071caf7d54b/litellm-1.83.0.tar.gz", hash = "sha256:860bebc76c4bb27b4cf90b4a77acd66dba25aced37e3db98750de8a1766bfb7a", size = 17333062, upload-time = "2026-03-31T05:08:25.331Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/19/2c/a670cc050fcd6f45c6199eb99e259c73aea92edba8d5c2fc1b3686d36217/litellm-1.83.0-py3-none-any.whl", hash = "sha256:88c536d339248f3987571493015784671ba3f193a328e1ea6780dbebaa2094a8", size = 15610306, upload-time = "2026-03-31T05:08:21.987Z" }, +] + [[package]] name = "llguidance" version = "1.3.0" @@ -2644,6 +2740,7 @@ wheels = [ name = "nemo-retriever" source = { editable = "." } dependencies = [ + { name = "aiohttp" }, { name = "debugpy" }, { name = "fastapi" }, { name = "httpx" }, @@ -2664,6 +2761,7 @@ dependencies = [ { name = "ray", extra = ["data", "serve"] }, { name = "requests" }, { name = "rich" }, + { name = "sqlglot" }, { name = "tqdm" }, { name = "typer" }, { name = "universal-pathlib" }, @@ -2712,6 +2810,11 @@ dev = [ { name = "build" }, { name = "pytest" }, ] +eval = [ + { name = "litellm" }, + { name = "pyyaml" }, + { name = "tenacity" }, +] local = [ { name = "accelerate" }, { name = "addict" }, @@ -2751,6 +2854,7 @@ stores = [ requires-dist = [ { name = "accelerate", marker = "extra == 'local'", specifier = "==1.12.0" }, { name = "addict", marker = "extra == 'local'" }, + { name = "aiohttp", specifier = ">=3.9.0" }, { name = "albumentations", marker = "extra == 'local'", specifier = "==2.0.8" }, { name = "build", marker = "extra == 'dev'", specifier = ">=1.2.2" }, { name = "cairosvg", marker = "extra == 'multimedia'", specifier = ">=2.7.0" }, @@ -2764,6 +2868,7 @@ requires-dist = [ { name = "httpx", specifier = ">=0.27.0" }, { name = "lancedb" }, { name = "langchain-nvidia-ai-endpoints", specifier = ">=0.3.0" }, + { name = "litellm", marker = "extra == 'eval'", specifier = ">=1.40.0" }, { name = "markitdown" }, { name = "nemo-retriever", extras = ["benchmarks", "local", "multimedia", "stores"], marker = "extra == 'all'" }, { name = "nemotron-graphic-elements-v1", marker = "extra == 'local'", specifier = ">=0.dev0", index = "https://test.pypi.org/simple/" }, @@ -2784,12 +2889,15 @@ requires-dist = [ { name = "pytest", marker = "extra == 'dev'", specifier = ">=8.0.2" }, { name = "python-multipart", specifier = ">=0.0.9" }, { name = "pyyaml", specifier = ">=6.0" }, + { name = "pyyaml", marker = "extra == 'eval'", specifier = ">=6.0" }, { name = "ray", extras = ["data", "serve"], specifier = ">=2.49.0" }, { name = "requests", specifier = ">=2.32.5" }, { name = "rich", specifier = ">=13.7.0" }, { name = "scikit-learn", marker = "extra == 'local'", specifier = ">=1.6.0" }, { name = "scipy", marker = "extra == 'multimedia'", specifier = ">=1.11.0" }, { name = "soundfile", marker = "extra == 'multimedia'", specifier = ">=0.12.0" }, + { name = "sqlglot", specifier = ">=30.0.0" }, + { name = "tenacity", marker = "extra == 'eval'", specifier = ">=8.0.0" }, { name = "timm", marker = "extra == 'local'", specifier = "==1.0.22" }, { name = "tokenizers", marker = "extra == 'local'", specifier = ">=0.20.3" }, { name = "torch", marker = "(sys_platform == 'linux' and extra == 'local') or (sys_platform == 'win32' and extra == 'local')", specifier = "~=2.9.1", index = "https://download.pytorch.org/whl/cu130" }, @@ -2806,7 +2914,7 @@ requires-dist = [ { name = "vllm", marker = "platform_machine != 'aarch64' and platform_machine != 'x86_64' and sys_platform == 'linux' and extra == 'local'", specifier = "==0.16.0" }, { name = "vllm", marker = "platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'local'", url = "https://github.com/vllm-project/vllm/releases/download/v0.16.0/vllm-0.16.0+cu130-cp38-abi3-manylinux_2_35_x86_64.whl" }, ] -provides-extras = ["local", "multimedia", "stores", "benchmarks", "dev", "all"] +provides-extras = ["local", "multimedia", "stores", "benchmarks", "eval", "dev", "all"] [[package]] name = "nemotron-graphic-elements-v1" @@ -3481,14 +3589,14 @@ name = "openai" version = "2.31.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "anyio", marker = "sys_platform == 'linux'" }, - { name = "distro", marker = "sys_platform == 'linux'" }, - { name = "httpx", marker = "sys_platform == 'linux'" }, - { name = "jiter", marker = "sys_platform == 'linux'" }, - { name = "pydantic", marker = "sys_platform == 'linux'" }, - { name = "sniffio", marker = "sys_platform == 'linux'" }, - { name = "tqdm", marker = "sys_platform == 'linux'" }, - { name = "typing-extensions", marker = "sys_platform == 'linux'" }, + { name = "anyio" }, + { name = "distro" }, + { name = "httpx" }, + { name = "jiter" }, + { name = "pydantic" }, + { name = "sniffio" }, + { name = "tqdm" }, + { name = "typing-extensions" }, ] sdist = { url = "https://files.pythonhosted.org/packages/94/fe/64b3d035780b3188f86c4f6f1bc202e7bb74757ef028802112273b9dcacf/openai-2.31.0.tar.gz", hash = "sha256:43ca59a88fc973ad1848d86b98d7fac207e265ebbd1828b5e4bdfc85f79427a5", size = 684772, upload-time = "2026-04-08T21:01:41.797Z" } wheels = [ @@ -5654,6 +5762,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e5/30/8519fdde58a7bdf155b714359791ad1dc018b47d60269d5d160d311fdc36/sqlalchemy-2.0.49-py3-none-any.whl", hash = "sha256:ec44cfa7ef1a728e88ad41674de50f6db8cfdb3e2af84af86e0041aaf02d43d0", size = 1942158, upload-time = "2026-04-03T16:53:44.135Z" }, ] +[[package]] +name = "sqlglot" +version = "30.4.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f8/9b/af8fcaca1b0821349ef4f88e2775e059bcd7640900bca6533832f1fb845d/sqlglot-30.4.3.tar.gz", hash = "sha256:3a4e9a1e1dd47f8e536ba822d77cb784681704da5e4a3e1a07d2ef86b6067826", size = 5827662, upload-time = "2026-04-13T17:05:15.456Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ee/19/7df8b292accba3bc0de92c611c1e89423b25c08c82c18b14ca1fdbcf6e44/sqlglot-30.4.3-py3-none-any.whl", hash = "sha256:58ea8e723444569da5cec91e4c8f16e385bce3f0ce0374b8c722c3088e1c1c7a", size = 670965, upload-time = "2026-04-13T17:05:13.128Z" }, +] + [[package]] name = "sse-starlette" version = "3.3.4" @@ -5816,31 +5933,46 @@ name = "tiktoken" version = "0.12.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "regex", marker = "sys_platform == 'linux'" }, - { name = "requests", marker = "sys_platform == 'linux'" }, + { name = "regex" }, + { name = "requests" }, ] sdist = { url = "https://files.pythonhosted.org/packages/7d/ab/4d017d0f76ec3171d469d80fc03dfbb4e48a4bcaddaa831b31d526f05edc/tiktoken-0.12.0.tar.gz", hash = "sha256:b18ba7ee2b093863978fcb14f74b3707cdc8d4d4d3836853ce7ec60772139931", size = 37806, upload-time = "2025-10-06T20:22:45.419Z" } wheels = [ + { url = "https://files.pythonhosted.org/packages/a4/85/be65d39d6b647c79800fd9d29241d081d4eeb06271f383bb87200d74cf76/tiktoken-0.12.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:b97f74aca0d78a1ff21b8cd9e9925714c15a9236d6ceacf5c7327c117e6e21e8", size = 1050728, upload-time = "2025-10-06T20:21:52.756Z" }, + { url = "https://files.pythonhosted.org/packages/4a/42/6573e9129bc55c9bf7300b3a35bef2c6b9117018acca0dc760ac2d93dffe/tiktoken-0.12.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:2b90f5ad190a4bb7c3eb30c5fa32e1e182ca1ca79f05e49b448438c3e225a49b", size = 994049, upload-time = "2025-10-06T20:21:53.782Z" }, { url = "https://files.pythonhosted.org/packages/66/c5/ed88504d2f4a5fd6856990b230b56d85a777feab84e6129af0822f5d0f70/tiktoken-0.12.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:65b26c7a780e2139e73acc193e5c63ac754021f160df919add909c1492c0fb37", size = 1129008, upload-time = "2025-10-06T20:21:54.832Z" }, { url = "https://files.pythonhosted.org/packages/f4/90/3dae6cc5436137ebd38944d396b5849e167896fc2073da643a49f372dc4f/tiktoken-0.12.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:edde1ec917dfd21c1f2f8046b86348b0f54a2c0547f68149d8600859598769ad", size = 1152665, upload-time = "2025-10-06T20:21:56.129Z" }, { url = "https://files.pythonhosted.org/packages/a3/fe/26df24ce53ffde419a42f5f53d755b995c9318908288c17ec3f3448313a3/tiktoken-0.12.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:35a2f8ddd3824608b3d650a000c1ef71f730d0c56486845705a8248da00f9fe5", size = 1194230, upload-time = "2025-10-06T20:21:57.546Z" }, { url = "https://files.pythonhosted.org/packages/20/cc/b064cae1a0e9fac84b0d2c46b89f4e57051a5f41324e385d10225a984c24/tiktoken-0.12.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:83d16643edb7fa2c99eff2ab7733508aae1eebb03d5dfc46f5565862810f24e3", size = 1254688, upload-time = "2025-10-06T20:21:58.619Z" }, + { url = "https://files.pythonhosted.org/packages/81/10/b8523105c590c5b8349f2587e2fdfe51a69544bd5a76295fc20f2374f470/tiktoken-0.12.0-cp312-cp312-win_amd64.whl", hash = "sha256:ffc5288f34a8bc02e1ea7047b8d041104791d2ddbf42d1e5fa07822cbffe16bd", size = 878694, upload-time = "2025-10-06T20:21:59.876Z" }, + { url = "https://files.pythonhosted.org/packages/00/61/441588ee21e6b5cdf59d6870f86beb9789e532ee9718c251b391b70c68d6/tiktoken-0.12.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:775c2c55de2310cc1bc9a3ad8826761cbdc87770e586fd7b6da7d4589e13dab3", size = 1050802, upload-time = "2025-10-06T20:22:00.96Z" }, + { url = "https://files.pythonhosted.org/packages/1f/05/dcf94486d5c5c8d34496abe271ac76c5b785507c8eae71b3708f1ad9b45a/tiktoken-0.12.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:a01b12f69052fbe4b080a2cfb867c4de12c704b56178edf1d1d7b273561db160", size = 993995, upload-time = "2025-10-06T20:22:02.788Z" }, { url = "https://files.pythonhosted.org/packages/a0/70/5163fe5359b943f8db9946b62f19be2305de8c3d78a16f629d4165e2f40e/tiktoken-0.12.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:01d99484dc93b129cd0964f9d34eee953f2737301f18b3c7257bf368d7615baa", size = 1128948, upload-time = "2025-10-06T20:22:03.814Z" }, { url = "https://files.pythonhosted.org/packages/0c/da/c028aa0babf77315e1cef357d4d768800c5f8a6de04d0eac0f377cb619fa/tiktoken-0.12.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:4a1a4fcd021f022bfc81904a911d3df0f6543b9e7627b51411da75ff2fe7a1be", size = 1151986, upload-time = "2025-10-06T20:22:05.173Z" }, { url = "https://files.pythonhosted.org/packages/a0/5a/886b108b766aa53e295f7216b509be95eb7d60b166049ce2c58416b25f2a/tiktoken-0.12.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:981a81e39812d57031efdc9ec59fa32b2a5a5524d20d4776574c4b4bd2e9014a", size = 1194222, upload-time = "2025-10-06T20:22:06.265Z" }, { url = "https://files.pythonhosted.org/packages/f4/f8/4db272048397636ac7a078d22773dd2795b1becee7bc4922fe6207288d57/tiktoken-0.12.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:9baf52f84a3f42eef3ff4e754a0db79a13a27921b457ca9832cf944c6be4f8f3", size = 1255097, upload-time = "2025-10-06T20:22:07.403Z" }, + { url = "https://files.pythonhosted.org/packages/8e/32/45d02e2e0ea2be3a9ed22afc47d93741247e75018aac967b713b2941f8ea/tiktoken-0.12.0-cp313-cp313-win_amd64.whl", hash = "sha256:b8a0cd0c789a61f31bf44851defbd609e8dd1e2c8589c614cc1060940ef1f697", size = 879117, upload-time = "2025-10-06T20:22:08.418Z" }, + { url = "https://files.pythonhosted.org/packages/ce/76/994fc868f88e016e6d05b0da5ac24582a14c47893f4474c3e9744283f1d5/tiktoken-0.12.0-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:d5f89ea5680066b68bcb797ae85219c72916c922ef0fcdd3480c7d2315ffff16", size = 1050309, upload-time = "2025-10-06T20:22:10.939Z" }, + { url = "https://files.pythonhosted.org/packages/f6/b8/57ef1456504c43a849821920d582a738a461b76a047f352f18c0b26c6516/tiktoken-0.12.0-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:b4e7ed1c6a7a8a60a3230965bdedba8cc58f68926b835e519341413370e0399a", size = 993712, upload-time = "2025-10-06T20:22:12.115Z" }, { url = "https://files.pythonhosted.org/packages/72/90/13da56f664286ffbae9dbcfadcc625439142675845baa62715e49b87b68b/tiktoken-0.12.0-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:fc530a28591a2d74bce821d10b418b26a094bf33839e69042a6e86ddb7a7fb27", size = 1128725, upload-time = "2025-10-06T20:22:13.541Z" }, { url = "https://files.pythonhosted.org/packages/05/df/4f80030d44682235bdaecd7346c90f67ae87ec8f3df4a3442cb53834f7e4/tiktoken-0.12.0-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:06a9f4f49884139013b138920a4c393aa6556b2f8f536345f11819389c703ebb", size = 1151875, upload-time = "2025-10-06T20:22:14.559Z" }, { url = "https://files.pythonhosted.org/packages/22/1f/ae535223a8c4ef4c0c1192e3f9b82da660be9eb66b9279e95c99288e9dab/tiktoken-0.12.0-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:04f0e6a985d95913cabc96a741c5ffec525a2c72e9df086ff17ebe35985c800e", size = 1194451, upload-time = "2025-10-06T20:22:15.545Z" }, { url = "https://files.pythonhosted.org/packages/78/a7/f8ead382fce0243cb625c4f266e66c27f65ae65ee9e77f59ea1653b6d730/tiktoken-0.12.0-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:0ee8f9ae00c41770b5f9b0bb1235474768884ae157de3beb5439ca0fd70f3e25", size = 1253794, upload-time = "2025-10-06T20:22:16.624Z" }, + { url = "https://files.pythonhosted.org/packages/93/e0/6cc82a562bc6365785a3ff0af27a2a092d57c47d7a81d9e2295d8c36f011/tiktoken-0.12.0-cp313-cp313t-win_amd64.whl", hash = "sha256:dc2dd125a62cb2b3d858484d6c614d136b5b848976794edfb63688d539b8b93f", size = 878777, upload-time = "2025-10-06T20:22:18.036Z" }, + { url = "https://files.pythonhosted.org/packages/72/05/3abc1db5d2c9aadc4d2c76fa5640134e475e58d9fbb82b5c535dc0de9b01/tiktoken-0.12.0-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:a90388128df3b3abeb2bfd1895b0681412a8d7dc644142519e6f0a97c2111646", size = 1050188, upload-time = "2025-10-06T20:22:19.563Z" }, + { url = "https://files.pythonhosted.org/packages/e3/7b/50c2f060412202d6c95f32b20755c7a6273543b125c0985d6fa9465105af/tiktoken-0.12.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:da900aa0ad52247d8794e307d6446bd3cdea8e192769b56276695d34d2c9aa88", size = 993978, upload-time = "2025-10-06T20:22:20.702Z" }, { url = "https://files.pythonhosted.org/packages/14/27/bf795595a2b897e271771cd31cb847d479073497344c637966bdf2853da1/tiktoken-0.12.0-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:285ba9d73ea0d6171e7f9407039a290ca77efcdb026be7769dccc01d2c8d7fff", size = 1129271, upload-time = "2025-10-06T20:22:22.06Z" }, { url = "https://files.pythonhosted.org/packages/f5/de/9341a6d7a8f1b448573bbf3425fa57669ac58258a667eb48a25dfe916d70/tiktoken-0.12.0-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:d186a5c60c6a0213f04a7a802264083dea1bbde92a2d4c7069e1a56630aef830", size = 1151216, upload-time = "2025-10-06T20:22:23.085Z" }, { url = "https://files.pythonhosted.org/packages/75/0d/881866647b8d1be4d67cb24e50d0c26f9f807f994aa1510cb9ba2fe5f612/tiktoken-0.12.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:604831189bd05480f2b885ecd2d1986dc7686f609de48208ebbbddeea071fc0b", size = 1194860, upload-time = "2025-10-06T20:22:24.602Z" }, { url = "https://files.pythonhosted.org/packages/b3/1e/b651ec3059474dab649b8d5b69f5c65cd8fcd8918568c1935bd4136c9392/tiktoken-0.12.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:8f317e8530bb3a222547b85a58583238c8f74fd7a7408305f9f63246d1a0958b", size = 1254567, upload-time = "2025-10-06T20:22:25.671Z" }, + { url = "https://files.pythonhosted.org/packages/80/57/ce64fd16ac390fafde001268c364d559447ba09b509181b2808622420eec/tiktoken-0.12.0-cp314-cp314-win_amd64.whl", hash = "sha256:399c3dd672a6406719d84442299a490420b458c44d3ae65516302a99675888f3", size = 921067, upload-time = "2025-10-06T20:22:26.753Z" }, + { url = "https://files.pythonhosted.org/packages/ac/a4/72eed53e8976a099539cdd5eb36f241987212c29629d0a52c305173e0a68/tiktoken-0.12.0-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:c2c714c72bc00a38ca969dae79e8266ddec999c7ceccd603cc4f0d04ccd76365", size = 1050473, upload-time = "2025-10-06T20:22:27.775Z" }, + { url = "https://files.pythonhosted.org/packages/e6/d7/0110b8f54c008466b19672c615f2168896b83706a6611ba6e47313dbc6e9/tiktoken-0.12.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:cbb9a3ba275165a2cb0f9a83f5d7025afe6b9d0ab01a22b50f0e74fee2ad253e", size = 993855, upload-time = "2025-10-06T20:22:28.799Z" }, { url = "https://files.pythonhosted.org/packages/5f/77/4f268c41a3957c418b084dd576ea2fad2e95da0d8e1ab705372892c2ca22/tiktoken-0.12.0-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:dfdfaa5ffff8993a3af94d1125870b1d27aed7cb97aa7eb8c1cefdbc87dbee63", size = 1129022, upload-time = "2025-10-06T20:22:29.981Z" }, { url = "https://files.pythonhosted.org/packages/4e/2b/fc46c90fe5028bd094cd6ee25a7db321cb91d45dc87531e2bdbb26b4867a/tiktoken-0.12.0-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:584c3ad3d0c74f5269906eb8a659c8bfc6144a52895d9261cdaf90a0ae5f4de0", size = 1150736, upload-time = "2025-10-06T20:22:30.996Z" }, { url = "https://files.pythonhosted.org/packages/28/c0/3c7a39ff68022ddfd7d93f3337ad90389a342f761c4d71de99a3ccc57857/tiktoken-0.12.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:54c891b416a0e36b8e2045b12b33dd66fb34a4fe7965565f1b482da50da3e86a", size = 1194908, upload-time = "2025-10-06T20:22:32.073Z" }, { url = "https://files.pythonhosted.org/packages/ab/0d/c1ad6f4016a3968c048545f5d9b8ffebf577774b2ede3e2e352553b685fe/tiktoken-0.12.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:5edb8743b88d5be814b1a8a8854494719080c28faaa1ccbef02e87354fe71ef0", size = 1253706, upload-time = "2025-10-06T20:22:33.385Z" }, + { url = "https://files.pythonhosted.org/packages/af/df/c7891ef9d2712ad774777271d39fdef63941ffba0a9d59b7ad1fd2765e57/tiktoken-0.12.0-cp314-cp314t-win_amd64.whl", hash = "sha256:f61c0aea5565ac82e2ec50a05e02a6c44734e91b51c10510b084ea1b8e633a71", size = 920667, upload-time = "2025-10-06T20:22:34.444Z" }, ] [[package]] From c9b31d9f4ae51cf4f64ed6f5cb674004c440a954 Mon Sep 17 00:00:00 2001 From: Jeremy Dyer Date: Thu, 16 Apr 2026 19:09:57 -0400 Subject: [PATCH 2/7] First pass at async ray operators --- .../nemo_retriever/graph/abstract_operator.py | 22 +++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/nemo_retriever/src/nemo_retriever/graph/abstract_operator.py b/nemo_retriever/src/nemo_retriever/graph/abstract_operator.py index 6a97a3d6e..89a8b5e2f 100644 --- a/nemo_retriever/src/nemo_retriever/graph/abstract_operator.py +++ b/nemo_retriever/src/nemo_retriever/graph/abstract_operator.py @@ -13,10 +13,32 @@ from nemo_retriever.graph.pipeline_graph import Graph, Node +_policy_checked = False + + +def _ensure_event_loop_policy() -> None: + """Reset to the default asyncio policy if the current one prevents loop creation. + + uvloop >= 0.22's ``get_event_loop()`` raises ``RuntimeError`` when no loop + is *running* (not merely *set*), which breaks Ray Data's async-actor + initialisation in freshly spawned worker processes. Falling back to the + default policy lets ``get_event_loop()`` create a loop implicitly. + """ + global _policy_checked + if _policy_checked: + return + _policy_checked = True + try: + asyncio.get_event_loop_policy().get_event_loop() + except RuntimeError: + asyncio.set_event_loop_policy(asyncio.DefaultEventLoopPolicy()) + + class AbstractOperator(ABC): """Base class for all pipeline operators.""" def __init__(self, **kwargs: Any) -> None: + _ensure_event_loop_policy() self._graph_init_kwargs = dict(kwargs) for key, value in kwargs.items(): setattr(self, key, value) From d4aff4c4ab69cc0d0ae1ebce3f551679531b80f8 Mon Sep 17 00:00:00 2001 From: Jeremy Dyer Date: Thu, 16 Apr 2026 19:30:30 -0400 Subject: [PATCH 3/7] First pass at async ray operators --- .../nemo_retriever/graph/abstract_operator.py | 43 +++++++++++-------- 1 file changed, 24 insertions(+), 19 deletions(-) diff --git a/nemo_retriever/src/nemo_retriever/graph/abstract_operator.py b/nemo_retriever/src/nemo_retriever/graph/abstract_operator.py index 89a8b5e2f..f1381289a 100644 --- a/nemo_retriever/src/nemo_retriever/graph/abstract_operator.py +++ b/nemo_retriever/src/nemo_retriever/graph/abstract_operator.py @@ -13,32 +13,32 @@ from nemo_retriever.graph.pipeline_graph import Graph, Node -_policy_checked = False - - -def _ensure_event_loop_policy() -> None: - """Reset to the default asyncio policy if the current one prevents loop creation. +def _ensure_event_loop() -> None: + """Guarantee an asyncio event loop exists for the current thread. uvloop >= 0.22's ``get_event_loop()`` raises ``RuntimeError`` when no loop - is *running* (not merely *set*), which breaks Ray Data's async-actor - initialisation in freshly spawned worker processes. Falling back to the - default policy lets ``get_event_loop()`` create a loop implicitly. + is *running*, which breaks Ray Data's async-actor initialisation in freshly + spawned worker processes. We fall back to the default asyncio policy (which + allows implicit loop creation) and explicitly create + set a loop so that + Ray's ``get_or_create_event_loop()`` finds one immediately. """ - global _policy_checked - if _policy_checked: - return - _policy_checked = True try: - asyncio.get_event_loop_policy().get_event_loop() + loop = asyncio.get_event_loop_policy().get_event_loop() + if loop.is_closed(): + raise RuntimeError("closed") except RuntimeError: asyncio.set_event_loop_policy(asyncio.DefaultEventLoopPolicy()) + asyncio.set_event_loop(asyncio.new_event_loop()) + + +# Run at import time so the loop exists before Ray's _init_async() runs. +_ensure_event_loop() class AbstractOperator(ABC): """Base class for all pipeline operators.""" def __init__(self, **kwargs: Any) -> None: - _ensure_event_loop_policy() self._graph_init_kwargs = dict(kwargs) for key, value in kwargs.items(): setattr(self, key, value) @@ -61,12 +61,17 @@ def run(self, data: Any, **kwargs: Any) -> Any: async def aprocess(self, data: Any, **kwargs: Any) -> Any: """Async version of :meth:`process`. - Default wraps the synchronous ``process`` in a thread so - compute-bound subclasses work without modification. I/O-bound - subclasses should override this with a proper ``await``-based - implementation. + The default calls ``process()`` synchronously on the event-loop + thread. This is intentional: most operators use C extensions + (pypdfium2, OpenCV, torch, …) that are **not** thread-safe, so + ``asyncio.to_thread`` would allow Ray Data to run multiple + batches in parallel threads inside the same actor, causing + memory corruption. + + I/O-bound subclasses that call remote endpoints should override + this with a proper ``await``-based implementation. """ - return await asyncio.to_thread(self.process, data, **kwargs) + return self.process(data, **kwargs) async def arun(self, data: Any, **kwargs: Any) -> Any: """Async version of :meth:`run`.""" From ba3bce6233352f47c2ea0f57a7cc2b08b56af051 Mon Sep 17 00:00:00 2001 From: Jeremy Dyer Date: Mon, 20 Apr 2026 18:11:33 -0400 Subject: [PATCH 4/7] Address Greptile issues --- .../src/nemo_retriever/chart/shared.py | 5 ++- .../nemo_retriever/graph/abstract_operator.py | 41 +++++++++++++++---- .../graph/operator_archetype.py | 4 +- .../nemo_retriever/text_embed/operators.py | 5 +-- nemo_retriever/tests/test_asr_actor.py | 2 +- .../tests/test_audio_chunk_actor.py | 2 +- .../tests/test_chart_graphic_elements.py | 2 +- nemo_retriever/tests/test_doc_to_pdf_actor.py | 2 +- nemo_retriever/tests/test_image_load.py | 2 +- .../tests/test_nemotron_rerank_v2.py | 2 +- nemo_retriever/tests/test_table_structure.py | 2 +- 11 files changed, 47 insertions(+), 22 deletions(-) diff --git a/nemo_retriever/src/nemo_retriever/chart/shared.py b/nemo_retriever/src/nemo_retriever/chart/shared.py index 0e19b761b..d9af167c0 100644 --- a/nemo_retriever/src/nemo_retriever/chart/shared.py +++ b/nemo_retriever/src/nemo_retriever/chart/shared.py @@ -8,9 +8,12 @@ import base64 import io +import logging import time import traceback +logger = logging.getLogger(__name__) + import pandas as pd from nemo_retriever.nim.nim import NIMClient, invoke_image_inference_batches from nemo_retriever.params import RemoteRetryParams @@ -783,7 +786,7 @@ def _run_local_ocr(): chart_items.append({"bbox_xyxy_norm": bbox, "text": text}) except BaseException as e: - print(f"Warning: graphic-elements+OCR failed: {type(e).__name__}: {e}") + logger.warning("graphic-elements+OCR failed: %s: %s", type(e).__name__, e, exc_info=True) row_error = { "stage": "graphic_elements_ocr_page_elements", "type": e.__class__.__name__, diff --git a/nemo_retriever/src/nemo_retriever/graph/abstract_operator.py b/nemo_retriever/src/nemo_retriever/graph/abstract_operator.py index f1381289a..53780ca6f 100644 --- a/nemo_retriever/src/nemo_retriever/graph/abstract_operator.py +++ b/nemo_retriever/src/nemo_retriever/graph/abstract_operator.py @@ -18,17 +18,31 @@ def _ensure_event_loop() -> None: uvloop >= 0.22's ``get_event_loop()`` raises ``RuntimeError`` when no loop is *running*, which breaks Ray Data's async-actor initialisation in freshly - spawned worker processes. We fall back to the default asyncio policy (which - allows implicit loop creation) and explicitly create + set a loop so that - Ray's ``get_or_create_event_loop()`` finds one immediately. + spawned worker processes. We try to create a loop under the *current* + policy first so that uvloop (or any other custom policy) is preserved. + The default-policy fallback only triggers if the installed policy itself + cannot create a usable loop. """ try: loop = asyncio.get_event_loop_policy().get_event_loop() if loop.is_closed(): raise RuntimeError("closed") + return except RuntimeError: - asyncio.set_event_loop_policy(asyncio.DefaultEventLoopPolicy()) - asyncio.set_event_loop(asyncio.new_event_loop()) + pass + + # Try to create a new loop under the current policy before falling back. + policy = asyncio.get_event_loop_policy() + try: + new_loop = policy.new_event_loop() + asyncio.set_event_loop(new_loop) + return + except Exception: + pass + + # Current policy cannot create a loop — fall back to the stdlib default. + asyncio.set_event_loop_policy(asyncio.DefaultEventLoopPolicy()) + asyncio.set_event_loop(asyncio.new_event_loop()) # Run at import time so the loop exists before Ray's _init_async() runs. @@ -80,9 +94,20 @@ async def arun(self, data: Any, **kwargs: Any) -> Any: data = self.postprocess(data, **kwargs) return data - async def __call__(self, data: Any, **kwargs: Any) -> Any: - """Make operators directly usable as Ray ``map_batches`` async callables.""" - return await self.arun(data, **kwargs) + def __call__(self, data: Any, **kwargs: Any) -> Any: + """Make operators directly usable as Ray ``map_batches`` async callables. + + When called from a synchronous context (no running event loop) the + operator executes synchronously via :meth:`run` so that existing + ``result = op(data)`` call-sites keep working. Inside an async + context the method returns the :meth:`arun` coroutine which the + caller (e.g. Ray Data) must ``await``. + """ + try: + asyncio.get_running_loop() + except RuntimeError: + return self.run(data, **kwargs) + return self.arun(data, **kwargs) def get_constructor_kwargs(self) -> dict[str, Any]: """Best-effort constructor kwargs for executor-side reconstruction.""" diff --git a/nemo_retriever/src/nemo_retriever/graph/operator_archetype.py b/nemo_retriever/src/nemo_retriever/graph/operator_archetype.py index c28e6e885..e9a3f43c8 100644 --- a/nemo_retriever/src/nemo_retriever/graph/operator_archetype.py +++ b/nemo_retriever/src/nemo_retriever/graph/operator_archetype.py @@ -77,9 +77,9 @@ async def aprocess(self, data: Any, **kwargs: Any) -> Any: async def arun(self, data: Any, **kwargs: Any) -> Any: return await self._resolve_delegate().arun(data, **kwargs) - async def __call__(self, data: Any, **kwargs: Any) -> Any: + def __call__(self, data: Any, **kwargs: Any) -> Any: delegate = self._resolve_delegate() - return await delegate(data, **kwargs) + return delegate(data, **kwargs) def _resolve_delegate(self, resources: ClusterResources | Resources | None = None) -> AbstractOperator: if not hasattr(self, "_resolved_delegate"): diff --git a/nemo_retriever/src/nemo_retriever/text_embed/operators.py b/nemo_retriever/src/nemo_retriever/text_embed/operators.py index 198ae57b2..79b88f9d1 100644 --- a/nemo_retriever/src/nemo_retriever/text_embed/operators.py +++ b/nemo_retriever/src/nemo_retriever/text_embed/operators.py @@ -22,7 +22,7 @@ description="Generates embeddings in batches using configurable embedding parameters", category_color="#e06cff", ) -class _BatchEmbedActor(ArchetypeOperator): +class BatchEmbedActor(ArchetypeOperator): """Graph-facing batch embedding archetype.""" @classmethod @@ -47,9 +47,6 @@ def __init__(self, params: Any) -> None: super().__init__(params=params) -BatchEmbedActor = _BatchEmbedActor - - def __getattr__(name: str): if name == "BatchEmbedCPUActor": from nemo_retriever.text_embed.cpu_operator import BatchEmbedCPUActor diff --git a/nemo_retriever/tests/test_asr_actor.py b/nemo_retriever/tests/test_asr_actor.py index 72baa6b75..07372654b 100644 --- a/nemo_retriever/tests/test_asr_actor.py +++ b/nemo_retriever/tests/test_asr_actor.py @@ -25,7 +25,7 @@ def _run(coro): """Run a coroutine synchronously in tests.""" - return asyncio.get_event_loop().run_until_complete(coro) + return asyncio.new_event_loop().run_until_complete(coro) def test_strip_pad_from_transcript(): diff --git a/nemo_retriever/tests/test_audio_chunk_actor.py b/nemo_retriever/tests/test_audio_chunk_actor.py index edd0b5502..ca8d3227d 100644 --- a/nemo_retriever/tests/test_audio_chunk_actor.py +++ b/nemo_retriever/tests/test_audio_chunk_actor.py @@ -22,7 +22,7 @@ def _run(coro): """Run a coroutine synchronously in tests.""" - return asyncio.get_event_loop().run_until_complete(coro) + return asyncio.new_event_loop().run_until_complete(coro) def _make_small_wav(path: Path, duration_sec: float = 0.5, sample_rate: int = 8000) -> None: diff --git a/nemo_retriever/tests/test_chart_graphic_elements.py b/nemo_retriever/tests/test_chart_graphic_elements.py index 2d73eb255..b703bc7ca 100644 --- a/nemo_retriever/tests/test_chart_graphic_elements.py +++ b/nemo_retriever/tests/test_chart_graphic_elements.py @@ -20,7 +20,7 @@ def _run(coro): """Run a coroutine synchronously in tests.""" - return asyncio.get_event_loop().run_until_complete(coro) + return asyncio.new_event_loop().run_until_complete(coro) def _can_import(mod: str) -> bool: diff --git a/nemo_retriever/tests/test_doc_to_pdf_actor.py b/nemo_retriever/tests/test_doc_to_pdf_actor.py index 32c159de7..6318659a2 100644 --- a/nemo_retriever/tests/test_doc_to_pdf_actor.py +++ b/nemo_retriever/tests/test_doc_to_pdf_actor.py @@ -14,7 +14,7 @@ def _run(coro): """Run a coroutine synchronously in tests.""" - return asyncio.get_event_loop().run_until_complete(coro) + return asyncio.new_event_loop().run_until_complete(coro) class TestConvertToPdfBytes: diff --git a/nemo_retriever/tests/test_image_load.py b/nemo_retriever/tests/test_image_load.py index a489c22be..c82718d82 100644 --- a/nemo_retriever/tests/test_image_load.py +++ b/nemo_retriever/tests/test_image_load.py @@ -27,7 +27,7 @@ def _run(coro): """Run a coroutine synchronously in tests.""" - return asyncio.get_event_loop().run_until_complete(coro) + return asyncio.new_event_loop().run_until_complete(coro) # -- Helpers ------------------------------------------------------------------ diff --git a/nemo_retriever/tests/test_nemotron_rerank_v2.py b/nemo_retriever/tests/test_nemotron_rerank_v2.py index 5ac665e1c..94c4b5280 100644 --- a/nemo_retriever/tests/test_nemotron_rerank_v2.py +++ b/nemo_retriever/tests/test_nemotron_rerank_v2.py @@ -21,7 +21,7 @@ def _run(coro): """Run a coroutine synchronously in tests.""" - return asyncio.get_event_loop().run_until_complete(coro) + return asyncio.new_event_loop().run_until_complete(coro) # --------------------------------------------------------------------------- diff --git a/nemo_retriever/tests/test_table_structure.py b/nemo_retriever/tests/test_table_structure.py index 58491da94..72f81cd90 100644 --- a/nemo_retriever/tests/test_table_structure.py +++ b/nemo_retriever/tests/test_table_structure.py @@ -20,7 +20,7 @@ def _run(coro): """Run a coroutine synchronously in tests.""" - return asyncio.get_event_loop().run_until_complete(coro) + return asyncio.new_event_loop().run_until_complete(coro) def _can_import(mod: str) -> bool: From d3d753af276464f49a360527b1a2740e7f09baae Mon Sep 17 00:00:00 2001 From: Jeremy Dyer Date: Mon, 20 Apr 2026 18:18:16 -0400 Subject: [PATCH 5/7] Unit test fixes --- nemo_retriever/tests/test_actor_operators.py | 8 +++++--- nemo_retriever/tests/test_asr_actor.py | 8 +++++--- nemo_retriever/tests/test_audio_chunk_actor.py | 8 +++++--- nemo_retriever/tests/test_chart_graphic_elements.py | 8 +++++--- nemo_retriever/tests/test_doc_to_pdf_actor.py | 8 +++++--- nemo_retriever/tests/test_image_load.py | 8 +++++--- nemo_retriever/tests/test_nemotron_rerank_v2.py | 8 +++++--- nemo_retriever/tests/test_table_structure.py | 8 +++++--- 8 files changed, 40 insertions(+), 24 deletions(-) diff --git a/nemo_retriever/tests/test_actor_operators.py b/nemo_retriever/tests/test_actor_operators.py index 26d907249..824efba17 100644 --- a/nemo_retriever/tests/test_actor_operators.py +++ b/nemo_retriever/tests/test_actor_operators.py @@ -14,11 +14,13 @@ from nemo_retriever.graph.abstract_operator import AbstractOperator -def _run(coro): - """Run a coroutine synchronously in tests.""" +def _run(coro_or_result): + """Run a coroutine synchronously in tests; pass through plain values.""" + if not asyncio.iscoroutine(coro_or_result): + return coro_or_result loop = asyncio.new_event_loop() try: - return loop.run_until_complete(coro) + return loop.run_until_complete(coro_or_result) finally: loop.close() diff --git a/nemo_retriever/tests/test_asr_actor.py b/nemo_retriever/tests/test_asr_actor.py index 07372654b..819efe1d4 100644 --- a/nemo_retriever/tests/test_asr_actor.py +++ b/nemo_retriever/tests/test_asr_actor.py @@ -23,9 +23,11 @@ from nemo_retriever.params import ASRParams -def _run(coro): - """Run a coroutine synchronously in tests.""" - return asyncio.new_event_loop().run_until_complete(coro) +def _run(coro_or_result): + """Run a coroutine synchronously in tests; pass through plain values.""" + if not asyncio.iscoroutine(coro_or_result): + return coro_or_result + return asyncio.new_event_loop().run_until_complete(coro_or_result) def test_strip_pad_from_transcript(): diff --git a/nemo_retriever/tests/test_audio_chunk_actor.py b/nemo_retriever/tests/test_audio_chunk_actor.py index ca8d3227d..00f5997c9 100644 --- a/nemo_retriever/tests/test_audio_chunk_actor.py +++ b/nemo_retriever/tests/test_audio_chunk_actor.py @@ -20,9 +20,11 @@ from nemo_retriever.params import AudioChunkParams -def _run(coro): - """Run a coroutine synchronously in tests.""" - return asyncio.new_event_loop().run_until_complete(coro) +def _run(coro_or_result): + """Run a coroutine synchronously in tests; pass through plain values.""" + if not asyncio.iscoroutine(coro_or_result): + return coro_or_result + return asyncio.new_event_loop().run_until_complete(coro_or_result) def _make_small_wav(path: Path, duration_sec: float = 0.5, sample_rate: int = 8000) -> None: diff --git a/nemo_retriever/tests/test_chart_graphic_elements.py b/nemo_retriever/tests/test_chart_graphic_elements.py index b703bc7ca..095f899f5 100644 --- a/nemo_retriever/tests/test_chart_graphic_elements.py +++ b/nemo_retriever/tests/test_chart_graphic_elements.py @@ -18,9 +18,11 @@ from nemo_retriever.utils.table_and_chart import join_graphic_elements_and_ocr_output -def _run(coro): - """Run a coroutine synchronously in tests.""" - return asyncio.new_event_loop().run_until_complete(coro) +def _run(coro_or_result): + """Run a coroutine synchronously in tests; pass through plain values.""" + if not asyncio.iscoroutine(coro_or_result): + return coro_or_result + return asyncio.new_event_loop().run_until_complete(coro_or_result) def _can_import(mod: str) -> bool: diff --git a/nemo_retriever/tests/test_doc_to_pdf_actor.py b/nemo_retriever/tests/test_doc_to_pdf_actor.py index 6318659a2..0db6d04c2 100644 --- a/nemo_retriever/tests/test_doc_to_pdf_actor.py +++ b/nemo_retriever/tests/test_doc_to_pdf_actor.py @@ -12,9 +12,11 @@ from nemo_retriever.utils.convert.to_pdf import DocToPdfConversionActor, convert_to_pdf_bytes, convert_batch_to_pdf -def _run(coro): - """Run a coroutine synchronously in tests.""" - return asyncio.new_event_loop().run_until_complete(coro) +def _run(coro_or_result): + """Run a coroutine synchronously in tests; pass through plain values.""" + if not asyncio.iscoroutine(coro_or_result): + return coro_or_result + return asyncio.new_event_loop().run_until_complete(coro_or_result) class TestConvertToPdfBytes: diff --git a/nemo_retriever/tests/test_image_load.py b/nemo_retriever/tests/test_image_load.py index c82718d82..ae7e92810 100644 --- a/nemo_retriever/tests/test_image_load.py +++ b/nemo_retriever/tests/test_image_load.py @@ -25,9 +25,11 @@ from nemo_retriever.image.ray_data import ImageLoadActor # noqa: E402 -def _run(coro): - """Run a coroutine synchronously in tests.""" - return asyncio.new_event_loop().run_until_complete(coro) +def _run(coro_or_result): + """Run a coroutine synchronously in tests; pass through plain values.""" + if not asyncio.iscoroutine(coro_or_result): + return coro_or_result + return asyncio.new_event_loop().run_until_complete(coro_or_result) # -- Helpers ------------------------------------------------------------------ diff --git a/nemo_retriever/tests/test_nemotron_rerank_v2.py b/nemo_retriever/tests/test_nemotron_rerank_v2.py index 94c4b5280..e6b0cee5c 100644 --- a/nemo_retriever/tests/test_nemotron_rerank_v2.py +++ b/nemo_retriever/tests/test_nemotron_rerank_v2.py @@ -19,9 +19,11 @@ import pytest -def _run(coro): - """Run a coroutine synchronously in tests.""" - return asyncio.new_event_loop().run_until_complete(coro) +def _run(coro_or_result): + """Run a coroutine synchronously in tests; pass through plain values.""" + if not asyncio.iscoroutine(coro_or_result): + return coro_or_result + return asyncio.new_event_loop().run_until_complete(coro_or_result) # --------------------------------------------------------------------------- diff --git a/nemo_retriever/tests/test_table_structure.py b/nemo_retriever/tests/test_table_structure.py index 72f81cd90..1f8246fe6 100644 --- a/nemo_retriever/tests/test_table_structure.py +++ b/nemo_retriever/tests/test_table_structure.py @@ -18,9 +18,11 @@ from nemo_retriever.utils.table_and_chart import join_table_structure_and_ocr_output -def _run(coro): - """Run a coroutine synchronously in tests.""" - return asyncio.new_event_loop().run_until_complete(coro) +def _run(coro_or_result): + """Run a coroutine synchronously in tests; pass through plain values.""" + if not asyncio.iscoroutine(coro_or_result): + return coro_or_result + return asyncio.new_event_loop().run_until_complete(coro_or_result) def _can_import(mod: str) -> bool: From f3885c220e0c20f651a43fbefd4cb61e34e0acf0 Mon Sep 17 00:00:00 2001 From: Jeremy Dyer Date: Tue, 21 Apr 2026 21:06:31 -0400 Subject: [PATCH 6/7] Address review comments --- .../nemo_retriever/graph/abstract_operator.py | 25 +----- .../src/nemo_retriever/recall/core.py | 76 +------------------ .../nemo_retriever/text_embed/operators.py | 2 + nemo_retriever/tests/test_actor_operators.py | 13 +--- nemo_retriever/tests/test_asr_actor.py | 9 +-- .../tests/test_audio_chunk_actor.py | 9 +-- .../tests/test_chart_graphic_elements.py | 9 +-- nemo_retriever/tests/test_doc_to_pdf_actor.py | 9 +-- nemo_retriever/tests/test_image_load.py | 9 +-- .../tests/test_nemotron_rerank_v2.py | 8 +- nemo_retriever/tests/test_table_structure.py | 9 +-- nemo_retriever/tests/testing_utils.py | 20 +++++ 12 files changed, 36 insertions(+), 162 deletions(-) create mode 100644 nemo_retriever/tests/testing_utils.py diff --git a/nemo_retriever/src/nemo_retriever/graph/abstract_operator.py b/nemo_retriever/src/nemo_retriever/graph/abstract_operator.py index 53780ca6f..3ffc95556 100644 --- a/nemo_retriever/src/nemo_retriever/graph/abstract_operator.py +++ b/nemo_retriever/src/nemo_retriever/graph/abstract_operator.py @@ -16,32 +16,15 @@ def _ensure_event_loop() -> None: """Guarantee an asyncio event loop exists for the current thread. - uvloop >= 0.22's ``get_event_loop()`` raises ``RuntimeError`` when no loop - is *running*, which breaks Ray Data's async-actor initialisation in freshly - spawned worker processes. We try to create a loop under the *current* - policy first so that uvloop (or any other custom policy) is preserved. - The default-policy fallback only triggers if the installed policy itself - cannot create a usable loop. + Ray Data spawns fresh worker processes that may not have an event loop. + ``asyncio.new_event_loop()`` delegates to the installed policy, so uvloop + (or any other custom policy) is automatically preserved. """ try: - loop = asyncio.get_event_loop_policy().get_event_loop() - if loop.is_closed(): - raise RuntimeError("closed") + asyncio.get_running_loop() return except RuntimeError: pass - - # Try to create a new loop under the current policy before falling back. - policy = asyncio.get_event_loop_policy() - try: - new_loop = policy.new_event_loop() - asyncio.set_event_loop(new_loop) - return - except Exception: - pass - - # Current policy cannot create a loop — fall back to the stdlib default. - asyncio.set_event_loop_policy(asyncio.DefaultEventLoopPolicy()) asyncio.set_event_loop(asyncio.new_event_loop()) diff --git a/nemo_retriever/src/nemo_retriever/recall/core.py b/nemo_retriever/src/nemo_retriever/recall/core.py index f70255024..58ced38e3 100644 --- a/nemo_retriever/src/nemo_retriever/recall/core.py +++ b/nemo_retriever/src/nemo_retriever/recall/core.py @@ -17,7 +17,6 @@ logger = logging.getLogger(__name__) AUDIO_MATCH_TOLERANCE_SECS = 2.0 -import numpy as np import pandas as pd @@ -230,77 +229,6 @@ def _normalize_query_df(df: pd.DataFrame, *, match_mode: str) -> pd.DataFrame: return df -def _resolve_embedding_endpoint(cfg: RecallConfig) -> Tuple[Optional[str], Optional[bool]]: - """ - Resolve which embedding endpoint to use. - - Returns (endpoint, use_grpc) where: - - endpoint is either an http(s) URL or a host:port string for gRPC - - use_grpc is True for gRPC, False for HTTP, None when no endpoint is configured - """ - http_ep = (cfg.embedding_http_endpoint or "").strip() if isinstance(cfg.embedding_http_endpoint, str) else None - grpc_ep = (cfg.embedding_grpc_endpoint or "").strip() if isinstance(cfg.embedding_grpc_endpoint, str) else None - single = (cfg.embedding_endpoint or "").strip() if isinstance(cfg.embedding_endpoint, str) else None - - if http_ep: - return http_ep, False - if grpc_ep: - return grpc_ep, True - if single: - # Infer protocol: if a URL scheme is present, treat as HTTP; otherwise gRPC. - return single, (not single.lower().startswith("http")) - - return None, None - - -def _embed_queries_nim( - queries: List[str], - *, - endpoint: str, - model: str, - api_key: str, - grpc: bool, -) -> List[List[float]]: - from nv_ingest_api.util.nim import infer_microservice - - # `infer_microservice` returns a list of embeddings. - embeddings = infer_microservice( - queries, - model_name=model, - embedding_endpoint=endpoint, - nvidia_api_key=(api_key or "").strip(), - grpc=bool(grpc), - input_type="query", - ) - # Some backends return numpy arrays; normalize to list-of-list floats. - out: List[List[float]] = [] - for e in embeddings: - if isinstance(e, np.ndarray): - out.append(e.astype("float32").tolist()) - else: - out.append(list(e)) - return out - - -def _embed_queries_local_hf( - queries: List[str], - *, - device: Optional[str], - cache_dir: Optional[str], - batch_size: int, - model_name: Optional[str] = None, -) -> List[List[float]]: - from nemo_retriever.model import create_local_embedder, is_vl_embed_model - - embedder = create_local_embedder(model_name, device=device, hf_cache_dir=cache_dir) - - if is_vl_embed_model(model_name): - vecs = embedder.embed_queries(queries, batch_size=int(batch_size)) - else: - vecs = embedder.embed(["query: " + q for q in queries], batch_size=int(batch_size)) - return vecs.detach().to("cpu").tolist() - - def _hits_to_keys(raw_hits: List[List[Dict[str, Any]]]) -> List[List[str]]: retrieved_keys: List[List[str]] = [] for hits in raw_hits: @@ -517,12 +445,12 @@ def retrieve_and_score( queries = df_query["query"].astype(str).tolist() gold = df_query["golden_answer"].astype(str).tolist() - endpoint, use_grpc = _resolve_embedding_endpoint(cfg) retriever = Retriever( lancedb_uri=cfg.lancedb_uri, lancedb_table=cfg.lancedb_table, embedder=cfg.embedding_model or VL_EMBED_MODEL, - embedding_http_endpoint=endpoint, + embedding_http_endpoint=cfg.embedding_http_endpoint, + embedding_endpoint=cfg.embedding_endpoint, embedding_api_key=cfg.embedding_api_key, top_k=cfg.top_k, nprobes=cfg.nprobes, diff --git a/nemo_retriever/src/nemo_retriever/text_embed/operators.py b/nemo_retriever/src/nemo_retriever/text_embed/operators.py index 79b88f9d1..a845b7fe3 100644 --- a/nemo_retriever/src/nemo_retriever/text_embed/operators.py +++ b/nemo_retriever/src/nemo_retriever/text_embed/operators.py @@ -56,6 +56,8 @@ def __getattr__(name: str): from nemo_retriever.text_embed.gpu_operator import BatchEmbedGPUActor return BatchEmbedGPUActor + if name == "_BatchEmbedActor": + return BatchEmbedActor if name == "_BatchEmbedCPUActor": from nemo_retriever.text_embed.cpu_operator import BatchEmbedCPUActor diff --git a/nemo_retriever/tests/test_actor_operators.py b/nemo_retriever/tests/test_actor_operators.py index 824efba17..0c8976c77 100644 --- a/nemo_retriever/tests/test_actor_operators.py +++ b/nemo_retriever/tests/test_actor_operators.py @@ -4,7 +4,6 @@ """Unit tests verifying all pipeline actors inherit from AbstractOperator.""" -import asyncio from pathlib import Path from unittest.mock import AsyncMock, patch @@ -12,17 +11,7 @@ import pytest from nemo_retriever.graph.abstract_operator import AbstractOperator - - -def _run(coro_or_result): - """Run a coroutine synchronously in tests; pass through plain values.""" - if not asyncio.iscoroutine(coro_or_result): - return coro_or_result - loop = asyncio.new_event_loop() - try: - return loop.run_until_complete(coro_or_result) - finally: - loop.close() +from nemo_retriever.tests.testing_utils import _run # --------------------------------------------------------------------------- diff --git a/nemo_retriever/tests/test_asr_actor.py b/nemo_retriever/tests/test_asr_actor.py index 819efe1d4..b98494232 100644 --- a/nemo_retriever/tests/test_asr_actor.py +++ b/nemo_retriever/tests/test_asr_actor.py @@ -10,7 +10,6 @@ into sys.modules so the real module is never loaded. """ -import asyncio import base64 import sys from unittest.mock import MagicMock @@ -21,13 +20,7 @@ from nemo_retriever.audio.asr_actor import ASRActor, ASRCPUActor from nemo_retriever.audio.asr_actor import apply_asr_to_df from nemo_retriever.params import ASRParams - - -def _run(coro_or_result): - """Run a coroutine synchronously in tests; pass through plain values.""" - if not asyncio.iscoroutine(coro_or_result): - return coro_or_result - return asyncio.new_event_loop().run_until_complete(coro_or_result) +from nemo_retriever.tests.testing_utils import _run def test_strip_pad_from_transcript(): diff --git a/nemo_retriever/tests/test_audio_chunk_actor.py b/nemo_retriever/tests/test_audio_chunk_actor.py index 00f5997c9..4874fab4e 100644 --- a/nemo_retriever/tests/test_audio_chunk_actor.py +++ b/nemo_retriever/tests/test_audio_chunk_actor.py @@ -6,7 +6,6 @@ Unit tests for nemo_retriever.audio: MediaChunkActor and audio_path_to_chunks_df. """ -import asyncio import wave from pathlib import Path @@ -18,13 +17,7 @@ from nemo_retriever.audio.chunk_actor import audio_path_to_chunks_df from nemo_retriever.audio.media_interface import is_media_available from nemo_retriever.params import AudioChunkParams - - -def _run(coro_or_result): - """Run a coroutine synchronously in tests; pass through plain values.""" - if not asyncio.iscoroutine(coro_or_result): - return coro_or_result - return asyncio.new_event_loop().run_until_complete(coro_or_result) +from nemo_retriever.tests.testing_utils import _run def _make_small_wav(path: Path, duration_sec: float = 0.5, sample_rate: int = 8000) -> None: diff --git a/nemo_retriever/tests/test_chart_graphic_elements.py b/nemo_retriever/tests/test_chart_graphic_elements.py index 095f899f5..c1dc6e0aa 100644 --- a/nemo_retriever/tests/test_chart_graphic_elements.py +++ b/nemo_retriever/tests/test_chart_graphic_elements.py @@ -6,7 +6,6 @@ from __future__ import annotations -import asyncio import base64 import importlib import io @@ -15,16 +14,10 @@ import pandas as pd import pytest +from nemo_retriever.tests.testing_utils import _run from nemo_retriever.utils.table_and_chart import join_graphic_elements_and_ocr_output -def _run(coro_or_result): - """Run a coroutine synchronously in tests; pass through plain values.""" - if not asyncio.iscoroutine(coro_or_result): - return coro_or_result - return asyncio.new_event_loop().run_until_complete(coro_or_result) - - def _can_import(mod: str) -> bool: return importlib.util.find_spec(mod) is not None diff --git a/nemo_retriever/tests/test_doc_to_pdf_actor.py b/nemo_retriever/tests/test_doc_to_pdf_actor.py index 0db6d04c2..4a41e9259 100644 --- a/nemo_retriever/tests/test_doc_to_pdf_actor.py +++ b/nemo_retriever/tests/test_doc_to_pdf_actor.py @@ -2,23 +2,16 @@ # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -import asyncio from unittest.mock import patch import pandas as pd import pytest from nemo_retriever.graph.abstract_operator import AbstractOperator +from nemo_retriever.tests.testing_utils import _run from nemo_retriever.utils.convert.to_pdf import DocToPdfConversionActor, convert_to_pdf_bytes, convert_batch_to_pdf -def _run(coro_or_result): - """Run a coroutine synchronously in tests; pass through plain values.""" - if not asyncio.iscoroutine(coro_or_result): - return coro_or_result - return asyncio.new_event_loop().run_until_complete(coro_or_result) - - class TestConvertToPdfBytes: @patch("nemo_retriever.utils.convert.to_pdf.shutil.which", return_value=None) def test_raises_when_libreoffice_missing(self, _mock_which): diff --git a/nemo_retriever/tests/test_image_load.py b/nemo_retriever/tests/test_image_load.py index ae7e92810..63805e1c6 100644 --- a/nemo_retriever/tests/test_image_load.py +++ b/nemo_retriever/tests/test_image_load.py @@ -6,7 +6,6 @@ from __future__ import annotations -import asyncio import base64 from io import BytesIO from unittest.mock import patch @@ -23,13 +22,7 @@ image_file_to_pages_df, ) from nemo_retriever.image.ray_data import ImageLoadActor # noqa: E402 - - -def _run(coro_or_result): - """Run a coroutine synchronously in tests; pass through plain values.""" - if not asyncio.iscoroutine(coro_or_result): - return coro_or_result - return asyncio.new_event_loop().run_until_complete(coro_or_result) +from nemo_retriever.tests.testing_utils import _run # noqa: E402 # -- Helpers ------------------------------------------------------------------ diff --git a/nemo_retriever/tests/test_nemotron_rerank_v2.py b/nemo_retriever/tests/test_nemotron_rerank_v2.py index e6b0cee5c..179a01a32 100644 --- a/nemo_retriever/tests/test_nemotron_rerank_v2.py +++ b/nemo_retriever/tests/test_nemotron_rerank_v2.py @@ -11,19 +11,13 @@ from __future__ import annotations -import asyncio import sys from types import ModuleType from unittest.mock import MagicMock, patch import pytest - -def _run(coro_or_result): - """Run a coroutine synchronously in tests; pass through plain values.""" - if not asyncio.iscoroutine(coro_or_result): - return coro_or_result - return asyncio.new_event_loop().run_until_complete(coro_or_result) +from nemo_retriever.tests.testing_utils import _run # --------------------------------------------------------------------------- diff --git a/nemo_retriever/tests/test_table_structure.py b/nemo_retriever/tests/test_table_structure.py index 1f8246fe6..e6c8a59ea 100644 --- a/nemo_retriever/tests/test_table_structure.py +++ b/nemo_retriever/tests/test_table_structure.py @@ -6,7 +6,6 @@ from __future__ import annotations -import asyncio import base64 import importlib import io @@ -15,16 +14,10 @@ import pandas as pd import pytest +from nemo_retriever.tests.testing_utils import _run from nemo_retriever.utils.table_and_chart import join_table_structure_and_ocr_output -def _run(coro_or_result): - """Run a coroutine synchronously in tests; pass through plain values.""" - if not asyncio.iscoroutine(coro_or_result): - return coro_or_result - return asyncio.new_event_loop().run_until_complete(coro_or_result) - - def _can_import(mod: str) -> bool: return importlib.util.find_spec(mod) is not None diff --git a/nemo_retriever/tests/testing_utils.py b/nemo_retriever/tests/testing_utils.py new file mode 100644 index 000000000..9989d245f --- /dev/null +++ b/nemo_retriever/tests/testing_utils.py @@ -0,0 +1,20 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024-25, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared helpers for unit tests.""" + +from __future__ import annotations + +import asyncio + + +def _run(coro_or_result): + """Run a coroutine synchronously in tests; pass through plain values.""" + if not asyncio.iscoroutine(coro_or_result): + return coro_or_result + loop = asyncio.new_event_loop() + try: + return loop.run_until_complete(coro_or_result) + finally: + loop.close() From c61c1f63cfb7efa00dc82798c63bf206bfc772f4 Mon Sep 17 00:00:00 2001 From: Jeremy Dyer Date: Tue, 21 Apr 2026 21:12:22 -0400 Subject: [PATCH 7/7] Address review comments --- .../nemo_retriever/graph/abstract_operator.py | 17 +++-------------- .../nemo_retriever/graph/operator_archetype.py | 4 ++-- 2 files changed, 5 insertions(+), 16 deletions(-) diff --git a/nemo_retriever/src/nemo_retriever/graph/abstract_operator.py b/nemo_retriever/src/nemo_retriever/graph/abstract_operator.py index 3ffc95556..e2fa11fbc 100644 --- a/nemo_retriever/src/nemo_retriever/graph/abstract_operator.py +++ b/nemo_retriever/src/nemo_retriever/graph/abstract_operator.py @@ -77,20 +77,9 @@ async def arun(self, data: Any, **kwargs: Any) -> Any: data = self.postprocess(data, **kwargs) return data - def __call__(self, data: Any, **kwargs: Any) -> Any: - """Make operators directly usable as Ray ``map_batches`` async callables. - - When called from a synchronous context (no running event loop) the - operator executes synchronously via :meth:`run` so that existing - ``result = op(data)`` call-sites keep working. Inside an async - context the method returns the :meth:`arun` coroutine which the - caller (e.g. Ray Data) must ``await``. - """ - try: - asyncio.get_running_loop() - except RuntimeError: - return self.run(data, **kwargs) - return self.arun(data, **kwargs) + async def __call__(self, data: Any, **kwargs: Any) -> Any: + """Make operators directly usable as Ray ``map_batches`` async callables.""" + return await self.arun(data, **kwargs) def get_constructor_kwargs(self) -> dict[str, Any]: """Best-effort constructor kwargs for executor-side reconstruction.""" diff --git a/nemo_retriever/src/nemo_retriever/graph/operator_archetype.py b/nemo_retriever/src/nemo_retriever/graph/operator_archetype.py index e9a3f43c8..c28e6e885 100644 --- a/nemo_retriever/src/nemo_retriever/graph/operator_archetype.py +++ b/nemo_retriever/src/nemo_retriever/graph/operator_archetype.py @@ -77,9 +77,9 @@ async def aprocess(self, data: Any, **kwargs: Any) -> Any: async def arun(self, data: Any, **kwargs: Any) -> Any: return await self._resolve_delegate().arun(data, **kwargs) - def __call__(self, data: Any, **kwargs: Any) -> Any: + async def __call__(self, data: Any, **kwargs: Any) -> Any: delegate = self._resolve_delegate() - return delegate(data, **kwargs) + return await delegate(data, **kwargs) def _resolve_delegate(self, resources: ClusterResources | Resources | None = None) -> AbstractOperator: if not hasattr(self, "_resolved_delegate"):