lmflow.pipeline.dpo_aligner

Module Contents

Classes

DPOAligner

Functions

get_paired_dataset(→ datasets.Dataset)

Load dataset and convert it to the necessary format.

lmflow.pipeline.dpo_aligner.get_paired_dataset(data_root: str, data_dir: str, sanity_check: bool = False, cache_dir: str | None = None, num_proc=24) → datasets.Dataset

Load dataset and convert it to the necessary format.

The dataset is converted to a dictionary with the following structure: {

‘prompt’: List[str], ‘chosen’: List[str], ‘rejected’: List[str],

}

Prompts are structured as follows:

“Question: “ + <prompt> + “

Answer: “

class lmflow.pipeline.dpo_aligner.DPOAligner(model_args, data_args, aligner_args)

Bases: lmflow.pipeline.base_aligner.BaseAligner

_initialize_trainer(model, tokenizer)

_load_dataset()

align(model, dataset, reward_model)