code/training/scripts/synthesize_think.py

#!/usr/bin/env python3
"""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()
    # check the negative phrase first so "not vulnerable" doesn't match "vulnerable"
    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:  # noqa: BLE001
                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:  # backfill -> keep the known-good answer, prepend the verified reasoning
            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())