UnpadOperation

功能

对于输入input_ids,把所有有效的token拼接在一起,并在最后补0。

使用场景

为了保证[batch,max_seq_len]的二维shape,某些batch的真实句子长度会被扩展成max_seq_len的长度,此算子目的为去除这些padding的部分。

使用示例:

假设有以下输入:

在这个例子中:

通过 unpad 算子处理后:

  1. 计算 cum_offsets_out:

    去除最后一个 batch 的数值:[[2]]

    整体后移一个 batch:[[0], [2]]

    第一个 batch 的值置为 0:[[0], [2]](cum_offsets_out的值)

  2. 移除填充 0:
    • 对于第一个 batch:
      • cum_offsets_out[0] 是 0,表示从位置 0 开始。
      • seq_len[0] 是 3,表示有 3 个有效 token。
      • 将 input_ids[0] 中的前 3 个有效 token 复制到 x_remove_padding 的前 3 个位置。
    • 对于第二个 batch:
      • cum_offsets_out[1] 是 2,表示从位置 2 开始。
      • seq_len[1] 是 2,表示有 2 个有效 token。
      • 将 input_ids[1] 中的前 2 个有效 token 复制到 x_remove_padding 的第 3 和第 4 个位置。
      • 根据 cum_offsets_out 和 seq_len,重新排列有效 token,得到 x_remove_padding:[1, 2, 3, 4, 5, 0, 0, 0, 0, 0]。

应用场景:

在自然语言处理(NLP)中,输入的句子长度通常不同。为了让它们在一个 batch 中处理,我们需要将它们填充(pad)到相同的长度。填充的方式通常是在句子的末尾添加零(0),直到所有句子长度相同。Unpad 用于在处理完这些填充的序列后,去除填充的 0,提取有效的 token 序列。

定义

struct UnpadParam {};

输入

参数

维度

数据类型

格式

描述

input_ids

[batch, max_seq_len]

int64

ND

输入tensor:经过pad(末尾填充0)之后,batch个token id序列。

cum_offsets_now

[batch, 1]

int32

ND

输入tensor:每个batch末尾填充0的数量组成的序列的前缀和。

token_num

[batch, 1]

int64

ND

输入tensor:有效token数量总和。

seq_len

[batch, 1]

int32

ND

输入tensor:每个batch中的有效token数量。

输出

参数

维度

数据类型

格式

描述

x_remove_padding

[1, batch * max_seq_len]

int64

ND

输出tensor:从input_ids中去除填充的0后得到的有效token序列(为了保持shape,在末尾填充0)。

cum_offsets_out

[batch, 1]

int32

ND

输出tensor:去除cum_offsets_now最后一个batch的数值,再将其整体后移一个batch,并将第一个batch的值置为0。

padding_offset

[1, batch * max_seq_len]

int32

ND

输出tensor:由seq_len和cum_offsets_out组合计算得到。