NVIDIA NIM

Plugin information

NVIDIA NIM – Generative AI Plugin

With this plugin, you can use LLMs with NVIDIA NIM on Dataiku Mesh.

Capabilities

• Custom LLM connection for LLMs hosted in NVIDIA Cloud or self-hosted NIM containers
• Connect to a deployed NVIDIA container chat completion endpoint to use in Dataiku Prompt Studios, LLM recipes, and via the LLM Mesh Python APIs
• Connect to a deployed NVIDIA container embedding endpoint to use in the Dataiku Embed recipe for Retrieval Augmented Generation (RAG)

Limitations

• Must use Dataiku >= v13.0.0
• Must have setup an NVIDIA Authentication API Key
• Must deploy and maintain NVIDIA container models

Setup access to NVIDIA NIM

An NVIDIA AI Enterprise license is required to get access to NIM Models. With a license you can access NVIDIA models via NVIDIA Cloud or by using Docker containers from the NGC Catalog. Follow the steps here to set up a NIM container.

Setup the plugin in Dataiku

Install the Plugin using the installation guide. Then navigate to Plugins → Installed → NVIDIA NIM Plugin → Settings → Add Preset → Define a Preset Key.

Every user can set up the preset under their Dataiku profile → API Credentials → Credentials → Plugin credentials. Click the Edit icon and paste your NIM API key.

Setup the NIM Connection in Dataiku

• Go to Administration → Connections → New Connection → Custom LLM (LLM Mesh)

• Provide a connection name and select NVIDIA NIM Connector in the Plugin dropdown
• To add models – Click Add Model
  • Id: Unique name to identify model within Dataiku
  • Capability: Chat Completion / Text Embedding
  • Type: NVIDIA NIM LLM Connector
  • Keys Preset: Select the preset name previously defined in the plugin settings
  • Endpoint URL: Provide complete URL (examples below)
    • For chat completion – https://hostname/v1/chat/completions
    • For text embedding – https://hostname/v1/retrieval/nvidia/embeddings
  • Model Key: Provide the model key accepted by NIM containers. E.g. ‘google/gemma-2b’ or ‘NV-Embed-QA’
  • Input Type (This property only applies to Embedding models. By default, it is set to query): query or passage

Once the setup is complete, you can access models both in LLM Powered Visual Recipes, Prompt Studios and using Python and REST LLM Mesh APIs

Plugin Demo

Advancing healthcare with Dataiku and Nvidia

Curious how Dataiku and Nvidia partnership is advancing healthcare? Checkout the following blog post:

Advancing healthcare with Dataiku and Nvidia