| |
|
|
| from __future__ import annotations |
|
|
| from typing import Any |
|
|
| from datasets import load_dataset, Dataset |
|
|
|
|
| def load_document_reasoning_sft( |
| tokenizer, |
| train_path: str, |
| validation_path: str | None = None, |
| stress_validation_path: str | None = None, |
| add_eos_to_completion: bool = True, |
| ): |
| """ |
| Loads JSONL made by DocumentsReasoningDataset.py and converts: |
| |
| messages = [system, user, assistant] |
| |
| into: |
| |
| prompt = rendered system+user chat |
| completion = assistant content |
| |
| This is safer than relying on assistant_only_loss=True because |
| prompt-completion SFT naturally trains only the completion side. |
| """ |
|
|
| data_files = {"train": train_path} |
|
|
| if validation_path: |
| data_files["validation"] = validation_path |
|
|
| if stress_validation_path: |
| data_files["validation_stress"] = stress_validation_path |
|
|
| raw = load_dataset("json", data_files=data_files) |
|
|
| def convert(example: dict[str, Any]) -> dict[str, str]: |
| messages = example["messages"] |
|
|
| if len(messages) < 2: |
| raise ValueError("Expected at least user + assistant messages.") |
|
|
| if messages[-1]["role"] != "assistant": |
| raise ValueError("Last message must be assistant.") |
|
|
| prompt_messages = messages[:-1] |
| assistant_content = messages[-1]["content"] |
|
|
| prompt = tokenizer.apply_chat_template( |
| prompt_messages, |
| tokenize=False, |
| add_generation_prompt=True, |
| ) |
|
|
| completion = assistant_content |
|
|
| if add_eos_to_completion and tokenizer.eos_token: |
| completion = completion + tokenizer.eos_token |
|
|
| return { |
| "prompt": prompt, |
| "completion": completion, |
| } |
|
|
| keep_cols = ["prompt", "completion"] |
|
|
| converted = raw.map( |
| convert, |
| remove_columns=raw["train"].column_names, |
| desc="Converting messages to prompt/completion", |
| ) |
|
|
| train_dataset = converted["train"] |
| eval_dataset = converted.get("validation", None) |
| stress_eval_dataset = converted.get("validation_stress", None) |
|
|
| return train_dataset, eval_dataset, stress_eval_dataset |
