lmflow.pipeline.raft_aligner

The Aligner class simplifies the process of running alignment.

Module Contents

Classes

RaftAligner

Initializes the RaftAligner class with given arguments.

Attributes

logger

lmflow.pipeline.raft_aligner.logger

class lmflow.pipeline.raft_aligner.RaftAligner(model_args, data_args, aligner_args, *args, **kwargs)

Bases: lmflow.pipeline.base_aligner.BaseAligner

Initializes the RaftAligner class with given arguments.

Parameters:

model_argsModelArguments object.

Contains the arguments required to load the model.

data_argsDatasetArguments object.

Contains the arguments required to load the dataset.

raft_aligner_argsRaftAlignerArguments object.

Contains the arguments required to perform alignment.

argsOptional.

Positional arguments.

kwargsOptional.

Keyword arguments.

_initialize_trainer(model, tokenizer, training_args)

This function takes the model and tokenizer as the input and initialize the trainer.

_load_dataset(selected_dataset, model, tokenizer, model_args, data_args, training_args)

This function prepares the dataset for every iteration.

_load_input_dataset(dataset, tokenizer)

Load input dataset (i.e. prompt/question dataset) for training.

Args:

dataset: A Dataset object.

The dataset to be loaded.

Returns:

dataloader (torch.utils.data.DataLoader):

The dataloader for the dataset.

_clean_text(text)

_discard_sample(text)

_get_batch_dataset_top(model, batch_input, alpha=0.2, iter_id=0, local_rank=0, output_min_length=16, output_max_length=48, infer_batch_size=8, generation_kwargs={}, tokenizer=None, training_args=None, reward_model=None, output_reward_path=None)

Parameters:

batch_input – input prompts

_get_batch_dataset_local(model, batch_input, K=8, iter_id=0, local_rank=0, output_min_length=16, output_max_length=48, infer_batch_size=8, generation_kwargs={}, tokenizer=None, training_args=None, reward_model=None, output_reward_path=None)

Parameters:

batch_input – input prompts

align(model, dataset, reward_model)

Perform alignment for a model

Parameters:

modelBaseModel object.

dataset: Dataset object.

Input dataset for model to generate outputs. The input and output

will then be feed into reward model to get the reward for alignment.

reward_model: RegressionModel object.