| |
| """Backfill <think> reasoning for rows that lack it (rejection sampling / STaR). |
| |
| Consumes the ``*.to_synthesize.jsonl`` queue emitted by ``build_sft_dataset.py`` |
| and produces ``*.synthesized.jsonl`` rows (with a verified ``<think>`` block) |
| that ``build_sft_dataset.py --include-synthesized`` folds back into the ready mix. |
| |
| Two verification modes (carried on each record's ``verify`` block): |
| |
| label -- rejection sampling. Sample N teacher completions; keep the first |
| whose final verdict matches the ground-truth label |
| (vulnerable / not_vulnerable). This is the STaR / RFT idea applied |
| to vuln detection: only correct reasoning survives. |
| |
| backfill -- the answer is fixed (secure-fix code, CVE write-up). Ask the teacher |
| to produce a reasoning trace that justifies the given answer; accept |
| a non-degenerate <think> block and keep the original answer. |
| |
| Teacher endpoint is any OpenAI-compatible chat-completions server (e.g. a vLLM |
| host). Configure via env or flags: |
| TEACHER_BASE_URL (default http://localhost:8000/v1) |
| TEACHER_MODEL |
| TEACHER_API_KEY (default EMPTY) |
| |
| Use ``--mock`` to run the whole pipeline offline (no endpoint) for testing: it |
| fabricates a deterministic, clearly-labelled trace so plumbing can be validated |
| without a GPU. Mock output must never be used for a real Stage 1 run. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import os |
| import sys |
| import urllib.request |
| from pathlib import Path |
| from typing import Any |
|
|
| THINK_INSTRUCTION = ( |
| "Think step by step inside a single <think>...</think> block, then give your " |
| "final answer after it. Always include the <think> block." |
| ) |
|
|
|
|
| def read_jsonl(path: Path) -> list[dict[str, Any]]: |
| rows: list[dict[str, Any]] = [] |
| with path.open("r", encoding="utf-8") as fh: |
| for line in fh: |
| line = line.strip() |
| if line: |
| rows.append(json.loads(line)) |
| return rows |
|
|
|
|
| def call_teacher(base_url: str, model: str, api_key: str, messages: list[dict[str, str]], |
| temperature: float, max_tokens: int, timeout: int = 120) -> str: |
| payload = { |
| "model": model, |
| "messages": messages, |
| "temperature": temperature, |
| "max_tokens": max_tokens, |
| } |
| req = urllib.request.Request( |
| base_url.rstrip("/") + "/chat/completions", |
| data=json.dumps(payload).encode("utf-8"), |
| headers={"Content-Type": "application/json", "Authorization": f"Bearer {api_key}"}, |
| method="POST", |
| ) |
| with urllib.request.urlopen(req, timeout=timeout) as resp: |
| body = json.loads(resp.read()) |
| return body["choices"][0]["message"]["content"] |
|
|
|
|
| def mock_completion(record: dict[str, Any], sample_idx: int) -> str: |
| """Deterministic offline stand-in for the teacher (testing only).""" |
| verify = record.get("verify", {}) or {} |
| mode = verify.get("mode", "backfill") |
| if mode == "label": |
| expected = verify.get("expected", "not_vulnerable") |
| cwe = ", ".join(verify.get("cwe", []) or []) or "no specific CWE" |
| verdict = "Vulnerable." if expected == "vulnerable" else "Not vulnerable." |
| return ( |
| f"<think>\n[mock trace] Inspecting the function for memory-safety and " |
| f"input-validation issues; the evidence points to {expected} ({cwe}).\n</think>\n\n" |
| f"{verdict} Most likely {cwe}." |
| ) |
| answer = verify.get("answer", record.get("reference_answer", "")) |
| return ( |
| "<think>\n[mock trace] Reasoning toward the provided secure answer: identify " |
| "the flaw, then justify why the answer resolves it.\n</think>\n\n" + answer |
| ) |
|
|
|
|
| def extract_think_and_body(text: str) -> tuple[str, str]: |
| if "<think>" in text and "</think>" in text: |
| start = text.index("<think>") |
| end = text.index("</think>") + len("</think>") |
| return text[start:end], text[end:].strip() |
| return "", text.strip() |
|
|
|
|
| def verdict_of(text: str) -> str | None: |
| low = text.lower() |
| |
| if "not vulnerable" in low or "no vulnerability" in low or "not a vulnerability" in low: |
| return "not_vulnerable" |
| if "vulnerable" in low or "vulnerability" in low: |
| return "vulnerable" |
| return None |
|
|
|
|
| def think_is_degenerate(think: str) -> bool: |
| inner = think.replace("<think>", "").replace("</think>", "").strip() |
| return len(inner) < 40 |
|
|
|
|
| def synthesize_one(record: dict[str, Any], args: argparse.Namespace) -> dict[str, Any] | None: |
| verify = record.get("verify", {}) or {} |
| mode = verify.get("mode", "backfill") |
| prompt_messages = record.get("prompt_messages", []) |
| teacher_messages = list(prompt_messages) |
| if teacher_messages and teacher_messages[0].get("role") == "system": |
| teacher_messages[0] = { |
| "role": "system", |
| "content": teacher_messages[0]["content"] + "\n\n" + THINK_INSTRUCTION, |
| } |
| else: |
| teacher_messages.insert(0, {"role": "system", "content": THINK_INSTRUCTION}) |
|
|
| for sample_idx in range(args.samples): |
| if args.mock: |
| completion = mock_completion(record, sample_idx) |
| else: |
| try: |
| completion = call_teacher( |
| args.base_url, args.model, args.api_key, teacher_messages, |
| args.temperature, args.max_tokens, |
| ) |
| except Exception as exc: |
| print(f"[warn] teacher call failed for {record.get('id')}: {exc!r}", file=sys.stderr) |
| continue |
|
|
| think, body = extract_think_and_body(completion) |
| if not think or think_is_degenerate(think): |
| continue |
|
|
| if mode == "label": |
| expected = verify.get("expected") |
| if verdict_of(body) != expected: |
| continue |
| assistant = completion if "<think>" in completion else f"{think}\n\n{body}" |
| else: |
| answer = verify.get("answer", record.get("reference_answer", "")) |
| assistant = f"{think}\n\n{answer}".strip() |
|
|
| messages = list(prompt_messages) + [{"role": "assistant", "content": assistant}] |
| return { |
| "id": record.get("id"), |
| "source": record.get("source"), |
| "license": record.get("license", "missing"), |
| "group": record.get("group"), |
| "messages": messages, |
| "metadata": { |
| **(record.get("metadata") or {}), |
| "think_status": "present", |
| "synthesized": True, |
| "synthesis_mode": mode, |
| "synthesis_mock": bool(args.mock), |
| }, |
| } |
| return None |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| parser.add_argument("--input", required=True) |
| parser.add_argument("--output", required=True) |
| parser.add_argument("--samples", type=int, default=4, help="Rollouts per row before giving up.") |
| parser.add_argument("--temperature", type=float, default=0.7) |
| parser.add_argument("--max-tokens", type=int, default=1024) |
| parser.add_argument("--limit", type=int, default=None, help="Only process the first N rows.") |
| parser.add_argument("--base-url", default=os.getenv("TEACHER_BASE_URL", "http://localhost:8000/v1")) |
| parser.add_argument("--model", default=os.getenv("TEACHER_MODEL", "teacher")) |
| parser.add_argument("--api-key", default=os.getenv("TEACHER_API_KEY", "EMPTY")) |
| parser.add_argument("--mock", action="store_true", help="Offline deterministic teacher (testing only).") |
| return parser.parse_args() |
|
|
|
|
| def main() -> int: |
| args = parse_args() |
| rows = read_jsonl(Path(args.input)) |
| if args.limit is not None: |
| rows = rows[: args.limit] |
|
|
| out_path = Path(args.output) |
| out_path.parent.mkdir(parents=True, exist_ok=True) |
|
|
| kept = 0 |
| failed = 0 |
| with out_path.open("w", encoding="utf-8") as out: |
| for record in rows: |
| result = synthesize_one(record, args) |
| if result is None: |
| failed += 1 |
| continue |
| out.write(json.dumps(result, ensure_ascii=False, sort_keys=True) + "\n") |
| kept += 1 |
|
|
| print(json.dumps( |
| { |
| "input": args.input, |
| "output": args.output, |
| "rows_in": len(rows), |
| "synthesized": kept, |
| "failed": failed, |
| "mock": bool(args.mock), |
| "samples_per_row": args.samples, |
| }, |
| indent=2, |
| )) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
