Skip to main content

Nvidia NIM

https://docs.api.nvidia.com/nim/reference/

tip

We support ALL Nvidia NIM models, just set model=nvidia_nim/<any-model-on-nvidia_nim> as a prefix when sending litellm requests

API Key​

# env variable
os.environ['NVIDIA_API_KEY']

Sample Usage​

from litellm import completion
import os

os.environ['NVIDIA_API_KEY'] = ""
response = completion(
    model="nvidia_nim/meta/llama3-70b-instruct",
    messages=[
        {
            "role": "user",
            "content": "What's the weather like in Boston today in Fahrenheit?",
        }
    ],
    temperature=0.2,        # optional
    top_p=0.9,              # optional
    frequency_penalty=0.1,  # optional
    presence_penalty=0.1,   # optional
    max_tokens=10,          # optional
    stop=["\n\n"],          # optional
)
print(response)

Sample Usage - Streaming​

from litellm import completion
import os

os.environ['NVIDIA_API_KEY'] = ""
response = completion(
    model="nvidia_nim/meta/llama3-70b-instruct",
    messages=[
        {
            "role": "user",
            "content": "What's the weather like in Boston today in Fahrenheit?",
        }
    ],
    stream=True,
    temperature=0.2,        # optional
    top_p=0.9,              # optional
    frequency_penalty=0.1,  # optional
    presence_penalty=0.1,   # optional
    max_tokens=10,          # optional
    stop=["\n\n"],          # optional
)

for chunk in response:
    print(chunk)

Function/Tool Calling​

from litellm import completion

# set env
os.environ['NVIDIA_API_KEY'] = ""

tools = [
    {
        "type": "function",
        "function": {
            "name": "get_current_weather",
            "description": "Get the current weather in a given location",
            "parameters": {
                "type": "object",
                "properties": {
                    "location": {
                        "type": "string",
                        "description": "The city and state, e.g. San Francisco, CA",
                    },
                    "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
                },
                "required": ["location"],
            },
        },
    }
]
messages = [{"role": "user", "content": "What's the weather like in Boston today?"}]

response = completion(
    model="nvidia/meta/llama-3.1-70b-instruct",
    messages=messages,
    tools=tools,
    tool_choice="auto",
)
# Add any assertions, here to check response args
print(response)
assert isinstance(response.choices[0].message.tool_calls[0].function.name, str)
assert isinstance(
    response.choices[0].message.tool_calls[0].function.arguments, str
)

Forcing Tool Use​

If you want LLM to use a specific tool to answer the user’s question

You can do this by specifying the tool in the tool_choice field like so:

response = completion(
    os.environ['NVIDIA_API_KEY'] = ""
    messages=messages,
    tools=tools,
    tool_choice={"type": "tool", "name": "get_weather"},
)

Usage - Vision​

from litellm import completion

# set env
os.environ['NVIDIA_API_KEY'] = ""

def encode_image(image_path):
    import base64

    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode("utf-8")


image_path = "nvidia-picasso.jpg"
# Getting the base64 string
base64_image = encode_image(image_path)
response = litellm.completion(
    model="nvidia/microsoft/phi-3-vision-128k-instruct",
    messages=[
        {
            "role": "user",
            "content": [
                {"type": "text", "text": "Whats in this image?"},
                {
                    "type": "image_url",
                    "image_url": {
                        "url": "data:image/jpeg;base64," + base64_image
                    },
                },
            ],
        }
    ],
)
print(f"\nResponse: {response}")

Usage - embedding​

import litellm
import os

response = litellm.embedding(
    model="nvidia_nim/nvidia/nv-embedqa-e5-v5",               # add `nvidia_nim/` prefix to model so litellm knows to route to Nvidia NIM
    input=["good morning from litellm"],
    encoding_format = "float", 
    user_id = "user-1234",

    # Nvidia NIM Specific Parameters
    input_type = "passage", # Optional
    truncate = "NONE" # Optional
)
print(response)

Usage - LiteLLM Proxy Server​

