Backport PR jupyterlab#579: Add nvidia provider

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, 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.
 

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'

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`

packages/jupyter-ai-magics/jupyter_ai_magics/__init__.py

@@ -22,6 +22,7 @@
     BedrockChatProvider,
     BedrockProvider,
     ChatAnthropicProvider,
+    ChatNVIDIAProvider,
     ChatOpenAIProvider,
     CohereProvider,
     GPT4AllProvider,

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`
@@ -809,3 +810,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")

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"