Compute variance of tensor with masked values.

MeanFlow sampler supporting both single-step and multi-step generation Based on Eq.(12): z_r = z_t - (t-r)u(z_t, r, t) For single-st

Merge two models using mergekit Args: config (`MergeConfig`): The merge configuration. out_path (`str`): The output path for

Compute the maximum value of a tensor, ignoring NaNs. This function only supports 1D tensors. Args: tensor (`torch.Tensor`): Input t

Compute the minimum value of a tensor, ignoring NaNs. This function only supports 1D tensors. Args: tensor (`torch.Tensor`): Input t

Compute the standard deviation of a tensor, ignoring NaNs. This function only supports 1D tensors. Args: tensor (`torch.Tensor`):

Context manager for handling null reference model (that is, peft adapter manipulation).

Convert a NumPy array of shape (batch, height, width, channels) to a list of PIL Images.

Convert a NumPy array of shape (batch, height, width, channels) to a list of PIL Images.

Compute ORPO's odds ratio (OR) loss for a batch of policy and reference model log probabilities. Args: policy_chosen_logps:

r""" Pack sequences in a dataset into chunks of size `seq_length`. Args: dataset (`Dataset` or `DatasetDict`): Dataset to

Parse request string. Expected format: <request><tool_name>query<call>

r""" Function invoked when calling the pipeline for generation. Args: prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guid

r""" Function to get RGB image with gradients attached to the model weights. Args: prompt (`str` or `list[str]`, *optional*, defaults

Prepare dataset example for training.

Prepare the text from a sample of the dataset.

Prepare the text from a sample of the dataset.

Print out a sample of model completions to the console with multiple reward metrics. This function creates a nicely formatted table showing

Prints the number of trainable parameters in the model.

Removes the optimizer hooks from a DeepSpeed ZeRO-3 model.

Run the environment on a list of queries. Args: queries (list[str]): A list of queries to run the model in the environme

Sample time steps (r, t). progress: float in [0,1] OR current_step/total_steps pair. If none provided, use start values (back

Save the content of this instance in JSON format inside `json_path`.

Overrides the parser's default values with those provided via keyword arguments, including for subparsers. Any argument with an upda

Shift input ids one token to the right, and pad with pad_token_id

Shuffles all sequence-like values in a dictionary along the first dimension in unison. Example: ```python >>> x = torch.arange(6).re

Splits `batch["pixel_values"]` into a list of tensors based on the product of each row in `batch["image_grid_thw"]`, while keeping other entr

Split the tokens into query and response tokens.

Splits a dictionary of tensors along the first dimension into `num_chunks` equal parts. Example: ```python >>> x = torch.arange(12).

Step the environment forward one turn. Args: history (`TextHistory`): The history to step forward.

Perform a single step of training. Args: epoch (int): The current epoch. global_step (int): The current glob

Drops labels to enable classifier-free guidance.

Tokenize a single row from a DPO specific dataset.

Train the model for a given number of epochs

Train the model for a given number of epochs

r""" Truncate sequences in a dataset to a specifed `max_length`. Args: dataset (`Dataset` or `DatasetDict`): Dataset to t

Truncate tensors to target length while preserving protected tokens. Args: ids (`torch.Tensor`): Input tensor of token I

r""" Unpair a preference dataset. Args: dataset (`Dataset` or `DatasetDict`): Preference dataset to unpair. The dataset m

Opposite of `split_pixel_values_by_grid`. Merges a list of tensors in `batch["pixel_values"]` back into a single tensor along the first dimen

Updates all parameters of the given model by calling `update_named_param` for each parameter in the model. Args: model (

Visualize the reward model logits prediction Args: num_print_samples (`int`, defaults to `4`): The numbe

Compute MeanFlow loss function with schedule support. Provide either training_progress (0..1) or current_step & total_steps.

Returns true if all generated sequences contain any of the stop strings.

Reroutes a method call through the `wrapper_module`'s `forward` method. Args: wrapper_module: The module that has `original_modul