Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
52 changes: 52 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,58 @@ Each config is a frozen dataclass with a `.from_env()` constructor and an `env_v

ai4RAG uses [`docling-core`](https://github.com/docling-project/docling-core) for document representation and chunking. Documents are represented as `DoclingDocument` instances, and the `DoclingChunker` leverages docling's `HybridChunker` for structure-aware, token-aware chunking. `docling-core`, `openai`, and the vector store clients (`pymilvus` with Milvus Lite, `pgvector`, `asyncpg`) are all installed automatically with `ai4rag`.

## Running on Disconnected Clusters

If you are running ai4rag on a **disconnected cluster** (no internet access), you must pre-download the required ML models before execution.

### What to Pre-Download

ai4rag depends on models from two sources:

| Component | Size | Purpose | Environment Variable |
|-----------|------|---------|----------------------|
| **Docling artifacts** | ~300-400 MB | Document text extraction and optional OCR | `DOCLING_ARTIFACTS_PATH` |
| **HuggingFace models** | Variable | Embeddings and foundation models | `HF_HOME` |

### Setup Instructions

1. **On an internet-connected machine**, download the artifacts:
```bash
pip install 'ai4rag[text-extraction]'

# Trigger Docling model download
python -c "from docling.document_converter import DocumentConverter; \
converter = DocumentConverter(); \
converter.convert_document_string('/tmp/test.txt')"

