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`           | N/A                             |
 | 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, etc. Navigate to Catalog > [AI Foundation Models](https://catalog.ngc.nvidia.com/ai-foundation-models) > (Model with API endpoint). Select the API option and click Generate Key > save as `NVIDIA_API_KEY`.
+
 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

@@ -7,6 +7,7 @@
     CohereEmbeddingsProvider,
     GPT4AllEmbeddingsProvider,
     HfHubEmbeddingsProvider,
+    NVIDIAEmbeddingsProvider,
     OpenAIEmbeddingsProvider,
     QianfanEmbeddingsEndpointProvider,
 )
@@ -22,6 +23,7 @@
     BedrockChatProvider,
     BedrockProvider,
     ChatAnthropicProvider,
+    ChatNVIDIAProvider,
     ChatOpenAIProvider,
     CohereProvider,
     GPT4AllProvider,

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

@@ -16,6 +16,7 @@
     OpenAIEmbeddings,
     QianfanEmbeddingsEndpoint,
 )
+from langchain_nvidia_ai_endpoints import NVIDIAEmbeddings
 
 
 class BaseEmbeddingsProvider(BaseModel):
@@ -139,3 +140,11 @@ class QianfanEmbeddingsEndpointProvider(
     model_id_key = "model"
     pypi_package_deps = ["qianfan"]
     auth_strategy = MultiEnvAuthStrategy(names=["QIANFAN_AK", "QIANFAN_SK"])
+
+
+class NVIDIAEmbeddingsProvider(BaseEmbeddingsProvider, NVIDIAEmbeddings):
+    id = "nvidia"
+    name = "NVIDIA"
+    models = ["playground_nvolveqa_40k"]
+    model_id_key = "model"
+    auth_strategy = EnvAuthStrategy(name="NVIDIA_API_KEY")

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

@@ -32,6 +32,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`
@@ -770,3 +771,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

@@ -27,6 +27,7 @@ dependencies = [
     "typing_extensions>=4.5.0",
     "click~=8.0",
     "jsonpath-ng>=1.5.3,<2",
+    "langchain_nvidia_ai_endpoints",
 ]
 
 [project.optional-dependencies]
@@ -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"
@@ -69,6 +71,7 @@ gpt4all = "jupyter_ai_magics:GPT4AllEmbeddingsProvider"
 huggingface_hub = "jupyter_ai_magics:HfHubEmbeddingsProvider"
 openai = "jupyter_ai_magics:OpenAIEmbeddingsProvider"
 qianfan = "jupyter_ai_magics:QianfanEmbeddingsEndpointProvider"
+nvidia = "jupyter_ai_magics:NVIDIAEmbeddingsProvider"
 
 [tool.hatch.version]
 source = "nodejs"