Format-adaptive tool rendering + resilient dataset loading

train_securecoder.py

@@ -170,26 +170,41 @@ def _normalise_parameters(params: Any) -> dict:
 
 
 def _normalise_tool_schema(tool: Any) -> dict | None:
-    """Force a tool schema into the shape Qwen's chat template expects:
-    {"type": "function", "function": {"name", "description", "parameters"}}."""
+    """Canonical **flat** tool schema: {"name", "description", "parameters"}.
+
+    Qwen3-Coder's chat template walks ``tool.parameters.properties`` directly, so
+    the flat form is what we store; ``_tools_for_style`` re-wraps it into the
+    OpenAI ``{"type": "function", "function": {...}}`` shape for templates that
+    want that instead.
+    """
     if not isinstance(tool, dict):
         return None
-    fn = tool.get("function") if "function" in tool else tool
+    fn = tool.get("function") if isinstance(tool.get("function"), dict) else tool
     if not isinstance(fn, dict) or not fn.get("name"):
         return None
     return {
-        "type": "function",
-        "function": {
-            "name": fn["name"],
-            "description": fn.get("description", ""),
-            "parameters": _normalise_parameters(fn.get("parameters")),
-        },
+        "name": fn["name"],
+        "description": fn.get("description", ""),
+        "parameters": _normalise_parameters(fn.get("parameters")),
     }
 
 
+def _tools_for_style(tools: list[dict], style: str) -> list[dict] | None:
+    if not tools:
+        return None
+    if style == "nested":
+        return [{"type": "function", "function": t} for t in tools]
+    return tools
+
+
 
 def _parse_calls(value: Any) -> list[dict] | None:
-    """Return OpenAI tool_calls if ``value`` is one or more JSON function calls."""
+    """Return OpenAI-style tool calls if ``value`` is one or more function calls.
+
+    ``arguments`` is kept as a **dict** plus a JSON-string copy, because templates
+    disagree: Qwen3-Coder iterates ``arguments | items`` (needs a mapping) while
+    others print a JSON string.
+    """
     if isinstance(value, str):
         text = value.strip()
         if not text.startswith(("{", "[")):
@@ -204,12 +219,21 @@ def _parse_calls(value: Any) -> list[dict] | None:
     calls = []
     for i, item in enumerate(items):
         arguments = item.get("arguments", item.get("parameters", {}))
-        if not isinstance(arguments, str):
-            arguments = json.dumps(arguments)
+        if isinstance(arguments, str):
+            try:
+                arguments = json.loads(arguments)
+            except json.JSONDecodeError:
+                arguments = {"value": arguments}
+        if not isinstance(arguments, dict):
+            arguments = {"value": arguments}
         calls.append({
             "id": f"call_{i}",
             "type": "function",
-            "function": {"name": item["name"], "arguments": arguments},
+            "function": {
+                "name": item["name"],
+                "arguments": arguments,
+                "arguments_json": json.dumps(arguments),
+            },
         })
     return calls
 
@@ -305,16 +329,55 @@ def _messages_from_any(row: dict, kind: str) -> tuple[list[dict], list[dict]]:
         messages.append({"role": "assistant", "content": str(row["solution"])})
         return messages, tools
 