# Pre-download HuggingFace models
export HF_HOME=/path/to/hf_cache
python -c "from transformers import AutoTokenizer; \
AutoTokenizer.from_pretrained('BAAI/bge-m3')"
```

2. **Transfer the artifacts to your disconnected cluster:**
```bash
rsync -av ~/.cache/docling/ cluster:/offline/docling/
rsync -av /path/to/hf_cache/ cluster:/offline/hf_cache/
```

3. **On the disconnected cluster, set environment variables:**
```bash
export DOCLING_ARTIFACTS_PATH=/offline/docling
export HF_HOME=/offline/hf_cache
export HF_HUB_OFFLINE=1 # Enforce offline mode
```

### Indexing Notebooks

The provided `maas_indexing_template.ipynb` includes a **"Prerequisites for Disconnected Clusters"** section with:
- Validation cells to verify artifact availability
- Configuration examples for custom model paths
- Step-by-step download instructions in the appendix

All indexing notebooks automatically detect and use pre-downloaded artifacts when `DOCLING_ARTIFACTS_PATH` is set.


## Quick start
1. [Prepare a MaaS client to integrate with your models.](#prepare-the-maas-client)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,103 @@
"outputs": [],
"source": "%pip install 'ai4rag[text-extraction]~={AI4RAG_VERSION}' | tail -n 1"
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"1## Prerequisites for Disconnected Clusters\n",
"\n",
"If you are running this notebook on a **disconnected cluster** (no internet access), you must pre-download the required ML models before execution.\n",
"\n",
"**Why:** This notebook uses **Docling** for document processing and **HuggingFace models** for embeddings. By default, these are downloaded on-the-fly from the internet. On disconnected clusters, you must prepare them offline.\n",
"\n",
"**What you need:**\n",
"- **Docling artifacts** (core + RapidOCR models) — ~300-400 MB\n",
"- **HuggingFace models** (if using custom embeddings) — size varies\n",
"\n",
"**Environment variables:**\n",
"- `DOCLING_ARTIFACTS_PATH`: Path to pre-downloaded docling artifacts\n",
"- `HF_HOME`: Directory containing pre-downloaded HuggingFace models\n",
"- `HF_HUB_OFFLINE`: Set to `\"1\"` to enforce offline mode\n",
"\n",
"👉 **Skip this section if you have internet access.** Jump to [Import required libraries](#Import-required-libraries).\n",
"\n",
"👉 **For setup instructions, see [Appendix: Downloading Models for Offline Use](#Appendix-Downloading-Models-for-Offline-Use) at the end of this notebook.**"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Configure Models for Disconnected Environments\n",
"\n",
"For disconnected clusters, configure the paths to your pre-downloaded model artifacts before proceeding.\n",
"\n",
"⚠️ **Required for disconnected clusters** | ✅ Skip if you have internet access"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"from pathlib import Path\n",
"\n",
"# For disconnected clusters: set these paths to your pre-downloaded artifacts\n",
"docling_path = os.getenv(\"DOCLING_ARTIFACTS_PATH\")\n",
"hf_home = os.getenv(\"HF_HOME\")\n",
"\n",
"if docling_path:\n",
" print(f\"✓ DOCLING_ARTIFACTS_PATH is set: {docling_path}\")\n",
" os.environ[\"HF_HUB_OFFLINE\"] = \"1\" # Enforce offline mode\n",
" print(\"✓ Offline mode enabled (HF_HUB_OFFLINE=1)\")\n",
"else:\n",
" print(\"ℹ️ DOCLING_ARTIFACTS_PATH not set. Models will be downloaded on-demand.\")\n",
" print(\" For disconnected clusters, see the appendix for setup instructions.\")\n",
"\n",
"if hf_home:\n",
" print(f\"✓ HF_HOME is set: {hf_home}\")\n",
"else:\n",
" print(\"ℹ️ HF_HOME not set. Will use default HuggingFace cache.\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Validate Offline Configuration\n",
"\n",
"If using a disconnected cluster, verify that model artifacts exist before starting extraction:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Validate artifacts (disconnected clusters only)\n",
"if docling_path:\n",
" artifacts_dir = Path(docling_path)\n",
" if not artifacts_dir.exists():\n",
" raise ValueError(f\"DOCLING_ARTIFACTS_PATH does not exist: {docling_path}\")\n",
" \n",
" print(f\"✓ Docling artifacts directory exists: {artifacts_dir}\")\n",
" \n",
" # Check for RapidOCR models (optional)\n",
" rapidocr_path = artifacts_dir / \"RapidOcr\" / \"onnx\" / \"PP-OCRv4\"\n",
" if rapidocr_path.exists():\n",
" det_models = list((rapidocr_path / \"det\").glob(\"*.onnx\"))\n",
" rec_models = list((rapidocr_path / \"rec\").glob(\"*.onnx\"))\n",
" print(f\"✓ RapidOCR models found: {len(det_models)} detection, {len(rec_models)} recognition\")\n",
" else:\n",
" print(\"ℹ️ RapidOCR models not found (OCR will be disabled if no internet)\")\n",
"else:\n",
" print(\"✓ Online mode: models will be downloaded as needed\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
Expand Down Expand Up @@ -156,7 +253,10 @@
"execution_count": null,
"metadata": {},
"outputs": [],
"source": "extracted_text_dir = step_output_dir / \"extracted_text\"\n\nextraction_result = extract_text(\n documents=result.to_dict()[\"documents\"],\n bucket=result.bucket,\n output_dir=extracted_text_dir,\n)\n\nprint(\n f\"Extracted {{extraction_result.processed_count}}/{{extraction_result.total_documents}} documents \"\n f\"({{extraction_result.error_count}} errors)\"\n)"
"source": [
"extracted_text_dir = step_output_dir / \"extracted_text\"extraction_result = extract_text( documents=result.to_dict()[\"documents\"], bucket=result.bucket, output_dir=extracted_text_dir, docling_artifacts_path=os.getenv(\"DOCLING_ARTIFACTS_PATH\"),\n",
")print( f\"Extracted {{extraction_result.processed_count}}/{{extraction_result.total_documents}} documents \" f\"({{extraction_result.error_count}} errors)\")"
]
},
{
"cell_type": "markdown",
Expand Down Expand Up @@ -280,6 +380,107 @@
"\n",
"This notebook successfully processed documents from S3 storage, extracted their text content using Docling, chunked the text into manageable pieces, and uploaded the embeddings to a vector store. The indexed documents are now ready for semantic search and retrieval in RAG applications."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"---\n",
"\n",
"## Appendix: Downloading Models for Offline Use\n",
"\n",
"This section guides you through pre-downloading and transferring ML models to a disconnected cluster.\n",
"\n",
"### Step 1: Download on an Internet-Connected Machine\n",
"\n",
"On a machine **with internet access**, install ai4rag and download artifacts:\n",
"\n",
"```bash\n",
"# Install ai4rag with text-extraction support\n",
"pip install 'ai4rag[text-extraction]'\n",
"```\n",
"\n",
"### Step 2: Download Docling Artifacts\n",
"\n",
"Create a Python script to trigger Docling model downloads:\n",
"\n",
"```python\n",
"from pathlib import Path\n",
"from docling.document_converter import DocumentConverter\n",
"import shutil\n",
"import tempfile\n",
"\n",
"# Create a temporary test file to trigger downloads\n",
"with tempfile.NamedTemporaryFile(suffix=\".txt\", delete=False, mode=\"w\") as f:\n",
" f.write(\"Test document for model download\")\n",
" test_file = f.name\n",
"\n",
"try:\n",
" # This will download all docling artifacts to ~/.cache/docling/\n",
" converter = DocumentConverter()\n",
" converter.convert_document_string(test_file)\n",
" print(\"✓ Docling artifacts downloaded to ~/.cache/docling/\")\n",
"finally:\n",
" Path(test_file).unlink()\n",
"\n",
"# Copy artifacts to your destination\n",
"import os\n",
"from_path = Path.home() / \".cache\" / \"docling\"\n",
"to_path = Path(\"/tmp/docling_artifacts\") # Change to your preferred location\n",
"to_path.mkdir(parents=True, exist_ok=True)\n",
"\n",
"if from_path.exists():\n",
" shutil.copytree(from_path, to_path / \"artifacts\", dirs_exist_ok=True)\n",
" print(f\"✓ Copied to {to_path / 'artifacts'}\")\n",
"```\n",
"\n",
"### Step 3: Download HuggingFace Models (If Needed)\n",
"\n",
"If using custom embedding models, pre-download them:\n",
"\n",
"```bash\n",
"export HF_HOME=/tmp/huggingface_cache\n",
"\n",
"# Download the model used in this notebook\n",
"python -c \"from transformers import AutoTokenizer, AutoModel; \\\n",
" tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-m3'); \\\n",
" model = AutoModel.from_pretrained('BAAI/bge-m3')\"\n",
"\n",
"# Or download any other model you're using\n",
"python -c \"from transformers import AutoTokenizer; \\\n",
" AutoTokenizer.from_pretrained('your-model-name')\"\n",
"```\n",
"\n",
"### Step 4: Transfer Artifacts to the Disconnected Cluster\n",
"\n",
"Copy the downloaded artifacts to your cluster:\n",
"\n",
"```bash\n",
"# Using rsync (recommended for large directories)\n",
"rsync -av /tmp/docling_artifacts/ user@cluster:/offline/models/docling/\n",
"rsync -av /tmp/huggingface_cache/ user@cluster:/offline/models/hf_cache/\n",
"\n",
"# Or using scp for smaller transfers\n",
"scp -r /tmp/docling_artifacts user@cluster:/offline/models/\n",
"scp -r /tmp/huggingface_cache user@cluster:/offline/models/\n",
"```\n",
"\n",
"### Step 5: Set Environment Variables on the Cluster\n",
"\n",
"On your **disconnected cluster**, set the environment variables before running the notebook:\n",
"\n",
"```bash\n",
"export DOCLING_ARTIFACTS_PATH=/offline/models/docling/artifacts\n",
"export HF_HOME=/offline/models/hf_cache\n",
"export HF_HUB_OFFLINE=1 # Enforce offline mode\n",
"```\n",
"\n",
"Or add these to your notebook environment (e.g., in OpenShift AI workbench settings).\n",
"\n",
"### Step 6: Verify and Run\n",
"\n",
"Run the notebook with the environment variables set. The validation cells will confirm that artifacts are accessible."
]
}
],
"metadata": {
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,95 @@
# -----------------------------------------------------------------------------
# Copyright IBM Corp. 2026
# SPDX-License-Identifier: Apache-2.0
# -----------------------------------------------------------------------------
"""Test that indexing notebook includes disconnected cluster documentation."""
import json
from pathlib import Path

import pytest


@pytest.fixture
def indexing_notebook() -> dict:
"""Load the MaaS indexing template notebook."""
notebook_path = Path(__file__).parents[4] / "ai4rag/assets_generator/notebook_templates/maas_indexing_template.ipynb"
with open(notebook_path) as f:
return json.load(f)


class TestDisconnectedClusterDocumentation:
"""Verify disconnected cluster support is documented in indexing notebook."""

def test_notebook_has_disconnected_cluster_prerequisites_section(self, indexing_notebook):
"""Notebook must include a 'Prerequisites for Disconnected Clusters' section."""
sources = [
"".join(cell["source"]) if isinstance(cell["source"], list) else cell["source"]
for cell in indexing_notebook["cells"]
if cell["cell_type"] == "markdown"
]

assert any("Prerequisites for Disconnected Clusters" in s for s in sources), \
"Notebook missing 'Prerequisites for Disconnected Clusters' section"

def test_notebook_explains_model_requirements(self, indexing_notebook):
"""Notebook must explain which models are required offline."""
sources = [
"".join(cell["source"]) if isinstance(cell["source"], list) else cell["source"]
for cell in indexing_notebook["cells"]
if cell["cell_type"] == "markdown"
]
full_text = "\n".join(sources)

assert "Docling" in full_text, "Notebook must mention Docling artifacts"
assert "DOCLING_ARTIFACTS_PATH" in full_text, "Notebook must mention DOCLING_ARTIFACTS_PATH env var"
assert "HF_HOME" in full_text, "Notebook must mention HF_HOME env var"

def test_notebook_includes_environment_setup_code(self, indexing_notebook):
"""Notebook must include code cells to set environment variables."""
code_sources = [
"".join(cell["source"]) if isinstance(cell["source"], list) else cell["source"]
for cell in indexing_notebook["cells"]
if cell["cell_type"] == "code"
]
full_code = "\n".join(code_sources)

assert "DOCLING_ARTIFACTS_PATH" in full_code, "Notebook must include code to configure DOCLING_ARTIFACTS_PATH"
assert "HF_HOME" in full_code, "Notebook must include code to configure HF_HOME"

def test_extract_text_includes_docling_artifacts_path(self, indexing_notebook):
"""The extract_text() call must include docling_artifacts_path parameter."""
code_sources = [
"".join(cell["source"]) if isinstance(cell["source"], list) else cell["source"]
for cell in indexing_notebook["cells"]
if cell["cell_type"] == "code"
]

extract_text_calls = [s for s in code_sources if "extract_text(" in s]
assert len(extract_text_calls) > 0, "Notebook must include an extract_text() call"

assert any("docling_artifacts_path" in call for call in extract_text_calls), \
"extract_text() call must include docling_artifacts_path parameter"

def test_notebook_includes_appendix_with_download_instructions(self, indexing_notebook):
"""Notebook must include appendix with offline download instructions."""
sources = [
"".join(cell["source"]) if isinstance(cell["source"], list) else cell["source"]
for cell in indexing_notebook["cells"]
if cell["cell_type"] == "markdown"
]
full_text = "\n".join(sources)

assert "Appendix" in full_text, "Notebook must include an appendix section"
assert "Download" in full_text, "Appendix must include download instructions"

def test_notebook_cells_are_valid_json(self, indexing_notebook):
"""All cells must be valid JSON with required fields."""
for i, cell in enumerate(indexing_notebook["cells"]):
assert "cell_type" in cell, f"Cell {i} missing cell_type"
assert cell["cell_type"] in ("code", "markdown"), f"Cell {i} has invalid type: {cell['cell_type']}"
assert "source" in cell, f"Cell {i} missing source"
assert "metadata" in cell, f"Cell {i} missing metadata"

if cell["cell_type"] == "code":
assert "execution_count" in cell, f"Code cell {i} missing execution_count"
assert "outputs" in cell, f"Code cell {i} missing outputs"
Loading