| |
| """Per-source schema adapters for the CyberGym SFT mix. |
| |
| Each raw Hugging Face dataset in ``training/configs/datasets.yaml`` has its own |
| column layout. The generic ``normalize_sft_jsonl.py`` only handles a few common |
| shapes (``messages`` / ``conversations`` / ``instruction|input|output`` / |
| ``system|user|assistant``) and silently rejects everything else. The Tier-1 |
| C/C++ detection tables (PrimeVul, DiverseVul) and the vuln/fix pair sets |
| (CrossVul) do **not** match any of those shapes, so they need dedicated |
| adapters. |
| |
| An adapter takes one raw row (a ``dict``) plus a small ``params`` dict from the |
| manifest and returns a list of :class:`Example` objects (0, 1, or 2 per row). |
| The orchestrator (``build_sft_dataset.py``) wraps each Example with provenance |
| (id/source/license) and routes it to the ``ready`` mix or the ``to_synthesize`` |
| queue based on ``think_status``. |
| |
| This module is deliberately stdlib-only so it can be unit-tested without the |
| ``datasets``/``torch`` stack (those only exist on the rented GPU host). |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| from dataclasses import dataclass, field |
| from typing import Any, Callable |
|
|
| DEFAULT_SYSTEM = "Authorized security research and education context." |
|
|
| |
| |
| DETECTION_SYSTEM = ( |
| "You are a senior security engineer reviewing source code for " |
| "memory-safety and other security vulnerabilities. Reason carefully, then " |
| "give a clear verdict. Authorized security research and education context." |
| ) |
|
|
|
|
| @dataclass |
| class Example: |
| """One normalized training example before provenance wrapping.""" |
|
|
| messages: list[dict[str, str]] |
| |
| |
| think_status: str = "needs_synthesis" |
| |
| |
| |
| verify: dict[str, Any] | None = None |
| metadata: dict[str, Any] = field(default_factory=dict) |
|
|
|
|
| |
| |
| |
| def coerce_list(value: Any) -> list[str]: |
| """Normalize a CWE-ish field into a clean list of strings. |
| |
| Handles native lists, JSON-encoded list strings (``'["CWE-310"]'`` as |
| returned by the datasets-server), comma/space separated strings, and the |
| various null sentinels seen in these datasets. |
| """ |
| if value is None: |
| return [] |
| if isinstance(value, list): |
| items = value |
| elif isinstance(value, str): |
| s = value.strip() |
| if not s or s.lower() in {"none", "null", "[]", "nan"}: |
| return [] |
| if s.startswith("["): |
| try: |
| parsed = json.loads(s) |
| items = parsed if isinstance(parsed, list) else [parsed] |
| except json.JSONDecodeError: |
| items = [p.strip() for p in s.strip("[]").split(",")] |
| else: |
| items = [p.strip() for p in s.replace(";", ",").split(",")] |
| else: |
| items = [value] |
| out: list[str] = [] |
| for item in items: |
| text = str(item).strip().strip("'\"") |
| if text and text.lower() not in {"none", "null", "nan"}: |
| out.append(text) |
| return out |
|
|
|
|
| def as_bool_label(value: Any) -> bool | None: |
| """Interpret a vulnerability label that may be 0/1, "0"/"1", or bool.""" |
| if value is None: |
| return None |
| if isinstance(value, bool): |
| return value |
| if isinstance(value, (int, float)): |
| return int(value) == 1 |
| text = str(value).strip().lower() |
| if text in {"1", "true", "vulnerable", "yes", "vuln"}: |
| return True |
| if text in {"0", "false", "not vulnerable", "no", "safe", "secure", "benign"}: |
| return False |
| return None |
|
|
|
|
| def get_first(row: dict[str, Any], keys: list[str]) -> Any: |
| for key in keys: |
| if key in row and row[key] not in (None, ""): |
| return row[key] |
| return None |
|
|
|
|
| def has_think(text: str) -> bool: |
| return "<think>" in text and "</think>" in text |
|
|
|
|
| def ensure_think_closed(think: str) -> str: |
| """Make sure a reasoning fragment is a well-formed <think> block.""" |
| t = think.strip() |
| if not t: |
| return "" |
| if "<think>" not in t: |
| t = "<think>\n" + t |
| if "</think>" not in t: |
| t = t + "\n</think>" |
| return t |
|
|
|
|
| def fence(code: str, language: str) -> str: |
| lang = (language or "").strip().lower() |
| alias = {"c++": "cpp", "c/c++": "cpp", "cpp": "cpp", "c": "c"}.get(lang, lang or "") |
| body = str(code).rstrip() |
| return f"```{alias}\n{body}\n```" |
|
|
|
|
| def _compose_grounded(task: str, label: str, cwes: list[str], desc: str | None, language: str) -> str: |
| """Compose a short, label-consistent <think> from ground-truth metadata. |
| |
| This is "grounded backfill": real reasoning derived from the dataset's own |
| labels/descriptions, so detection data is training-ready with a <think> block |
| without a slow teacher/rejection-sampling pass. The reasoning always agrees |
| with the gold answer (never teaches a wrong conclusion). |
| """ |
| cwe_txt = ", ".join(cwes) if cwes else "" |
| lang = language or "code" |
| if task == "secure_fix": |
| cp = f" (it is affected by {cwe_txt})" if cwe_txt else "" |
| return ( |
| "<think>\n" |
| f"The original {lang} code contains a security defect{cp}. I rewrite the unsafe " |
| "operation with proper bounds and input validation while preserving the intended behavior.\n" |
| "</think>" |
| ) |
| if task == "secure_code": |
| focus = (desc or "a common security weakness").rstrip(".") |
| return ( |
| "<think>\n" |
| f"This task touches {focus}. I implement it so the weakness cannot occur — validating " |
| "inputs and using safe, bounded APIs.\n" |
| "</think>" |
| ) |
| |
| if label == "vulnerable": |
| lead = (desc.strip() + " ") if desc else ( |
| "Tracing the data flow, an unsafe operation is reachable with attacker-influenced input. " |
| ) |
| tail = ( |
| f"This is consistent with {cwe_txt}, so the function is vulnerable." |
| if cwe_txt else "An exploitable weakness is present, so the function is vulnerable." |
| ) |
| return f"<think>\nI review this {lang} function for memory-safety and input-validation defects. {lead}{tail}\n</think>" |
| return ( |
| "<think>\n" |
| f"I review this {lang} function for common weaknesses (out-of-bounds access, integer overflow, " |
| "use-after-free, format-string, unvalidated input). The operations are bounded and inputs are " |
| "handled safely, so I find no security-relevant defect.\n" |
| "</think>" |
| ) |
|
|
|
|
| def _sys_user_assistant(system: str, user: str, assistant: str) -> list[dict[str, str]]: |
| return [ |
| {"role": "system", "content": system}, |
| {"role": "user", "content": user}, |
| {"role": "assistant", "content": assistant}, |
| ] |
|
|
|
|
| |
| |
| |
| def adapt_detection_func_target(row: dict[str, Any], params: dict[str, Any]) -> list[Example]: |
| """PrimeVul / DiverseVul: a function + a binary `target` (1 == vulnerable).""" |
| code_field = params.get("code_field", "func") |
| label_field = params.get("label_field", "target") |
| func = get_first(row, [code_field, "func", "function", "code"]) |
| label = as_bool_label(row.get(label_field, row.get("target"))) |
| if not func or label is None: |
| return [] |
|
|
| language = params.get("language", "C/C++") |
| cwes = coerce_list(get_first(row, ["cwe", "cwe_ids", "cwe_id"])) |
| project = get_first(row, ["project", "repo_name", "repo"]) |
| cve = get_first(row, ["cve", "cve_id"]) |
| cve_desc = get_first(row, ["cve_desc", "cve_description"]) |
|
|
| ctx = f"Project: {project}\n\n" if project else "" |
| user = ( |
| f"Review the following {language} function for security vulnerabilities.\n\n" |
| f"{ctx}{fence(func, language)}\n\n" |
| "Does this function contain a security vulnerability? Answer " |
| "'Vulnerable' or 'Not vulnerable'. If vulnerable, name the most likely " |
| "CWE and explain the root cause; if not, briefly justify why." |
| ) |
|
|
| if label: |
| verdict = "Vulnerable." |
| if cwes: |
| verdict += f" Most likely {', '.join(cwes)}." |
| if cve: |
| verdict += f" (Associated CVE: {cve}.)" |
| if cve_desc: |
| verdict += f"\n\n{cve_desc}" |
| expected = "vulnerable" |
| else: |
| verdict = "Not vulnerable. No security-relevant defect is evident in this function." |
| expected = "not_vulnerable" |
|
|
| if params.get("grounded_think"): |
| think = _compose_grounded("vuln_detection", expected, cwes, cve_desc, language) |
| assistant, status = f"{think}\n\n{verdict}", "present" |
| else: |
| assistant, status = verdict, "needs_synthesis" |
|
|
| return [ |
| Example( |
| messages=_sys_user_assistant(DETECTION_SYSTEM, user, assistant), |
| think_status=status, |
| verify={"mode": "label", "expected": expected, "cwe": cwes}, |
| metadata={ |
| "task": "vuln_detection", |
| "language": language, |
| "label": expected, |
| "cwe": cwes, |
| "cve": cve, |
| "project": project, |
| }, |
| ) |
| ] |
|
|
|
|
| def adapt_vuln_fix_pair(row: dict[str, Any], params: dict[str, Any]) -> list[Example]: |
| """CrossVul: vulnerable_code / fixed_code pairs across many languages.""" |
| vuln = get_first(row, ["vulnerable_code", "vuln_code", "before"]) |
| fixed = get_first(row, ["fixed_code", "fix_code", "after"]) |
| if not vuln: |
| return [] |
| language = get_first(row, ["language", "lang"]) or params.get("language", "code") |
| cwes = coerce_list(get_first(row, ["cwe_id", "cwe", "cwe_ids"])) |
| cwe_desc = get_first(row, ["cwe_description", "cwe_desc"]) |
| tasks = params.get("tasks", ["detect", "fix"]) |
|
|
| out: list[Example] = [] |
|
|
| grounded = bool(params.get("grounded_think")) |
|
|
| if "detect" in tasks: |
| user = ( |
| f"Review the following {language} code for security vulnerabilities.\n\n" |
| f"{fence(vuln, language)}\n\n" |
| "Is this code vulnerable? If so, identify the CWE and the root cause." |
| ) |
| verdict = "Vulnerable." |
| if cwes: |
| verdict += f" Most likely {', '.join(cwes)}." |
| if cwe_desc: |
| verdict += f"\n\n{cwe_desc}" |
| if grounded: |
| think = _compose_grounded("vuln_detection", "vulnerable", cwes, cwe_desc, language) |
| d_assistant, d_status = f"{think}\n\n{verdict}", "present" |
| else: |
| d_assistant, d_status = verdict, "needs_synthesis" |
| out.append( |
| Example( |
| messages=_sys_user_assistant(DETECTION_SYSTEM, user, d_assistant), |
| think_status=d_status, |
| verify={"mode": "label", "expected": "vulnerable", "cwe": cwes}, |
| metadata={"task": "vuln_detection", "language": language, "label": "vulnerable", "cwe": cwes}, |
| ) |
| ) |
|
|
| if "fix" in tasks and fixed and str(fixed).strip() != str(vuln).strip(): |
| cwe_hint = f" (it is affected by {', '.join(cwes)})" if cwes else "" |
| user = ( |
| f"The following {language} code contains a security vulnerability{cwe_hint}.\n\n" |
| f"{fence(vuln, language)}\n\n" |
| "Rewrite it to remove the vulnerability while preserving behavior. " |
| "Explain what was wrong and how your fix addresses it." |
| ) |
| answer = f"{fence(fixed, language)}" |
| if grounded: |
| think = _compose_grounded("secure_fix", "vulnerable", cwes, cwe_desc, language) |
| f_assistant, f_status = f"{think}\n\n{answer}", "present" |
| else: |
| f_assistant, f_status = answer, "needs_synthesis" |
| out.append( |
| Example( |
| messages=_sys_user_assistant(DETECTION_SYSTEM, user, f_assistant), |
| think_status=f_status, |
| verify={"mode": "backfill", "answer": answer}, |
| metadata={"task": "secure_fix", "language": language, "cwe": cwes}, |
| ) |
| ) |
|
|
| return out |
|
|
|
|
| def adapt_instruction_io(row: dict[str, Any], params: dict[str, Any]) -> list[Example]: |
| """MegaVul-style instruction / input / output rows (analysis already written).""" |
| instruction = get_first(row, ["instruction", "Instruction"]) |
| input_text = get_first(row, ["input", "Input"]) |
| output = get_first(row, ["output", "completion", "answer", "Answer"]) |
| if output is None or (instruction is None and input_text is None): |
| return [] |
|
|
| user_parts = [p for p in (instruction, input_text) if p] |
| user = "\n\n".join(str(p) for p in user_parts) |
| assistant = str(output) |
|
|
| is_vuln = as_bool_label(row.get("is_vulnerable")) |
| cwes = coerce_list(get_first(row, ["cwe_ids", "cwe", "cwe_id"])) |
| verify: dict[str, Any] | None = None |
| if has_think(assistant): |
| status = "present" |
| elif params.get("grounded_think") and is_vuln is not None: |
| label = "vulnerable" if is_vuln else "not_vulnerable" |
| think = _compose_grounded("vuln_detection", label, cwes, None, "C/C++") |
| assistant = f"{think}\n\n{assistant}" |
| status = "present" |
| else: |
| status = "needs_synthesis" |
| if is_vuln is not None: |
| verify = { |
| "mode": "label", |
| "expected": "vulnerable" if is_vuln else "not_vulnerable", |
| "cwe": cwes, |
| } |
| else: |
| verify = {"mode": "backfill", "answer": assistant} |
|
|
| system = get_first(row, ["system", "System"]) or DETECTION_SYSTEM |
| return [ |
| Example( |
| messages=_sys_user_assistant(str(system), user, assistant), |
| think_status=status, |
| verify=verify, |
| metadata={"task": "vuln_detection", "cwe": cwes}, |
| ) |
| ] |
|
|
|
|
| def adapt_system_user_assistant(row: dict[str, Any], params: dict[str, Any]) -> list[Example]: |
| """AlicanKiraz0 CVE records: System / User / Assistant triples.""" |
| system = get_first(row, ["system", "System"]) or DEFAULT_SYSTEM |
| user = get_first(row, ["user", "User", "question", "Question", "prompt"]) |
| assistant = get_first(row, ["assistant", "Assistant", "answer", "Answer", "output", "response"]) |
| if not user or assistant is None: |
| return [] |
| assistant = str(assistant) |
| status = "present" if has_think(assistant) else "needs_synthesis" |
| verify = None if status == "present" else {"mode": "backfill", "answer": assistant} |
| return [ |
| Example( |
| messages=_sys_user_assistant(str(system), str(user), assistant), |
| think_status=status, |
| verify=verify, |
| metadata={"task": "cve_knowledge"}, |
| ) |
| ] |
|
|
|
|
| def adapt_reasoning_field(row: dict[str, Any], params: dict[str, Any]) -> list[Example]: |
| """Datasets where the reasoning lives in a separate column or a `generations` list. |
| |
| Covers OpenR1-Math (`generations` already wrap <think>), Code-Reasoning (`r1_generation`), |
| and column-split sets (thinking/solution, thought/answer, claude_thinking_trajectory/...). |
| |
| params: question_field(s), reasoning_field(s), answer_field(s), task. If the reasoning text |
| already contains <think>...</think> it is used as-is; otherwise it is wrapped, then the |
| answer (if a separate field) is appended. |
| """ |
| qf = params.get("question_fields", ["question", "problem", "prompt", "instruction", "query"]) |
| rf = params.get("reasoning_fields", ["generations", "r1_generation", "thinking", "reasoning", |
| "thought", "claude_thinking_trajectory", "reasoning_content"]) |
| af = params.get("answer_fields", ["solution", "answer", "response", "output", "claude_attempt", "completion"]) |
| task = params.get("task", "reasoning") |
|
|
| user = get_first(row, qf) |
| |
| if not user and isinstance(row.get("messages"), list): |
| for m in row["messages"]: |
| if isinstance(m, dict) and (m.get("role") == "user" or m.get("from") == "human"): |
| user = m.get("content") or m.get("value") |
| break |
| if not user: |
| return [] |
|
|
| reasoning = get_first(row, rf) |
| if isinstance(reasoning, list): |
| reasoning = next((str(x) for x in reasoning if x and str(x).strip()), None) |
| answer = get_first(row, af) |
|
|
| if reasoning and has_think(str(reasoning)): |
| assistant = str(reasoning) |
| if answer and str(answer).strip() and str(answer) not in assistant: |
| assistant = f"{assistant}\n\n{answer}" |
| elif reasoning: |
| assistant = f"<think>\n{str(reasoning).strip()}\n</think>" |
| if answer and str(answer).strip(): |
| assistant += f"\n\n{answer}" |
| elif answer and has_think(str(answer)): |
| assistant = str(answer) |
| else: |
| return [] |
|
|
| system = get_first(row, ["system", "System"]) or DEFAULT_SYSTEM |
| return [ |
| Example( |
| messages=_sys_user_assistant(str(system), str(user), assistant), |
| think_status="present" if has_think(assistant) else "needs_synthesis", |
| metadata={"task": task}, |
| ) |
| ] |
|
|
|
|
| def adapt_dpo_to_sft(row: dict[str, Any], params: dict[str, Any]) -> list[Example]: |
| """CyberNative DPO: use the `chosen` (secure) answer for SFT. |
| |
| The rejected/chosen pair is retained separately for the Stage-2 DPO run; for |
| SFT we only learn the secure answer. |
| """ |
| question = get_first(row, ["question", "prompt", "instruction"]) |
| chosen = get_first(row, ["chosen", "output", "answer"]) |
| if not question or not chosen: |
| return [] |
| system = get_first(row, ["system", "System"]) or DEFAULT_SYSTEM |
| vulnerability = get_first(row, ["vulnerability"]) |
| lang = get_first(row, ["lang", "language"]) |
| user = str(question) |
| if vulnerability: |
| user += f"\n\n(Security focus: {vulnerability})" |
| chosen = str(chosen) |
| verify = None |
| if has_think(chosen): |
| status = "present" |
| elif params.get("grounded_think"): |
| think = _compose_grounded("secure_code", "", [], vulnerability, lang) |
| chosen = f"{think}\n\n{chosen}" |
| status = "present" |
| else: |
| status = "needs_synthesis" |
| verify = {"mode": "backfill", "answer": chosen} |
| return [ |
| Example( |
| messages=_sys_user_assistant(str(system) or DEFAULT_SYSTEM, user, chosen), |
| think_status=status, |
| verify=verify, |
| metadata={"task": "secure_code", "language": lang}, |
| ) |
| ] |
|
|
|
|
| def adapt_mcq_think(row: dict[str, Any], params: dict[str, Any]) -> list[Example]: |
| """theelderemo pentesting-explanations: MCQ + explanation + ready <think>.""" |
| prompt = get_first(row, ["prompt"]) or DEFAULT_SYSTEM |
| question = get_first(row, ["question"]) |
| choices = row.get("choices") |
| think = get_first(row, ["think"]) or "" |
| response = get_first(row, ["response", "explanation"]) or "" |
|
|
| |
| |
| if question: |
| user = str(question) |
| if isinstance(choices, list) and choices: |
| letters = [chr(ord("A") + i) for i in range(len(choices))] |
| rendered = "\n".join(f"{letters[i]}. {c}" for i, c in enumerate(choices)) |
| user += "\n\n" + rendered |
| assistant_parts = [] |
| block = ensure_think_closed(str(think)) |
| if block: |
| assistant_parts.append(block) |
| if response: |
| assistant_parts.append(str(response).strip()) |
| assistant = "\n\n".join(assistant_parts).strip() |
| if not assistant: |
| return [] |
| status = "present" if has_think(assistant) else "needs_synthesis" |
| system = str(prompt) if prompt else DEFAULT_SYSTEM |
| return [ |
| Example( |
| messages=_sys_user_assistant(system, user, assistant), |
| think_status=status, |
| metadata={"task": "offensive_mcq"}, |
| ) |
| ] |
|
|
| msgs = row.get("messages") |
| if isinstance(msgs, list) and msgs: |
| return adapt_chatml_conversations(row, params) |
| return [] |
|
|
|
|
| ROLE_MAP = { |
| "human": "user", |
| "user": "user", |
| "gpt": "assistant", |
| "assistant": "assistant", |
| "system": "system", |
| "tool": "tool", |
| "observation": "tool", |
| "function": "tool", |
| } |
|
|
|
|
| def adapt_chatml_conversations(row: dict[str, Any], params: dict[str, Any]) -> list[Example]: |
| """interstellarninja / generic multi-turn: conversations or messages lists.""" |
| raw = row.get("messages") if isinstance(row.get("messages"), list) else row.get("conversations") |
| if not isinstance(raw, list) or not raw: |
| return [] |
| messages: list[dict[str, str]] = [] |
| for item in raw: |
| if not isinstance(item, dict): |
| continue |
| role = item.get("role", item.get("from", item.get("speaker", ""))) |
| content = item.get("content", item.get("value", item.get("text", ""))) |
| if not role or content is None: |
| continue |
| messages.append({"role": ROLE_MAP.get(str(role).strip().lower(), str(role).strip().lower()), "content": str(content)}) |
|
|
| if not any(m["role"] == "assistant" for m in messages): |
| return [] |
|
|
| |
| |
| tools = row.get("tools") |
| if params.get("inline_tools") and tools: |
| tools_text = tools if isinstance(tools, str) else json.dumps(tools) |
| sys_msg = f"{DEFAULT_SYSTEM}\n\nAvailable tools:\n{tools_text}" |
| if messages and messages[0]["role"] == "system": |
| messages[0]["content"] = messages[0]["content"] + "\n\nAvailable tools:\n" + tools_text |
| else: |
| messages.insert(0, {"role": "system", "content": sys_msg}) |
| elif messages[0]["role"] != "system": |
| messages.insert(0, {"role": "system", "content": DEFAULT_SYSTEM}) |
|
|
| status = "present" if any( |
| m["role"] == "assistant" and has_think(m["content"]) for m in messages |
| ) else "needs_synthesis" |
| return [ |
| Example( |
| messages=messages, |
| think_status=status, |
| metadata={"task": "agentic_tool_loop"}, |
| ) |
| ] |
|
|
|
|
| ADAPTERS: dict[str, Callable[[dict[str, Any], dict[str, Any]], list[Example]]] = { |
| "detection_func_target": adapt_detection_func_target, |
| "vuln_fix_pair": adapt_vuln_fix_pair, |
| "instruction_io": adapt_instruction_io, |
| "system_user_assistant": adapt_system_user_assistant, |
| "reasoning_field": adapt_reasoning_field, |
| "dpo_to_sft": adapt_dpo_to_sft, |
| "mcq_think": adapt_mcq_think, |
| "chatml_conversations": adapt_chatml_conversations, |
| } |
|
|
|
|
| def apply_adapter(name: str, row: dict[str, Any], params: dict[str, Any] | None = None) -> list[Example]: |
| if name not in ADAPTERS: |
| raise KeyError(f"Unknown adapter {name!r}. Known: {sorted(ADAPTERS)}") |
| return ADAPTERS[name](row, params or {}) |
|
|