Here's how to call an Nvidia NIM Endpoint with the LiteLLM Proxy Server

  1. Modify the config.yaml 

    model_list:
      - model_name: my-model
        litellm_params:
          model: nvidia_nim/<your-model-name>  # add nvidia_nim/ prefix to route as Nvidia NIM provider
          api_key: api-key                 # api key to send your model

  1. Start the proxy 

    $ litellm --config /path/to/config.yaml

  2. Send Request to LiteLLM Proxy Server

    import openai
    client = openai.OpenAI(
        api_key="sk-1234",             # pass litellm proxy key, if you're using virtual keys
        base_url="http://0.0.0.0:4000" # litellm-proxy-base url
    )

    response = client.chat.completions.create(
        model="my-model",
        messages = [
            {
                "role": "user",
                "content": "what llm are you"
            }
        ],
    )

    print(response)

Supported Models - 💥 ALL Nvidia NIM Models Supported!​

We support ALL nvidia_nim models, just set nvidia_nim/ as a prefix when sending completion requests

Model NameFunction Call
nvidia/nemotron-4-340b-rewardcompletion(model="nvidia_nim/nvidia/nemotron-4-340b-reward", messages)
01-ai/yi-largecompletion(model="nvidia_nim/01-ai/yi-large", messages)
aisingapore/sea-lion-7b-instructcompletion(model="nvidia_nim/aisingapore/sea-lion-7b-instruct", messages)
databricks/dbrx-instructcompletion(model="nvidia_nim/databricks/dbrx-instruct", messages)
google/gemma-7bcompletion(model="nvidia_nim/google/gemma-7b", messages)
google/gemma-2bcompletion(model="nvidia_nim/google/gemma-2b", messages)
google/codegemma-1.1-7bcompletion(model="nvidia_nim/google/codegemma-1.1-7b", messages)
google/codegemma-7bcompletion(model="nvidia_nim/google/codegemma-7b", messages)
google/recurrentgemma-2bcompletion(model="nvidia_nim/google/recurrentgemma-2b", messages)
ibm/granite-34b-code-instructcompletion(model="nvidia_nim/ibm/granite-34b-code-instruct", messages)
ibm/granite-8b-code-instructcompletion(model="nvidia_nim/ibm/granite-8b-code-instruct", messages)
mediatek/breeze-7b-instructcompletion(model="nvidia_nim/mediatek/breeze-7b-instruct", messages)
meta/codellama-70bcompletion(model="nvidia_nim/meta/codellama-70b", messages)
meta/llama2-70bcompletion(model="nvidia_nim/meta/llama2-70b", messages)
meta/llama3-8bcompletion(model="nvidia_nim/meta/llama3-8b", messages)
meta/llama3-70bcompletion(model="nvidia_nim/meta/llama3-70b", messages)
microsoft/phi-3-medium-4k-instructcompletion(model="nvidia_nim/microsoft/phi-3-medium-4k-instruct", messages)
microsoft/phi-3-mini-128k-instructcompletion(model="nvidia_nim/microsoft/phi-3-mini-128k-instruct", messages)
microsoft/phi-3-mini-4k-instructcompletion(model="nvidia_nim/microsoft/phi-3-mini-4k-instruct", messages)
microsoft/phi-3-small-128k-instructcompletion(model="nvidia_nim/microsoft/phi-3-small-128k-instruct", messages)
microsoft/phi-3-small-8k-instructcompletion(model="nvidia_nim/microsoft/phi-3-small-8k-instruct", messages)
mistralai/codestral-22b-instruct-v0.1completion(model="nvidia_nim/mistralai/codestral-22b-instruct-v0.1", messages)
mistralai/mistral-7b-instructcompletion(model="nvidia_nim/mistralai/mistral-7b-instruct", messages)
mistralai/mistral-7b-instruct-v0.3completion(model="nvidia_nim/mistralai/mistral-7b-instruct-v0.3", messages)
mistralai/mixtral-8x7b-instructcompletion(model="nvidia_nim/mistralai/mixtral-8x7b-instruct", messages)
mistralai/mixtral-8x22b-instructcompletion(model="nvidia_nim/mistralai/mixtral-8x22b-instruct", messages)
mistralai/mistral-largecompletion(model="nvidia_nim/mistralai/mistral-large", messages)
nvidia/nemotron-4-340b-instructcompletion(model="nvidia_nim/nvidia/nemotron-4-340b-instruct", messages)
seallms/seallm-7b-v2.5completion(model="nvidia_nim/seallms/seallm-7b-v2.5", messages)
snowflake/arcticcompletion(model="nvidia_nim/snowflake/arctic", messages)
upstage/solar-10.7b-instructcompletion(model="nvidia_nim/upstage/solar-10.7b-instruct", messages)