fix: update nim api key (#262)

<!-- SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION
& AFFILIATES. All rights reserved. -->
<!-- SPDX-License-Identifier: Apache-2.0 -->

<!-- Thank you for contributing to Safe Synthesizer! -->

# Summary
<!-- Brief description of changes -->
- Standardizing `NIM_API_KEY` to `NSS_INFERENCE_KEY` and
`NIM_ENDPOINT_URL` with `NSS_INFERENCE_ENDPOINT ` in NSS to match what
it's called on the basic build docs.
- Set `NSS_INFERENCE_ENDPOINT = https://integrate.api.nvidia.com/v1` as
the default.
- Update the docs to clarify that `NSS_INFERENCE_KEY` and optional
`NSS_INFERENCE_ENDPOINT` need to be set when using CLI and SDK.
- Address feedbacks from the notebook:
   -  Remove addition of the inference end point since a default is set.
   - Remove "next steps" section.
- Remove `with_replace_pii`, and only add a comment about how it
disables PII replacement.

## Pre-Review Checklist

<!-- These checks should be completed before a PR is reviewed, -->
<!-- but you can submit a draft early to indicate that the issue is
being worked on. -->

Ensure that the following pass:

- [x] `make format && make check` or via prek validation.
- [x] `make test` passes locally
- [ ] `make test-e2e` passes locally
- [ ] `make test-ci-container` passes locally (recommended)
- [ ] GPU CI status check passes -- comment `/sync` on this PR to
trigger a run (auto-triggers on ready-for-review)

## Pre-Merge Checklist

<!-- These checks need to be completed before a PR is merged, -->
<!-- but as PRs often change significantly during review, -->
<!-- it's OK for them to be incomplete when review is first requested.
-->

- [ ] New or updated tests for any fix or new behavior
- [ ] Updated documentation for new features and behaviors, including
docstrings for API docs.

## Other Notes

<!-- Please add the issue number that should be closed when this PR is
merged. -->
- Closes #<issue>

---------

Signed-off-by: memadi <memadi@nvidia.com>
Signed-off-by: Marjan Emadi <memadi@nvidia.com>
Co-authored-by: Nina Xu <19981858+nina-xu@users.noreply.github.com>

Author: memadi-nv

README.md

@@ -235,17 +235,19 @@ Common values: `FLASHINFER`, `FLASH_ATTN`, `TORCH_SDPA`, `TRITON_ATTN`, `FLEX_AT
 ## NIM Integration
 
 Column classification uses a NIM/OpenAI-compatible endpoint to detect entity types
-in your data. The endpoint is configured via `NIM_ENDPOINT_URL`; if it is unset,
-classification is skipped and the pipeline falls back to default entity detection,
-logging an error and falling back rather than raising it to the user.
+in your data. `NSS_INFERENCE_ENDPOINT` defaults to `https://integrate.api.nvidia.com/v1`;
+override it to use a different endpoint.
+
+When using the CLI or Python SDK, set `NSS_INFERENCE_KEY` (and `NSS_INFERENCE_ENDPOINT` only if not
+using the default) so column classification can run.
 
 ### Local Endpoint
 
 To point to a locally hosted LLM:
 
 ```bash
-export NIM_ENDPOINT_URL="https://your-local-nim-endpoint"
-export NIM_API_KEY="your-api-key"  # pragma: allowlist secret
+export NSS_INFERENCE_ENDPOINT="https://your-local-nim-endpoint"
+export NSS_INFERENCE_KEY="your-api-key"  # pragma: allowlist secret
 ```
 
 ### Disable Classification

containers/Dockerfile.cuda

@@ -26,7 +26,7 @@
 #   -v HOST:CONTAINER   bind-mount data and HF cache (use absolute paths)
 #   -e HF_HOME=...      persist model downloads across container runs
 #   -e HF_TOKEN=...     HF token for gated models (Llama, Mistral, etc.)
-#   -e NIM_API_KEY=...  NIM endpoint key for PII classification (optional)
+#   -e NSS_INFERENCE_KEY=...  inference API key for PII column classification (optional)
 #
 # Interactive shell:
 #   docker run -it --gpus all --shm-size=1g \

containers/README.md

@@ -72,7 +72,7 @@ Key flags:
 - `-v HOST:CONTAINER` -- bind-mount data and HF cache; Docker requires absolute paths (use `$(pwd)` to expand relative ones)
 - `-e HF_HOME=...` -- persist model downloads across container runs
 - `-e HF_TOKEN=...` -- Hugging Face token for gated models (Llama, Mistral, etc.)
-- `-e NIM_API_KEY=...` -- NIM endpoint key for PII classification (optional, only when `NIM_ENDPOINT_URL` is set)
+- `-e NSS_INFERENCE_KEY=...` -- inference API key for PII column classification (optional; `NSS_INFERENCE_ENDPOINT` defaults to the NVIDIA integrate URL if unset)
 - `-e WANDB_API_KEY=...` -- WandB API key for experiment tracking (optional)
 - `--user "$(id -u):$(id -g)"` -- match host uid if you get "Permission denied" writing artifacts
 

docs/tutorials/safe-synthesizer-101.ipynb

@@ -59,11 +59,9 @@
       "metadata": {},
       "source": [
         "\n",
-        "### 🔑 Set the NIM API key and configure column classification\n",
+        "### 🔑 Set the inference API key for column classification\n",
         "\n",
-        "Setting `NIM_API_KEY` is optional but strongly recommended.\n",
-        "\n",
-        "NeMo Safe Synthesizer uses an LLM‑based column classifier to automatically infer column types and improve PII detection accuracy. To enable this feature, you must set both `NIM_ENDPOINT_URL` and `NIM_API_KEY`. You can obtain an API key from [build.nvidia.com](https://build.nvidia.com/settings/api-keys)\n"
+        "NeMo Safe Synthesizer uses an LLM‑based column classifier to automatically infer column types and improve PII detection accuracy. To enable this feature, set `NSS_INFERENCE_KEY` (the inference endpoint defaults to the NVIDIA integrate URL. You can obtain an API key from [build.nvidia.com](https://build.nvidia.com/settings/api-keys)). Setting this value is optional but strongly recommended.\n"
       ]
     },
     {
@@ -76,18 +74,14 @@
         "import os\n",
         "import getpass\n",
         "\n",
-        "# Set the NIM endpoint URL\n",
-        "os.environ[\"NIM_ENDPOINT_URL\"] = \"https://integrate.api.nvidia.com/v1\"\n",
-        "print(\"NIM_ENDPOINT_URL is set.\")\n",
-        "\n",
-        "# Setting NIM_API_KEY is optional but strongly recommended for PII replacement.\n",
-        "if \"NIM_API_KEY\" not in os.environ:\n",
-        "    os.environ[\"NIM_API_KEY\"] = getpass.getpass(\"Paste NIM API key (or press Enter to skip): \")\n",
-        "if os.environ.get(\"NIM_API_KEY\"):\n",
-        "    print(\"NIM_API_KEY is set\")\n",
+        "# Setting NSS_INFERENCE_KEY is optional but strongly recommended for PII replacement.\n",
+        "if \"NSS_INFERENCE_KEY\" not in os.environ:\n",
+        "    os.environ[\"NSS_INFERENCE_KEY\"] = getpass.getpass(\"Paste inference API key (or press Enter to skip): \")\n",
+        "if os.environ.get(\"NSS_INFERENCE_KEY\"):\n",
+        "    print(\"NSS_INFERENCE_KEY is set\")\n",
         "else:\n",
         "    print(\n",
-        "        \"NIM_API_KEY is not set. \"\n",
+        "        \"NSS_INFERENCE_KEY is not set. \"\n",
         "        \"We strongly recommend setting a key.\"\n",
         "    )"
       ]
