115 lines
4.0 KiB
Python
115 lines
4.0 KiB
Python
"""Build a fine-tuning JSONL dataset from batch request + output files.
|
|
|
|
Joins the original request JSONL (system + user messages) with the batch
|
|
output JSONL (assistant completions) by custom_id to produce a ChatML-style
|
|
messages JSONL suitable for fine-tuning.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
from pathlib import Path
|
|
from typing import Annotated
|
|
|
|
import typer
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
HTTP_OK = 200
|
|
|
|
|
|
def load_requests(path: Path) -> dict[str, list[dict]]:
|
|
"""Parse request JSONL into {custom_id: messages}."""
|
|
results: dict[str, list[dict]] = {}
|
|
with path.open(encoding="utf-8") as handle:
|
|
for raw_line in handle:
|
|
stripped = raw_line.strip()
|
|
if not stripped:
|
|
continue
|
|
record = json.loads(stripped)
|
|
custom_id = record["custom_id"]
|
|
messages = record["body"]["messages"]
|
|
results[custom_id] = messages
|
|
return results
|
|
|
|
|
|
def load_completions(path: Path) -> dict[str, str]:
|
|
"""Parse batch output JSONL into {custom_id: assistant_content}."""
|
|
results: dict[str, str] = {}
|
|
with path.open(encoding="utf-8") as handle:
|
|
for line_number, raw_line in enumerate(handle, 1):
|
|
stripped = raw_line.strip()
|
|
if not stripped:
|
|
continue
|
|
record = json.loads(stripped)
|
|
custom_id = record["custom_id"]
|
|
response = record.get("response", {})
|
|
if response.get("status_code") != HTTP_OK:
|
|
logger.warning("Skipping %s (line %d): status %s", custom_id, line_number, response.get("status_code"))
|
|
continue
|
|
body = response.get("body", {})
|
|
choices = body.get("choices", [])
|
|
if not choices:
|
|
logger.warning("Skipping %s (line %d): no choices", custom_id, line_number)
|
|
continue
|
|
content = choices[0].get("message", {}).get("content", "")
|
|
if not content:
|
|
logger.warning("Skipping %s (line %d): empty content", custom_id, line_number)
|
|
continue
|
|
results[custom_id] = content
|
|
return results
|
|
|
|
|
|
def main(
|
|
requests_path: Annotated[Path, typer.Option("--requests", help="Batch request JSONL")] = Path(
|
|
"output/openai_batch/requests.jsonl",
|
|
),
|
|
batch_output: Annotated[Path, typer.Option("--batch-output", help="Batch output JSONL")] = Path(
|
|
"batch_69d84558d91c819091d53f08d78f9fd6_output.jsonl",
|
|
),
|
|
output_path: Annotated[Path, typer.Option("--output", help="Fine-tuning JSONL output")] = Path(
|
|
"output/finetune_dataset.jsonl",
|
|
),
|
|
log_level: Annotated[str, typer.Option(help="Log level")] = "INFO",
|
|
) -> None:
|
|
"""Build fine-tuning dataset by joining request and output JSONL files."""
|
|
logging.basicConfig(level=log_level, format="%(asctime)s %(levelname)s %(name)s: %(message)s")
|
|
|
|
logger.info("Loading requests from %s", requests_path)
|
|
requests = load_requests(requests_path)
|
|
logger.info("Loaded %d requests", len(requests))
|
|
|
|
logger.info("Loading completions from %s", batch_output)
|
|
completions = load_completions(batch_output)
|
|
logger.info("Loaded %d completions", len(completions))
|
|
|
|
output_path.parent.mkdir(parents=True, exist_ok=True)
|
|
matched = 0
|
|
skipped = 0
|
|
|
|
with output_path.open("w", encoding="utf-8") as handle:
|
|
for custom_id, messages in requests.items():
|
|
assistant_content = completions.get(custom_id)
|
|
if assistant_content is None:
|
|
skipped += 1
|
|
continue
|
|
|
|
example = {
|
|
"messages": [*messages, {"role": "assistant", "content": assistant_content}],
|
|
}
|
|
handle.write(json.dumps(example, ensure_ascii=False))
|
|
handle.write("\n")
|
|
matched += 1
|
|
|
|
logger.info("Wrote %d examples to %s (skipped %d unmatched)", matched, output_path, skipped)
|
|
|
|
|
|
def cli() -> None:
|
|
"""Typer entry point."""
|
|
typer.run(main)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
cli()
|