Add nvidia provider (jupyterlab#579)

* Add NVIDIA chat and embeddings providers * Fix typo * Some NVIDIA provider documentation * Update docs and naming for NVIDIA * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update docs/source/users/index.md Co-authored-by: Jason Weill <[email protected]> * Update index.md - add note about api key access * Remove nvidia embeddings model * Remove nvidia embedding provider * Remove nvidia embeddings provider * Mentions conda install instructions in docs * Uninstallation with conda * Makes NVIDIA package dependency optional, updates docs * Revert "Uninstallation with conda" This reverts commit 8082fa4. * Revert "Mentions conda install instructions in docs" This reverts commit 348ae2e. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Jason Weill <[email protected]> Co-authored-by: Jason Weill <[email protected]>
dbelgrod · Feb 1, 2024 · 7b767c8 · 7b767c8
1 parent 6f3dc39
commit 7b767c8
Show file tree

Hide file tree

Showing 6 changed files with 30 additions and 2 deletions.
diff --git a/README.md b/README.md
@@ -10,7 +10,7 @@ in JupyterLab and the Jupyter Notebook. More specifically, Jupyter AI offers:
   This works anywhere the IPython kernel runs (JupyterLab, Jupyter Notebook, Google Colab, Kaggle, VSCode, etc.).
 * A native chat UI in JupyterLab that enables you to work with generative AI as a conversational assistant.
 * Support for a wide range of generative model providers, including AI21, Anthropic, AWS, Cohere,
-  Hugging Face, and OpenAI.
+  Hugging Face, NVIDIA, and OpenAI.
 * Local model support through GPT4All, enabling use of generative AI models on consumer grade machines
   with ease and privacy.
 

diff --git a/docs/source/index.md b/docs/source/index.md
@@ -8,7 +8,7 @@ in JupyterLab and the Jupyter Notebook. More specifically, Jupyter AI offers:
   This works anywhere the IPython kernel runs (JupyterLab, Jupyter Notebook, Google Colab, VSCode, etc.).
 * A native chat UI in JupyterLab that enables you to work with generative AI as a conversational assistant.
 * Support for a wide range of generative model providers and models
-  (AI21, Anthropic, Cohere, Hugging Face, OpenAI, SageMaker, etc.).
+  (AI21, Anthropic, Cohere, Hugging Face, OpenAI, SageMaker, NVIDIA, etc.).
 
 <img src="_static/jupyter-ai-screenshot.png"
     alt='A screenshot of Jupyter AI showing the chat interface and the magic commands'

diff --git a/docs/source/users/index.md b/docs/source/users/index.md
@@ -133,6 +133,7 @@ Jupyter AI supports the following model providers:
 | ERNIE-Bot           | `qianfan`            | `QIANFAN_AK`, `QIANFAN_SK` | `qianfan`                       |
 | GPT4All             | `gpt4all`            | N/A                        | `gpt4all`                       |
 | Hugging Face Hub    | `huggingface_hub`    | `HUGGINGFACEHUB_API_TOKEN` | `huggingface_hub`, `ipywidgets`, `pillow` |
+| NVIDIA              | `nvidia-chat`        | `NVIDIA_API_KEY`           | `langchain_nvidia_ai_endpoints` |
 | OpenAI              | `openai`             | `OPENAI_API_KEY`           | `openai`                        |
 | OpenAI (chat)       | `openai-chat`        | `OPENAI_API_KEY`           | `openai`                        |
 | SageMaker           | `sagemaker-endpoint` | N/A                        | `boto3`                         |
@@ -150,6 +151,8 @@ You need the `pillow` Python package to use Hugging Face Hub's text-to-image mod
 
 You can find a list of Hugging Face's models at [https://huggingface.co/models](https://huggingface.co/models).
 
+To use NVIDIA models, create a free account with the [NVIDIA NGC service](https://catalog.ngc.nvidia.com/), which hosts AI solution catalogs, containers, models, and more. Navigate to Catalog > [AI Foundation Models](https://catalog.ngc.nvidia.com/ai-foundation-models), and select a model with an API endpoint. Click "API" on the model's detail page, and click "Generate Key". Save this key, and set it as the environment variable `NVIDIA_API_KEY` to access any of the model endpoints.
+
 SageMaker endpoint names are created when you deploy a model. For more information, see
 ["Create your endpoint and deploy your model"](https://docs.aws.amazon.com/sagemaker/latest/dg/realtime-endpoints-deployment.html)
 in the SageMaker documentation.
@@ -515,6 +518,7 @@ We currently support the following language model providers:
 - `bedrock-chat`
 - `cohere`
 - `huggingface_hub`
+- `nvidia-chat`
 - `openai`
 - `openai-chat`
 - `sagemaker-endpoint`

diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/__init__.py b/packages/jupyter-ai-magics/jupyter_ai_magics/__init__.py
@@ -22,6 +22,7 @@
     BedrockChatProvider,
     BedrockProvider,
     ChatAnthropicProvider,
+    ChatNVIDIAProvider,
     ChatOpenAIProvider,
     CohereProvider,
     GPT4AllProvider,

diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py b/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py
@@ -38,6 +38,7 @@
     OpenAI,
     SagemakerEndpoint,
 )
+from langchain_nvidia_ai_endpoints import ChatNVIDIA
 
 # this is necessary because `langchain.pydantic_v1.main` does not include
 # `ModelMetaclass`, as it is not listed in `__all__` by the `pydantic.main`
@@ -858,3 +859,23 @@ class QianfanProvider(BaseProvider, QianfanChatEndpoint):
     model_id_key = "model_name"
     pypi_package_deps = ["qianfan"]
     auth_strategy = MultiEnvAuthStrategy(names=["QIANFAN_AK", "QIANFAN_SK"])
+
+
+class ChatNVIDIAProvider(BaseProvider, ChatNVIDIA):
+    id = "nvidia-chat"
+    name = "NVIDIA"
+    models = [
+        "playground_llama2_70b",
+        "playground_nemotron_steerlm_8b",
+        "playground_mistral_7b",
+        "playground_nv_llama2_rlhf_70b",
+        "playground_llama2_13b",
+        "playground_steerlm_llama_70b",
+        "playground_llama2_code_13b",
+        "playground_yi_34b",
+        "playground_mixtral_8x7b",
+        "playground_neva_22b",
+        "playground_llama2_code_34b",
+    ]
+    model_id_key = "model"
+    auth_strategy = EnvAuthStrategy(name="NVIDIA_API_KEY")
diff --git a/packages/jupyter-ai-magics/pyproject.toml b/packages/jupyter-ai-magics/pyproject.toml
@@ -41,6 +41,7 @@ all = [
     "gpt4all",
     "huggingface_hub",
     "ipywidgets",
+    "langchain_nvidia_ai_endpoints",
     "pillow",
     "openai~=1.6.1",
     "boto3",
@@ -61,6 +62,7 @@ amazon-bedrock = "jupyter_ai_magics:BedrockProvider"
 anthropic-chat = "jupyter_ai_magics:ChatAnthropicProvider"
 amazon-bedrock-chat = "jupyter_ai_magics:BedrockChatProvider"
 qianfan = "jupyter_ai_magics:QianfanProvider"
+nvidia-chat = "jupyter_ai_magics:ChatNVIDIAProvider"
 
 [project.entry-points."jupyter_ai.embeddings_model_providers"]
 bedrock = "jupyter_ai_magics:BedrockEmbeddingsProvider"