`nemo_automodel.components.datasets.llm.agent_chat`#

Multi-turn agent SFT dataset adapter.

Loads function-calling chat datasets where each example contains tool definitions and a multi-turn message list that includes tool calls and tool responses, then renders them through the tokenizer’s chat template with answer_only_loss_mask=True so that user and tool tokens are excluded from the loss.

Two input schemas are accepted:

Swift / chatml messages schema (used by AI-ModelScope/function-calling-chatml)::

{
    "tools": "[{...openai tool schema...}]",
    "messages": [
        {"role": "user",          "content": "..."},
        {"role": "tool_call",     "content": "{\"name\": ..., \"arguments\": ...}"},
        {"role": "tool_response", "content": "..."},
        {"role": "assistant",     "content": "..."}
    ]
}

ShareGPT conversations schema (used by llamafactory/glaive_toolcall_en and similar)::

{
    "tools": "[...]",
    "conversations": [
        {"from": "human",         "value": "..."},
        {"from": "function_call", "value": "..."},
        {"from": "observation",   "value": "..."},
        {"from": "gpt",           "value": "..."}
    ]
}

Consecutive tool_call entries are merged into a single assistant message with parallel tool_calls; the following tool_response entries are paired with those calls in order.

Module Contents#

Functions#

`_json_load_if_str`	Parse `value` as JSON if it is a string, otherwise return as-is.
`_sharegpt_to_chatml`	Convert ShareGPT `{from, value}` turns to chatml `{role, content}`.
`_convert_messages`	Convert chatml-style agent messages to OpenAI chat-completions format.
`_mask_labels_to_last_turn`	Restrict the loss to the final assistant turn (`mask_history`).
`_format_example`	Render one agent example into tokenized `input_ids` / `labels`.
`make_agent_chat_dataset`	Load a multi-turn function-calling SFT dataset.

Data#

`logger`
`_VALID_MESSAGE_ROLES`
`_SHAREGPT_ROLE_MAP`

API#

nemo_automodel.components.datasets.llm.agent_chat.logger#: ‘getLogger(…)’

nemo_automodel.components.datasets.llm.agent_chat._VALID_MESSAGE_ROLES#: None

nemo_automodel.components.datasets.llm.agent_chat._SHAREGPT_ROLE_MAP#: None

nemo_automodel.components.datasets.llm.agent_chat._json_load_if_str(value: Any) → Any#: Parse value as JSON if it is a string, otherwise return as-is.

nemo_automodel.components.datasets.llm.agent_chat._sharegpt_to_chatml( conversations: List[Dict[str, Any]], ) → List[Dict[str, Any]]#: Convert ShareGPT {from, value} turns to chatml {role, content}.

nemo_automodel.components.datasets.llm.agent_chat._convert_messages( messages: List[Dict[str, Any]], example_id: Optional[Union[int, str]] = None, ) → List[Dict[str, Any]]#

Convert chatml-style agent messages to OpenAI chat-completions format.

Consecutive tool_call entries collapse into one assistant message with parallel tool_calls. When the preceding emitted turn is an assistant message without tool_calls (e.g. an assistant content that reasons before calling tools), the tool_calls attach to that message instead of creating a second consecutive assistant turn — this preserves the natural single-turn shape the model produces at inference. tool_response (or tool) entries that follow are paired with those tool_call ids in order.

Parameters:

messages – chatml-style turns with roles in _VALID_MESSAGE_ROLES.
example_id – Optional identifier used to derive unique tool_call ids.

Returns:

A list of OpenAI-format messages suitable for apply_chat_template.

nemo_automodel.components.datasets.llm.agent_chat._mask_labels_to_last_turn( labels: List[int], ignore_index: int = -100, ) → List[int]#

Restrict the loss to the final assistant turn (mask_history).

labels come from :func:format_chat_template with every assistant turn supervised; non-assistant tokens are already ignore_index. Because the chat template renders each assistant message as a single contiguous span, supervised tokens form one maximal run per assistant turn separated by ignore_index runs. This keeps only the last such run and masks every earlier supervised token in place.

Parameters:

labels – per-token labels (ignore_index marks unsupervised tokens).
ignore_index – the value marking unsupervised tokens.

Returns:

The same list, mutated so only the final supervised run is kept.

nemo_automodel.components.datasets.llm.agent_chat._format_example( example: Dict[str, Any], tokenizer, eos_token_id: int, pad_token_id: int, seq_length: Optional[int] = None, padding: Union[str, bool] = False, truncation: Union[str, bool] = False, mask_reasoning_content: bool = False, train_on_last_turn_only: bool = False, ) → Dict[str, List[int]]#: Render one agent example into tokenized input_ids / labels.

nemo_automodel.components.datasets.llm.agent_chat.make_agent_chat_dataset( tokenizer, *, dataset_name: Optional[str] = None, path: Optional[Union[str, List[str]]] = None, split: str = 'train', seq_length: Optional[int] = None, limit_dataset_samples: Optional[int] = None, padding: Union[str, bool] = False, truncation: Union[str, bool] = False, mask_reasoning_content: bool = False, train_on_last_turn_only: bool = False, ) → nemo_automodel.components.datasets.lazy_mapped_dataset.LazyMappedDataset#

Load a multi-turn function-calling SFT dataset.

Exactly one of dataset_name (HuggingFace Hub id) or path (local JSON/JSONL file or list of files) must be provided. The loaded examples are lazily rendered through the tokenizer’s chat template; tool/user tokens are excluded from the loss via answer_only_loss_mask=True.

Parameters:

tokenizer – HuggingFace tokenizer with a chat template.
dataset_name – HF Hub dataset id, e.g. llamafactory/glaive_toolcall_en.
path – Local JSON/JSONL file path or list of paths.
split – Dataset split (only used with dataset_name).
seq_length – Optional max sequence length for the tokenizer.
limit_dataset_samples – If set, keep only the first N examples.
padding – Padding strategy forwarded to the tokenizer.
truncation – Truncation strategy forwarded to the tokenizer.
mask_reasoning_content – If True, exclude assistant reasoning_content (thinking) tokens from the loss while still rendering them into the prompt. Requires a chat template that emits reasoning_content. Defaults to False, which trains on reasoning tokens like any other assistant content.
train_on_last_turn_only – If True, supervise only the final assistant turn of each dialogue (mask_history); all earlier assistant turns are excluded from the loss. Defaults to False, which supervises every assistant turn.

Returns:

A LazyMappedDataset yielding dicts with input_ids, labels and attention_mask ready for the trainer collator.

nemo_automodel.components.datasets.llm.agent_chat#

Module Contents#

Functions#

Data#

API#

`nemo_automodel.components.datasets.llm.agent_chat`#