快捷方式

History

class torchrl.data.llm.History(role: 'str', content: 'str | ContentBase', is_complete: 'bool' = True, tool_calls: 'list[dict] | None' = None, tool_responses: 'list[str] | None' = None, *, batch_size, device=None, names=None)[源代码]
append(history: History, *, inplace: bool = True, dim: int = - 1) History[源代码]

将新的历史附加到当前历史。

参数:
  • history (History) – 要附加的新历史。

  • inplace (bool, optional) – 是否原地执行操作。默认为 True

  • dim (int, optional) – 附加的维度。默认为 -1。

返回:

附加后的 History 对象。

返回类型:

History

apply_chat_template(*, tokenizer: transformers.AutoTokenizer | transformers.AutoProcessor, add_generation_prompt: bool = True, chat_template: str | None = None, chat_template_name: str | None = None, continue_final_message: bool = False, tokenize: bool | None = None, padding: bool | str = False, truncation: bool | str = False, return_tensors: str | None = None, return_dict: bool | None = None, return_assistant_tokens_mask: bool = False, **kwargs) str | list[str] | TensorDict[源代码]

应用聊天模板到历史记录。

关键字参数:
  • tokenizer (transformers.PreTrainedTokenizer | transformers.AutoProcessor) – 要使用的分词器。

  • add_generation_prompt (bool, optional) – 是否添加生成提示(例如 “<|im_start|>assistant”)。默认为 True

  • chat_template (str, optional) – 要使用的聊天模板。默认为分词器的默认模板。

  • chat_template_name (str, optional) – 要使用的聊天模板的名称。优先于 tokenizer.chat_template。如果为 None,则方法将自动检测模型系列并使用适当的模板。默认为 None

  • continue_final_message (bool, optional) – 是否继续最后一条消息。默认为 False

  • tokenize (bool, optional) – 是否对输出进行分词。默认为 False

  • padding (bool | str, optional) – 要使用的填充策略。默认为 False

  • truncation (bool | str, optional) – 要使用的截断策略。默认为 False

  • return_tensors (str | None, optional) – 要返回的张量类型。默认为 “pt”。

  • return_dict (bool, optional) – 是否返回字典。默认为 False

  • return_assistant_tokens_mask (bool, optional) –

    是否返回助手生成的 token 的掩码。如果为 True,则掩码将写入 assistant_masks 键。对于助手生成的 token,掩码将包含 1。对于用户和系统 token,掩码将包含 0。此功能仅对支持它的聊天模板(通过 {% generation %} 关键字)可用。默认为 False

    注意

    助手 token 掩码支持多种模型系列: - Qwen 系列:使用自定义模板,支持完整的工具调用 - DialoGPT 系列:使用自定义模板进行对话格式 - Falcon 系列:使用自定义模板进行指令格式 - DeepSeek 系列:使用自定义模板,支持原生格式 - 其他模型:使用默认的 chatml_format 模板

    该方法会自动检测模型系列并选择合适的模板。

  • **kwargs – 传递给分词器 apply_chat_template 方法的其他关键字参数。

返回:

格式化后的历史记录。

classmethod default_spec(shape=(- 1,))[源代码]

在返回 History 对象的转换/环境中使用的默认规范。

参数:

shape (torch.Size, optional) – 返回的 History 规范的形状。默认为 (-1)(在时间维度上可变长度)。

示例

>>> import tensordict
>>> from torchrl.data import History
>>> tensordict.set_list_to_stack(True).set()
>>>
>>> history = History(role=["system", "user"], content=["a message", "another message"], batch_size=(2,))
>>> spec = history.default_spec()
>>> print(spec)
Composite(
    role: NonTensor(
        shape=torch.Size([-1]),
        space=None,
        device=None,
        dtype=None,
        domain=None,
        example_data=foo),
    content: NonTensor(
        shape=torch.Size([-1]),
        space=None,
        device=None,
        dtype=None,
        domain=None,
        example_data=foo),
    device=None,
    shape=torch.Size([-1]))
>>> print(spec.zero())
History(
    content=NonTensorData(data=foo, batch_size=torch.Size([1]), device=None),
    role=NonTensorData(data=foo, batch_size=torch.Size([1]), device=None),
    batch_size=torch.Size([1]),
    device=None,
    is_shared=False)
property device: device

检索张量类的设备类型。

dumps(prefix: str | None = None, copy_existing: bool = False, *, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False) T

将tensordict保存到磁盘。

此函数是 memmap() 的代理。

classmethod fields()

返回一个描述此数据类的字段的元组。字段类型为 Field。

接受一个数据类或其实例。元组元素为 Field 类型。

classmethod from_chats(chats: list[list[dict]]) History[源代码]

从聊天列表创建 History 对象。

参数:

chats (list[list[dict]]) – 聊天列表,其中每个聊天是一个字典列表。

classmethod from_tensordict(tensordict, non_tensordict=None, safe=True)

用于实例化新张量类对象的张量类包装器。

参数:
  • tensordict (TensorDict) – 张量类型的字典

  • non_tensordict (dict) – 包含非张量和嵌套张量类对象的字典

