nemo-retriever 26.3.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- nemo_retriever-26.3.0/PKG-INFO +455 -0
- nemo_retriever-26.3.0/README.md +393 -0
- nemo_retriever-26.3.0/pyproject.toml +114 -0
- nemo_retriever-26.3.0/setup.cfg +4 -0
- nemo_retriever-26.3.0/src/nemo_retriever/__init__.py +46 -0
- nemo_retriever-26.3.0/src/nemo_retriever/__main__.py +9 -0
- nemo_retriever-26.3.0/src/nemo_retriever/_build_info.py +4 -0
- nemo_retriever-26.3.0/src/nemo_retriever/adapters/__init__.py +5 -0
- nemo_retriever-26.3.0/src/nemo_retriever/adapters/cli/__init__.py +7 -0
- nemo_retriever-26.3.0/src/nemo_retriever/adapters/cli/main.py +62 -0
- nemo_retriever-26.3.0/src/nemo_retriever/adapters/service/__init__.py +7 -0
- nemo_retriever-26.3.0/src/nemo_retriever/adapters/service/app.py +101 -0
- nemo_retriever-26.3.0/src/nemo_retriever/adapters/service/cli.py +127 -0
- nemo_retriever-26.3.0/src/nemo_retriever/api/__init__.py +7 -0
- nemo_retriever-26.3.0/src/nemo_retriever/api/ingestor.py +9 -0
- nemo_retriever-26.3.0/src/nemo_retriever/application/__init__.py +7 -0
- nemo_retriever-26.3.0/src/nemo_retriever/application/modes/__init__.py +20 -0
- nemo_retriever-26.3.0/src/nemo_retriever/application/modes/executor.py +21 -0
- nemo_retriever-26.3.0/src/nemo_retriever/application/modes/factory.py +38 -0
- nemo_retriever-26.3.0/src/nemo_retriever/application/modes/run_batch.py +18 -0
- nemo_retriever-26.3.0/src/nemo_retriever/application/modes/run_fused.py +18 -0
- nemo_retriever-26.3.0/src/nemo_retriever/application/modes/run_inprocess.py +18 -0
- nemo_retriever-26.3.0/src/nemo_retriever/application/modes/run_online.py +18 -0
- nemo_retriever-26.3.0/src/nemo_retriever/application/pipeline/__init__.py +17 -0
- nemo_retriever-26.3.0/src/nemo_retriever/application/pipeline/build_plan.py +70 -0
- nemo_retriever-26.3.0/src/nemo_retriever/application/pipeline/contracts.py +31 -0
- nemo_retriever-26.3.0/src/nemo_retriever/application/pipeline/run_plan.py +28 -0
- nemo_retriever-26.3.0/src/nemo_retriever/application/pipeline/stage_registry.py +42 -0
- nemo_retriever-26.3.0/src/nemo_retriever/audio/__init__.py +31 -0
- nemo_retriever-26.3.0/src/nemo_retriever/audio/__main__.py +12 -0
- nemo_retriever-26.3.0/src/nemo_retriever/audio/asr_actor.py +245 -0
- nemo_retriever-26.3.0/src/nemo_retriever/audio/chunk_actor.py +133 -0
- nemo_retriever-26.3.0/src/nemo_retriever/audio/cli.py +26 -0
- nemo_retriever-26.3.0/src/nemo_retriever/audio/media_interface.py +240 -0
- nemo_retriever-26.3.0/src/nemo_retriever/audio/stage.py +297 -0
- nemo_retriever-26.3.0/src/nemo_retriever/chart/__init__.py +21 -0
- nemo_retriever-26.3.0/src/nemo_retriever/chart/__main__.py +16 -0
- nemo_retriever-26.3.0/src/nemo_retriever/chart/chart_detection.py +637 -0
- nemo_retriever-26.3.0/src/nemo_retriever/chart/commands.py +137 -0
- nemo_retriever-26.3.0/src/nemo_retriever/chart/config.py +49 -0
- nemo_retriever-26.3.0/src/nemo_retriever/chart/processor.py +50 -0
- nemo_retriever-26.3.0/src/nemo_retriever/chart/stage.py +19 -0
- nemo_retriever-26.3.0/src/nemo_retriever/config/__init__.py +7 -0
- nemo_retriever-26.3.0/src/nemo_retriever/config/loader.py +37 -0
- nemo_retriever-26.3.0/src/nemo_retriever/config_utils.py +17 -0
- nemo_retriever-26.3.0/src/nemo_retriever/examples/batch_pipeline.py +845 -0
- nemo_retriever-26.3.0/src/nemo_retriever/examples/common.py +51 -0
- nemo_retriever-26.3.0/src/nemo_retriever/examples/fused_pipeline.py +368 -0
- nemo_retriever-26.3.0/src/nemo_retriever/examples/inprocess_pipeline.py +408 -0
- nemo_retriever-26.3.0/src/nemo_retriever/examples/online_pipeline.py +256 -0
- nemo_retriever-26.3.0/src/nemo_retriever/harness/__init__.py +7 -0
- nemo_retriever-26.3.0/src/nemo_retriever/harness/__main__.py +11 -0
- nemo_retriever-26.3.0/src/nemo_retriever/harness/artifacts.py +81 -0
- nemo_retriever-26.3.0/src/nemo_retriever/harness/cli.py +22 -0
- nemo_retriever-26.3.0/src/nemo_retriever/harness/config.py +391 -0
- nemo_retriever-26.3.0/src/nemo_retriever/harness/nightly.py +122 -0
- nemo_retriever-26.3.0/src/nemo_retriever/harness/parsers.py +66 -0
- nemo_retriever-26.3.0/src/nemo_retriever/harness/recall_adapters.py +98 -0
- nemo_retriever-26.3.0/src/nemo_retriever/harness/reporting.py +260 -0
- nemo_retriever-26.3.0/src/nemo_retriever/harness/run.py +581 -0
- nemo_retriever-26.3.0/src/nemo_retriever/harness/slack.py +371 -0
- nemo_retriever-26.3.0/src/nemo_retriever/html/__init__.py +21 -0
- nemo_retriever-26.3.0/src/nemo_retriever/html/__main__.py +107 -0
- nemo_retriever-26.3.0/src/nemo_retriever/html/convert.py +196 -0
- nemo_retriever-26.3.0/src/nemo_retriever/html/ray_data.py +52 -0
- nemo_retriever-26.3.0/src/nemo_retriever/image/__init__.py +21 -0
- nemo_retriever-26.3.0/src/nemo_retriever/image/load.py +160 -0
- nemo_retriever-26.3.0/src/nemo_retriever/image/ray_data.py +59 -0
- nemo_retriever-26.3.0/src/nemo_retriever/infographic/__init__.py +32 -0
- nemo_retriever-26.3.0/src/nemo_retriever/infographic/commands.py +76 -0
- nemo_retriever-26.3.0/src/nemo_retriever/infographic/config.py +36 -0
- nemo_retriever-26.3.0/src/nemo_retriever/infographic/infographic_detection.py +782 -0
- nemo_retriever-26.3.0/src/nemo_retriever/infographic/processor.py +50 -0
- nemo_retriever-26.3.0/src/nemo_retriever/infographic/ray_data.py +36 -0
- nemo_retriever-26.3.0/src/nemo_retriever/infographic/stage.py +10 -0
- nemo_retriever-26.3.0/src/nemo_retriever/ingest_config.py +144 -0
- nemo_retriever-26.3.0/src/nemo_retriever/ingest_modes/__init__.py +24 -0
- nemo_retriever-26.3.0/src/nemo_retriever/ingest_modes/batch.py +1150 -0
- nemo_retriever-26.3.0/src/nemo_retriever/ingest_modes/fused.py +315 -0
- nemo_retriever-26.3.0/src/nemo_retriever/ingest_modes/gpu_pool.py +448 -0
- nemo_retriever-26.3.0/src/nemo_retriever/ingest_modes/inprocess.py +1787 -0
- nemo_retriever-26.3.0/src/nemo_retriever/ingest_modes/lancedb_utils.py +235 -0
- nemo_retriever-26.3.0/src/nemo_retriever/ingest_modes/online.py +258 -0
- nemo_retriever-26.3.0/src/nemo_retriever/ingest_modes/pipeline.py +43 -0
- nemo_retriever-26.3.0/src/nemo_retriever/ingest_modes/serve.py +9 -0
- nemo_retriever-26.3.0/src/nemo_retriever/ingestor.py +213 -0
- nemo_retriever-26.3.0/src/nemo_retriever/io/__init__.py +17 -0
- nemo_retriever-26.3.0/src/nemo_retriever/io/dataframe.py +72 -0
- nemo_retriever-26.3.0/src/nemo_retriever/io/markdown.py +410 -0
- nemo_retriever-26.3.0/src/nemo_retriever/io/stage_files.py +19 -0
- nemo_retriever-26.3.0/src/nemo_retriever/local/__init__.py +7 -0
- nemo_retriever-26.3.0/src/nemo_retriever/local/__main__.py +32 -0
- nemo_retriever-26.3.0/src/nemo_retriever/local/stages/__init__.py +12 -0
- nemo_retriever-26.3.0/src/nemo_retriever/local/stages/stage1_pdf_extraction.py +28 -0
- nemo_retriever-26.3.0/src/nemo_retriever/local/stages/stage2_infographic_extraction.py +105 -0
- nemo_retriever-26.3.0/src/nemo_retriever/local/stages/stage3_table_extractor.py +150 -0
- nemo_retriever-26.3.0/src/nemo_retriever/local/stages/stage4_chart_extractor.py +109 -0
- nemo_retriever-26.3.0/src/nemo_retriever/local/stages/stage5_text_embeddings.py +27 -0
- nemo_retriever-26.3.0/src/nemo_retriever/local/stages/stage6_vdb_upload.py +26 -0
- nemo_retriever-26.3.0/src/nemo_retriever/local/stages/stage7_vdb_query.py +27 -0
- nemo_retriever-26.3.0/src/nemo_retriever/local/stages/stage999_post_mortem_analysis.py +2223 -0
- nemo_retriever-26.3.0/src/nemo_retriever/model/__init__.py +75 -0
- nemo_retriever-26.3.0/src/nemo_retriever/model/local/__init__.py +54 -0
- nemo_retriever-26.3.0/src/nemo_retriever/model/local/llama_nemotron_embed_1b_v2_embedder.py +126 -0
- nemo_retriever-26.3.0/src/nemo_retriever/model/local/llama_nemotron_embed_vl_1b_v2_embedder.py +151 -0
- nemo_retriever-26.3.0/src/nemo_retriever/model/local/nemotron_graphic_elements_v1.py +160 -0
- nemo_retriever-26.3.0/src/nemo_retriever/model/local/nemotron_ocr_v1.py +268 -0
- nemo_retriever-26.3.0/src/nemo_retriever/model/local/nemotron_page_elements_v3.py +238 -0
- nemo_retriever-26.3.0/src/nemo_retriever/model/local/nemotron_parse_v1_2.py +191 -0
- nemo_retriever-26.3.0/src/nemo_retriever/model/local/nemotron_rerank_v2.py +210 -0
- nemo_retriever-26.3.0/src/nemo_retriever/model/local/nemotron_table_structure_v1.py +143 -0
- nemo_retriever-26.3.0/src/nemo_retriever/model/local/parakeet_ctc_1_1b_asr.py +180 -0
- nemo_retriever-26.3.0/src/nemo_retriever/model/model.py +104 -0
- nemo_retriever-26.3.0/src/nemo_retriever/nim/__init__.py +5 -0
- nemo_retriever-26.3.0/src/nemo_retriever/nim/nim.py +218 -0
- nemo_retriever-26.3.0/src/nemo_retriever/ocr/__init__.py +7 -0
- nemo_retriever-26.3.0/src/nemo_retriever/ocr/ocr.py +1140 -0
- nemo_retriever-26.3.0/src/nemo_retriever/online/__init__.py +5 -0
- nemo_retriever-26.3.0/src/nemo_retriever/online/__main__.py +9 -0
- nemo_retriever-26.3.0/src/nemo_retriever/page_elements/__init__.py +18 -0
- nemo_retriever-26.3.0/src/nemo_retriever/page_elements/page_elements.py +781 -0
- nemo_retriever-26.3.0/src/nemo_retriever/params/__init__.py +53 -0
- nemo_retriever-26.3.0/src/nemo_retriever/params/models.py +306 -0
- nemo_retriever-26.3.0/src/nemo_retriever/params/utils.py +43 -0
- nemo_retriever-26.3.0/src/nemo_retriever/pdf/__init__.py +24 -0
- nemo_retriever-26.3.0/src/nemo_retriever/pdf/__main__.py +16 -0
- nemo_retriever-26.3.0/src/nemo_retriever/pdf/config.py +46 -0
- nemo_retriever-26.3.0/src/nemo_retriever/pdf/extract.py +384 -0
- nemo_retriever-26.3.0/src/nemo_retriever/pdf/io.py +43 -0
- nemo_retriever-26.3.0/src/nemo_retriever/pdf/split.py +197 -0
- nemo_retriever-26.3.0/src/nemo_retriever/pdf/stage.py +705 -0
- nemo_retriever-26.3.0/src/nemo_retriever/recall/__init__.py +16 -0
- nemo_retriever-26.3.0/src/nemo_retriever/recall/__main__.py +20 -0
- nemo_retriever-26.3.0/src/nemo_retriever/recall/core.py +380 -0
- nemo_retriever-26.3.0/src/nemo_retriever/recall/main_recall.py +142 -0
- nemo_retriever-26.3.0/src/nemo_retriever/recall/vdb_recall.py +317 -0
- nemo_retriever-26.3.0/src/nemo_retriever/rerank/__init__.py +24 -0
- nemo_retriever-26.3.0/src/nemo_retriever/rerank/rerank.py +375 -0
- nemo_retriever-26.3.0/src/nemo_retriever/retriever.py +297 -0
- nemo_retriever-26.3.0/src/nemo_retriever/table/__init__.py +17 -0
- nemo_retriever-26.3.0/src/nemo_retriever/table/commands.py +150 -0
- nemo_retriever-26.3.0/src/nemo_retriever/table/config.py +49 -0
- nemo_retriever-26.3.0/src/nemo_retriever/table/processor.py +50 -0
- nemo_retriever-26.3.0/src/nemo_retriever/table/stage.py +19 -0
- nemo_retriever-26.3.0/src/nemo_retriever/table/table_detection.py +552 -0
- nemo_retriever-26.3.0/src/nemo_retriever/text_embed/__init__.py +32 -0
- nemo_retriever-26.3.0/src/nemo_retriever/text_embed/commands.py +248 -0
- nemo_retriever-26.3.0/src/nemo_retriever/text_embed/config.py +20 -0
- nemo_retriever-26.3.0/src/nemo_retriever/text_embed/main_text_embed.py +705 -0
- nemo_retriever-26.3.0/src/nemo_retriever/text_embed/processor.py +112 -0
- nemo_retriever-26.3.0/src/nemo_retriever/text_embed/ray_data.py +40 -0
- nemo_retriever-26.3.0/src/nemo_retriever/text_embed/stage.py +10 -0
- nemo_retriever-26.3.0/src/nemo_retriever/text_embed/text_embed.py +220 -0
- nemo_retriever-26.3.0/src/nemo_retriever/txt/__init__.py +25 -0
- nemo_retriever-26.3.0/src/nemo_retriever/txt/__main__.py +108 -0
- nemo_retriever-26.3.0/src/nemo_retriever/txt/ray_data.py +75 -0
- nemo_retriever-26.3.0/src/nemo_retriever/txt/split.py +293 -0
- nemo_retriever-26.3.0/src/nemo_retriever/utils/__init__.py +5 -0
- nemo_retriever-26.3.0/src/nemo_retriever/utils/benchmark/__init__.py +11 -0
- nemo_retriever-26.3.0/src/nemo_retriever/utils/benchmark/__main__.py +21 -0
- nemo_retriever-26.3.0/src/nemo_retriever/utils/benchmark/all_actor.py +106 -0
- nemo_retriever-26.3.0/src/nemo_retriever/utils/benchmark/audio_extract_actor.py +188 -0
- nemo_retriever-26.3.0/src/nemo_retriever/utils/benchmark/common.py +206 -0
- nemo_retriever-26.3.0/src/nemo_retriever/utils/benchmark/extract_actor.py +75 -0
- nemo_retriever-26.3.0/src/nemo_retriever/utils/benchmark/ocr_actor.py +82 -0
- nemo_retriever-26.3.0/src/nemo_retriever/utils/benchmark/page_elements_actor.py +69 -0
- nemo_retriever-26.3.0/src/nemo_retriever/utils/benchmark/split_actor.py +65 -0
- nemo_retriever-26.3.0/src/nemo_retriever/utils/compare/__init__.py +11 -0
- nemo_retriever-26.3.0/src/nemo_retriever/utils/compare/__main__.py +18 -0
- nemo_retriever-26.3.0/src/nemo_retriever/utils/compare/compare_json.py +321 -0
- nemo_retriever-26.3.0/src/nemo_retriever/utils/compare/compare_results.py +871 -0
- nemo_retriever-26.3.0/src/nemo_retriever/utils/convert/__init__.py +17 -0
- nemo_retriever-26.3.0/src/nemo_retriever/utils/convert/to_pdf.py +151 -0
- nemo_retriever-26.3.0/src/nemo_retriever/utils/detection_summary.py +251 -0
- nemo_retriever-26.3.0/src/nemo_retriever/utils/hf_cache.py +24 -0
- nemo_retriever-26.3.0/src/nemo_retriever/utils/hf_model_registry.py +59 -0
- nemo_retriever-26.3.0/src/nemo_retriever/utils/image/__init__.py +15 -0
- nemo_retriever-26.3.0/src/nemo_retriever/utils/image/__main__.py +16 -0
- nemo_retriever-26.3.0/src/nemo_retriever/utils/image/render.py +241 -0
- nemo_retriever-26.3.0/src/nemo_retriever/utils/input_files.py +34 -0
- nemo_retriever-26.3.0/src/nemo_retriever/utils/ray_resource_hueristics.py +494 -0
- nemo_retriever-26.3.0/src/nemo_retriever/utils/remote_auth.py +11 -0
- nemo_retriever-26.3.0/src/nemo_retriever/utils/table_and_chart.py +668 -0
- nemo_retriever-26.3.0/src/nemo_retriever/vector_store/__init__.py +19 -0
- nemo_retriever-26.3.0/src/nemo_retriever/vector_store/__main__.py +20 -0
- nemo_retriever-26.3.0/src/nemo_retriever/vector_store/lancedb_store.py +331 -0
- nemo_retriever-26.3.0/src/nemo_retriever/vector_store/stage.py +85 -0
- nemo_retriever-26.3.0/src/nemo_retriever/version.py +162 -0
- nemo_retriever-26.3.0/src/nemo_retriever.egg-info/PKG-INFO +455 -0
- nemo_retriever-26.3.0/src/nemo_retriever.egg-info/SOURCES.txt +225 -0
- nemo_retriever-26.3.0/src/nemo_retriever.egg-info/dependency_links.txt +1 -0
- nemo_retriever-26.3.0/src/nemo_retriever.egg-info/entry_points.txt +2 -0
- nemo_retriever-26.3.0/src/nemo_retriever.egg-info/requires.txt +52 -0
- nemo_retriever-26.3.0/src/nemo_retriever.egg-info/top_level.txt +1 -0
- nemo_retriever-26.3.0/tests/test_asr_actor.py +198 -0
- nemo_retriever-26.3.0/tests/test_audio_benchmark.py +42 -0
- nemo_retriever-26.3.0/tests/test_audio_chunk_actor.py +89 -0
- nemo_retriever-26.3.0/tests/test_audio_pipeline_batch.py +311 -0
- nemo_retriever-26.3.0/tests/test_audio_stage.py +100 -0
- nemo_retriever-26.3.0/tests/test_batch_ingestor.py +58 -0
- nemo_retriever-26.3.0/tests/test_batch_pipeline.py +14 -0
- nemo_retriever-26.3.0/tests/test_chart_graphic_elements.py +416 -0
- nemo_retriever-26.3.0/tests/test_create_local_embedder.py +110 -0
- nemo_retriever-26.3.0/tests/test_embed_params.py +95 -0
- nemo_retriever-26.3.0/tests/test_factory.py +72 -0
- nemo_retriever-26.3.0/tests/test_harness_config.py +307 -0
- nemo_retriever-26.3.0/tests/test_harness_nightly.py +230 -0
- nemo_retriever-26.3.0/tests/test_harness_parsers.py +57 -0
- nemo_retriever-26.3.0/tests/test_harness_recall_adapters.py +36 -0
- nemo_retriever-26.3.0/tests/test_harness_reporting.py +102 -0
- nemo_retriever-26.3.0/tests/test_harness_run.py +527 -0
- nemo_retriever-26.3.0/tests/test_html_convert.py +91 -0
- nemo_retriever-26.3.0/tests/test_image_load.py +216 -0
- nemo_retriever-26.3.0/tests/test_ingest_config.py +78 -0
- nemo_retriever-26.3.0/tests/test_ingest_interface.py +38 -0
- nemo_retriever-26.3.0/tests/test_io_dataframe.py +81 -0
- nemo_retriever-26.3.0/tests/test_io_markdown.py +195 -0
- nemo_retriever-26.3.0/tests/test_lancedb_utils.py +244 -0
- nemo_retriever-26.3.0/tests/test_multimodal_embed.py +332 -0
- nemo_retriever-26.3.0/tests/test_nemotron_rerank_v2.py +595 -0
- nemo_retriever-26.3.0/tests/test_params_utils.py +55 -0
- nemo_retriever-26.3.0/tests/test_pdf_render_scale.py +226 -0
- nemo_retriever-26.3.0/tests/test_recall_core.py +29 -0
- nemo_retriever-26.3.0/tests/test_resource_heuristics.py +221 -0
- nemo_retriever-26.3.0/tests/test_retriever_queries.py +372 -0
- nemo_retriever-26.3.0/tests/test_table_structure.py +506 -0
- nemo_retriever-26.3.0/tests/test_txt_split.py +85 -0
|
@@ -0,0 +1,455 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: nemo-retriever
|
|
3
|
+
Version: 26.3.0
|
|
4
|
+
Summary: A modern RAG ingestion pipeline from Nvidia
|
|
5
|
+
Author-email: Jeremy Dyer <jdyer@nvidia.com>
|
|
6
|
+
License-Expression: Apache-2.0
|
|
7
|
+
Classifier: Programming Language :: Python :: 3
|
|
8
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
9
|
+
Classifier: Operating System :: OS Independent
|
|
10
|
+
Requires-Python: >=3.12
|
|
11
|
+
Description-Content-Type: text/markdown
|
|
12
|
+
Requires-Dist: ray[data,serve]>=2.49.0
|
|
13
|
+
Requires-Dist: pypdfium2==4.30.0
|
|
14
|
+
Requires-Dist: pandas<3,>=2.0
|
|
15
|
+
Requires-Dist: tqdm>=4.66.0
|
|
16
|
+
Requires-Dist: typer>=0.12.0
|
|
17
|
+
Requires-Dist: pyyaml>=6.0
|
|
18
|
+
Requires-Dist: lancedb
|
|
19
|
+
Requires-Dist: nv-ingest
|
|
20
|
+
Requires-Dist: nv-ingest-api
|
|
21
|
+
Requires-Dist: nv-ingest-client
|
|
22
|
+
Requires-Dist: fastapi>=0.114.0
|
|
23
|
+
Requires-Dist: uvicorn[standard]>=0.30.0
|
|
24
|
+
Requires-Dist: httpx>=0.27.0
|
|
25
|
+
Requires-Dist: requests>=2.32.5
|
|
26
|
+
Requires-Dist: pydantic>=2.8.0
|
|
27
|
+
Requires-Dist: typer>=0.12.0
|
|
28
|
+
Requires-Dist: rich>=13.7.0
|
|
29
|
+
Requires-Dist: pillow==12.1.1
|
|
30
|
+
Requires-Dist: numpy>=1.26.0
|
|
31
|
+
Requires-Dist: nltk==3.9.3
|
|
32
|
+
Requires-Dist: urllib3==2.6.3
|
|
33
|
+
Requires-Dist: debugpy>=1.8.0
|
|
34
|
+
Requires-Dist: python-multipart>=0.0.9
|
|
35
|
+
Requires-Dist: transformers>=5.0.0
|
|
36
|
+
Requires-Dist: tokenizers>=0.20.3
|
|
37
|
+
Requires-Dist: accelerate>=1.1.0
|
|
38
|
+
Requires-Dist: torch~=2.9.1
|
|
39
|
+
Requires-Dist: torchvision<0.25,>=0.24
|
|
40
|
+
Requires-Dist: einops
|
|
41
|
+
Requires-Dist: easydict
|
|
42
|
+
Requires-Dist: addict
|
|
43
|
+
Requires-Dist: nemotron-page-elements-v3>=0.dev0
|
|
44
|
+
Requires-Dist: nemotron-graphic-elements-v1>=0.dev0
|
|
45
|
+
Requires-Dist: nemotron-table-structure-v1>=0.dev0
|
|
46
|
+
Requires-Dist: nemotron-ocr>=0.dev0
|
|
47
|
+
Requires-Dist: markitdown
|
|
48
|
+
Requires-Dist: timm==1.0.22
|
|
49
|
+
Requires-Dist: tqdm
|
|
50
|
+
Requires-Dist: accelerate==1.12.0
|
|
51
|
+
Requires-Dist: albumentations==2.0.8
|
|
52
|
+
Requires-Dist: scikit-learn>=1.6.0
|
|
53
|
+
Requires-Dist: open-clip-torch==3.2.0
|
|
54
|
+
Requires-Dist: soundfile>=0.12.0
|
|
55
|
+
Requires-Dist: scipy>=1.11.0
|
|
56
|
+
Requires-Dist: nvidia-ml-py
|
|
57
|
+
Provides-Extra: svg
|
|
58
|
+
Requires-Dist: cairosvg>=2.7.0; extra == "svg"
|
|
59
|
+
Provides-Extra: dev
|
|
60
|
+
Requires-Dist: build>=1.2.2; extra == "dev"
|
|
61
|
+
Requires-Dist: pytest>=8.0.2; extra == "dev"
|
|
62
|
+
|
|
63
|
+
# Quick Start for NeMo Retriever Library
|
|
64
|
+
|
|
65
|
+
NeMo Retriever Library is a retrieval-augmented generation (RAG) ingestion pipeline for documents that can parse text, tables, charts, and infographics. NeMo Retriever Library parses documents, creates embeddings, optionally stores embeddings in LanceDB, and performs recall evaluation.
|
|
66
|
+
|
|
67
|
+
This quick start guide shows how to run NeMo Retriever Library as a library all within local Python processes without containers. NeMo Retriever Library supports two inference options:
|
|
68
|
+
- Pull and run [Nemotron RAG models from Hugging Face](https://huggingface.co/collections/nvidia/nemotron-rag) on your local GPU(s).
|
|
69
|
+
- Make over the network inference calls to build.nvidia.com hosted or locally deployed NeMo Retriever NIM endpoints.
|
|
70
|
+
|
|
71
|
+
You’ll set up a CUDA 13–compatible environment, install the library and its dependencies, and run GPU‑accelerated ingestion pipelines that convert PDFs, HTML, plain text, audio, or video into vector embeddings stored in LanceDB (on local disk), with Ray‑based scaling and built‑in recall benchmarking.
|
|
72
|
+
|
|
73
|
+
## Prerequisites
|
|
74
|
+
|
|
75
|
+
<<<<<<< HEAD
|
|
76
|
+
Before you start, make sure your system meets the following requirements:
|
|
77
|
+
=======
|
|
78
|
+
Before starting, make sure your system meets the following requirements:
|
|
79
|
+
>>>>>>> 15b2bc05 (Updating the nemo_retriever quickstart README (#1632))
|
|
80
|
+
|
|
81
|
+
- The host is running CUDA 13.x so that `libcudart.so.13` is available.
|
|
82
|
+
- Your GPUs are visible to the system and compatible with CUDA 13.x.
|
|
83
|
+
|
|
84
|
+
If optical character recognition (OCR) fails with a `libcudart.so.13` error, install the CUDA 13 runtime for your platform and update `LD_LIBRARY_PATH` to include the CUDA lib64 directory, then rerun the pipeline.
|
|
85
|
+
|
|
86
|
+
For example, the following command can be used to update the `LD_LIBRARY_PATH` value.
|
|
87
|
+
|
|
88
|
+
```bash
|
|
89
|
+
export LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:/usr/local/cuda/lib64
|
|
90
|
+
```
|
|
91
|
+
|
|
92
|
+
## Setup your environment
|
|
93
|
+
|
|
94
|
+
Complete the following steps to setup your environment. You will create and activate isolated Python and project virtual environments, install the NeMo Retriever Library and its dependencies, and then run the provided ingestion snippets to validate your setup.
|
|
95
|
+
|
|
96
|
+
1. Create and activate the NeMo Retriever Library environment
|
|
97
|
+
|
|
98
|
+
Before installing NeMo Retriever Library, create an isolated Python environment so its dependencies do not conflict with other projects on your system. In this step, you set up a new virtual environment and activate it so that all subsequent installs are scoped to NeMo Retriever Library.
|
|
99
|
+
|
|
100
|
+
In your terminal, run the following commands from any location.
|
|
101
|
+
|
|
102
|
+
```bash
|
|
103
|
+
uv venv retriever --python 3.12
|
|
104
|
+
source retriever/bin/activate
|
|
105
|
+
uv pip install -i https://test.pypi.org/simple/ --extra-index-url https://pypi.org/simple nemo-retriever==26.3.0rc2 nv-ingest-client==26.3.0rc2 nv-ingest==26.3.0rc2 nv-ingest-api==26.3.0rc2
|
|
106
|
+
```
|
|
107
|
+
This creates a dedicated Python environment and installs the `nemo-retriever` PyPI package, the canonical distribution for the NeMo Retriever Library.
|
|
108
|
+
|
|
109
|
+
2. Install CUDA 13 builds of Torch and Torchvision
|
|
110
|
+
|
|
111
|
+
To ensure NeMo Retriever Library’s OCR and GPU‑accelerated components run correctly on your system, you need PyTorch and TorchVision builds that are compiled for CUDA 13. In this step, you uninstall any existing Torch/TorchVision packages and reinstall them from a dedicated CUDA 13.0 wheel index so they link against the same CUDA runtime as the rest of your pipeline.
|
|
112
|
+
|
|
113
|
+
Use the CUDA 13.0 wheels from the dedicated index by running the following command.
|
|
114
|
+
|
|
115
|
+
```bash
|
|
116
|
+
uv pip uninstall torch torchvision
|
|
117
|
+
uv pip install torch==2.9.1 torchvision -i https://download.pytorch.org/whl/cu130
|
|
118
|
+
```
|
|
119
|
+
This ensures the OCR and GPU‑accelerated components in NeMo Retriever Library run against the right CUDA runtime.
|
|
120
|
+
|
|
121
|
+
## Run the pipeline
|
|
122
|
+
|
|
123
|
+
The [test PDF](../data/multimodal_test.pdf) contains text, tables, charts, and images. Additional test data resides [here](../data/).
|
|
124
|
+
|
|
125
|
+
> **Note:** `batch` is the primary intended run_mode of operation for this library. Other modes are experimental and subject to change or removal.
|
|
126
|
+
|
|
127
|
+
### Ingest a test pdf
|
|
128
|
+
```python
|
|
129
|
+
from nemo_retriever import create_ingestor
|
|
130
|
+
from nemo_retriever.io import to_markdown, to_markdown_by_page
|
|
131
|
+
from pathlib import Path
|
|
132
|
+
|
|
133
|
+
documents = [str(Path("../data/multimodal_test.pdf"))]
|
|
134
|
+
ingestor = create_ingestor(run_mode="batch")
|
|
135
|
+
|
|
136
|
+
# ingestion tasks are chainable and defined lazily
|
|
137
|
+
ingestor = (
|
|
138
|
+
ingestor.files(documents)
|
|
139
|
+
.extract(
|
|
140
|
+
# below are the default values, but content types can be controlled
|
|
141
|
+
extract_text=True,
|
|
142
|
+
extract_charts=True,
|
|
143
|
+
extract_tables=True,
|
|
144
|
+
extract_infographics=True
|
|
145
|
+
)
|
|
146
|
+
.embed()
|
|
147
|
+
.vdb_upload()
|
|
148
|
+
)
|
|
149
|
+
|
|
150
|
+
# ingestor.ingest() actually executes the pipeline
|
|
151
|
+
# results are returned as a ray dataset and inspectable as chunks
|
|
152
|
+
ray_dataset = ingestor.ingest()
|
|
153
|
+
chunks = ray_dataset.get_dataset().take_all()
|
|
154
|
+
```
|
|
155
|
+
|
|
156
|
+
### Inspect extracts
|
|
157
|
+
You can inspect how recall accuracy optimized text chunks for various content types were extracted into text representations:
|
|
158
|
+
```python
|
|
159
|
+
# page 1 raw text:
|
|
160
|
+
>>> chunks[0]["text"]
|
|
161
|
+
'TestingDocument\r\nA sample document with headings and placeholder text\r\nIntroduction\r\nThis is a placeholder document that can be used for any purpose...'
|
|
162
|
+
|
|
163
|
+
# markdown formatted table from the first page
|
|
164
|
+
>>> chunks[1]["text"]
|
|
165
|
+
'| Table | 1 |\n| This | table | describes | some | animals, | and | some | activities | they | might | be | doing | in | specific |\n| locations. |\n| Animal | Activity | Place |\n| Giraffe | Driving | a | car | At | the | beach |\n| Lion | Putting | on | sunscreen | At | the | park |\n| Cat | Jumping | onto | a | laptop | In | a | home | office |\n| Dog | Chasing | a | squirrel | In | the | front | yard |\n| Chart | 1 |'
|
|
166
|
+
|
|
167
|
+
# a chart from the first page
|
|
168
|
+
>>> chunks[2]["text"]
|
|
169
|
+
'Chart 1\nThis chart shows some gadgets, and some very fictitious costs.\nGadgets and their cost\n$160.00\n$140.00\n$120.00\n$100.00\nDollars\n$80.00\n$60.00\n$40.00\n$20.00\n$-\nPowerdrill\nBluetooth speaker\nMinifridge\nPremium desk fan\nHammer\nCost'
|
|
170
|
+
|
|
171
|
+
# markdown formatting for full pages or documents:
|
|
172
|
+
>>> to_markdown_by_page(chunks).keys()
|
|
173
|
+
dict_keys([1, 2, 3])
|
|
174
|
+
|
|
175
|
+
>>> to_markdown_by_page(chunks)[1]
|
|
176
|
+
'## Page 1\n\nTestingDocument\r\nA sample document with headings and placeholder text\r\nIntroduction\r\nThis is a placeholder document that can be used for any purpose. It contains some \r\nheadings and some placeholder text to fill the space. The text is not important and contains \r\nno real value, but it is useful for testing. Below, we will have some simple tables and charts \r\nthat we can use to confirm Ingest is working as expected.\r\nTable 1\r\nThis table describes some animals, and some activities they might be doing in specific \r\nlocations.\r\nAnimal Activity Place\r\nGira@e Driving a car At the beach\r\nLion Putting on sunscreen At the park\r\nCat Jumping onto a laptop In a home o@ice\r\nDog Chasing a squirrel In the front yard\r\nChart 1\r\nThis chart shows some gadgets, and some very fictitious costs.\n\n| This | table | describes | some | animals, | and | some | activities | they | might | be | doing | in | specific |\n| locations. |\n| Animal | Activity | Place |\n| Giraffe | Driving | a | car | At | the | beach |\n| Lion | Putting | on | sunscreen | At | the | park |\n| Cat | Jumping | onto | a | laptop | In | a | home | office |\n| Dog | Chasing | a | squirrel | In | the | front | yard |\n| Chart | 1 |\n\nChart 1 This chart shows some gadgets, and some very fictitious costs...'
|
|
177
|
+
|
|
178
|
+
# full document markdown
|
|
179
|
+
>>> to_markdown(chunks)
|
|
180
|
+
```
|
|
181
|
+
|
|
182
|
+
Since the ingestion job automatically populated a lancedb table with all these chunks, you can use queries to retrieve semantically relevant chunks for feeding directly into an LLM:
|
|
183
|
+
|
|
184
|
+
### Run a recall query
|
|
185
|
+
|
|
186
|
+
```python
|
|
187
|
+
from nemo_retriever.retriever import Retriever
|
|
188
|
+
|
|
189
|
+
retriever = Retriever(
|
|
190
|
+
# default values
|
|
191
|
+
lancedb_uri="lancedb",
|
|
192
|
+
lancedb_table="nv-ingest",
|
|
193
|
+
embedder="nvidia/llama-3.2-nv-embedqa-1b-v2",
|
|
194
|
+
top_k=5,
|
|
195
|
+
reranker=False
|
|
196
|
+
)
|
|
197
|
+
|
|
198
|
+
query = "Given their activities, which animal is responsible for the typos in my documents?"
|
|
199
|
+
|
|
200
|
+
# you can also submit a list with retriever.queries[...]
|
|
201
|
+
hits = retriever.query(query)
|
|
202
|
+
```
|
|
203
|
+
|
|
204
|
+
```python
|
|
205
|
+
# retrieved text from the first page
|
|
206
|
+
>>> hits[0]
|
|
207
|
+
{'text': 'TestingDocument\r\nA sample document with headings and placeholder text\r\nIntroduction\r\nThis is a placeholder document that can be used for any purpose. It contains some \r\nheadings and some placeholder text to fill the space. The text is not important and contains \r\nno real value, but it is useful for testing. Below, we will have some simple tables and charts \r\nthat we can use to confirm Ingest is working as expected.\r\nTable 1\r\nThis table describes some animals, and some activities they might be doing in specific \r\nlocations.\r\nAnimal Activity Place\r\nGira@e Driving a car At the beach\r\nLion Putting on sunscreen At the park\r\nCat Jumping onto a laptop In a home o@ice\r\nDog Chasing a squirrel In the front yard\r\nChart 1\r\nThis chart shows some gadgets, and some very fictitious costs.', 'metadata': '{"page_number": 1, "pdf_page": "multimodal_test_1", "page_elements_v3_num_detections": 9, "page_elements_v3_counts_by_label": {"table": 1, "chart": 1, "title": 3, "text": 4}, "ocr_table_detections": 1, "ocr_chart_detections": 1, "ocr_infographic_detections": 0}', 'source': '{"source_id": "/home/dev/projects/NeMo-Retriever/data/multimodal_test.pdf"}', 'page_number': 1, '_distance': 1.5822279453277588}
|
|
208
|
+
|
|
209
|
+
# retrieved text of the table from the first page
|
|
210
|
+
>>> hits[1]
|
|
211
|
+
{'text': '| Table | 1 |\n| This | table | describes | some | animals, | and | some | activities | they | might | be | doing | in | specific |\n| locations. |\n| Animal | Activity | Place |\n| Giraffe | Driving | a | car | At | the | beach |\n| Lion | Putting | on | sunscreen | At | the | park |\n| Cat | Jumping | onto | a | laptop | In | a | home | office |\n| Dog | Chasing | a | squirrel | In | the | front | yard |\n| Chart | 1 |', 'metadata': '{"page_number": 1, "pdf_page": "multimodal_test_1", "page_elements_v3_num_detections": 9, "page_elements_v3_counts_by_label": {"table": 1, "chart": 1, "title": 3, "text": 4}, "ocr_table_detections": 1, "ocr_chart_detections": 1, "ocr_infographic_detections": 0}', 'source': '{"source_id": "/home/dev/projects/NeMo-Retriever/data/multimodal_test.pdf"}', 'page_number': 1, '_distance': 1.614684820175171}
|
|
212
|
+
```
|
|
213
|
+
|
|
214
|
+
### Generate a query answer using an LLM
|
|
215
|
+
The above retrieval results are often feedable directly to an LLM for answer generation.
|
|
216
|
+
|
|
217
|
+
To do so, first install the openai client and set your [build.nvidia.com](https://build.nvidia.com/) API key:
|
|
218
|
+
```bash
|
|
219
|
+
uv pip install -y openai
|
|
220
|
+
export NVIDIA_API_KEY=nvapi-...
|
|
221
|
+
```
|
|
222
|
+
|
|
223
|
+
```python
|
|
224
|
+
from openai import OpenAI
|
|
225
|
+
import os
|
|
226
|
+
|
|
227
|
+
client = OpenAI(
|
|
228
|
+
base_url = "https://integrate.api.nvidia.com/v1",
|
|
229
|
+
api_key = os.environ.get("NVIDIA_API_KEY")
|
|
230
|
+
)
|
|
231
|
+
|
|
232
|
+
hit_texts = [hit["text"] for hit in hits]
|
|
233
|
+
prompt = f"""
|
|
234
|
+
Given the following retrieved documents, answer the question: {query}
|
|
235
|
+
|
|
236
|
+
Documents:
|
|
237
|
+
{hit_texts}
|
|
238
|
+
"""
|
|
239
|
+
|
|
240
|
+
completion = client.chat.completions.create(
|
|
241
|
+
model="nvidia/nemotron-3-super-120b-a12b",
|
|
242
|
+
messages=[{"role":"user","content":prompt}],
|
|
243
|
+
stream=False
|
|
244
|
+
)
|
|
245
|
+
|
|
246
|
+
answer = completion.choices[0].message.content
|
|
247
|
+
print(answer)
|
|
248
|
+
```
|
|
249
|
+
|
|
250
|
+
Answer:
|
|
251
|
+
```
|
|
252
|
+
Cat is the animal whose activity (jumping onto a laptop) matches the location of the typos, so the cat is responsible for the typos in the documents.
|
|
253
|
+
```
|
|
254
|
+
|
|
255
|
+
### Ingest other types of content:
|
|
256
|
+
|
|
257
|
+
For PowerPoint and Docx files, ensure libeoffice is installed by your system's package manager. This is required to make their pages renderable as images for our [page-elements content classifier](https://huggingface.co/nvidia/nemotron-page-elements-v3).
|
|
258
|
+
|
|
259
|
+
For example, with apt-get on Ubuntu:
|
|
260
|
+
```bash
|
|
261
|
+
sudo apt install -y libreoffice
|
|
262
|
+
```
|
|
263
|
+
|
|
264
|
+
Example usage:
|
|
265
|
+
```python
|
|
266
|
+
# docx and pptx files
|
|
267
|
+
documents = [str(Path(f"../data/*{ext}")) for ext in [".pptx", ".docx"]]
|
|
268
|
+
# mixed types of images
|
|
269
|
+
images = [str(Path(f"../data/*{ext}")) for ext in [".png", ".jpeg", ".bmp"]]
|
|
270
|
+
ingestor = (
|
|
271
|
+
# above file types can be combined into a single job
|
|
272
|
+
ingestor.files(documents + images)
|
|
273
|
+
.extract()
|
|
274
|
+
)
|
|
275
|
+
```
|
|
276
|
+
|
|
277
|
+
*Note:* the `split()` task uses a tokenizer to split texts by a max_token length
|
|
278
|
+
|
|
279
|
+
<<<<<<< HEAD
|
|
280
|
+
### Render results as markdown
|
|
281
|
+
|
|
282
|
+
If you want a readable markdown view of extracted results, pass the full in-process result list
|
|
283
|
+
to `nemo_retriever.io.to_markdown`. The helper now returns a `dict[str, str]` keyed by input
|
|
284
|
+
filename, where each value is the document collapsed into one markdown string without per-page
|
|
285
|
+
headers, so both single-document and multi-document runs follow the same contract.
|
|
286
|
+
=======
|
|
287
|
+
PDF text is split at the page level.
|
|
288
|
+
|
|
289
|
+
HTML and .txt files have no natural page delimiters, so they almost always need to be paired with the `.split()` task.
|
|
290
|
+
>>>>>>> 15b2bc05 (Updating the nemo_retriever quickstart README (#1632))
|
|
291
|
+
|
|
292
|
+
```python
|
|
293
|
+
# html and text files - include a split task to prevent texts from exceeding the embedder's max sequence length
|
|
294
|
+
documents = [str(Path(f"../data/*{ext}")) for ext in [".txt", ".html"]]
|
|
295
|
+
ingestor = (
|
|
296
|
+
ingestor.files(documents)
|
|
297
|
+
.extract()
|
|
298
|
+
.split(max_tokens=5) #1024 by default, set low here to demonstrate chunking
|
|
299
|
+
)
|
|
300
|
+
<<<<<<< HEAD
|
|
301
|
+
results = ingestor.ingest()
|
|
302
|
+
markdown_docs = to_markdown(results)
|
|
303
|
+
print(markdown_docs["multimodal_test.pdf"])
|
|
304
|
+
```
|
|
305
|
+
|
|
306
|
+
Use `to_markdown_by_page(results)` when you want a nested
|
|
307
|
+
`dict[str, dict[int, str]]` instead, where each filename maps to its per-page markdown strings.
|
|
308
|
+
=======
|
|
309
|
+
```
|
|
310
|
+
|
|
311
|
+
For audio and video files, ensure ffmpeg is installed by your system's package manager.
|
|
312
|
+
>>>>>>> 15b2bc05 (Updating the nemo_retriever quickstart README (#1632))
|
|
313
|
+
|
|
314
|
+
For example, with apt-get on Ubuntu:
|
|
315
|
+
```bash
|
|
316
|
+
sudo apt install -y ffmpeg
|
|
317
|
+
```
|
|
318
|
+
|
|
319
|
+
```python
|
|
320
|
+
ingestor = create_ingestor(run_mode="batch")
|
|
321
|
+
ingestor = ingestor.files([str(INPUT_AUDIO)]).extract_audio()
|
|
322
|
+
```
|
|
323
|
+
|
|
324
|
+
### Explore Different Pipeline Options:
|
|
325
|
+
|
|
326
|
+
You can use the [Nemotron RAG VL Embedder](https://huggingface.co/nvidia/llama-nemotron-embed-vl-1b-v2)
|
|
327
|
+
|
|
328
|
+
```python
|
|
329
|
+
ingestor = (
|
|
330
|
+
ingestor.files(documents)
|
|
331
|
+
.extract()
|
|
332
|
+
.embed(
|
|
333
|
+
model_name="nvidia/llama-nemotron-embed-vl-1b-v2",
|
|
334
|
+
#works with plain "text"s, "image"s, and "text_image" pairs
|
|
335
|
+
embed_modality="text_image"
|
|
336
|
+
)
|
|
337
|
+
)
|
|
338
|
+
```
|
|
339
|
+
|
|
340
|
+
You can use a different ingestion pipeline based on [Nemotron-Parse](https://huggingface.co/nvidia/NVIDIA-Nemotron-Parse-v1.2) combined with the default embedder:
|
|
341
|
+
```python
|
|
342
|
+
ingestor = ingestor.files(documents).extract(method="nemotron_parse")
|
|
343
|
+
```
|
|
344
|
+
|
|
345
|
+
## Run with remote inference, no local GPU required:
|
|
346
|
+
|
|
347
|
+
For build.nvidia.com hosted inference, make sure you have NVIDIA_API_KEY set as an environment variable.
|
|
348
|
+
|
|
349
|
+
```python
|
|
350
|
+
ingestor = (
|
|
351
|
+
ingestor.files(documents)
|
|
352
|
+
.extract(
|
|
353
|
+
# for self hosted NIMs, your URLs will depend on your NIM container DNS settings
|
|
354
|
+
page_elements_invoke_url="https://ai.api.nvidia.com/v1/cv/nvidia/nemotron-page-elements-v3",
|
|
355
|
+
graphic_elements_invoke_url="https://ai.api.nvidia.com/v1/cv/nvidia/nemotron-graphic-elements-v1",
|
|
356
|
+
ocr_invoke_url="https://ai.api.nvidia.com/v1/cv/nvidia/nemoretriever-ocr-v1",
|
|
357
|
+
table_structure_invoke_url="https://ai.api.nvidia.com/v1/cv/nvidia/nemotron-table-structure-v1"
|
|
358
|
+
)
|
|
359
|
+
.embed(
|
|
360
|
+
embed_invoke_url="https://integrate.api.nvidia.com/v1/embeddings",
|
|
361
|
+
model_name="nvidia/llama-nemotron-embed-1b-v2",
|
|
362
|
+
embed_modality="text",
|
|
363
|
+
)
|
|
364
|
+
.vdb_upload()
|
|
365
|
+
)
|
|
366
|
+
```
|
|
367
|
+
|
|
368
|
+
## Ray cluster setup
|
|
369
|
+
|
|
370
|
+
NeMo Retriever Library uses Ray Data for distributed ingestion and benchmarking. [NeMo Ray run guide](https://docs.nvidia.com/nemo/run/latest/guides/ray.html)
|
|
371
|
+
|
|
372
|
+
### Local Ray cluster with dashboard
|
|
373
|
+
|
|
374
|
+
To start a Ray cluster with the dashboard on a single machine use the following command.
|
|
375
|
+
|
|
376
|
+
```bash
|
|
377
|
+
ray start --head
|
|
378
|
+
```
|
|
379
|
+
|
|
380
|
+
Open `http://127.0.0.1:8265` in your browser for the Ray Dashboard, and run your NeMo Retriever Library pipeline on the same machine with `--ray-address auto` to attach to this cluster. [Connecting to a remote Ray cluster on Kubernetes](https://discuss.ray.io/t/connecting-to-remote-ray-cluster-on-k8s/7460)
|
|
381
|
+
|
|
382
|
+
### Single‑GPU cluster on multi‑GPU nodes
|
|
383
|
+
|
|
384
|
+
To restrict Ray to a single GPU on a multi‑GPU node use the following command.
|
|
385
|
+
|
|
386
|
+
```bash
|
|
387
|
+
CUDA_VISIBLE_DEVICES=0 ray start --head --num-gpus=1
|
|
388
|
+
```
|
|
389
|
+
Then run your pipeline as before with `--ray-address auto` so it connects to this single‑GPU Ray cluster. [NeMo Ray run guide](https://docs.nvidia.com/nemo/run/latest/guides/ray.html)
|
|
390
|
+
|
|
391
|
+
## Running multiple NIM instances on multi‑GPU hosts
|
|
392
|
+
|
|
393
|
+
### Resource heuristics (batch mode)
|
|
394
|
+
|
|
395
|
+
By default, batch mode computes resources using this order:
|
|
396
|
+
|
|
397
|
+
1. Auto-detected resources (Ray cluster if connected, otherwise local machine)
|
|
398
|
+
2. Environment variables
|
|
399
|
+
3. Explicit function arguments (highest precedence)
|
|
400
|
+
|
|
401
|
+
This means defaults are deterministic but easy to override when you need fixed behavior.
|
|
402
|
+
|
|
403
|
+
### Default behavior
|
|
404
|
+
|
|
405
|
+
- `cpu_count` / `gpu_count` are detected from Ray (`cluster_resources`) or local host.
|
|
406
|
+
- Worker heuristics:
|
|
407
|
+
- `page_elements_workers = gpu_count * page_elements_per_gpu`
|
|
408
|
+
- `detect_workers = gpu_count * ocr_per_gpu`
|
|
409
|
+
- `embed_workers = gpu_count * embed_per_gpu`
|
|
410
|
+
- minimum of `1` per stage
|
|
411
|
+
- Stage GPU defaults:
|
|
412
|
+
- If `gpu_count >= 2` and `concurrent_gpu_stage_count == 3`, uses high-overlap values for page-elements/OCR/embed.
|
|
413
|
+
- Otherwise uses `min(max_gpu_per_stage, gpu_count / concurrent_gpu_stage_count)`.
|
|
414
|
+
|
|
415
|
+
### Override variables
|
|
416
|
+
|
|
417
|
+
| Variable | Where to set | Meaning |
|
|
418
|
+
|---|---|---|
|
|
419
|
+
| `override_cpu_count`, `override_gpu_count` | function args | Highest-priority CPU/GPU override |
|
|
420
|
+
|
|
421
|
+
### Running multiple NIM service instances on multi-GPU hosts
|
|
422
|
+
|
|
423
|
+
### Start two stacks on separate GPUs
|
|
424
|
+
|
|
425
|
+
```bash
|
|
426
|
+
# GPU 0 stack
|
|
427
|
+
GPU_ID=0 \
|
|
428
|
+
PAGE_ELEMENTS_HTTP_PORT=8000 PAGE_ELEMENTS_GRPC_PORT=8001 PAGE_ELEMENTS_METRICS_PORT=8002 \
|
|
429
|
+
OCR_HTTP_PORT=8019 OCR_GRPC_PORT=8010 OCR_METRICS_PORT=8011 \
|
|
430
|
+
docker compose -p retriever-gpu0 up -d page-elements ocr
|
|
431
|
+
|
|
432
|
+
# GPU 1 stack
|
|
433
|
+
GPU_ID=1 \
|
|
434
|
+
PAGE_ELEMENTS_HTTP_PORT=8100 PAGE_ELEMENTS_GRPC_PORT=8101 PAGE_ELEMENTS_METRICS_PORT=8102 \
|
|
435
|
+
OCR_HTTP_PORT=8119 OCR_GRPC_PORT=8110 OCR_METRICS_PORT=8111 \
|
|
436
|
+
docker compose -p retriever-gpu1 up -d page-elements ocr
|
|
437
|
+
```
|
|
438
|
+
|
|
439
|
+
The `-p` project names create isolated stacks, `GPU_ID` pins each stack to a specific physical GPU, and distinct host ports avoid collisions between the services.
|
|
440
|
+
|
|
441
|
+
### Check and tear down stacks
|
|
442
|
+
|
|
443
|
+
To verify that both stacks are running use the following command.
|
|
444
|
+
|
|
445
|
+
```bash
|
|
446
|
+
docker compose -p retriever-gpu0 ps
|
|
447
|
+
docker compose -p retriever-gpu1 ps
|
|
448
|
+
```
|
|
449
|
+
|
|
450
|
+
To stop and remove both stacks use the following command.
|
|
451
|
+
|
|
452
|
+
```bash
|
|
453
|
+
docker compose -p retriever-gpu0 down
|
|
454
|
+
docker compose -p retriever-gpu1 down
|
|
455
|
+
```
|