-    # --- SecOps reasoning: goal/command -> unified_interpretation ---------
-    if row.get("unified_interpretation"):
-        ask = [f"Tool: {row.get('tool', 'shell')}", f"Goal: {row.get('goal', '')}"]
-        for key in ("command", "command_sequence", "nmap_context"):
-            if row.get(key):
-                ask.append(f"{key}: {row[key]}")
-        ask.append("Explain what the output means, what it tells you about the target, "
-                   "and what the next step should be.")
+    # --- SecOps reasoning: goal/command -> interpretation ------------------
+    # dpevzner's rows carry goal + command_sequence + interpretation +
+    # classification + safety_and_scope (unified_interpretation is empty in the
+    # published revision). Scope metadata goes into the prompt so the model
+    # learns the framed, authorised-use context alongside the command knowledge.
+    if row.get("goal") and row.get("interpretation"):
+        seq = row.get("command_sequence") or {}
+        if isinstance(seq, str):
+            seq = {"command": seq}
+        interp = row.get("interpretation") or {}
+        if isinstance(interp, str):
+            interp = {"what_it_means": [interp]}
+        scope = row.get("safety_and_scope") or {}
+        if isinstance(scope, str):
+            scope = {}
+
+        tool = row.get("tool") or {}
+        if isinstance(tool, str):
+            tool = {"name": tool}
+        ask = [f"Environment: {tool.get('name', 'shell')} ({tool.get('platform', 'unknown')})"]
+        if seq.get("command"):
+            ask.append(f"Command: {seq['command']}")
+        ask.append(f"Goal: {row['goal']}")
+        if isinstance(scope, dict) and scope:
+            ask.append("Scope: " + ", ".join(f"{k}={v}" for k, v in list(scope.items())[:4]))
+        if row.get("classification"):
+            cls = row["classification"]
+            if isinstance(cls, dict):
+                ask.append("Context: " + ", ".join(f"{k}={v}" for k, v in list(cls.items())[:3]))
+        ask.append("Explain what this command does, what its output means, what it tells you "
+                   "about the target, and the next step in an authorised assessment.")
+
+        answer = []
+        if seq.get("description"):
+            answer.append(f"**What it does.** {seq['description']}")
+        if seq.get("expected_output_pattern"):
+            answer.append("**Expected output.** " + ", ".join(map(str, seq["expected_output_pattern"])))
+        for item in interp.get("what_it_means", []) if isinstance(interp, dict) else []:
+            answer.append(f"**What it means.** {item}")
+        for item in interp.get("risk_indicators", []) if isinstance(interp, dict) else []:
+            answer.append(f"**Risk indicators.** {item}")
+        if row.get("ambiguity_analysis"):
+            answer.append(f"**Ambiguity.** {row['ambiguity_analysis']}")
+        if isinstance(scope, dict) and scope.get("authorization_required"):
+            answer.append("**Scope.** Only run this against systems you are authorised to test.")
+        if len(answer) < 2:
+            return [], tools
         messages.append({"role": "user", "content": "\n".join(str(a) for a in ask)})
-        messages.append({"role": "assistant", "content": str(row["unified_interpretation"])})
+        messages.append({"role": "assistant", "content": "\n\n".join(answer)})
         return messages, tools
 
     # --- generic prompt/completion fallback -------------------------------
@@ -331,34 +394,118 @@ def _messages_from_any(row: dict, kind: str) -> tuple[list[dict], list[dict]]:
 # --------------------------------------------------------------------------
 def load_source(src: Source, token: str | None, progress: bool = False) -> list[dict]:
     """Pull up to ``limit`` rows from one Hub dataset, streaming so we never
-    download more than we need."""
-    from datasets import load_dataset
+    download more than we need.
 