@@ -144,7 +138,9 @@
       "source": [
         "from nemo_safe_synthesizer.sdk.library_builder import SafeSynthesizer\n",
         "\n",
-        "builder = SafeSynthesizer().with_data_source(df).with_replace_pii()\n",
+        "\n",
+        "# To disable PII replacement for the run, chain `.with_replace_pii(enable=False)` on the builder before `run()`.\n",
+        "builder = SafeSynthesizer().with_data_source(df)\n",
         "\n",
         "builder.run()\n",
         "results = builder.results"
@@ -209,30 +205,6 @@
         "        f.write(results.evaluation_report_html)\n",
         "    print(f\"The HTML evaluation report is saved in {report_path}.\")"
       ]
-    },
-    {
-      "cell_type": "markdown",
-      "id": "e9a19fcc",
-      "metadata": {},
-      "source": [
-        "### ➡️ Next Steps\n",
-        "\n",
-        "Now that you've completed your first Safe Synthesizer job, explore more advanced features:\n",
-        "\n",
-        "### Advanced Tutorials\n",
-        "\n",
-        "- [Differential Privacy Tutorial](https://aire.gitlab-master-pages.nvidia.com/microservices/nmp/latest/nemo-microservices/latest/safe-synthesizer/tutorials/differential-privacy.html) - Apply mathematical privacy guarantees\n",
-        "\n",
-        "- [PII Replacement Tutorial](https://aire.gitlab-master-pages.nvidia.com/microservices/nmp/latest/nemo-microservices/latest/safe-synthesizer/tutorials/pii-replacement.html) - Advanced PII detection and replacement\n",
-        "\n",
-        "\n",
-        "### Try These Next\n",
-        "\n",
-        "1. **Customize PII replacement**: Configure specific entity types and replacement strategies\n",
-        "2. **Enable differential privacy**: Add formal privacy guarantees with epsilon and delta parameters\n",
-        "3. **Tune generation parameters**: Experiment with temperature and sampling to understand how they impact quality and privacy scores. More on generation parameters [here](https://github.com/NVIDIA-NeMo/Safe-Synthesizer/blob/main/docs/user-guide/configuration.md#generation)\n",
-        "4. **Use your own data**: Replace the sample dataset with your sensitive data\n"
-      ]
     }
   ],
   "metadata": {

docs/user-guide/configuration.md

@@ -149,7 +149,7 @@ Key config parameters:
 
 | Field | Default | Description | Guidance |
 |-------|---------|-------------|----------|
-| `replace_pii.globals.classify.enable_classify` | `true` | Enable LLM-based column classification | Requires `NIM_ENDPOINT_URL`; set to `false` if no LLM endpoint is available |
+| `replace_pii.globals.classify.enable_classify` | `true` | Enable LLM-based PII column classification | When using the CLI, set `NSS_INFERENCE_KEY` (and optionally `NSS_INFERENCE_ENDPOINT`); set to `false` if no LLM endpoint is available |
 | `replace_pii.globals.classify.entities` | (see default list) | Entity types used for LLM-based column classification. Defaults to 15 types covering names, addresses, phone numbers, emails, SSN, national/tax IDs, and credit/debit cards -- see [PII Replacement](../product-overview/pii_replacement.md) and [`PiiReplacerConfig`][nemo_safe_synthesizer.config.replace_pii.PiiReplacerConfig] | Override to add or remove entity types from classification |
 | `replace_pii.globals.ner.ner_threshold` | `0.3` | GLiNER confidence threshold for NER detection | Lower to catch more entities (more false positives); raise to reduce false positives |
 

docs/user-guide/docker.md

@@ -123,8 +123,8 @@ docker run --gpus all --shm-size=1g \
 | Variable | Required | Purpose |
 |----------|----------|---------|
 | `HF_TOKEN` | For gated models | Hugging Face token for downloading gated models (Llama, Mistral, etc.). Get one at [hf.co/settings/tokens](https://huggingface.co/settings/tokens) |
-| `NIM_API_KEY` | For PII classification | API key for the NIM endpoint used by PII column classification. Only needed when `NIM_ENDPOINT_URL` is set |
-| `NIM_ENDPOINT_URL` | For PII classification | NIM/OpenAI-compatible endpoint URL for PII column classification |
+| `NSS_INFERENCE_KEY` | For PII classification | API key for `NSS_INFERENCE_ENDPOINT`. Set when using the CLI/SDK for column classification |
+| `NSS_INFERENCE_ENDPOINT` | For PII classification | NIM/OpenAI-compatible endpoint URL (default: `https://integrate.api.nvidia.com/v1`). Override for a custom endpoint |
 | `WANDB_API_KEY` | For experiment tracking | WandB API key. Only needed when `--wandb-mode online` is used |
 
 If `HF_TOKEN` is already stored in your HF cache (`~/.cache/huggingface/token`),

docs/user-guide/environment.md

@@ -27,8 +27,8 @@ are cached, and which network endpoints are used.
 | `NSS_DATASET_REGISTRY` | `--dataset-registry` | Dataset registry YAML path/URL |
 | `NSS_WANDB_MODE` | `--wandb-mode` | WandB mode (alias for `WANDB_MODE`) |
 | `NSS_WANDB_PROJECT` | `--wandb-project` | WandB project name (alias for `WANDB_PROJECT`) |
-| `NIM_ENDPOINT_URL` | -- | LLM endpoint for PII column classification |
-| `NIM_API_KEY` | -- | API key (optional -- only for direct endpoints) |
+| `NSS_INFERENCE_ENDPOINT` | -- | LLM endpoint for PII column classification (default: `https://integrate.api.nvidia.com/v1`) |
+| `NSS_INFERENCE_KEY` | -- | API key for the `NSS_INFERENCE_ENDPOINT` is required for column classification in both CLI and SDK. |
 | `NIM_MODEL_ID` | -- | Column classification model ID |
 | `LOCAL_FILES_ONLY` | -- | Set to `true` for offline mode (Unsloth, GLiNER) |
 | `SAFE_SYNTHESIZER_CPU_COUNT` | -- | NER CPU processes |
@@ -162,17 +162,19 @@ Common values: `FLASHINFER`, `FLASH_ATTN`, `TORCH_SDPA`, `TRITON_ATTN`,
 
 NIM endpoint, API keys, and CPU parallelism for PII detection.
 
-### `NIM_ENDPOINT_URL`
+### `NSS_INFERENCE_ENDPOINT`
 
-The NIM/OpenAI-compatible endpoint used for PII column classification. When
-unset, an error is logged and the pipeline falls back to NER-only detection.
-Set this to enable LLM-based column classification:
+The NIM/OpenAI-compatible endpoint used for PII column classification. Defaults
+to `https://integrate.api.nvidia.com/v1` when unset. Override for a custom endpoint:
 
 ```bash
-export NIM_ENDPOINT_URL="https://your-local-nim-endpoint"
-export NIM_API_KEY="your-api-key"  # pragma: allowlist secret
+export NSS_INFERENCE_ENDPOINT="https://your-llm-inference-endpoint"
+export NSS_INFERENCE_KEY="your-api-key"  # pragma: allowlist secret
 ```
 
+When using the CLI or SDK: for column classification to work, set `NSS_INFERENCE_KEY` (and
+`NSS_INFERENCE_ENDPOINT` only if you are not using the default URL).
+
 To disable column classification entirely instead of pointing it at a local
 endpoint, use the `replace_pii.globals.classify.enable_classify` config option.
 PII classify config is deeply nested -- use YAML or SDK:
@@ -201,10 +203,10 @@ PII classify config is deeply nested -- use YAML or SDK:
     )
     ```
 
-### `NIM_API_KEY`
+### `NSS_INFERENCE_KEY`
 
-API key for the NIM endpoint. Required when `NIM_ENDPOINT_URL` points to an
-authenticated endpoint.
+API key for the NSS inference endpoint. Required for PII column classification when using the
+CLI and SDK (with the default or custom `NSS_INFERENCE_ENDPOINT`).
 
 ### `NIM_MODEL_ID`
 

docs/user-guide/evaluating-data.md

@@ -89,8 +89,10 @@ or
 Could not perform classify, falling back to default entities.
 ```
 
-Fix: set entity types explicitly in your config, or check that `NIM_ENDPOINT_URL`
-is reachable. PII classify config is deeply nested -- use YAML or SDK:
+When `NSS_INFERENCE_KEY` is not set, the same log line is followed by guidance to set it (and a note that `NSS_INFERENCE_ENDPOINT` is optional with the default API). When the key is set, a traceback may be included to show the underlying API error.
+
+Fix: set entity types explicitly in your config, or when using the CLI ensure
+`NSS_INFERENCE_KEY` is set (and `NSS_INFERENCE_ENDPOINT` if not using the default). PII classify config is deeply nested -- use YAML or SDK:
 
 === "Config reference"
 

docs/user-guide/running.md

@@ -441,20 +441,22 @@ default in both the CLI and SDK. PII on by default means no config flag is neede
 
 ### LLM Column Classification
 
-To enable LLM-based PII column classification (optional), set the endpoint
-before running the pipeline. Any OpenAI-compatible inference endpoint
-works -- not just NVIDIA NIM:
+To enable LLM-based PII column classification (optional), set the API key
+before running the pipeline. The endpoint defaults to
+`https://integrate.api.nvidia.com/v1`; override `NSS_INFERENCE_ENDPOINT` for a
+custom OpenAI-compatible endpoint.
 
-```bash
-export NIM_ENDPOINT_URL="https://integrate.api.nvidia.com/v1"  # or your own OpenAI-compatible endpoint
+When using the CLI, set both for column classification:
 
-export NIM_API_KEY="your-api-key"  # pragma: allowlist secret  (optional -- only needed for direct endpoints, not inference gateways)
+```bash
+export NSS_INFERENCE_ENDPOINT="https://integrate.api.nvidia.com/v1"  # optional; this is the default
+export NSS_INFERENCE_KEY="your-api-key"  # pragma: allowlist secret  (required for column classification with the inference endpoint)
 ```
 
-When `NIM_ENDPOINT_URL` is unset, the classification step is attempted but
+PII column classification requires `NSS_INFERENCE_KEY` (and optionally `NSS_INFERENCE_ENDPOINT` if not using the default).
+When `NSS_INFERENCE_KEY` is unset, the classification step is attempted but
 falls back to NER-only detection (with an error log). No environment
-variables are required for NER-only PII replacement; column classification
-requires `NIM_ENDPOINT_URL`.
+variables are required for NER-only PII replacement.
 
 See [Configuration Reference -- Replacing PII](configuration.md#replacing-pii) for the full parameter reference.