data_preprocessing/common.py

from __future__ import annotations

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


BESPOKE_TOKEN_REPLACEMENTS = {
    "<|begin_of_thought|>": "<think>",
    "<|end_of_thought|>": "</think>",
    "<|begin_of_solution|>": "<answer>",
    "<|end_of_solution|>": "</answer>",
}


def clean_text(value: Any) -> str:
    if value is None:
        return ""

    return str(value).strip()


def normalize_reasoning_tokens(value: Any) -> str:
    text = clean_text(value)

    for old, new in BESPOKE_TOKEN_REPLACEMENTS.items():
        text = text.replace(old, new)

    return text.strip()


def make_user_assistant_record(
    user: Any,
    assistant: Any,
    *,
    source: str | None = None,
    include_source: bool = False,
) -> dict[str, str] | None:
    user_text = clean_text(user)
    assistant_text = normalize_reasoning_tokens(assistant)

    if not user_text or not assistant_text:
        return None

    record = {
        "user": user_text,
        "assistant": assistant_text,
    }

    if include_source and source:
        record["source"] = source

    return record


def read_jsonl(path: str | Path) -> Iterator[tuple[int, dict[str, Any]]]:
    jsonl_path = Path(path)

    with jsonl_path.open("r", encoding="utf-8") as f:
        for row_index, line in enumerate(f):
            line = line.strip()

            if not line:
                continue

            yield row_index, json.loads(line)


def write_jsonl(records: Iterable[dict[str, Any]], path: str | Path) -> None:
    out_path = Path(path)
    out_path.parent.mkdir(parents=True, exist_ok=True)

    with out_path.open("w", encoding="utf-8") as f:
        for record in records:
            f.write(json.dumps(record, ensure_ascii=False))
            f.write("\n")


def resolve_paths(paths: str | Path | Iterable[str | Path]) -> list[Path]:
    if isinstance(paths, (str, Path)):
        return [Path(paths)]

    return [Path(path) for path in paths]


def render_records(records: Iterable[dict[str, Any]]) -> str:
    return json.dumps(list(records), ensure_ascii=False, indent=2)