| |
| """Normalize raw JSONL rows into the training ChatML contract.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| from pathlib import Path |
| from typing import Any |
|
|
|
|
| DEFAULT_SYSTEM = "Authorized security research and education context." |
|
|
|
|
| ROLE_MAP = { |
| "human": "user", |
| "user": "user", |
| "gpt": "assistant", |
| "assistant": "assistant", |
| "system": "system", |
| "tool": "tool", |
| } |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--input", required=True) |
| parser.add_argument("--output", required=True) |
| parser.add_argument("--source", required=True) |
| parser.add_argument("--license", default="missing") |
| parser.add_argument("--split", default="unsplit") |
| parser.add_argument("--system", default=DEFAULT_SYSTEM) |
| parser.add_argument( |
| "--missing-think-policy", |
| choices=["reject", "wrap", "allow"], |
| default="reject", |
| help="How to handle assistant rows that do not contain <think> blocks.", |
| ) |
| return parser.parse_args() |
|
|
|
|
| def read_jsonl(path: Path) -> list[dict[str, Any]]: |
| rows = [] |
| with path.open("r", encoding="utf-8") as fh: |
| for line_no, line in enumerate(fh, start=1): |
| line = line.strip() |
| if not line: |
| continue |
| try: |
| row = json.loads(line) |
| except json.JSONDecodeError as exc: |
| raise ValueError(f"Invalid JSON in {path}:{line_no}: {exc}") from exc |
| if not isinstance(row, dict): |
| raise ValueError(f"Expected object row in {path}:{line_no}") |
| rows.append(row) |
| return rows |
|
|
|
|
| def stable_value(row: dict[str, Any], keys: list[str]) -> str | None: |
| for key in keys: |
| value = row.get(key) |
| if value not in (None, ""): |
| return str(value) |
| return None |
|
|
|
|
| def normalize_role(role: str) -> str: |
| return ROLE_MAP.get(role.strip().lower(), role.strip().lower()) |
|
|
|
|
| def normalize_message_list(items: list[Any]) -> list[dict[str, str]]: |
| messages = [] |
| for item in items: |
| 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": normalize_role(str(role)), "content": str(content)}) |
| return messages |
|
|
|
|
| def row_to_messages(row: dict[str, Any], system_prompt: str) -> list[dict[str, str]] | None: |
| if isinstance(row.get("messages"), list): |
| messages = normalize_message_list(row["messages"]) |
| elif isinstance(row.get("conversations"), list): |
| messages = normalize_message_list(row["conversations"]) |
| else: |
| system = stable_value(row, ["system", "System"]) or system_prompt |
| instruction = stable_value(row, ["instruction", "Instruction"]) |
| prompt = stable_value(row, ["prompt", "question", "Question", "user", "User"]) |
| input_text = stable_value(row, ["input", "Input"]) |
| answer = stable_value(row, ["output", "completion", "answer", "Answer", "assistant", "Assistant", "positive"]) |
|
|
| user_parts = [] |
| if instruction: |
| user_parts.append(instruction) |
| if prompt and prompt not in user_parts: |
| user_parts.append(prompt) |
| if input_text: |
| user_parts.append(input_text) |
|
|
| if not user_parts or answer is None: |
| return None |
|
|
| messages = [ |
| {"role": "system", "content": system}, |
| {"role": "user", "content": "\n\n".join(user_parts)}, |
| {"role": "assistant", "content": answer}, |
| ] |
|
|
| if not messages: |
| return None |
| if messages[0]["role"] != "system": |
| messages.insert(0, {"role": "system", "content": system_prompt}) |
| return messages |
|
|
|
|
| def has_think(content: str) -> bool: |
| return "<think>" in content and "</think>" in content |
|
|
|
|
| def enforce_think(messages: list[dict[str, str]], policy: str) -> bool: |
| keep = True |
| for message in messages: |
| if message["role"] != "assistant": |
| continue |
| if has_think(message["content"]): |
| continue |
| if policy == "reject": |
| keep = False |
| elif policy == "wrap": |
| message["content"] = ( |
| "<think>\n" |
| "No source reasoning trace was provided; this row is for pipeline smoke testing only.\n" |
| "</think>\n\n" |
| + message["content"] |
| ) |
| return keep |
|
|
|
|
| def main() -> int: |
| args = parse_args() |
| input_path = Path(args.input) |
| output_path = Path(args.output) |
| rows = read_jsonl(input_path) |
| output_path.parent.mkdir(parents=True, exist_ok=True) |
|
|
| kept = 0 |
| rejected_no_schema = 0 |
| rejected_no_think = 0 |
|
|
| with output_path.open("w", encoding="utf-8") as out: |
| for index, row in enumerate(rows): |
| messages = row_to_messages(row, args.system) |
| if messages is None: |
| rejected_no_schema += 1 |
| continue |
| if not enforce_think(messages, args.missing_think_policy): |
| rejected_no_think += 1 |
| continue |
| row_id = stable_value(row, ["id", "_id", "key", "cve_id", "sha"]) or str(index) |
| normalized = { |
| "id": f"{args.source}:{row_id}", |
| "source": args.source, |
| "license": args.license, |
| "messages": messages, |
| "metadata": { |
| "split": args.split, |
| "row_index": index, |
| }, |
| } |
| out.write(json.dumps(normalized, sort_keys=True) + "\n") |
| kept += 1 |
|
|
| print( |
| json.dumps( |
| { |
| "input": str(input_path), |
| "output": str(output_path), |
| "rows_in": len(rows), |
| "rows_out": kept, |
| "rejected_no_schema": rejected_no_schema, |
| "rejected_no_think": rejected_no_think, |
| "missing_think_policy": args.missing_think_policy, |
| }, |
| indent=2, |
| ) |
| ) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
