File size: 1,876 Bytes
a797f9a | 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 | from __future__ import annotations
from typing import Any
from datasets import load_dataset
import utils.configs_loader as ProjectConfigs
def load_document_grounding_dpo(tokenizer, dataset_key: str = "DOCUMENT_GROUNDING_DPO"):
cfg = ProjectConfigs.get(("DATASET", dataset_key))
ds = load_dataset(
"json",
data_files=cfg.Path,
split=cfg.Split,
)
test_size = getattr(cfg, "TestSize", 0.0)
seed = getattr(cfg, "Seed", 42)
if test_size and test_size > 0:
split = ds.train_test_split(
test_size=test_size,
seed=seed,
)
train_ds = split["train"]
eval_ds = split["test"]
else:
train_ds = ds
eval_ds = None
def convert(example: dict[str, Any]) -> dict[str, str]:
prompt_messages = example["prompt"]
chosen_messages = example["chosen"]
rejected_messages = example["rejected"]
prompt_text = tokenizer.apply_chat_template(
prompt_messages,
tokenize=False,
add_generation_prompt=True,
)
chosen_text = chosen_messages[0]["content"]
rejected_text = rejected_messages[0]["content"]
if tokenizer.eos_token:
chosen_text += tokenizer.eos_token
rejected_text += tokenizer.eos_token
return {
"prompt": prompt_text,
"chosen": chosen_text,
"rejected": rejected_text,
}
remove_cols = train_ds.column_names
train_ds = train_ds.map(
convert,
remove_columns=remove_cols,
desc="Converting document grounding DPO train set",
)
if eval_ds is not None:
eval_ds = eval_ds.map(
convert,
remove_columns=remove_cols,
desc="Converting document grounding DPO eval set",
)
return train_ds, eval_ds |