get(key: NestedKey, *args, **kwargs)

获取输入键对应的存储值。

参数:
  • key (str, tuple of str) – 要查询的键。如果是字符串元组,则等同于链式调用 getattr。

  • default – 如果在张量类中找不到键,则返回默认值。

返回:

存储在输入键下的值

classmethod load(prefix: str | Path, *args, **kwargs) T

从磁盘加载 tensordict。

此类方法是 load_memmap() 的代理。

load_(prefix: str | Path, *args, **kwargs)

在当前 tensordict 中从磁盘加载 tensordict。

此类方法是 load_memmap_() 的代理。

classmethod load_memmap(prefix: str | Path, device: torch.device | None = None, non_blocking: bool = False, *, out: TensorDictBase | None = None) T

从磁盘加载内存映射的 tensordict。

参数:
  • prefix (strPath to folder) – 应获取已保存 tensordict 的文件夹路径。

  • device (torch.device等效项, optional) – 如果提供,数据将异步转换为该设备。支持 “meta” 设备,在这种情况下,数据不会被加载,而是创建一组空的 “meta” 张量。这有助于了解模型的总大小和结构,而无需实际打开任何文件。

  • non_blocking (bool, optional) – 如果为 True,则在将张量加载到设备后不会进行同步。默认为 False

  • out (TensorDictBase, optional) – 数据应写入的可选 tensordict。

示例

>>> from tensordict import TensorDict
>>> td = TensorDict.fromkeys(["a", "b", "c", ("nested", "e")], 0)
>>> td.memmap("./saved_td")
>>> td_load = TensorDict.load_memmap("./saved_td")
>>> assert (td == td_load).all()

此方法还允许加载嵌套的 tensordicts。

示例

>>> nested = TensorDict.load_memmap("./saved_td/nested")
>>> assert nested["e"] == 0

tensordict 也可以在“meta”设备上加载,或者作为假张量加载。

示例

>>> import tempfile
>>> td = TensorDict({"a": torch.zeros(()), "b": {"c": torch.zeros(())}})
>>> with tempfile.TemporaryDirectory() as path:
...     td.save(path)
...     td_load = TensorDict.load_memmap(path, device="meta")
...     print("meta:", td_load)
...     from torch._subclasses import FakeTensorMode
...     with FakeTensorMode():
...         td_load = TensorDict.load_memmap(path)
...         print("fake:", td_load)
meta: TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=meta, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([]), device=meta, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([]),
            device=meta,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=meta,
    is_shared=False)