-    kwargs: dict[str, Any] = {"split": src.split, "streaming": True}
-    if src.config:
-        kwargs["name"] = src.config
-    if token:
-        kwargs["token"] = token
+    Datasets move: builder configs get renamed (a config called ``default`` last
+    week is ``chatml`` today) and splits get added. Each candidate is tried in
+    turn so one rename cannot silently empty a slice of the mix.
+    """
+    from datasets import load_dataset
 
-    ds = load_dataset(src.repo, **kwargs)
-    rows = []
-    for i, row in enumerate(ds):
-        if i >= src.limit:
-            break
-        rows.append(dict(row))
-        if progress and i and i % 2500 == 0:
-            log.info("    %s: %d rows...", source_name(src), i)
-    return rows
+    candidates = [
+        (src.config, src.split),
+        (src.config, "train"),
+        (src.config, "test"),
+        (None, src.split),
+        (None, "train"),
+        (None, "test"),
+    ]
+    seen: set = set()
+    last_exc: Exception | None = None
+
+    for config, split in candidates:
+        if (config, split) in seen:
+            continue
+        seen.add((config, split))
+        kwargs: dict[str, Any] = {"split": split, "streaming": True}
+        if config:
+            kwargs["name"] = config
+        if token:
+            kwargs["token"] = token
+        try:
+            ds = load_dataset(src.repo, **kwargs)
+            rows = []
+            for i, row in enumerate(ds):
+                if i >= src.limit:
+                    break
+                rows.append(dict(row))
+                if progress and i and i % 2500 == 0:
+                    log.info("    %s: %d rows...", source_name(src), i)
+            if not rows:
+                last_exc = ValueError(f"config={config} split={split} streamed 0 rows")
+                continue
+            if (config, split) != (src.config, src.split):
+                log.info("  %s: fell back to config=%s split=%s",
+                         src.repo, config, split)
+            return rows
+        except Exception as exc:  # noqa: BLE001 - try the next candidate
+            last_exc = exc
+
+    raise last_exc if last_exc else RuntimeError(f"could not load {src.repo}")
+
+
+_RENDER_STYLE: str | None = None
+STYLE_ATTEMPTS = (("flat", "dict"), ("nested", "dict"), ("flat", "string"), ("nested", "string"))
+
+
+def _apply_arg_style(messages: list[dict], arg_style: str) -> list[dict]:
+    """Copy messages, swapping tool-call arguments between dict and JSON string."""
+    if arg_style != "string":
+        return messages
+    out = []
+    for message in messages:
+        if message.get("tool_calls"):
+            message = dict(message)
+            message["tool_calls"] = [
+                {
+                    "id": call["id"],
+                    "type": "function",
+                    "function": {
+                        "name": call["function"]["name"],
+                        "arguments": call["function"].get(
+                            "arguments_json", json.dumps(call["function"]["arguments"])
+                        ),
+                    },
+                }
+                for call in message["tool_calls"]
+            ]
+        out.append(message)
+    return out
 
 
 def render_record(tokenizer, messages: list[dict], tools: list[dict] | None = None) -> str:
-    """Render to the model's native chat format (Qwen3 emits <tool_call> blocks)."""
-    return tokenizer.apply_chat_template(
-        messages,
-        tools=tools or None,
-        tokenize=False,
-        add_generation_prompt=False,
-    )
+    """Render with the model's native chat template.
+
+    Tool templates disagree in two independent ways: whether tool schemas are
+    flat (`{"name", "parameters"}`) or OpenAI-nested (`{"type": "function", ...}`),
+    and whether tool-call arguments are a mapping or a JSON string. Qwen3-Coder
+    renders ``<function=NAME><parameter=...>`` blocks and iterates
+    ``arguments | items``, so a JSON string there is a hard error. Rather than
+    hard-code one convention, detect it once and reuse it for the rest of the
+    run - with the model, the mix and the template all free to change.
+    """
+    global _RENDER_STYLE
+
+    order = []
+    if _RENDER_STYLE:
+        order.append(tuple(_RENDER_STYLE.split("+")))
+    order += [style for style in STYLE_ATTEMPTS if style not in order]
+
+    last_exc: Exception | None = None
+    for tool_style, arg_style in order:
+        try:
+            text = tokenizer.apply_chat_template(
+                _apply_arg_style(messages, arg_style),
+                tools=_tools_for_style(tools, tool_style),
+                tokenize=False,
+                add_generation_prompt=False,
+            )
+            _RENDER_STYLE = f"{tool_style}+{arg_style}"
+            return text
+        except Exception as exc:  # noqa: BLE001 - try the next convention
+            last_exc = exc
+    raise last_exc if last_exc else RuntimeError("render failed")
+
 
 
 def build_dataset(tokenizer, sources: list[Source], token: str | None, validate: bool):