diff --git a/src/post_training/chat_templates/templates/olmo3-instruct-sft.jinja b/src/post_training/chat_templates/templates/olmo3-instruct-sft.jinja index 5d9bd97..49036b0 100644 --- a/src/post_training/chat_templates/templates/olmo3-instruct-sft.jinja +++ b/src/post_training/chat_templates/templates/olmo3-instruct-sft.jinja @@ -6,7 +6,7 @@ You are a helpful function-calling AI assistant. ' -}}{%- if tools is none or (t ' -}}{%- elif message['role'] == 'user' -%}{{- '<|im_start|>user ' + message['content'] + '<|im_end|> ' -}}{%- elif message['role'] == 'assistant' -%}{{- '<|im_start|>assistant -' -}}{%- generation -%}{%- if message.get('content', none) is not none -%}{{- message['content'] -}}{%- endif -%}{%- if message.get('function_calls', none) is not none -%}{{- '' + message['function_calls'] + '' -}}{% elif message.get('tool_calls', none) is not none %}{{- '' -}}{%- for tool_call in message['tool_calls'] %}{%- if tool_call is mapping and tool_call.get('function', none) is not none %}{%- set args = tool_call['function']['arguments'] -%}{%- set ns = namespace(arguments_list=[]) -%}{%- for key, value in args.items() -%}{%- set ns.arguments_list = ns.arguments_list + [key ~ '=' ~ (value | tojson)] -%}{%- endfor -%}{%- set arguments = ns.arguments_list | join(', ') -%}{{- tool_call['function']['name'] + '(' + arguments + ')' -}}{%- if not loop.last -%}{{ ' +' -}}{%- generation -%}{%- if message.get('content', none) is not none -%}{{- message['content'] -}}{%- endif -%}{%- if message.get('function_calls', none) is not none -%}{{- '' + message['function_calls'] + '' -}}{% elif message.get('tool_calls', none) is not none %}{{- '' -}}{%- for tool_call in message['tool_calls'] %}{%- if tool_call is mapping and tool_call.get('function', none) is not none %}{%- set args = tool_call['function']['arguments'] -%}{%- if args is string -%}{{- raise_exception('olmo3-instruct-sft needs tool_call arguments as a dict, not a JSON string. For Dolci rows, use transform: null; this template renders their function_calls directly.') -}}{%- endif -%}{%- set ns = namespace(arguments_list=[]) -%}{%- for key, value in args.items() -%}{%- set ns.arguments_list = ns.arguments_list + [key ~ '=' ~ (value | tojson)] -%}{%- endfor -%}{%- set arguments = ns.arguments_list | join(', ') -%}{{- tool_call['function']['name'] + '(' + arguments + ')' -}}{%- if not loop.last -%}{{ ' ' }}{%- endif -%}{% else %}{{- tool_call -}}{%- endif %}{%- endfor %}{{- '' -}}{%- endif -%}{%- if not loop.last -%}{{- '<|im_end|>' + ' ' -}}{%- else -%}{{- eos_token -}}{%- endif -%}{%- endgeneration -%}{%- elif message['role'] == 'environment' -%}{{- '<|im_start|>environment ' + message['content'] + '<|im_end|> diff --git a/src/post_training/data/loader.py b/src/post_training/data/loader.py index 4d38677..5bec1be 100644 --- a/src/post_training/data/loader.py +++ b/src/post_training/data/loader.py @@ -14,7 +14,7 @@ from pathlib import Path from typing import TYPE_CHECKING -from datasets import Dataset, Features, concatenate_datasets, load_dataset +from datasets import Dataset, Features, List, concatenate_datasets, load_dataset from post_training.data.transforms import get_transform @@ -99,6 +99,47 @@ def _resample_to_size(ds: Dataset, target_n: int, seed: int) -> Dataset: return concatenate_datasets(copies) +def _align_message_fields(datasets: list[Dataset], num_proc: int) -> list[Dataset]: + """Give every dataset the same message fields, so that they concatenate. + + A transformed entry has the message fields of the cast features (e.g. + ``tool_calls``); a native entry keeps its own (e.g. only ``role`` and + ``content``, or Dolci's ``functions`` and ``function_calls``). Each dataset + gets the union of all fields, with null for a field it does not have, so no + data is lost. If two datasets give one field different types, nothing + changes, and ``concatenate_datasets`` reports the conflict. + """ + structs = [ + ds.features["messages"].feature + for ds in datasets + if isinstance(ds.features.get("messages"), List) + and isinstance(ds.features["messages"].feature, dict) + ] + union: dict = {} + for struct in structs: + for name, feature in struct.items(): + if union.setdefault(name, feature) != feature: + return datasets + if all(list(struct) == list(union) for struct in structs): + return datasets + + aligned = [] + for ds in datasets: + messages = ds.features.get("messages") + if ( + not isinstance(messages, List) + or not isinstance(messages.feature, dict) + or list(messages.feature) == list(union) + ): + aligned.append(ds) + continue + features = ds.features.copy() + features["messages"] = List(dict(union)) + # The identity map re-encodes the rows: datasets fills missing fields with null. + aligned.append(ds.map(lambda row: row, features=features, num_proc=num_proc)) + return aligned + + def load_and_mix_datasets( config: DataConfig, dataset_filter_fn: Callable[[Dataset, int], Dataset] | None = None, @@ -239,6 +280,6 @@ def load_and_mix_datasets( if len(resampled_datasets) == 1: mixed = resampled_datasets[0] else: - mixed = concatenate_datasets(resampled_datasets) + mixed = concatenate_datasets(_align_message_fields(resampled_datasets, num_proc)) return mixed.shuffle(seed=seed) diff --git a/src/post_training/data/transforms.py b/src/post_training/data/transforms.py index a26ed75..7ecef91 100644 --- a/src/post_training/data/transforms.py +++ b/src/post_training/data/transforms.py @@ -12,9 +12,14 @@ from __future__ import annotations +import ast +import io import json +import keyword import re -from collections.abc import Callable +import tokenize +import warnings +from collections.abc import Callable, Iterator from typing import Any # --------------------------------------------------------------------------- @@ -555,3 +560,350 @@ def helpsteer3_principle(example: dict[str, Any]) -> dict[str, Any]: {"role": "assistant", "content": example["response"]}, ] } + + +# --------------------------------------------------------------------------- +# Dolci tool use +# --------------------------------------------------------------------------- + +# Dolci tool rows open with this system prompt. It describes the OLMo 3 +# / format, which the transform replaces. +_DOLCI_TOOL_PROMPT = ( + "You are a helpful function-calling AI assistant. You are provided with function signatures " + "within XML tags. You may call one or more functions to assist with " + "the user query. Output any function calls within XML tags. " + "Don't make assumptions about what values to plug into functions." +) +_DOLCI_PYTHONIC_HINT = " Format calls like: tool_name(query=...)" + +_JSON_NAMES = {"true": True, "false": False, "null": None} +_CALL_RE = re.compile(r"\s*([^()]+?)\((.*)\)\s*", re.DOTALL) +# Python keywords used as argument names (``from="EUR"``) are a syntax error; +# they are renamed before parsing and restored afterwards. +_KEYWORD_PREFIX = "_kw_" +# Some blocks lost a "\" before a call name: "\reverse(...)" was decoded as a +# carriage return + "everse(...)". The line break and the letter are restored. +_ESCAPE_LETTERS = {"\r": "r", "\t": "t", "\x0c": "f", "\x08": "b", "\x0b": "v", "\x07": "a"} +_LOST_ESCAPE_RE = re.compile(r"\)([\r\t\x0c\x08\x0b\x07])(?=[\w.]+\()") +# Nodes allowed in an argument value that is an expression over literals, +# e.g. ``12*60``, ``[0.03]*15`` or ``[c for c in "noon"]``. A bare name that is +# not bound by a comprehension is a placeholder and fails with NameError. +_SAFE_NODES = ( + ast.Expression, ast.Constant, ast.List, ast.Tuple, ast.Dict, ast.Set, ast.BinOp, ast.UnaryOp, + ast.Add, ast.Sub, ast.Mult, ast.Div, ast.FloorDiv, ast.Mod, ast.USub, ast.UAdd, ast.ListComp, + ast.comprehension, ast.Name, ast.Call, ast.Subscript, ast.Slice, ast.Load, ast.Store, + ast.Compare, ast.cmpop, +) # fmt: skip +# Most items a value may build with repetition or range(), e.g. [0.5] * 12 or +# list(range(60)); a dataset value must not allocate unbounded memory. +_MAX_ITEMS = 10_000 + + +class _DolciCallError(ValueError): + """A Dolci call line cannot be converted to a JSON tool call.""" + + +class _JsonNames(ast.NodeTransformer): + """Replace JSON ``true`` / ``false`` / ``null`` names with Python constants.""" + + def visit_Name(self, node: ast.Name) -> ast.AST: + if node.id in _JSON_NAMES: + # compile() needs a source location on every node. + return ast.copy_location(ast.Constant(_JSON_NAMES[node.id]), node) + return node + + +def _is_number_expr(node: ast.expr) -> bool: + """True if ``node`` is arithmetic over number constants only, e.g. ``12 * 60``.""" + return all( + isinstance(n, ast.BinOp | ast.UnaryOp | ast.operator | ast.unaryop) + or (isinstance(n, ast.Constant) and isinstance(n.value, int | float)) + for n in ast.walk(node) + ) + + +def _small_int(node: ast.expr) -> int | None: + """The int value of ``node`` if it is in ``[-_MAX_ITEMS, _MAX_ITEMS]``, else None. + + ``node`` may be arithmetic over numbers, e.g. ``30 * 12``. Without ``**`` + and ``<<`` its value grows only with the length of the expression. + """ + if not _is_number_expr(node): + return None + code = compile(ast.Expression(node), "", "eval") + try: + value = eval(code, {"__builtins__": {}}) # noqa: S307 + except ArithmeticError: + return None + return value if type(value) is int and abs(value) <= _MAX_ITEMS else None + + +def _within_size_limit(node: ast.expr) -> bool: + """True if every repetition and range() in ``node`` has a small constant size. + + The product of all repetition counts and range() lengths must be at most + ``_MAX_ITEMS``, so nested and chained repetition is bounded too. A + comprehension may have one ``for`` only. + """ + items = 1 + for n in ast.walk(node): + if isinstance(n, ast.BinOp) and isinstance(n.op, ast.Mult): + if _is_number_expr(n.left) and _is_number_expr(n.right): + continue # arithmetic, e.g. 12 * 60 + count = _small_int(n.right) if not _is_number_expr(n.left) else _small_int(n.left) + if count is None: + return False + items *= max(count, 1) + elif isinstance(n, ast.Call) and isinstance(n.func, ast.Name) and n.func.id == "range": + bounds = [_small_int(a) for a in n.args] + if n.keywords or not 1 <= len(bounds) <= 3 or None in bounds or bounds[2:] == [0]: + return False + items *= max(len(range(*bounds)), 1) + elif isinstance(n, ast.ListComp) and len(n.generators) > 1: + return False + return items <= _MAX_ITEMS + + +def _dolci_value(node: ast.expr) -> Any: + """Evaluate one argument value: a literal, or an expression over literals only.""" + node = _JsonNames().visit(node) + try: + return ast.literal_eval(node) + except (ValueError, TypeError): # TypeError: an unhashable key, e.g. {[]: 1} + pass + if not all( + isinstance(n, _SAFE_NODES) + and not (isinstance(n, ast.Subscript) and isinstance(n.ctx, ast.Store)) + and not ( + isinstance(n, ast.BinOp) + and isinstance(n.op, ast.Mod) + and not (_is_number_expr(n.left) and _is_number_expr(n.right)) + ) + for n in ast.walk(node) + ) or not _within_size_limit(node): + raise _DolciCallError(f"unsupported value: {ast.unparse(node)}") + # New globals for each value, so that no value can change how a later row evaluates. + safe_globals = {"__builtins__": {}, "range": range, "list": list} + try: + return eval(compile(ast.Expression(node), "", "eval"), safe_globals) # noqa: S307 + except Exception as exc: + raise _DolciCallError(f"cannot evaluate value: {ast.unparse(node)}") from exc + + +def _tokens(text: str) -> list[tuple[tokenize.TokenInfo, int, int]] | None: + """Python tokens of ``text`` with their start and end offsets, or None if it does not tokenize. + + The tokenizer knows strings and comments, so code built on it never + mistakes a ``(``, ``#`` or ``from=`` inside a quoted value for syntax. Names + with spaces (``Rates API``) still tokenize, unlike with ``ast.parse``. + """ + if "\r" in text: + return None # the tokenizer counts "\r" as a line break; the offsets below do not + # The last entry is for the NEWLINE / ENDMARKER tokens after the last line. + line_starts = [0] + [i + 1 for i, c in enumerate(text) if c == "\n"] + [len(text)] + try: + with warnings.catch_warnings(): + warnings.simplefilter("ignore", SyntaxWarning) + tokens = list(tokenize.generate_tokens(io.StringIO(text).readline)) + except (tokenize.TokenError, SyntaxError): + return None + if any(tok.type == tokenize.ERRORTOKEN for tok in tokens): + return None + return [ + ( + tok, + line_starts[tok.start[0] - 1] + tok.start[1], + line_starts[tok.end[0] - 1] + tok.end[1], + ) + for tok in tokens + ] + + +def _rename_keyword_args(arg_text: str) -> str: + """Prefix Python keywords used as argument names (``from="EUR"``), which are a syntax error.""" + tokens = _tokens(arg_text) + if tokens is None: + return arg_text # the parse fails next and reports the error + for (tok, start, _), (nxt, _, _) in reversed(list(zip(tokens, tokens[1:]))): + if ( + tok.type == tokenize.NAME + and keyword.iskeyword(tok.string) + and tok.string not in ("True", "False", "None") + and nxt.string == "=" + ): + arg_text = f"{arg_text[:start]}{_KEYWORD_PREFIX}{arg_text[start:]}" + return arg_text + + +def _dolci_call(text: str) -> dict[str, Any]: + """Convert one Python-style call, ``name(k=v, ...)``, to an OpenAI-style tool call. + + The name may contain spaces, ``-``, ``.`` or ``/``, so it is split off with a + regex. ``arguments`` is a JSON string, serialized like the HF ``tojson`` + filter and the vLLM ``hermes`` parser do, so a call the model writes and + the client sends back renders to the same text as in training. + """ + match = _CALL_RE.fullmatch(text) + if not match: + raise _DolciCallError(f"not a call: {text!r}") + name = match[1].strip() + if '"' in name or "\\" in name: + raise _DolciCallError(f"name breaks the rendered JSON: {name!r}") + arg_text = _rename_keyword_args(match[2]) + try: + # Dolci strings hold regex escapes such as '\.'; they parse fine but warn. + with warnings.catch_warnings(): + warnings.simplefilter("ignore", SyntaxWarning) + call = ast.parse(f"_f({arg_text})", mode="eval").body + except SyntaxError as exc: + raise _DolciCallError(f"syntax error: {text!r}") from exc + if not isinstance(call, ast.Call) or call.args or any(kw.arg is None for kw in call.keywords): + raise _DolciCallError(f"not a keyword-only call: {text!r}") + + args = {} + for kw in call.keywords: + unprefixed = kw.arg.removeprefix(_KEYWORD_PREFIX) + args[unprefixed if keyword.iskeyword(unprefixed) else kw.arg] = _dolci_value(kw.value) + try: + arguments = json.dumps(args, ensure_ascii=False) + round_trip_ok = json.loads(arguments) == args + except (TypeError, ValueError): + round_trip_ok = False + if not round_trip_ok: + # Tuples, sets, non-string keys or NaN would come back from the parser changed. + raise _DolciCallError(f"arguments do not survive a JSON round trip: {text!r}") + return {"type": "function", "function": {"name": name, "arguments": arguments}} + + +def _restore_lost_escapes(block: str) -> str: + """Apply the ``_LOST_ESCAPE_RE`` repair, but not inside string literals. + + The string positions come from the tokenizer, with the control characters + read as spaces (the same length, so the offsets stay correct). If the block + does not tokenize, the repair applies everywhere. + """ + tokens = _tokens(block.translate({ord(c): " " for c in _ESCAPE_LETTERS})) + strings = [(start, end) for tok, start, end in tokens or [] if tok.type == tokenize.STRING] + + def repair(match: re.Match[str]) -> str: + if any(start <= match.start() < end for start, end in strings): + return match[0] + return f")\n{_ESCAPE_LETTERS[match[1]]}" + + return _LOST_ESCAPE_RE.sub(repair, block) + + +def _dolci_call_texts(block: str) -> Iterator[str]: + """Split a Dolci ``function_calls`` block into one text per call. + + Skips ``null`` (Dolci's marker for "no call"), blank lines and ``#`` + comments, splits on a literal ``\\n`` between calls, and joins a call that + spans several lines. + + The split follows Python's tokenizer: a call ends at a NEWLINE token (or a + ``;``), which the tokenizer only emits outside brackets and strings. Thus + ``(``, ``)`` and ``#`` inside strings are safe, and names with spaces + (``Get Rounds(id=1)``) work. The literal ``\\n`` and ``\\r`` line breaks are + only replaced if the block does not tokenize as it is, so a ``)\\n`` inside + a string is kept. + + A block that does not tokenize at all, e.g. a name with a quote + (``Get Player's Media(id=1)``), falls back to a line split that counts + parentheses. A wrong join there gives a syntax error, so the row is + dropped, not changed. + """ + block = _restore_lost_escapes(block) + fixed = re.sub(r"\)\\n(?=\S)", ")\n", block.replace("\r\n", "\n").replace("\r", "\n")) + for text in (block, fixed): + tokens = _tokens(text) + if tokens is None: + continue + call: list[tuple[int, int]] = [] + for tok, start, end in tokens: + if tok.type in (tokenize.NEWLINE, tokenize.ENDMARKER) or tok.string == ";": + if call and text[call[0][0] : call[-1][1]] != "null": + yield text[call[0][0] : call[-1][1]] + call = [] + elif tok.type not in (tokenize.COMMENT, tokenize.NL, tokenize.INDENT, tokenize.DEDENT): + call.append((start, end)) + return + + pending = "" + for line in fixed.splitlines(): + if line.strip() in ("", "null") or line.lstrip().startswith("#"): + continue + pending = f"{pending}\n{line}" if pending else line + if pending.count("(") > pending.count(")"): + continue + yield pending + pending = "" + if pending: + yield pending + + +@register_transform("dolci_tool_use") +def dolci_tool_use(example: dict[str, Any]) -> dict[str, Any]: + """Transform allenai/Dolci-Instruct-SFT tool rows to the standard tool format. + + Dolci stores tools in the OLMo 3 layout: the tool list as a JSON string in + ``messages[0]["functions"]``, each call as a Python-style line in an + assistant ``function_calls`` field, and each result as role + ``environment``. This transform produces what chat templates and the vLLM + ``hermes`` parser expect: + + - ``messages[0]["functions"]`` -> the ``tools`` column (JSON string) + - assistant ``function_calls`` -> ``tool_calls`` with JSON ``arguments`` + - role ``environment`` -> role ``tool`` + + It also removes the Dolci system prompt text about ```` tags + (keeping any task instructions after it) and drops assistant turns with + no content and no calls. Rows without a tool list pass through unchanged. + + The row is dropped (empty ``messages``) if any call cannot be converted: + placeholders such as ``loan_amount=X``, calls nested in calls, positional + arguments, or values that change in a JSON round trip. + """ + messages = example["messages"] + if not messages or not messages[0].get("functions"): + return { + "messages": [{"role": m["role"], "content": m["content"]} for m in messages], + "tools": None, + } + + tool_names = {t.get("function", {}).get("name") for t in json.loads(messages[0]["functions"])} + out: list[dict[str, Any]] = [] + system = ( + messages[0]["content"] + .removeprefix(_DOLCI_TOOL_PROMPT) + .replace(_DOLCI_PYTHONIC_HINT, "") + .strip() + ) + if system: + out.append({"role": "system", "content": system}) + + for message in messages[1:]: + role, content = message["role"], message["content"] or "" + if role == "environment": + out.append({"role": "tool", "content": content}) + elif role == "assistant": + try: + tool_calls = [ + _dolci_call(t) for t in _dolci_call_texts(message.get("function_calls") or "") + ] + except _DolciCallError: + # Every row must return the same keys, or the datasets writer fails. + return {"messages": [], "tools": None} + for call in tool_calls: + # A "\" before a name that starts with "n" became a plain line + # break (see _LOST_ESCAPE_RE), so only the tool list can show it. + name = call["function"]["name"] + if name not in tool_names and f"n{name}" in tool_names: + call["function"]["name"] = f"n{name}" + if tool_calls: + out.append({"role": "assistant", "content": content, "tool_calls": tool_calls}) + elif content.strip(): + out.append({"role": "assistant", "content": content}) + else: + out.append({"role": role, "content": content}) + + return {"messages": out, "tools": messages[0]["functions"]} diff --git a/src/post_training/methods/sft.py b/src/post_training/methods/sft.py index e49257a..90547e6 100644 --- a/src/post_training/methods/sft.py +++ b/src/post_training/methods/sft.py @@ -104,6 +104,16 @@ def _classify_row( { "content": Value("string"), "role": Value("string"), + # OpenAI-style calls. `arguments` is a JSON string, which the + # chat template inserts verbatim; it avoids declaring each + # tool's argument schema. datasets fills it with null for + # messages that have no calls. + "tool_calls": List( + { + "type": Value("string"), + "function": {"name": Value("string"), "arguments": Value("string")}, + } + ), } ), # A JSON string, matching what TRL reads: it does diff --git a/tests/test_data_loader.py b/tests/test_data_loader.py index 8cabfc6..8ab4bee 100644 --- a/tests/test_data_loader.py +++ b/tests/test_data_loader.py @@ -230,6 +230,61 @@ def test_tools_column_survives_mixing_and_is_null_where_absent(monkeypatch): assert None in tools +def test_native_and_transformed_entries_with_different_message_fields_mix(monkeypatch): + """A transformed entry gets the cast message fields (here `tool_calls`); native + entries keep their own. Each dataset gets the union of the fields, with null + where a field is absent, so `concatenate_datasets` does not fail and no field + is lost.""" + from post_training.methods.sft import MESSAGES_FEATURES + + tools = '[{"type": "function", "function": {"name": "wx"}}]' + dolci_row = { + "messages": [ + {"role": "system", "content": "", "functions": tools, "function_calls": None}, + {"role": "user", "content": "Weather?", "functions": None, "function_calls": None}, + { + "role": "assistant", + "content": "", + "functions": None, + "function_calls": 'wx(q="Oslo")', + }, + ] + } + _patch_load_dataset( + monkeypatch, + { + "plain": _dataset("plain", 1).select_columns(["messages"]), + "dolci-native": Dataset.from_list([dolci_row]), + "dolci": Dataset.from_list([dolci_row]), + }, + ) + + mixed = loader.load_and_mix_datasets( + _config( + DatasetEntry(name="plain", path="plain"), + DatasetEntry(name="dolci-native", path="dolci-native"), + DatasetEntry(name="dolci", path="dolci", transform="dolci_tool_use"), + ), + columns_to_keep=["messages", "tools"], + features=MESSAGES_FEATURES, + ) + + assert set(mixed.features["messages"].feature) == { + "role", "content", "tool_calls", "functions", "function_calls" + } # fmt: skip + plain = next(r["messages"] for r in mixed if r["messages"][0]["content"] == "plain-0") + native = next(r["messages"] for r in mixed if r["messages"][0]["role"] == "system") + converted = next( + r["messages"] for r in mixed if r["tools"] and r["messages"][0]["role"] == "user" + ) + + assert plain[0]["tool_calls"] is None and plain[0]["functions"] is None + assert native[2]["function_calls"] == 'wx(q="Oslo")' # kept for the olmo3 template + assert native[2]["tool_calls"] is None + assert converted[1]["tool_calls"][0]["function"]["name"] == "wx" + assert converted[1]["function_calls"] is None + + def test_native_plain_chat_dataset_skips_schema_cast(monkeypatch): _patch_load_dataset(monkeypatch, {"dataset-a": _dataset("a", 2)}) features = Features( diff --git a/tests/test_transforms.py b/tests/test_transforms.py index 11a05e7..0bba954 100644 --- a/tests/test_transforms.py +++ b/tests/test_transforms.py @@ -1,6 +1,21 @@ """Tests for selected dataset transforms.""" -from post_training.data.transforms import aegis_safety, helpsteer2, helpsteer3_preference +import json + +import pytest +from datasets import Dataset + +from post_training.chat_templates.registry import get_chat_template +from post_training.config import DataConfig, DatasetEntry +from post_training.data import loader +from post_training.data.transforms import ( + _DOLCI_TOOL_PROMPT, + aegis_safety, + dolci_tool_use, + helpsteer2, + helpsteer3_preference, +) +from post_training.methods.sft import MESSAGES_FEATURES def test_aegis_safety_keeps_unsafe_non_redacted_response(): @@ -47,3 +62,214 @@ def test_helpsteer3_preference_selects_response(): helpsteer3_preference({**example, "overall_preference": 0})["messages"][-1]["content"] == "first" ) + + +_DOLCI_TOOLS = json.dumps( + [ + { + "type": "function", + "function": {"name": "wx.get", "parameters": {"properties": {"q": {"type": "string"}}}}, + } + ] +) + + +def _dolci_row(function_calls, content="", system=_DOLCI_TOOL_PROMPT): + return { + "messages": [ + { + "role": "system", + "content": system, + "functions": _DOLCI_TOOLS, + "function_calls": None, + }, + {"role": "user", "content": "Weather?", "functions": None, "function_calls": "null"}, + { + "role": "assistant", + "content": content, + "functions": None, + "function_calls": function_calls, + }, + { + "role": "environment", + "content": '{"temp": 5}', + "functions": None, + "function_calls": None, + }, + {"role": "assistant", "content": "It is 5.", "functions": None, "function_calls": None}, + ] + } + + +def _arguments(out): + return [json.loads(c["function"]["arguments"]) for c in out["messages"][1]["tool_calls"]] + + +def test_dolci_tool_use_converts_to_standard_format(): + out = dolci_tool_use( + _dolci_row('wx.get(q="Oslo", metric=true)\nwx.get(q="Bergen", days=[1, null])') + ) + + assert out["tools"] == _DOLCI_TOOLS + assert [m["role"] for m in out["messages"]] == ["user", "assistant", "tool", "assistant"] + assert out["messages"][1]["tool_calls"][0] == { + "type": "function", + "function": {"name": "wx.get", "arguments": '{"q": "Oslo", "metric": true}'}, + } + assert _arguments(out)[1] == {"q": "Bergen", "days": [1, None]} + assert out["messages"][2] == {"role": "tool", "content": '{"temp": 5}'} + + +@pytest.mark.parametrize( + ("function_calls", "expected"), + [ + ("wx.get(minutes=12*60, rates=[0.5]*2)", [{"minutes": 720, "rates": [0.5, 0.5]}]), + ("wx.get(chars=[c for c in 'ab'])", [{"chars": ["a", "b"]}]), + ('Rates API(from="EUR", to="USD")', [{"from": "EUR", "to": "USD"}]), + ("# look up both\nwx.get(q='a')\n\nwx.get(q='b')", [{"q": "a"}, {"q": "b"}]), + ("wx.get(q='a')\\nwx.get(q='b')", [{"q": "a"}, {"q": "b"}]), + ("wx.get(q='a',\n days=2)", [{"q": "a", "days": 2}]), + ('wx.get(q="(()")\nwx.get(q="())()")', [{"q": "(()"}, {"q": "())()"}]), + ("wx.get(q='a') # (may fail)", [{"q": "a"}]), + ("wx.get(q='#a')\n# skip\nwx.get(q='b')", [{"q": "#a"}, {"q": "b"}]), + ('wx.get(q="a)\\nb")', [{"q": "a)\nb"}]), + ("wx.get(q='a')\nnull", [{"q": "a"}]), + ( + 'Rates API(from="EUR", q="(")\nRates API(from="USD", q=")")', + [{"from": "EUR", "q": "("}, {"from": "USD", "q": ")"}], + ), + ( + 'wx.get(html="
", q="from=x")', + [{"html": "
", "q": "from=x"}], + ), + ("wx.get(flags=[true, false]*2)", [{"flags": [True, False, True, False]}]), + ("wx.get(q='a')\r\nwx.get(q='b')", [{"q": "a"}, {"q": "b"}]), + ("Get Player's Media(id=1)\nGet Player's Media(id=2)", [{"id": 1}, {"id": 2}]), + ('wx.get(q=")\tget(")', [{"q": ")\tget("}]), + ("wx.get(q=list(range(3)), r=[[0] * 2] * 2)", [{"q": [0, 1, 2], "r": [[0, 0], [0, 0]]}]), + ("wx.get(q=[1] * (2 * 2))", [{"q": [1, 1, 1, 1]}]), + ], +) +def test_dolci_tool_use_repairs_call_lines(function_calls, expected): + assert _arguments(dolci_tool_use(_dolci_row(function_calls))) == expected + + +@pytest.mark.parametrize( + "function_calls", + [ + "wx.get(q=X)", + "wx.get(q=[random numbers])", + "wx.get(q=other(q=1))", + "min([2, 4, 6])", + "wx.get(q=(1, 2))", + "wx.get(q={1: 'a'})", + "wx.get(q={[]: 1})", + "wx.get(q=[0]*1000000000)", + "wx.get(q=list(range(1000000000)))", + "wx.get(q=[[0] * 10000] * 10000)", + "wx.get(q=1 << 100)", + "wx.get(q=[x for x in 'ab' for y in 'cd'])", + "wx.get(q=2**10)", + 'say"hi(q=1)', + "wx.get", + ], +) +def test_dolci_tool_use_drops_row_with_unconvertible_call(function_calls): + assert dolci_tool_use(_dolci_row(function_calls)) == {"messages": [], "tools": None} + + +def test_dolci_tool_use_value_cannot_change_how_later_rows_evaluate(): + assert not dolci_tool_use(_dolci_row("wx.get(q=[0 for __builtins__['X'] in ['Oslo']])"))[ + "messages" + ] + assert not dolci_tool_use(_dolci_row("wx.get(q=X)"))["messages"] + + +def test_olmo3_instruct_sft_rejects_json_string_arguments_with_a_clear_error(): + chat_template_utils = pytest.importorskip("transformers.utils.chat_template_utils") + compiled = chat_template_utils._compile_jinja_template(get_chat_template("olmo3-instruct-sft")) + out = dolci_tool_use(_dolci_row('wx.get(q="Oslo")')) + with pytest.raises(Exception, match="use transform: null"): + chat_template_utils._render_with_assistant_indices( + compiled, out["messages"], json.loads(out["tools"]), None, False + ) + + +def test_dolci_tool_use_restores_names_that_lost_an_escaped_letter(): + # Dolci decoded "\reverse" as "\r" + "everse"; "\next" became a plain line break. + row = _dolci_row("wx.get(q='a')\reverse(s='x')\x0clatten(x=1)\next(n=2)") + row["messages"][0]["functions"] = json.dumps( + [ + {"type": "function", "function": {"name": n}} + for n in ("wx.get", "reverse", "flatten", "next") + ] + ) + calls = dolci_tool_use(row)["messages"][1]["tool_calls"] + assert [c["function"]["name"] for c in calls] == ["wx.get", "reverse", "flatten", "next"] + + +def test_dolci_tool_use_null_call_keeps_refusal_and_drops_empty_turn(): + refusal = dolci_tool_use(_dolci_row("null", content="I can't do that.")) + assert refusal["messages"][1] == {"role": "assistant", "content": "I can't do that."} + + empty = dolci_tool_use(_dolci_row("null")) + assert [m["role"] for m in empty["messages"]] == ["user", "tool", "assistant"] + + +def test_dolci_tool_use_keeps_task_instructions_in_system_prompt(): + system = ( + _DOLCI_TOOL_PROMPT + + " You do research.\n\n- Use tools. Format calls like: tool_name(query=...) (below)." + ) + out = dolci_tool_use(_dolci_row("wx.get(q='a')", system=system)) + + assert out["messages"][0] == { + "role": "system", + "content": "You do research.\n\n- Use tools. (below).", + } + + +def test_dolci_tool_use_passes_rows_without_tools(): + row = { + "messages": [{"role": "user", "content": "Hi", "functions": None, "function_calls": None}] + } + + assert dolci_tool_use(row) == {"messages": [{"role": "user", "content": "Hi"}], "tools": None} + + +def test_dolci_tool_use_survives_features_cast_and_renders_qwen3_tool_call(monkeypatch): + """The transform output must keep `tool_calls` through the loader's + MESSAGES_FEATURES cast and render as a Qwen3 / hermes ``.""" + plain = { + "messages": [{"role": "user", "content": "Hi", "functions": None, "function_calls": None}] + } + rows = [_dolci_row('wx.get(q="Oslo")'), plain, _dolci_row("wx.get(q=X)")] + monkeypatch.setattr(loader, "load_dataset", lambda path, **kwargs: Dataset.from_list(rows)) + + ds = loader.load_and_mix_datasets( + DataConfig( + num_proc=1, + datasets=[DatasetEntry(name="dolci", path="dolci", transform="dolci_tool_use")], + ), + columns_to_keep=["messages", "tools"], + features=MESSAGES_FEATURES, + ) + tool_row = next(r for r in ds if r["tools"]) + plain_row = next(r for r in ds if not r["tools"] and r["messages"]) + + assert sum(not r["messages"] for r in ds) == 1 # the X placeholder row + + assert plain_row["messages"] == [{"role": "user", "content": "Hi", "tool_calls": None}] + chat_template_utils = pytest.importorskip("transformers.utils.chat_template_utils") + compiled = chat_template_utils._compile_jinja_template(get_chat_template("qwen3")) + rendered, indices = chat_template_utils._render_with_assistant_indices( + compiled, tool_row["messages"], json.loads(tool_row["tools"]), None, False + ) + spans = [rendered[start:end] for start, end in indices] + + assert '\n{"type": "function", "function": {"name": "wx.get"' in rendered + assert '\n{"temp": 5}\n' in rendered + assert ( + spans[0] + == '\n{"name": "wx.get", "arguments": {"q": "Oslo"}}\n<|im_end|>' + )