data_preprocessing/process_reasoning_dpo_jsonl.py

from __future__ import annotations

import argparse
from pathlib import Path
from typing import Iterable, Iterator

from .common import (
    make_user_assistant_record,
    read_jsonl,
    render_records,
    resolve_paths,
    write_jsonl,
)


def iter_records(
    paths: str | Path | Iterable[str | Path],
    *,
    assistant_column: str = "chosen",
    include_source: bool = False,
) -> Iterator[dict[str, str]]:
    for path in resolve_paths(paths):
        for _, row in read_jsonl(path):
            record = make_user_assistant_record(
                row.get("prompt"),
                row.get(assistant_column),
                source=str(path),
                include_source=include_source,
            )

            if record is not None:
                yield record


def load_records(
    paths: str | Path | Iterable[str | Path],
    *,
    assistant_column: str = "chosen",
    include_source: bool = False,
) -> list[dict[str, str]]:
    return list(
        iter_records(
            paths,
            assistant_column=assistant_column,
            include_source=include_source,
        )
    )


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Normalize ReasoningDPO JSONL files to user/assistant rows."
    )
    parser.add_argument("inputs", nargs="+", help="Input JSONL path(s).")
    parser.add_argument("--output", help="Optional output JSONL path.")
    parser.add_argument("--take", type=int, default=0, help="Print the first N rows.")
    parser.add_argument(
        "--assistant-column",
        default="chosen",
        choices=("chosen", "rejected"),
        help="Which DPO response column to use as assistant text.",
    )
    parser.add_argument("--include-source", action="store_true")
    args = parser.parse_args()

    records = load_records(
        args.inputs,
        assistant_column=args.assistant_column,
        include_source=args.include_source,
    )

    if args.output:
        write_jsonl(records, args.output)

    if args.take:
        print(render_records(records[: args.take]))


if __name__ == "__main__":
    main()