UnpadOperation

Description

Concatenates all valid tokens and pads 0s for input_ids.

Application Scenarios

To ensure the two-dimensional shape of [batch, max_seq_len], the actual sentence length of some batches is extended to the length of max_seq_len. This operator is used to remove the padding.

Example:

Input:

input_ids: [[1, 2, 3, 0, 0], [4, 5, 0, 0, 0]]
cum_offsets_now: [[2], [5]]
seq_len: [[3], [2]]

In this example:

The first batch has two padded 0s. Therefore, the first value of cum_offsets_now is 2.
The second batch has five padded 0s (including the two padded 0s of the first batch). Therefore, the second value of cum_offsets_now is 5.

After the unpad operator is used for processing:

Calculate cum_offsets_out:
Remove the value of the last batch: [[2]]

Move backward by one batch: [[0], [2]]

The value of the first batch is set to 0: [[0], [2]] (value of cum_offsets_out).
Remove padded 0s:
- For the first batch:
  - cum_offsets_out[0] is 0, indicating that the offset starts from position 0.
  - seq_len[0] is 3, indicating that there are three valid tokens.
  - Copy the first three valid tokens in input_ids[0] to the first three positions of x_remove_padding.
- For the second batch:
  - cum_offsets_out[1] is 2, indicating that the offset starts from position 2.
  - seq_len[1] is 2, indicating that there are two valid tokens.
  - Copy the first two valid tokens in input_ids[1] to the third and fourth positions of x_remove_padding.
  - Re-arrange valid tokens based on cum_offsets_out and seq_len to obtain x_remove_padding: [1, 2, 3, 4, 5, 0, 0, 0, 0, 0].

Application scenarios:

In natural language processing (NLP), the length of an input sentence is usually different. To process them in a batch, we need to pad them to the same length. A padding method is usually adding zeros (0) at the end of a sentence until all sentences have a same length. Unpad is used to remove the padded 0s and extract a valid token sequence after these padded sequences are processed.

Definition

struct UnpadParam {
    uint8_t rsv[8] = {0};
};

Parameters

Member	Type	Default Value	Description
rsv[8]	uint8_t	{0}	Reserved

Input

Parameter	Dimension	Data Type	Format	Description
input_ids	[batch, max_seq_len]	int64	ND	Input tensor: token ID sequences (number determined by batch) after padding (after 0 is padded at the end).
cum_offsets_now	[batch, 1]	int32	ND	Input tensor: prefix sum of the sequence consisting of the number of 0s padded at the end of each batch.
token_num	[1, 1]	int64	ND	Input tensor: total number of valid tokens.
seq_len	[batch, 1]	int32	ND	Input tensor: number of valid tokens in each batch.

Output

Parameter	Dimension	Data Type	Format	Description
x_remove_padding	[1, batch * max_seq_len]	int64	ND	Output tensor: valid token sequence obtained after the padded 0s are removed from input_ids (0s are padded at the end to maintain the shape).
cum_offsets_out	[batch, 1]	int32	ND	Output tensor: removes the value of the last batch of cum_offsets_now, moves it backward by one batch, and sets the value of the first batch to 0.
padding_offset	[1, batch * max_seq_len]	int32	ND	Output tensor: computed based on the combination of seq_len and cum_offsets_out.

Parent topic: atb/infer_op_params.h