feat: Support passing in tool calls with OpenAI chat format when doing SFT - #1181
Conversation
Signed-off-by: Jiaqi Zeng <jiaqiz@nvidia.com>
ð WalkthroughWalkthroughAdds optional tool-calling support across SFT: propagates tools from datasets through preprocessing into chat template formatting; preserves heterogeneous sample keys via a new PreservingDataset and JSON fallback loader; supports chat template loading from .jinja files; extends config with OpenAI-format options and observability settings. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
actor User
participant Loader as Dataset Loader
participant Preproc as sft_preprocessor
participant Formatter as get_formatted_message_log
participant Tok as Tokenizer.apply_chat_template
rect rgb(245,248,255)
note over Loader: Load samples (HF path â fallback JSONL)
User->>Loader: Configure OpenAI-format dataset (tool_key optional)
Loader-->>User: Samples (may include tools)
end
rect rgb(245,255,245)
User->>Preproc: Sample datum
Preproc->>Formatter: messages, roles, tools=datum.tools?
note right of Formatter: Builds template kwargs incl. tools (if provided)
Formatter->>Tok: apply_chat_template(messages, kwargs...)
Tok-->>Formatter: formatted text
Formatter-->>Preproc: formatted conversation
end
sequenceDiagram
autonumber
participant Config as Config (sft.yaml)
participant Utils as get_tokenizer
participant FS as Filesystem
Config->>Utils: chat_template value
alt chat_template endswith ".jinja"
Utils->>FS: Read template file
FS-->>Utils: Template string
Utils-->>Config: Tokenizer with file-based template
else default/inline/None
Utils-->>Config: Tokenizer with existing handling
end
Estimated code review effortðŊ 3 (Moderate) | âąïļ ~25 minutes Pre-merge checks and finishing touchesâ Failed checks (3 warnings)
â Passed checks (3 passed)
âĻ Finishing touches
𧊠Generate unit tests
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and canât be posted inline due to platform limitations.
â ïļ Outside diff range comments (3)
nemo_rl/data/datasets/response_datasets/__init__.py (1)
58-65: Do not set config defaults in code; pass through tool_key without a fallback.Per nemo_rl guidelines, avoid non-None defaults in code. Using data_config.get("tool_key", "tools") sets a hidden default. Let YAML (examples/configs) carry the default, or pass None and let the dataset handle it explicitly.
Apply this diff:
- base_dataset = OpenAIFormatDataset( + base_dataset = OpenAIFormatDataset( data_config["train_data_path"], data_config["val_data_path"], data_config["chat_key"], data_config["system_key"], data_config["system_prompt"], - data_config.get("tool_key", "tools"), + data_config.get("tool_key"), # no in-code default )nemo_rl/data/datasets/response_datasets/oai_format_dataset.py (2)
104-118: Do not encode a default for tool_key in library code.In nemo_rl, defaults belong in YAML, not code. tool_key should default to None here and be set in config.
- tool_key: str | None = "tools", + tool_key: str | None = None,Also remove any code that assumes "tools" when None; treat None as âno tools columnâ.
172-190: Use explicit validation instead of assert for data checks.assert can be stripped with -O. Raise a clear error instead.
- assert messages[-1]["role"] == "assistant" + if not messages or messages[-1].get("role") != "assistant": + raise ValueError("OpenAIFormatDataset expects the last message to be from 'assistant'.")
ð§đ Nitpick comments (4)
nemo_rl/algorithms/utils.py (1)
242-254: Support .jinja file path more robustly (existence, encoding) and update docs.
- Guard path handling, read with UTFâ8, and raise on missing file.
- Also document file-path support in the function docstring.
Apply this diff:
elif tokenizer_config["chat_template"].lower() == "default": print("Using tokenizer's default chat template") - elif tokenizer_config["chat_template"].endswith(".jinja"): - # Load template from file - template_path = tokenizer_config["chat_template"] - print(f"Loading chat template from file: {template_path}") - with open(template_path, "r") as f: - tokenizer.chat_template = f.read() + elif isinstance(tokenizer_config["chat_template"], str) and tokenizer_config["chat_template"].endswith(".jinja"): + # Load template from file + from pathlib import Path + template_path = Path(tokenizer_config["chat_template"]).expanduser() + if not template_path.is_file(): + raise FileNotFoundError(f"Chat template file not found: {template_path}") + print(f"Loading chat template from file: {template_path}") + tokenizer.chat_template = template_path.read_text(encoding="utf-8")And amend the docstring section describing chat_template (Lines 166-171) to include âor a path to a .jinja fileâ.
examples/configs/sft.yaml (1)
159-167: Surface the recommended default for tool_key in YAML instead of code.To avoid hidden defaults in code, show the default in exemplar configs. Consider uncommenting tool_key with "tools".
- ## OpenAI format specific configs + ## OpenAI format specific configs # train_data_path: "/path/to/train.jsonl" # Path to training data # val_data_path: "/path/to/val.jsonl" # Path to validation data # chat_key: "messages" # Key for messages in the data # system_key: null # Key for system message (optional) # system_prompt: null # Default system prompt (optional) - # tool_key: null # Key for tools in the data (optional, defaults to "tools" if not specified) + tool_key: "tools" # Key for tools in the data (optional; recommended default)nemo_rl/data/datasets/response_datasets/oai_format_dataset.py (2)
23-84: PreservingDataset polish: remove unused args; minor Ruff nits.
- map(..., *args) unused; drop it or name as _ to silence ARG002.
- Docstrings are good; keep them concise.
- def map(self, function: Callable, *args, **kwargs) -> "PreservingDataset": + def map(self, function: Callable, **kwargs) -> "PreservingDataset": """Apply a function to each sample in the dataset.
118-132: Narrow exception handling; add UTFâ8 when reading JSONL; reduce noisy prints.
- Avoid catching bare Exception (BLE001). Catch specific errors from datasets and JSON.
- Use encoding="utf-8".
- Consider logging at debug level instead of unconditional prints.
- except (TypeError, ValueError, Exception) as e: + except (TypeError, ValueError) as e: + # Fall back when schema heterogeneity breaks HF dataset assumptions # Fallback to custom loading for heterogeneous schemas @@ - with open(train_ds_path, "r") as f: + with open(train_ds_path, "r", encoding="utf-8") as f: train_data = [json.loads(line) for line in f] @@ - with open(val_ds_path, "r") as f: + with open(val_ds_path, "r", encoding="utf-8") as f: val_data = [json.loads(line) for line in f]Optionally, import datasets as ds and also catch (OSError, json.JSONDecodeError). Replace print(...) with logging.debug(...) if a logger is available.
Also applies to: 132-143, 144-161
ð Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
ð Files selected for processing (6)
examples/configs/sft.yaml(2 hunks)examples/run_sft.py(1 hunks)nemo_rl/algorithms/utils.py(1 hunks)nemo_rl/data/datasets/response_datasets/__init__.py(1 hunks)nemo_rl/data/datasets/response_datasets/oai_format_dataset.py(3 hunks)nemo_rl/data/llm_message_utils.py(6 hunks)
𧰠Additional context used
ð Path-based instructions (3)
**/*.py
ð CodeRabbit inference engine (CODING_GUIDELINES.md)
**/*.py: Follow the Google Python Style Guide for all Python code
Target Python 3.12+ for all Python code in NeMo-RL
Indent Python code with 4 spaces; do not use tabs
Python filenames should be snake_case (e.g., some_file.py)
Class names should be PascalCase
Function and method names should be snake_case
Local variable names should be snake_case; if starting with a number, prefix with k (e.g., k_99th_percentile)
Global variables should be UPPER_SNAKE_CASE and prefixed with G_ (e.g., G_MY_GLOBAL)
Constants should be UPPER_SNAKE_CASE
Avoid shadowing variables declared in an outer scope
Initialize all externally visible members of a class in the constructor
For public interfaces used outside a file, prefer docstrings over comments
Use comments mainly for code within a function or interfaces local to a file
Commented-out code must include a nearby comment explaining usage and why it is commented out; otherwise remove before merging
Use Google-style docstrings for classes and functions (Sphinx-parseable)
Avoid using reflection when functionality can be easily achieved without it
Limit except clauses to the smallest specific set of exceptions possible
For duck-typing via try/except, keep the try body minimal and use else for main logic
Add the NVIDIA copyright header (with current year) at the top of all Python files, excluding tests/ and test-only scripts
Files:
nemo_rl/algorithms/utils.pynemo_rl/data/llm_message_utils.pynemo_rl/data/datasets/response_datasets/oai_format_dataset.pynemo_rl/data/datasets/response_datasets/__init__.pyexamples/run_sft.py
nemo_rl/**/*.py
ð CodeRabbit inference engine (CODING_GUIDELINES.md)
nemo_rl/**/*.py: Do not set non-None configuration defaults in code; YAML is the single source of truth for defaults
Access required config attributes directly (e.g., policy_cfg["precision"]) and assume presence; do not introduce hidden defaults
Express configuration optionality via TypedDict using typing.NotRequired
When adding a new config key to a TypedDict subclass, document the keyâs purpose, valid values/types, and recommended default in code
For any class or function decorated with @ray.remote, add '# pragma: no cover' on the class/def line (and on remote functions)
Files:
nemo_rl/algorithms/utils.pynemo_rl/data/llm_message_utils.pynemo_rl/data/datasets/response_datasets/oai_format_dataset.pynemo_rl/data/datasets/response_datasets/__init__.py
examples/configs/*.yaml
ð CodeRabbit inference engine (CODING_GUIDELINES.md)
examples/configs/*.yaml: Exemplar configs under examples/configs/.yaml must include documented defaults
When adding a new config key, reflect its recommended default in exemplar YAMLs under examples/configs/.yaml
Files:
examples/configs/sft.yaml
𧎠Code graph analysis (2)
nemo_rl/algorithms/utils.py (2)
tests/unit/models/generation/test_vllm_generation.py (1)
tokenizer(236-239)tests/unit/utils/test_native_checkpoint.py (1)
tokenizer(113-116)
nemo_rl/data/datasets/response_datasets/oai_format_dataset.py (2)
nemo_rl/data/interfaces.py (1)
TaskDataSpec(53-86)nemo_rl/data/datasets/response_datasets/response_dataset.py (1)
add_messages_key(72-80)
ðŠ Ruff (0.13.1)
nemo_rl/data/datasets/response_datasets/oai_format_dataset.py
55-57: Avoid specifying long messages outside the exception class
(TRY003)
62-64: Avoid specifying long messages outside the exception class
(TRY003)
69-69: Unused method argument: args
(ARG002)
132-132: Do not catch blind exception: Exception
(BLE001)
â° Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (3)
- GitHub Check: Lint check
- GitHub Check: Post automodel integration comment / Comment on PR
- GitHub Check: Post submodule check comment / Comment on PR
ð Additional comments (5)
examples/configs/sft.yaml (1)
28-29: Good: inline note about .jinja path support.Comment clarifies new behavior and aligns with utils.get_tokenizer. No action needed.
examples/run_sft.py (1)
69-77: LGTM: tools are forwarded to formatting safely.Passing tools via datum_dict.get("tools") keeps behavior unchanged when absent.
Please confirm your chat templates (default or custom) actually reference the âtoolsâ variable. If not, passing it is harmless, but older Transformers may reject unknown kwargs. See follow-up suggestion in llm_message_utils to add a safe fallback.
nemo_rl/data/llm_message_utils.py (3)
427-433: LGTM: safe handling for None content in multimodal/tool-call turns.Returning [] for None content avoids errors when extracting images.
533-537: Double-check add_generation_prompt for 'tool' role.Typically add_generation_prompt is only for user turns. Including tool turns might change formatting for some templates.
Would you like to keep add_generation_prompt on 'tool' messages? If unintentional, restrict to user only:
- "add_generation_prompt": add_generation_prompt - and message["role"] in ["user", "tool"], + "add_generation_prompt": add_generation_prompt and message["role"] == "user",
531-543: Verify transformers version; add fallback only if supporting <4.45transformers.PreTrainedTokenizerBase.apply_chat_template accepts a "tools" kwarg starting in transformers v4.45 (used in Llama/Qwen/Gemma docs). If this repo pins transformers >= 4.45 you can leave the code as-is; if you must support older transformers or thirdâparty tokenizers that may lack the kwarg, apply the TypeError retry below.
Location: nemo_rl/data/llm_message_utils.py (lines 531â543)
- formatted_message: str = tokenizer.apply_chat_template( # type: ignore - message_log_strs[: i + 1], **template_kwargs - ) + try: + formatted_message: str = tokenizer.apply_chat_template( # type: ignore + message_log_strs[: i + 1], **template_kwargs + ) + except TypeError as e: + # Retry without tools for older tokenizers that don't accept this kwarg + if "tools" in template_kwargs: + template_kwargs.pop("tools", None) + formatted_message = tokenizer.apply_chat_template( # type: ignore + message_log_strs[: i + 1], **template_kwargs + ) + else: + raise
|
@yuki-97 could you help review |
yuki-97
left a comment
yuki-97
left a comment
There was a problem hiding this comment.
thanks @HeyyyyyyG for adding this! overall LGTM, left some comments.
Signed-off-by: Jiaqi Zeng <jiaqiz@nvidia.com>
Signed-off-by: Jiaqi Zeng <jiaqiz@nvidia.com>
github-actions
Bot
added
the
Documentation
|
Made changes according to comments, and updated docs. @yuki-97 @terrykong could you take a look again? |
yuki-97
left a comment
yuki-97
left a comment
There was a problem hiding this comment.
thanks @HeyyyyyyG , LGTM! @terrykong can you also take a review?
terrykong
left a comment
terrykong
left a comment
There was a problem hiding this comment.
lgtm @HeyyyyyyG . just one minor comment
Co-authored-by: Terry Kong <terrycurtiskong@gmail.com> Signed-off-by: Jiaqi Zeng <49757268+HeyyyyyyG@users.noreply.github.com>
|
@HeyyyyyyG to resolve the lint issue, you need to update your branch |
Co-authored-by: Terry Kong <terrycurtiskong@gmail.com> Signed-off-by: Jiaqi Zeng <49757268+HeyyyyyyG@users.noreply.github.com>
Signed-off-by: Jiaqi Zeng <jiaqiz@nvidia.com>
|
@terrykong fixed |
|
great, thanks @jiemingz . enqueued |
âĶg SFT (NVIDIA-NeMo#1181) Signed-off-by: Jiaqi Zeng <jiaqiz@nvidia.com> Signed-off-by: Jiaqi Zeng <49757268+HeyyyyyyG@users.noreply.github.com> Co-authored-by: Terry Kong <terrycurtiskong@gmail.com>
âĶg SFT (NVIDIA-NeMo#1181) Signed-off-by: Jiaqi Zeng <jiaqiz@nvidia.com> Signed-off-by: Jiaqi Zeng <49757268+HeyyyyyyG@users.noreply.github.com> Co-authored-by: Terry Kong <terrycurtiskong@gmail.com> Signed-off-by: yuanhangs <yuanhangs@nvidia.com>
4 participants
What does this PR do ?
Support passing in tool calls with OpenAI chat format when doing SFT.
Support passing in a .jinja file as the chat template.
Issues
List issues that this PR closes (syntax):
close #1002
Usage
# Add a code snippet demonstrating how to use thisBefore your PR is "Ready for review"
Pre checks:
Additional Information
Summary by CodeRabbit
New Features
Bug Fixes
Documentation