数据实用工具

is_conversational

trl.is_conversational

< source >

( example: dict[str, typing.Any] ) → bool

参数

example (dict[str, Any]) — 数据集的单条数据项。根据数据集类型的不同，该数据项可能包含不同的键。

返回值

bool

如果数据为对话格式，则为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

maybe_convert_to_chatml

trl.maybe_convert_to_chatml

< source >

( example: dict[str, list] ) → dict[str, list]

参数

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

返回值

dict[str, list]

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

将包含 from 和 value 字段的对话数据集转换为 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

< source >

( example: dict[str, collections.abc.Sequence] ) → dict[str, list]

参数

example (dict[str, list]) — 表示偏好数据集中单个数据条目的字典。它必须包含键 "chosen" 和 "rejected"，其中每个值均为对话式或标准格式（str）。

返回值

dict[str, list]

包含以下内容的字典：

"prompt"：“chosen”和“rejected”补全之间最长的公共前缀。
"chosen"：移除提示词后的“chosen”补全的剩余部分。
"rejected"：移除提示词后，“rejected”补全的剩余部分。

从偏好数据示例中提取共享提示词，其中提示词隐含在选中的和拒绝的补全内容中。

该函数会识别“chosen”和“rejected”补全之间对话轮次的最长公共序列（前缀），并将其提取为提示词。随后，它会从对应的“chosen”和“rejected”补全中移除该提示词。

示例：

>>> 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.'}]}

或者，使用 Dataset 的 map 方法：

>>> 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

< source >

( dataset: typing.Union[~DatasetType, ~IterableDatasetType]**map_kwargs )

参数

dataset (Dataset 或 DatasetDict 或 IterableDataset 或 IterableDatasetDict) — 用于拆分的偏好数据集。该数据集必须包含 "chosen"、"rejected" 列，并可选包含 "prompt" 列。
**map_kwargs (dict, 可选) — 在拆分偏好时传递给数据集 map 方法的额外关键字参数。

取消配对偏好数据集。

输出包含 "prompt"、"completion" 和 "label"，以及任何额外的列，这些列会为每个 chosen 和 rejected 行进行重复。

示例：

>>> 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}

Update on GitHub

TRL

数据实用工具

is_conversational

trl.is_conversational

maybe_convert_to_chatml

trl.maybe_convert_to_chatml

extract_prompt

trl.extract_prompt

unpair_preference_dataset

trl.unpair_preference_dataset