Compute mean of tensor with a masked values.

Delete all items inside the given dictionaries.

If regular alignment, return the HALO-defined reward for the sequence (log [policy(y|x)/reference(y|x)]). Args:

Print, but only on rank 0.

r""" Instantiates a new model from a pretrained model from `transformers`. The pretrained model is loaded using the `from_pretrained`

r""" Return the state_dict of the pretrained model.

Takes a list of examples and returns a batch of examples with consistent padding across all processes. Uses a fixed maximum length fo

Remove double spaces in the prompt and generations to standardize spacing.

Tokenize a single batch element and truncate if prompt + generation is too long. Batch element is turned into Pytorch tensors in sel

Run evaluation on all the examples in the test data and return the metrics from get_batch_metrics. This is close-ended evaluation and

Compute the token-level log probabilities of the given labels under the given logits.

Sample from the given model. NOTE: If the policy is being trained with FSDP, then sampling from it directly will produce gibberish.

Save tokenizer, policy model, optimizer, scheduler state to disk.

r""" Save the pretrained model to a directory. This method is a wrapper around `transformers.PreTrainedModel.save_pretrained`. Please

Close the JSON array and the file.

Format a dictionary for printing.

Return the subset of data to be processed in the current process.

Compute the DPO loss for a batch of policy and reference model token-level log probabilities.

Create feedback item with 50% chance of swapping pairs

Run the given model on the given batch of inputs. Returns: chosen_logps: log probabilities of chosen examples

Run the given model on the given batch of inputs, concatenating the chosen and rejected inputs together. Return two tensors of shape (batch

Run the given model on the given batch of inputs. Args: model: model to run forward pass on batch: input batch (forwa

Compute variance of tensor with masked values.

Separate the kwargs from the arguments that we support inside `supported_args` and the ones that we don't.

Concatenate the chosen and rejected inputs into a single tensor. The first half is chosen outputs, the second half is rejected. Args:

Run the given model on the given batch of inputs. Args: - model: the model to use for the forward pass - batc

Run the given model on the given batch of inputs. Args: - model: the model to use for the forward pass - batc

r""" A simple wrapper around the `generate` method of the wrapped model. Please refer to the [`generate`](https://huggingface.co/docs/

Return the state dict for the base model given the state dict for a lora-wrapped AutoModelForCausalLM, merging the lora weights as needed.

Compute the loss and other metrics for the given batch of inputs. Arg: batch: dictionary of inputs for the batch (what is

Given a batch that has been processed in the outer loop of PPO, return the batch statistics and the loss.

Return the flattened data for the given prompts across all processes as a list of tuples. The first element of this tuple should alwa

Get the processed dict for the entire batch. Args: batch: dictionary containing batch data (shoud have keys 'values', 'r

Load the Anthropic Helpful-Harmless dataset from Huggingface and convert it into to a Dataset. For this dataset, the SFT text is the preferre

Return the probability ratio under the policy vs the reference [policy(y|x)/reference(y|x)]. Apply humanline if specified. A

Enable gradient chekpointing.

Process a single log file and extract top-level model metrics.

Resize the vocabulary size of the language model.

Update the reference model to have the policy weights.

Train with PPO.

Write a single item to the JSON array.

Compute the token-level log probabilities of the given labels under the given logits.

r""" Initializes the weights of the value head. The default initialization strategy is random. Users can pass a different initializati

Calculate the log probabilities of every input-output sequence in every iterator in self.iterators. Save these in self.logprobs as th

Estimate the advantages and rewards for every token taken. Args: values: the estimated values of the tokens. Should alre

Convert samples to binary feedback format.

Convert samples to pairwise feedback format.

Calculate entropy from logits. Args: logits: tensor of shape (batch_size, sequence length, vocab) mask: tensor of shape (batc

Run evaluation on all the examples in the test data and return the metrics from get_batch_metrics. This is close-ended evaluation and

Extract AlpacaEval metrics from the log text.

Extract only top-level metric values from the log text.

Get logits for both chosen and rejected examples. Args: model: The Bradley-Terry model batch: Dictionary cont

Return the cached log probabilities for the given input ids.

Get completion from API with retry logic.

Get list of available datasets by finding all get_* functions in dataloader.py

Get the number of training steps (needed for learning rate scheduler).

Get reward scores either from reward model or binary labels. Args: batch: Dictionary containing batch data

Compute the Kahneman-Tversky loss for a batch of policy and reference model log probabilities. If generation y ~ p_desirable, we have the 'de

Compute the GRPO loss. Args: policy_logps: log probability of the output under the policy (microbatch_size, sequence_len

Given the batch statistics and the current episode's values, calculate the loss and return some loss statistics. Args: b

Compute Bradley-Terry loss given the logits for chosen and rejected examples. The Bradley-Terry model predicts P(A > B) = sigmoid(sco

Main entry point for training. Validates config, creates/initializes model(s), and starts training.

r""" Post initialization method. This method is called after the model is instantiated and loaded from a checkpoint. It can be used to

Prepare the Accelerator.

Process a batch through the reward model using the already tokenized sequences.

Process all log files in a directory.

Process all samples through the API.

r""" Push the pretrained model to the hub. This method is a wrapper around `transformers.PreTrainedModel.push_to_hub`. Please refer to

Calculate the row-wise product over all the elements that have not been masked out. Args: mat: tensor of shape (batch_size, sequence

Save tokenizer, policy model, optimizer, scheduler state to disk.

Validate that all requested datasets have corresponding get_* functions

Get an attribute of the Example by its key. Args: key: The attribute name to retrieve Retur

A trainer for a language model, supporting either SFT, HALO, or offline PPO training.

r""" Initializes the model. Args: pretrained_model (`transformers.PreTrainedModel`): The model to wrap. I

Args: - reference_accelerator: accelerator that should be used for caching (different from main accelerator) - refere

Create a flat version of the data and yield batches.