diff --git a/README.md b/README.md index 601eb58e..2a7a01c4 100644 --- a/README.md +++ b/README.md @@ -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) diff --git a/ai4rag/assets_generator/notebook_templates/maas_indexing_template.ipynb b/ai4rag/assets_generator/notebook_templates/maas_indexing_template.ipynb index d1b5bb60..7c682a9d 100644 --- a/ai4rag/assets_generator/notebook_templates/maas_indexing_template.ipynb +++ b/ai4rag/assets_generator/notebook_templates/maas_indexing_template.ipynb @@ -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": {}, @@ -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", @@ -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": { diff --git a/tests/unit/ai4rag/assets_generator/test_notebook_disconnected_clusters.py b/tests/unit/ai4rag/assets_generator/test_notebook_disconnected_clusters.py new file mode 100644 index 00000000..df42a73b --- /dev/null +++ b/tests/unit/ai4rag/assets_generator/test_notebook_disconnected_clusters.py @@ -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"