Instructions to use nvidia/Nemotron-Labs-Diffusion-8B-Base with libraries, inference providers, notebooks, and local apps. Follow these links to get started.

Libraries

How to use nvidia/Nemotron-Labs-Diffusion-8B-Base with Transformers:

# Use a pipeline as a high-level helper
from transformers import pipeline

pipe = pipeline("text-generation", model="nvidia/Nemotron-Labs-Diffusion-8B-Base", trust_remote_code=True)
messages = [
    {"role": "user", "content": "Who are you?"},
]
pipe(messages)

# Load model directly
from transformers import AutoModel
model = AutoModel.from_pretrained("nvidia/Nemotron-Labs-Diffusion-8B-Base", trust_remote_code=True, dtype="auto")

Notebooks
Google Colab
Kaggle
Local Apps

vLLM

How to use nvidia/Nemotron-Labs-Diffusion-8B-Base with vLLM:

Install from pip and serve model

# Install vLLM from pip:
pip install vllm
# Start the vLLM server:
vllm serve "nvidia/Nemotron-Labs-Diffusion-8B-Base"
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:8000/v1/chat/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "nvidia/Nemotron-Labs-Diffusion-8B-Base",
		"messages": [
			{
				"role": "user",
				"content": "What is the capital of France?"
			}
		]
	}'

Use Docker

docker model run hf.co/nvidia/Nemotron-Labs-Diffusion-8B-Base

SGLang

How to use nvidia/Nemotron-Labs-Diffusion-8B-Base with SGLang:

Install from pip and serve model

# Install SGLang from pip:
pip install sglang
# Start the SGLang server:
python3 -m sglang.launch_server \
    --model-path "nvidia/Nemotron-Labs-Diffusion-8B-Base" \
    --host 0.0.0.0 \
    --port 30000
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:30000/v1/chat/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "nvidia/Nemotron-Labs-Diffusion-8B-Base",
		"messages": [
			{
				"role": "user",
				"content": "What is the capital of France?"
			}
		]
	}'

Use Docker images

docker run --gpus all \
    --shm-size 32g \
    -p 30000:30000 \
    -v ~/.cache/huggingface:/root/.cache/huggingface \
    --env "HF_TOKEN=<secret>" \
    --ipc=host \
    lmsysorg/sglang:latest \
    python3 -m sglang.launch_server \
        --model-path "nvidia/Nemotron-Labs-Diffusion-8B-Base" \
        --host 0.0.0.0 \
        --port 30000
# Call the server using curl (OpenAI-compatible API):
curl -X POST "http://localhost:30000/v1/chat/completions" \
	-H "Content-Type: application/json" \
	--data '{
		"model": "nvidia/Nemotron-Labs-Diffusion-8B-Base",
		"messages": [
			{
				"role": "user",
				"content": "What is the capital of France?"
			}
		]
	}'

Docker Model Runner
How to use nvidia/Nemotron-Labs-Diffusion-8B-Base with Docker Model Runner:
```
docker model run hf.co/nvidia/Nemotron-Labs-Diffusion-8B-Base
```

YongganFu commited on Feb 2

Commit

0a9534a

verified ·

1 Parent(s): 0820ac4

Update chat_template.jinja

Browse files

Files changed (1) hide show

chat_template.jinja +75 -85

chat_template.jinja CHANGED Viewed

