File size: 2,218 Bytes
a822d67 | 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 | # training/document_reasoning_loader.py
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 |