TokenizedDatasetLoader¶
- class torchrl.data.TokenizedDatasetLoader(split, max_length, dataset_name, tokenizer_fn: Type[TensorDictTokenizer], pre_tokenization_hook=None, root_dir=None, from_disk=False, valid_size: int = 2000, num_workers: Optional[int] = None, tokenizer_class=None, tokenizer_model_name=None)[source]¶
加载并标记数据集，并缓存其内存映射副本。
- 参数:
split (str) –
"train"或"valid"之一。max_length (int) – 最大序列长度。
dataset_name (str) – 数据集的名称。
tokenizer_fn (callable) – 标记化方法构造函数，例如
torchrl.data.rlhf.TensorDictTokenizer。 调用时，它应返回tensordict.TensorDict实例或包含标记化数据的类字典结构。pre_tokenization_hook (callable, optional) – 在标记化之前在数据集上调用。 它应返回修改后的 Dataset 对象。 预期用途是执行需要修改整个数据集的任务，而不是修改单个数据点，例如根据特定条件丢弃某些数据点。 数据的标记化和其他“逐元素”操作由进程函数执行，该函数映射到数据集。
root_dir (path, optional) – 存储数据集的路径。 默认为
"$HOME/.cache/torchrl/data"from_disk (bool, optional) – 如果为
True，将使用datasets.load_from_disk()。 否则，将使用datasets.load_dataset()。 默认为False。valid_size (int, optional) – 验证数据集的大小（如果 split 以
"valid"开头）将被截断为此值。 默认为 2000 项。num_workers (int, optional) – 用于
datasets.dataset.map()的工作进程数，在标记化期间调用。 默认为max(os.cpu_count() // 2, 1)。tokenizer_class (type, optional) – 标记器类，例如
AutoTokenizer(默认)。tokenizer_model_name (str, optional) – 应从中收集词汇表的模型。 默认为
"gpt2"。
数据集将存储在
<root_dir>/<split>/<max_length>/中。示例
>>> from torchrl.data.rlhf import TensorDictTokenizer >>> from torchrl.data.rlhf.reward import pre_tokenization_hook >>> split = "train" >>> max_length = 550 >>> dataset_name = "CarperAI/openai_summarize_comparisons" >>> loader = TokenizedDatasetLoader( ... split, ... max_length, ... dataset_name, ... TensorDictTokenizer, ... pre_tokenization_hook=pre_tokenization_hook, ... ) >>> dataset = loader.load() >>> print(dataset) TensorDict( fields={ attention_mask: MemoryMappedTensor(shape=torch.Size([185068, 550]), device=cpu, dtype=torch.int64, is_shared=False), input_ids: MemoryMappedTensor(shape=torch.Size([185068, 550]), device=cpu, dtype=torch.int64, is_shared=False)}, batch_size=torch.Size([185068]), device=None, is_shared=False)
- static dataset_to_tensordict(dataset: 'datasets.Dataset' | TensorDict, data_dir: Path, prefix: NestedKey = None, features: Sequence[str] = None, batch_dims=1, valid_mask_key=None)[source]¶
将数据集转换为内存映射的 TensorDict。
如果数据集已经是
TensorDict实例，则它仅被转换为内存映射的 TensorDict。 否则，数据集应具有features属性，该属性是一个字符串序列，指示可以在数据集中找到的特征。 如果没有，则必须将features显式传递给此函数。- 参数:
dataset (datasets.Dataset, TensorDict 或等效项) – 要转换为内存映射 TensorDict 的数据集。 如果
features为None，则它必须具有features属性，其中包含要写入 tensordict 的键列表。data_dir (Path 或等效项) – 应在其中写入数据的目录。
prefix (NestedKey, optional) – 数据集位置的前缀。 这可以用于区分同一数据集的多个副本，这些副本经历了不同的预处理。
features (sequence of str, optional) – 指示可以在数据集中找到的特征的字符串序列。
batch_dims (int, optional) – 数据的批次维度数（即 tensordict 可以沿其索引的维度数）。 默认为 1。
valid_mask_key (NestedKey, optional) – 如果提供，则将尝试收集此条目并用于过滤数据。 默认为
None(即，无过滤器键)。
返回：包含带有数据集的内存映射张量的 TensorDict。
示例
>>> from datasets import Dataset >>> import tempfile >>> data = Dataset.from_dict({"tokens": torch.randint(20, (10, 11)), "labels": torch.zeros(10, 11)}) >>> with tempfile.TemporaryDirectory() as tmpdir: ... data_memmap = TokenizedDatasetLoader.dataset_to_tensordict( ... data, data_dir=tmpdir, prefix=("some", "prefix"), features=["tokens", "labels"] ... ) ... print(data_memmap) TensorDict( fields={ some: TensorDict( fields={ prefix: TensorDict( fields={ labels: MemoryMappedTensor(shape=torch.Size([10, 11]), device=cpu, dtype=torch.float32, is_shared=False), tokens: MemoryMappedTensor(shape=torch.Size([10, 11]), device=cpu, dtype=torch.int64, is_shared=False)}, batch_size=torch.Size([10]), device=None, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)