@@ -1,99 +1,89 @@
-{#- Default system message if no system prompt is passed. #}
-{%- set default_system_message = '' %}
-{#- Begin of sequence token. #}
-{{- bos_token }}
-{#- Handle system prompt if it exists. #}
-{#- System prompt supports text content or text chunks. #}
-{%- if messages[0]['role'] == 'system' %}
-    {{- '[SYSTEM_PROMPT]' -}}
-    {%- if messages[0]['content'] is string %}
-        {{- messages[0]['content'] -}}
-    {%- else %}
-        {%- for block in messages[0]['content'] %}
-            {%- if block['type'] == 'text' %}
-                {{- block['text'] }}
-            {%- else %}
-                {{- raise_exception('Only text chunks are supported in system message contents.') }}
-            {%- endif %}
-        {%- endfor %}
     {%- endif %}
-    {{- '[/SYSTEM_PROMPT]' -}}
-    {%- set loop_messages = messages[1:] %}
 {%- else %}
-    {%- set loop_messages = messages %}
-    {%- if default_system_message != '' %}
-        {{- '[SYSTEM_PROMPT]' + default_system_message + '[/SYSTEM_PROMPT]' }}
     {%- endif %}
 {%- endif %}
-{#- Tools definition #}
-{%- set tools_definition = '' %}
-{%- set has_tools = false %}
-{%- if tools is defined and tools is not none and tools|length > 0 %}
-    {%- set has_tools = true %}
-    {%- set tools_definition = '[AVAILABLE_TOOLS]' + (tools| tojson) + '[/AVAILABLE_TOOLS]' %}
-    {{- tools_definition }}
-{%- endif %}
-{#- Handle conversation messages. #}
-{%- for message in loop_messages %}
-    {#- User messages supports text content or text and image chunks. #}
-    {%- if message['role'] == 'user' %}
-        {%- if message['content'] is string %}
-            {{- '[INST]' + message['content'] + '[/INST]' }}
-        {%- elif message['content'] | length > 0 %}
-            {{- '[INST]' }}
-            {%- if message['content'] | length == 2 %}
-                {%- set blocks = message['content'] | sort(attribute='type') %}
             {%- else %}
-                {%- set blocks = message['content'] %}
             {%- endif %}
-            {%- for block in blocks %}
-                {%- if block['type'] == 'text' %}
-                    {{- block['text'] }}
-                {%- elif block['type'] in ['image', 'image_url'] %}
-                    {{- '[IMG]' }}
-                {%- else %}
-                    {{- raise_exception('Only text, image and image_url chunks are supported in user message content.') }}
-                {%- endif %}
-            {%- endfor %}
-            {{- '[/INST]' }}
         {%- else %}
-            {{- raise_exception('User message must have a string or a list of chunks in content') }}
-        {%- endif %}
-    {#- Assistant messages supports text content or text and image chunks. #}
-    {%- elif message['role'] == 'assistant' %}
-        {%- if (message['content'] is none or message['content'] == '' or message['content']|length == 0) and (message['tool_calls'] is not defined or message['tool_calls'] is none or message['tool_calls']|length == 0) %}
-            {{- raise_exception('Assistant message must have a string or a list of chunks in content or a list of tool calls.') }}
         {%- endif %}
-        {%- if message['content'] is string %}
-            {{- message['content'] }}
-        {%- elif message['content'] | length > 0 %}
-            {%- for block in message['content'] %}
-                {%- if block['type'] == 'text' %}
-                    {{- block['text'] }}
                 {%- else %}
-                    {{- raise_exception('Only text chunks are supported in assistant message contents.') }}
                 {%- endif %}
             {%- endfor %}
         {%- endif %}
-        {%- if message['tool_calls'] is defined and message['tool_calls'] is not none and message['tool_calls']|length > 0 %}
-            {%- for tool in message['tool_calls'] %}
-                {%- set arguments = tool['function']['arguments'] %}
-                {%- if arguments is not string %}
-                    {%- set arguments = arguments|tojson|safe %}
-                {%- elif arguments == '' %}
-                    {%- set arguments = '{}' %}
-                {%- endif %}
-                {{- '[TOOL_CALLS]' + tool['function']['name'] + '[ARGS]' + arguments }}
-            {%- endfor %}
         {%- endif %}
-        {#- End of sequence token for each assistant messages. #}
-        {{- eos_token }}
-    {#- Tool messages only supports text content. #}
-    {%- elif message['role'] == 'tool' %}
-        {{- '[TOOL_RESULTS]' + message['content']|string + '[/TOOL_RESULTS]' }}
-    {#- Raise exception for unsupported roles. #}
-    {%- else %}
-        {{- raise_exception('Only user, assistant and tool roles are supported, got ' + message['role'] + '.') }}
     {%- endif %}
-{%- endfor %}

+{%- if tools %}
+    {{- '<s>system\n' }}
+    {%- if messages[0].role == 'system' %}
+        {{- messages[0].content + '\n\n' }}
     {%- endif %}
+    {{- "# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within [AVAILABLE_TOOLS][/AVAILABLE_TOOLS] XML tags:\n[AVAILABLE_TOOLS]" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n[/AVAILABLE_TOOLS]\n\nFor each function call, return a json object with function name and arguments within [TOOL_CALLS][SPECIAL_10] XML tags:\n[TOOL_CALLS]\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n[SPECIAL_10]</s>\n" }}
 {%- else %}
+    {%- if messages[0].role == 'system' %}
+        {{- '<s>system\n' + messages[0].content + '</s>\n' }}
     {%- endif %}
 {%- endif %}
+{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
+{%- for message in messages[::-1] %}
+    {%- set index = (messages|length - 1) - loop.index0 %}
+    {%- if ns.multi_step_tool and message.role == "user" and message.content is string and not(message.content.startswith('[TOOL_RESULTS]') and message.content.endswith('[/TOOL_RESULTS]')) %}
+        {%- set ns.multi_step_tool = false %}
+        {%- set ns.last_query_index = index %}
+    {%- endif %}
+{%- endfor %}
+{%- for message in messages %}
+    {%- if message.content is string %}
+        {%- set content = message.content %}
+    {%- else %}
+        {%- set content = '' %}
+    {%- endif %}
+    {%- if (message.role == "user") or (message.role == "system" and not loop.first) %}
+        {{- '<s>' + message.role + '\n' + content + '</s>' + '\n' }}
+    {%- elif message.role == "assistant" %}
+        {%- set reasoning_content = '' %}
+        {%- if message.reasoning_content is string %}
+            {%- set reasoning_content = message.reasoning_content %}
+        {%- else %}
+            {%- if '<SPECIAL_12>' in content %}
+                {%- set reasoning_content = content.split('<SPECIAL_12>')[0].rstrip('\n').split('<SPECIAL_11>')[-1].lstrip('\n') %}
+                {%- set content = content.split('<SPECIAL_12>')[-1].lstrip('\n') %}
+            {%- endif %}
+        {%- endif %}
+        {%- if loop.index0 > ns.last_query_index %}
+            {%- if loop.last or (not loop.last and reasoning_content) %}
+                {{- '<s>' + message.role + '\n<SPECIAL_11>\n' + reasoning_content.strip('\n') + '\n<SPECIAL_12>\n\n' + content.lstrip('\n') }}
             {%- else %}
+                {{- '<s>' + message.role + '\n' + content }}
             {%- endif %}
         {%- else %}
+            {{- '<s>' + message.role + '\n' + content }}
         {%- endif %}
+        {%- if message.tool_calls %}
+            {%- for tool_call in message.tool_calls %}
+                {%- if (loop.first and content) or (not loop.first) %}
+                    {{- '\n' }}
+                {%- endif %}
+                {%- if tool_call.function %}
+                    {%- set tool_call = tool_call.function %}
+                {%- endif %}
+                {{- '[TOOL_CALLS]\n{"name": "' }}
+                {{- tool_call.name }}
+                {{- '", "arguments": ' }}
+                {%- if tool_call.arguments is string %}
+                    {{- tool_call.arguments }}
                 {%- else %}
+                    {{- tool_call.arguments | tojson }}
                 {%- endif %}
+                {{- '}\n[SPECIAL_10]' }}
             {%- endfor %}
         {%- endif %}
+        {{- '</s>\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}
+            {{- '<s>user' }}
+        {%- endif %}
+        {{- '\n[TOOL_RESULTS]\n' }}
+        {{- content }}
+        {{- '\n[/TOOL_RESULTS]' }}
+        {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
+            {{- '</s>\n' }}
         {%- endif %}
     {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<s>assistant\n' }}
+    {%- if enable_thinking is defined and enable_thinking is false %}
+        {{- '<SPECIAL_11>\n\n<SPECIAL_12>\n\n' }}
+    {%- endif %}
+{%- endif %}