File size: 2,103 Bytes
61449ba | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 | 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)
|