fake: TensorDict(
    fields={
        a: FakeTensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: FakeTensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([]),
            device=cpu,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=cpu,
    is_shared=False)
load_state_dict(state_dict: dict[str, Any], strict=True, assign=False, from_flatten=False)

尝试将 state_dict 加载到目标张量类中(原地)。

memmap(prefix: str | None = None, copy_existing: bool = False, *, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False, existsok: bool = True) T

将所有张量写入内存映射的 Tensor 中,并放入新的 tensordict。

参数:
  • prefix (str) – 内存映射张量将存储的目录前缀。目录树结构将模仿 tensordict 的结构。

  • copy_existing (bool) – 如果为 False(默认值),则如果 tensordict 中的条目已经是存储在磁盘上的张量且具有关联文件,但未按 prefix 正确保存,则会引发异常。如果为 True,则任何现有的张量都将被复制到新位置。

关键字参数:
  • num_threads (int, optional) – 用于写入 memmap 张量的线程数。默认为 0

  • return_early (bool, optional) – 如果为 Truenum_threads>0,则该方法将返回一个 tensordict 的 future。

  • share_non_tensor (bool, optional) – 如果为 True,则非张量数据将在进程之间共享,并且在单个节点内的任何工作进程上的写入操作(例如原地更新或设置)将更新所有其他工作进程上的值。如果非张量叶节点数量很高(例如,共享大量非张量数据),这可能会导致 OOM 或类似错误。默认为 False

  • existsok (bool, optional) – 如果为 False,则如果张量已存在于同一路径中,则会引发异常。默认为 True

然后 tensordict 被锁定,这意味着任何非原地写入操作都将引发异常(例如,重命名、设置或删除条目)。一旦 tensordict 解锁,内存映射属性将被设置为 False,因为跨进程标识不再有保证。

返回:

一个包含存储在磁盘上的张量的新 tensordict(如果 return_early=False),否则为 TensorDictFuture 实例。

注意

以这种方式序列化对于深度嵌套的 tensordicts 来说可能很慢,因此不建议在训练循环中调用此方法。

memmap_(prefix: str | None = None, copy_existing: bool = False, *, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False, existsok: bool = True) T

将所有张量原地写入相应的内存映射张量。

参数:
  • prefix (str) – 内存映射张量将存储的目录前缀。目录树结构将模仿 tensordict 的结构。

  • copy_existing (bool) – 如果为 False(默认值),则如果 tensordict 中的条目已经是存储在磁盘上的张量且具有关联文件,但未按 prefix 正确保存,则会引发异常。如果为 True,则任何现有的张量都将被复制到新位置。

关键字参数:
  • num_threads (int, optional) – 用于写入 memmap 张量的线程数。默认为 0

  • return_early (bool, optional) – 如果为 Truenum_threads>0,则该方法将返回一个 tensordict 的 future。可以通过 future.result() 查询结果 tensordict。

  • share_non_tensor (bool, optional) – 如果为 True,则非张量数据将在进程之间共享,并且在单个节点内的任何工作进程上的写入操作(例如原地更新或设置)将更新所有其他工作进程上的值。如果非张量叶节点数量很高(例如,共享大量非张量数据),这可能会导致 OOM 或类似错误。默认为 False

  • existsok (bool, optional) – 如果为 False,则如果张量已存在于同一路径中,则会引发异常。默认为 True

然后 tensordict 被锁定,这意味着任何非原地写入操作都将引发异常(例如,重命名、设置或删除条目)。一旦 tensordict 解锁,内存映射属性将被设置为 False,因为跨进程标识不再有保证。

返回:

如果 return_early=False,则返回 self,否则返回 TensorDictFuture 实例。

注意

以这种方式序列化对于深度嵌套的 tensordicts 来说可能很慢,因此不建议在训练循环中调用此方法。

memmap_like(prefix: str | None = None, copy_existing: bool = False, *, existsok: bool = True, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False) T

创建一个无内容的内存映射 tensordict,其形状与原始 tensordict 相同。

参数:
  • prefix (str) – 内存映射张量将存储的目录前缀。目录树结构将模仿 tensordict 的结构。

  • copy_existing (bool) – 如果为 False(默认值),则如果 tensordict 中的条目已经是存储在磁盘上的张量且具有关联文件,但未按 prefix 正确保存,则会引发异常。如果为 True,则任何现有的张量都将被复制到新位置。

关键字参数:
  • num_threads (int, optional) – 用于写入 memmap 张量的线程数。默认为 0

  • return_early (bool, optional) – 如果为 Truenum_threads>0,则该方法将返回一个 tensordict 的 future。

  • share_non_tensor (bool, optional) – 如果为 True,则非张量数据将在进程之间共享,并且在单个节点内的任何工作进程上的写入操作(例如原地更新或设置)将更新所有其他工作进程上的值。如果非张量叶节点数量很高(例如,共享大量非张量数据),这可能会导致 OOM 或类似错误。默认为 False

  • existsok (bool, optional) – 如果为 False,则如果张量已存在于同一路径中,则会引发异常。默认为 True

然后 tensordict 被锁定,这意味着任何非原地写入操作都将引发异常(例如,重命名、设置或删除条目)。一旦 tensordict 解锁,内存映射属性将被设置为 False,因为跨进程标识不再有保证。

返回:

一个包含数据作为内存映射张量的新 TensorDict 实例(如果 return_early=False),否则为 TensorDictFuture 实例。

注意

这是将一组大型缓冲区写入磁盘的推荐方法,因为 memmap_() 会复制信息,这对于大型内容来说可能很慢。

示例

>>> td = TensorDict({
...     "a": torch.zeros((3, 64, 64), dtype=torch.uint8),
...     "b": torch.zeros(1, dtype=torch.int64),
... }, batch_size=[]).expand(1_000_000)  # expand does not allocate new memory
>>> buffer = td.memmap_like("/path/to/dataset")
memmap_refresh_()

如果内存映射的 tensordict 具有 saved_path,则刷新其内容。

如果没有任何路径与之关联,此方法将引发异常。

save(prefix: str | None = None, copy_existing: bool = False, *, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False) T

将tensordict保存到磁盘。

此函数是 memmap() 的代理。

set(key: NestedKey, value: Any, inplace: bool = False, non_blocking: bool = False)

设置一个新的键值对。

参数:
  • key (str, tuple of str) – 要设置的键的名称。如果是字符串元组,则等同于链式调用 getattr,然后进行最终的 setattr。

  • value (Any) – 要存储在张量类中的值

  • inplace (bool, optional) – 如果为 True,则 set 会尝试进行原地更新。如果为 False 或键不存在,则值将简单地写入其目标位置。

返回:

self

state_dict(destination=None, prefix='', keep_vars=False, flatten=False) dict[str, Any]

返回一个 state_dict 字典,可用于保存和加载张量类的数据。

to_tensordict(*, retain_none: bool | None = None) TensorDict

将张量类转换为常规 TensorDict。

复制所有条目。内存映射和共享内存张量将被转换为常规张量。

参数:

retain_none (bool) – 如果为 True,则 None 值将被写入 tensordict。否则将被丢弃。默认为 True

返回:

包含与张量类相同值的新的 TensorDict 对象。

unbind(dim: int)

返回沿指定维度解绑的索引张量类实例的元组。

结果张量类实例将共享初始张量类实例的存储。

文档

访问全面的 PyTorch 开发者文档

查看文档

教程

为初学者和高级开发者提供深入的教程

查看教程

资源

查找开发资源并让您的问题得到解答

查看资源