22ae21a6c2
* add kto support * test cleanup * fix outdated comment * fix llama3 ultra * chore: lint * update to use rl_beta instead of dpo_beta --------- Co-authored-by: Wing Lian <wing.lian@gmail.com>
40 lines
1.4 KiB
Python
40 lines
1.4 KiB
Python
"""
|
|
User-defined KTO strategies
|
|
"""
|
|
# pylint: disable=duplicate-code
|
|
|
|
|
|
def default(cfg, dataset_idx=0, **kwargs): # pylint: disable=unused-argument
|
|
ds_cfg = cfg["datasets"][dataset_idx]["type"]
|
|
if not isinstance(ds_cfg, dict):
|
|
raise ValueError(
|
|
f"User-defined dataset type must be a dictionary. Got: {ds_cfg}"
|
|
)
|
|
field_prompt = ds_cfg.get("field_prompt", "prompt")
|
|
field_system = ds_cfg.get("field_system", "system")
|
|
field_completion = ds_cfg.get("field_completion", "completion")
|
|
field_label = ds_cfg.get("field_label", "label")
|
|
prompt_format = ds_cfg.get("prompt_format")
|
|
if not prompt_format:
|
|
prompt_format = "{" + field_prompt + "}"
|
|
completion_format = ds_cfg.get("completion_format")
|
|
if not completion_format:
|
|
chosen_format = "{" + field_completion + "}"
|
|
|
|
def transform_fn(sample):
|
|
if (
|
|
"{" + field_system + "}" in prompt_format
|
|
and field_system in sample
|
|
and sample[field_system]
|
|
):
|
|
sample["prompt"] = prompt_format.format(
|
|
system=sample[field_system], prompt=sample[field_prompt]
|
|
)
|
|
else:
|
|
sample["prompt"] = prompt_format.format(prompt=sample["prompt"])
|
|
sample["completion"] = chosen_format.format(chosen=sample[field_completion])
|
|
sample["label"] = sample[field_label]
|
|
return sample
|
|
|
|
return transform_fn
|