data_preprocessing/process_complex_reasoning_bespoke_parquet.py

from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any, Iterator

from .common import make_user_assistant_record, render_records, write_jsonl


ROLE_ALIASES = {
    "human": "user",
    "user": "user",
    "assistant": "assistant",
    "gpt": "assistant",
    "bot": "assistant",
    "system": "system",
}


def _missing_parquet_reader_error() -> ImportError:
    return ImportError(
        "ComplexReasoningBespoke.parquet needs a parquet reader. "
        "Install one of: pyarrow, polars, or pandas with parquet support. "
        "Recommended: pip install pyarrow"
    )


def _iter_rows_pyarrow(path: Path, *, batch_size: int) -> Iterator[dict[str, Any]]:
    import pyarrow.parquet as pq

    parquet_file = pq.ParquetFile(path)

    for batch in parquet_file.iter_batches(batch_size=batch_size):
        for row in batch.to_pylist():
            yield row


def _iter_rows_polars(path: Path) -> Iterator[dict[str, Any]]:
    import polars as pl

    frame = pl.read_parquet(path)

    for row in frame.iter_rows(named=True):
        yield row


def _iter_rows_pandas(path: Path) -> Iterator[dict[str, Any]]:
    import pandas as pd

    frame = pd.read_parquet(path)

    for row in frame.to_dict(orient="records"):
        yield row


def iter_parquet_rows(
    path: str | Path,
    *,
    batch_size: int = 512,
) -> Iterator[dict[str, Any]]:
    parquet_path = Path(path)
    errors: list[Exception] = []

    for loader in (_iter_rows_pyarrow, _iter_rows_polars, _iter_rows_pandas):
        try:
            if loader is _iter_rows_pyarrow:
                yield from loader(parquet_path, batch_size=batch_size)
            else:
                yield from loader(parquet_path)

            return
        except ModuleNotFoundError as exc:
            errors.append(exc)
        except ImportError as exc:
            errors.append(exc)

    raise _missing_parquet_reader_error() from (errors[-1] if errors else None)


def _coerce_conversations(value: Any) -> list[dict[str, Any]]:
    if value is None:
        return []

    if isinstance(value, str):
        try:
            value = json.loads(value)
        except json.JSONDecodeError:
            return []

    if not isinstance(value, list):
        return []

    return [turn for turn in value if isinstance(turn, dict)]


def _role(turn: dict[str, Any]) -> str:
    raw_role = turn.get("from", turn.get("role", ""))
    return ROLE_ALIASES.get(str(raw_role).strip().lower(), str(raw_role).strip().lower())


def _content(turn: dict[str, Any]) -> str:
    return str(turn.get("value", turn.get("content", ""))).strip()


def _iter_user_assistant_pairs(
    conversations: list[dict[str, Any]],
    *,
    system_prompt: str = "",
    include_system: bool = False,
) -> Iterator[tuple[str, str]]:
    pending_user = ""

    for turn in conversations:
        role = _role(turn)
        content = _content(turn)

        if not content:
            continue

        if role == "user":
            pending_user = content
            continue

        if role == "assistant" and pending_user:
            user = pending_user

            if include_system and system_prompt:
                user = f"System:\n{system_prompt.strip()}\n\nUser:\n{user}"

            yield user, content
            pending_user = ""


def iter_records(
    path: str | Path,
    *,
    include_system: bool = False,
    include_source: bool = False,
    limit: int | None = None,
    batch_size: int = 512,
) -> Iterator[dict[str, str]]:
    emitted = 0
    parquet_path = Path(path)

    for row in iter_parquet_rows(parquet_path, batch_size=batch_size):
        system_prompt = str(row.get("system", "") or "")
        conversations = _coerce_conversations(row.get("conversations"))

        for user, assistant in _iter_user_assistant_pairs(
            conversations,
            system_prompt=system_prompt,
            include_system=include_system,
        ):
            record = make_user_assistant_record(
                user,
                assistant,
                source=str(parquet_path),
                include_source=include_source,
            )

            if record is None:
                continue

            yield record
            emitted += 1

            if limit is not None and emitted >= limit:
                return


def load_records(
    path: str | Path,
    *,
    include_system: bool = False,
    include_source: bool = False,
    limit: int | None = None,
    batch_size: int = 512,
) -> list[dict[str, str]]:
    return list(
        iter_records(
            path,
            include_system=include_system,
            include_source=include_source,
            limit=limit,
            batch_size=batch_size,
        )
    )


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Normalize ComplexReasoningBespoke parquet to user/assistant rows."
    )
    parser.add_argument("input", help="Input parquet path.")
    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("--limit", type=int, help="Stop after N emitted rows.")
    parser.add_argument("--batch-size", type=int, default=512)
    parser.add_argument("--include-system", action="store_true")
    parser.add_argument("--include-source", action="store_true")
    args = parser.parse_args()

    records = load_records(
        args.input,
        include_system=args.include_system,
        include_source=args.include_source,
        limit=args.limit,
        batch_size=args.batch_size,
    )

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

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


if __name__ == "__main__":
    main()