TRL 文档

数据工具

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

数据工具

is_conversational

trl.is_conversational

< >

( example: dict ) bool

参数

  • example (dict[str, Any]) — 数据集的单个数据条目。根据数据集类型的不同,该示例可以有不同的键。

返回

布尔值

如果数据是对话格式,则为 True,否则为 False

检查示例是否为对话格式。

示例

>>> example = {"prompt": [{"role": "user", "content": "What color is the sky?"}]}
>>> is_conversational(example)
True

>>> example = {"prompt": "The sky is"}
>>> is_conversational(example)
False

apply_chat_template

trl.apply_chat_template

< >

( example: dict tokenizer: PreTrainedTokenizerBase tools: typing.Optional[list[typing.Union[dict, typing.Callable]]] = None **template_kwargs )

将聊天模板应用于对话示例,同时应用 tools 中函数列表的模式。

更多详情,请参见 maybe_apply_chat_template()

maybe_apply_chat_template

trl.maybe_apply_chat_template

< >

( example: dict tokenizer: PreTrainedTokenizerBase tools: typing.Optional[list[typing.Union[dict, typing.Callable]]] = None **template_kwargs: typing.Any ) dict[str, str]

参数

  • example (dict[str, list[dict[str, str]]) — 表示对话数据集单个数据条目的字典。每个数据条目可以有不同的键,具体取决于数据集类型。支持的数据集类型有:

    • 语言建模数据集:"messages"
    • 仅提示数据集:"prompt"
    • 提示-完成数据集:"prompt""completion"
    • 偏好数据集:"prompt""chosen""rejected"
    • 带隐式提示的偏好数据集:"chosen""rejected"
    • 未配对偏好数据集:"prompt""completion""label"

    对于键 "messages""prompt""chosen""rejected""completion",其值是消息列表,其中每个消息是一个包含键 "role""content" 的字典。

  • tokenizer (PreTrainedTokenizerBase) — 用于应用聊天模板的分词器。
  • tools (list[Union[dict, Callable]]None可选,默认为 None) — 模型可访问的工具列表(可调用函数)。如果模板不支持函数调用,此参数将无效。
  • **template_kwargs (Any可选) — 传递给模板渲染器的额外关键字参数。聊天模板将可以访问这些参数。

返回

dict[str, str]

应用了聊天模板的格式化示例。

如果示例是对话格式,则对其应用聊天模板。

备注

  • 此函数不改变键,但语言建模数据集除外,其中 "messages" 会被替换为 "text"

  • 对于仅提示的数据,如果最后一个角色是 "user",则将生成提示添加到提示中。否则,如果最后一个角色是 "assistant",则继续最后一条消息。

示例

>>> from transformers import AutoTokenizer

>>> tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-3-mini-128k-instruct")
>>> example = {
...     "prompt": [{"role": "user", "content": "What color is the sky?"}],
...     "completion": [{"role": "assistant", "content": "It is blue."}],
... }
>>> apply_chat_template(example, tokenizer)
{'prompt': '<|user|>\nWhat color is the sky?<|end|>\n<|assistant|>\n', 'completion': 'It is blue.<|end|>\n<|endoftext|>'}

maybe_convert_to_chatml

trl.maybe_convert_to_chatml

< >

( example: dict ) dict[str, list]

参数

  • example (dict[str, list]) — 包含消息列表的单个数据条目。

返回

dict[str, list]

重新格式化为 ChatML 风格的示例。

将包含字段 fromvalue 的对话数据集转换为 ChatML 格式。

此函数修改对话数据以符合 OpenAI 的 ChatML 格式

  • 将消息字典中的键 "from" 替换为 "role"
  • 将消息字典中的键 "value" 替换为 "content"
  • 为了与 ChatML 保持一致,将 "conversations" 重命名为 "messages"

示例

>>> from trl import maybe_convert_to_chatml

>>> example = {
...     "conversations": [
...         {"from": "user", "value": "What color is the sky?"},
...         {"from": "assistant", "value": "It is blue."},
...     ]
... }
>>> maybe_convert_to_chatml(example)
{'messages': [{'role': 'user', 'content': 'What color is the sky?'},
              {'role': 'assistant', 'content': 'It is blue.'}]}

extract_prompt

trl.extract_prompt

< >

( example: dict )

从偏好数据示例中提取共享提示,其中提示隐含在“选择的”和“拒绝的”两个完成中。

更多详情,请参见 maybe_extract_prompt()

maybe_extract_prompt

trl.maybe_extract_prompt

< >

( example: dict ) dict[str, list]

参数

  • example (dict[str, list]) — 代表偏好数据集中单个数据条目的字典。它必须包含键 "chosen""rejected",其中每个值既可以是对话式,也可以是标准 (str) 格式。

返回

dict[str, list]

一个包含以下内容的字典:

  • "prompt":“选择的”和“拒绝的”两个完成之间最长的公共前缀。
  • "chosen":“选择的”完成中去除提示后的剩余部分。
  • "rejected":“拒绝的”完成中去除提示后的剩余部分。

从偏好数据示例中提取共享提示,其中提示隐含在“选择的”和“拒绝的”两个完成中。

如果示例已经包含 "prompt" 键,函数将按原样返回该示例。否则,函数会识别“选择的”和“拒绝的”两个完成之间最长的公共对话轮次序列(前缀),并将其提取为提示。然后,它会从各自的“选择的”和“拒绝的”完成中移除这个提示。

示例

>>> example = {
...     "chosen": [
...         {"role": "user", "content": "What color is the sky?"},
...         {"role": "assistant", "content": "It is blue."},
...     ],
...     "rejected": [
...         {"role": "user", "content": "What color is the sky?"},
...         {"role": "assistant", "content": "It is green."},
...     ],
... }
>>> extract_prompt(example)
{'prompt': [{'role': 'user', 'content': 'What color is the sky?'}],
 'chosen': [{'role': 'assistant', 'content': 'It is blue.'}],
 'rejected': [{'role': 'assistant', 'content': 'It is green.'}]}

或者,使用 datasets.Datasetmap 方法

>>> from trl import extract_prompt
>>> from datasets import Dataset

>>> dataset_dict = {
...     "chosen": [
...         [
...             {"role": "user", "content": "What color is the sky?"},
...             {"role": "assistant", "content": "It is blue."},
...         ],
...         [
...             {"role": "user", "content": "Where is the sun?"},
...             {"role": "assistant", "content": "In the sky."},
...         ],
...     ],
...     "rejected": [
...         [
...             {"role": "user", "content": "What color is the sky?"},
...             {"role": "assistant", "content": "It is green."},
...         ],
...         [
...             {"role": "user", "content": "Where is the sun?"},
...             {"role": "assistant", "content": "In the sea."},
...         ],
...     ],
... }
>>> dataset = Dataset.from_dict(dataset_dict)
>>> dataset = dataset.map(extract_prompt)
>>> dataset[0]
{'prompt': [{'role': 'user', 'content': 'What color is the sky?'}],
 'chosen': [{'role': 'assistant', 'content': 'It is blue.'}],
 'rejected': [{'role': 'assistant', 'content': 'It is green.'}]}

unpair_preference_dataset

trl.unpair_preference_dataset

< >

( dataset: ~DatasetType num_proc: typing.Optional[int] = None desc: typing.Optional[str] = None ) Dataset

参数

  • dataset (DatasetDatasetDict) — 要取消配对的偏好数据集。数据集必须包含 "chosen""rejected" 列,以及可选的 "prompt" 列。
  • num_proc (intNone可选,默认为 None) — 用于处理数据集的进程数。
  • desc (strNone可选,默认为 None) — 在映射示例时,与进度条一起显示的描述性文字。

返回

数据集

未配对的偏好数据集。

取消偏好数据集的配对。

示例

>>> from datasets import Dataset

>>> dataset_dict = {
...     "prompt": ["The sky is", "The sun is"],
...     "chosen": [" blue.", "in the sky."],
...     "rejected": [" green.", " in the sea."],
... }
>>> dataset = Dataset.from_dict(dataset_dict)
>>> dataset = unpair_preference_dataset(dataset)
>>> dataset
Dataset({
    features: ['prompt', 'completion', 'label'],
    num_rows: 4
})

>>> dataset[0]
{'prompt': 'The sky is', 'completion': ' blue.', 'label': True}

maybe_unpair_preference_dataset

trl.maybe_unpair_preference_dataset

< >

( dataset: ~DatasetType num_proc: typing.Optional[int] = None desc: typing.Optional[str] = None ) DatasetDatasetDict

参数

  • dataset (DatasetDatasetDict) — 要取消配对的偏好数据集。数据集必须包含 "chosen""rejected" 列,以及可选的 "prompt" 列。
  • num_proc (intNone可选,默认为 None) — 用于处理数据集的进程数。
  • desc (strNone可选,默认为 None) — 在映射示例时,与进度条一起显示的描述性文字。

返回

DatasetDatasetDict

如果偏好数据集已配对,则返回未配对的数据集,否则返回原始数据集。

如果偏好数据集已配对,则取消配对。

示例

>>> from datasets import Dataset

>>> dataset_dict = {
...     "prompt": ["The sky is", "The sun is"],
...     "chosen": [" blue.", "in the sky."],
...     "rejected": [" green.", " in the sea."],
... }
>>> dataset = Dataset.from_dict(dataset_dict)
>>> dataset = unpair_preference_dataset(dataset)
>>> dataset
Dataset({
    features: ['prompt', 'completion', 'label'],
    num_rows: 4
})

>>> dataset[0]
{'prompt': 'The sky is', 'completion': ' blue.', 'label': True}

pack_dataset

trl.pack_dataset

< >

( dataset: ~DatasetType seq_length: int strategy: str = 'bfd' map_kwargs: typing.Optional[dict[str, typing.Any]] = None ) DatasetDatasetDict

参数

  • dataset (DatasetDatasetDict) — 要打包的数据集
  • seq_length (int) — 目标打包序列长度。
  • strategy (str可选,默认为 "bfd") — 使用的打包策略。可以是:

    • "bfd" (最佳拟合递减法):较慢但保留序列边界。序列永远不会在中间被切断。
    • "wrapped":更快但更具侵略性。忽略序列边界,会为了完全填充每个打包序列而切断序列。
  • map_kwargs (dictNone可选,默认为 None) — 在打包示例时传递给数据集的 map 方法的额外关键字参数。

返回

DatasetDatasetDict

打包了序列的数据集。由于序列被合并,示例数量可能会减少。

将数据集中的序列打包成大小为 seq_length 的块。

示例

>>> from datasets import Dataset
>>> from trl import pack_dataset

>>> examples = {
...     "input_ids": [[1, 2, 3], [4, 5], [6, 7, 8], [9]],
...     "attention_mask": [[1, 1, 0], [1, 0], [1, 0, 0], [1]],
... }
>>> dataset = Dataset.from_dict(examples)
>>> packed_dataset = pack_dataset(dataset, seq_length=4, strategy="bfd")
>>> packed_dataset[:]
{'input_ids': [[1, 2, 3, 9], [6, 7, 8, 4, 5]],
 'attention_mask': [[1, 1, 0, 1], [1, 0, 0, 1, 0]]}

truncate_dataset

trl.truncate_dataset

< >

( dataset: ~DatasetType max_length: int map_kwargs: typing.Optional[dict[str, typing.Any]] = None ) DatasetDatasetDict

参数

  • dataset (DatasetDatasetDict) — 要截断的数据集。
  • seq_length (int) — 要截断到的最大序列长度。
  • map_kwargs (dictNone可选,默认为 None) — 在截断示例时传递给数据集的 map 方法的额外关键字参数。

返回

DatasetDatasetDict

序列被截断的数据集。

将数据集中的序列截断到指定的 max_length

示例

>>> from datasets import Dataset

>>> examples = {
...     "input_ids": [[1, 2, 3], [4, 5, 6, 7], [8]],
...     "attention_mask": [[0, 1, 1], [0, 0, 1, 1], [1]],
... }
>>> dataset = Dataset.from_dict(examples)
>>> truncated_dataset = truncate_dataset(dataset, max_length=2)
>>> truncated_dataset[:]
{'input_ids': [[1, 2], [4, 5], [8]],
 'attention_mask': [[0, 1], [0, 0], [1]]}
< > 在 GitHub 上更新