lmflow.pipeline.raft_aligner#

The Aligner class simplifies the process of running alignment.

Attributes#

logger

Classes#

RaftAligner

Initializes the RaftAligner class with given arguments.

Module Contents#

lmflow.pipeline.raft_aligner.logger[source]#

class lmflow.pipeline.raft_aligner.RaftAligner(model_args, data_args, aligner_args, *args, **kwargs)[source]#

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.

model_args[source]#

data_args[source]#

aligner_args[source]#

INF = 888888888[source]#

_initialize_trainer(model, tokenizer, training_args)[source]#: 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)[source]#: This function prepares the dataset for every iteration.

_load_input_dataset(dataset, tokenizer)[source]#

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)[source]#

_discard_sample(text)[source]#

_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)[source]#

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)[source]#

Parameters:: batch_input – input prompts

align(model, dataset, reward_model)[source]#

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.