scripts/filter_sft_quality.py

#!/usr/bin/env python3
"""Quality gate for SFT message rows — reject bad assistant targets."""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
from sft_pipeline_utils import (
    asks_for_code,
    assistant_repeats_user_prompt,
    has_code,
    verification_is_vague,
)

ROOT = Path(__file__).resolve().parent.parent
DEFAULT_INPUT = ROOT / "data" / "train" / "mythos_sft_messages.jsonl"
DEFAULT_CLEAN = ROOT / "data" / "train" / "mythos_sft_messages_clean.jsonl"
DEFAULT_REJECTED = ROOT / "data" / "rejected" / "sft_rejected.jsonl"


def reject_reasons(user: str, assistant: str) -> list[str]:
    reasons: list[str] = []
    a = str(assistant or "")
    u = str(user or "")

    if asks_for_code(u) and not has_code(a):
        reasons.append("code_request_no_code")

    if assistant_repeats_user_prompt(u, a):
        reasons.append("repeats_user_prompt")

    if "additional targeted edits" in a.lower():
        reasons.append("additional_targeted_edits")

    if a.startswith("Diagnosis:") and asks_for_code(u):
        reasons.append("code_request_starts_with_diagnosis")

    if "Diagnosis:" in a and "Implementation:" in a:
        impl = a.split("Implementation:", 1)[1].split("Verification:", 1)[0].strip()
        if not impl or (not has_code(impl) and len(impl) < 40):
            reasons.append("diagnosis_without_implementation")
    elif a.startswith("Here is the complete code:"):
        code_block = a.split("Here is the complete code:", 1)[1].split("Verification:", 1)[0].strip()
        if not code_block or not has_code(code_block):
            reasons.append("code_header_without_code")
    elif not has_code(a) and asks_for_code(u):
        reasons.append("code_request_no_code")

    ver_section = ""
    if "Verification:" in a:
        ver_section = a.split("Verification:", 1)[1].strip()
    if asks_for_code(u) and verification_is_vague(ver_section) and has_code(a):
        # soft — only reject if verification completely empty
        if not ver_section:
            reasons.append("missing_verification")

    return reasons


def main() -> int:
    parser = argparse.ArgumentParser(description="Filter low-quality SFT rows.")
    parser.add_argument("--input", default=str(DEFAULT_INPUT))
    parser.add_argument("--clean-output", default=str(DEFAULT_CLEAN))
    parser.add_argument("--rejected-output", default=str(DEFAULT_REJECTED))
    args = parser.parse_args()

    input_path = Path(args.input)
    clean_path = Path(args.clean_output)
    rejected_path = Path(args.rejected_output)
    clean_path.parent.mkdir(parents=True, exist_ok=True)
    rejected_path.parent.mkdir(parents=True, exist_ok=True)

    if not input_path.exists():
        print(f"Error: input not found: {input_path}", file=sys.stderr)
        return 1

    clean: list[dict] = []
    rejected: list[dict] = []

    with input_path.open("r", encoding="utf-8") as f:
        for line_num, line in enumerate(f, 1):
            if not line.strip():
                continue
            row = json.loads(line)
            user = row["messages"][1]["content"]
            assistant = row["messages"][2]["content"]
            reasons = reject_reasons(user, assistant)
            if reasons:
                rejected.append({
                    "line": line_num,
                    "reasons": reasons,
                    "user_prompt": user,
                    "assistant_preview": assistant[:600],
                    "row": row,
                })
            else:
                clean.append(row)

    with clean_path.open("w", encoding="utf-8") as f:
        for row in clean:
            f.write(json.dumps(row, ensure_ascii=False) + "\n")

    with rejected_path.open("w", encoding="utf-8") as f:
        for item in rejected:
            f.write(json.dumps(item, ensure_ascii=False) + "\n")

    print(f"Input:    {sum(1 for _ in input_path.open(encoding='utf-8') if _.strip())}")
    print(f"Clean:    {len(clean)} -> {clean_path}")
    print(f"Rejected: {len(rejected)} -> {rejected_path}")
    if rejected:
        from collections import Counter
        c = Counter(r for item in rejected for r in item["reasons"])
        print("Rejection reasons:", dict(c))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())