PromptTensorDictTokenizer

class torchrl.data.PromptTensorDictTokenizer(tokenizer, max_length, key='prompt', padding='max_length', truncation=True, return_tensordict=True, device=None)

提示数据集的 Tokenization 配方。

返回一个 tokenizer 函数，该函数读取包含提示和标签的示例并对其进行 tokenization。

参数:

• tokenizer (来自 transformers 库的 tokenizer) – 要使用的 tokenizer。

• max_length (int) – 序列的最大长度。

• key (str, optional) – 查找文本的键。默认为 "prompt"。

• padding (str, optional) – 填充类型。默认为 "max_length"。

• truncation (bool, optional) – 是否应将序列截断为 max_length。

• return_tensordict (bool, optional) – 如果为 True，则返回 TensoDict。否则，将返回原始数据。

• device (torch.device, optional) – 存储数据的设备。如果 return_tensordict=False，则忽略此选项。

此类 __call__() 方法将执行以下操作

• 读取与 label 字符串联系的 prompt 字符串并对其进行 tokenization。结果将存储在 "input_ids" TensorDict 条目中。

• 写入 "prompt_rindex" 条目，其中包含来自提示的最后一个有效 token 的索引。

• 写入 "valid_sample"，用于标识 tensordict 中哪个条目具有足够的 token 以满足 max_length 标准。

• 返回带有 tokenized 输入的 tensordict.TensorDict 实例。

tensordict 批大小将与输入的批大小匹配。

示例

>>> from transformers import AutoTokenizer
>>> tokenizer = AutoTokenizer.from_pretrained("gpt2")
>>> tokenizer.pad_token = tokenizer.eos_token
>>> example = {
...     "prompt": ["This prompt is long enough to be tokenized.", "this one too!"],
...     "label": ["Indeed it is.", 'It might as well be.'],
... }
>>> fn = PromptTensorDictTokenizer(tokenizer, 50)
>>> print(fn(example))
TensorDict(
    fields={
        attention_mask: Tensor(shape=torch.Size([2, 50]), device=cpu, dtype=torch.int64, is_shared=False),
        input_ids: Tensor(shape=torch.Size([2, 50]), device=cpu, dtype=torch.int64, is_shared=False),
        prompt_rindex: Tensor(shape=torch.Size([2]), device=cpu, dtype=torch.int64, is_shared=False),
        valid_sample: Tensor(shape=torch.Size([2]), device=cpu, dtype=torch.bool, is_shared=False)},
    batch_size=torch.Size([2]),
    device=None,
    is_shared=False)