bridge.data.datasets.packing_utils

Module Contents

Functions

find_first_bin_that_fits	Finds the first bin in a list of bins that has enough space to fit a sequence of size ‘s’.
first_fit	Packs sequences of varying lengths into bins using the First-Fit algorithm.
first_fit_decreasing	Packs sequences of varying lengths into bins using the First-Fit Decreasing algorithm.
first_fit_shuffle	Packs sequences of varying lengths into bins using the First-Fit with Shuffling algorithm.
create_hist	Creates a histogram of sequence lengths from a tokenized dataset.
create_packing_strategy	Packs sequences into bins using the specified packing algorithm.
fill_packing_strategy	Fills the packing strategy with actual sequence data based on assignments and sequence information.

Data

PACKING_ALGOS
logger

API

bridge.data.datasets.packing_utils.PACKING_ALGOS

[‘first_fit_decreasing’, ‘first_fit_shuffle’]

bridge.data.datasets.packing_utils.logger

‘getLogger(…)’

bridge.data.datasets.packing_utils.find_first_bin_that_fits(

bins: List[List[int]],

s: int,

bin_size: int,

) → int

Finds the first bin in a list of bins that has enough space to fit a sequence of size ‘s’.

Parameters:

• bins – A list of lists, where each inner list represents a bin and contains the current elements in that bin.

• s – The size of the sequence to be placed in a bin.

• bin_size – The maximum capacity of each bin.

Returns:

The index of the first bin that can fit the sequence ‘s’, or -1 if no such bin exists.

bridge.data.datasets.packing_utils.first_fit(

seqlens: List[int],

pack_size: int,

) → List[List[int]]

Packs sequences of varying lengths into bins using the First-Fit algorithm.

Parameters:

• seqlens – A list of integers, representing the lengths of the sequences to be packed.

• pack_size – The maximum capacity of each bin.

Returns:

A list of lists, where each inner list represents a bin and contains the indices of the sequences assigned to that bin.

bridge.data.datasets.packing_utils.first_fit_decreasing(

seqlens: List[int],

pack_size: int,

) → List[List[int]]

Packs sequences of varying lengths into bins using the First-Fit Decreasing algorithm.

This is a variation of the First-Fit algorithm where the sequences are sorted by decreasing length before packing.

Parameters:

• seqlens – A list of integers, representing the lengths of the sequences to be packed.

• pack_size – The maximum capacity of each bin.

Returns:

A list of lists, similar to the output of the ‘first_fit’ function.

bridge.data.datasets.packing_utils.first_fit_shuffle(

seqlens: List[int],

pack_size: int,

) → List[List[int]]

Packs sequences of varying lengths into bins using the First-Fit with Shuffling algorithm.

This variation shuffles the order of the sequences before applying the First-Fit algorithm.

Parameters:

• seqlens – A list of integers, representing the lengths of the sequences to be packed.

• pack_size – The maximum capacity of each bin.

Returns:

A list of lists, similar to the output of the ‘first_fit’ function.

bridge.data.datasets.packing_utils.create_hist(dataset: numpy.array, truncate_seq_len: int)

Creates a histogram of sequence lengths from a tokenized dataset.

This function analyzes the tokenized dataset and creates a histogram showing the distribution of sequence lengths.

Parameters:

• dataset – A NumPy array containing the tokenized sequences. Each element is a dictionary that contains at minimum the key input_ids.

• truncate_seq_len – The maximum sequence length to consider in the histogram.

Returns:

A dictionary where keys are sequence lengths and values are lists of corresponding sequences from the dataset. histogram: A list representing the histogram data (number of sequences for each length).

Return type:

sequences

bridge.data.datasets.packing_utils.create_packing_strategy(

histogram: List[int],

pack_size: int,

packing_algorithm: str = 'first_fit',

) → List[List[int]]

Packs sequences into bins using the specified packing algorithm.

This function takes the histogram of sequence lengths, desired pack size, and a string representing the packing algorithm to use. It then calls the corresponding function (e.g., ‘first_fit_decreasing’) and performs the packing process using only sequence lengths as input (without the actual sequences).

Parameters:

• histogram – A list representing the histogram data (number of sequences for each length).

• pack_size – The maximum capacity of each bin.

• packing_algorithm – One of the supported packing algorithms from [‘first_fit_decreasing’, ‘first_fit_shuffle’]

Returns:

A list of lists, where each inner list represents a bin and contains the indices of the sequence lengths assigned to that bin. pack_metadata: A dict that records packing metadata, for instance the max number of samples per bin.

Return type:

assignments

bridge.data.datasets.packing_utils.fill_packing_strategy(

assignments: List[List[int]],

sequences: Dict[int, List[Dict]],

pack_size: int,

pad_id: int,

) → List[Dict]

Fills the packing strategy with actual sequence data based on assignments and sequence information.

This function takes the assignments generated by the packing algorithm (containing sequence length indices), the original sequences data, and the pack size. It iterates through the assignments, retrieves the corresponding sequences from the sequences dictionary, and constructs the final output data structure with input IDs, loss masks (if available), and starting indices for each sequence in a packed sequence.

Parameters:

• assignments – A list of lists, where each inner list represents a bin and contains the indices of the sequence lengths assigned to that bin (output of ‘create_packing_strategy’).

• sequences – A dictionary where keys are sequence lengths and values are lists of corresponding sequences from the dataset (output of ‘create_hist’).

• pack_size – The maximum capacity of each bin.

• pad_id – The tokenizer’s padding token.

Returns:

A list of dictionaries, where each dictionary represents a packed sequence with its input IDs, loss mask (if available), and starting indices.

Return type:

output_data