TensorDictTokenizer

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

用於在文字範例上套用 tokenizers 的處理函數的工廠。

參數:

• tokenizer (來自 transformers library 的 tokenizer) – 要使用的 tokenizer。

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

• key (str, optional) – 尋找文字的鍵。預設為 "text"。

• padding (str, optional) – 填補類型。預設為 "max_length"。

• truncation (bool, optional) – 是否應將序列截斷為 max_length。

• return_tensordict (bool, optional) – 如果 True，則傳回 TensoDict。否則，將傳回原始資料。

• device (torch.device, optional) – 儲存資料的裝置。如果 return_tensordict=False，則會忽略此選項。

有關 tokenizers 的更多資訊，請參閱 transformers library

填補和截斷：https://huggingface.co/docs/transformers/pad_truncation

傳回：一個 tensordict.TensorDict 實例，其批次大小與輸入資料相同。

範例

>>> from transformers import AutoTokenizer
>>> tokenizer = AutoTokenizer.from_pretrained("gpt2")
>>> tokenizer.pad_token = 100
>>> process = TensorDictTokenizer(tokenizer, max_length=10)
>>> # example with a single input
>>> example = {"text": "I am a little worried"}
>>> process(example)
TensorDict(
    fields={
        attention_mask: Tensor(shape=torch.Size([10]), device=cpu, dtype=torch.int64, is_shared=False),
        input_ids: Tensor(shape=torch.Size([10]), device=cpu, dtype=torch.int64, is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)
>>> # example with a multiple inputs
>>> example = {"text": ["Let me reassure you", "It will be ok"]}
>>> process(example)
TensorDict(
    fields={
        attention_mask: Tensor(shape=torch.Size([2, 10]), device=cpu, dtype=torch.int64, is_shared=False),
        input_ids: Tensor(shape=torch.Size([2, 10]), device=cpu, dtype=torch.int64, is_shared=False)},
    batch_size=torch.Size([2]),
    device=None,
    is_shared=False)