in trl/data_utils.py [0:0]
def maybe_extract_prompt(example: dict[str, list]) -> dict[str, list]:
r"""
Extracts the shared prompt from a preference data example, where the prompt is implicit within both the chosen and
rejected completions.
If the example already contains a `"prompt"` key, the function returns the example as is. Else, the function
identifies the longest common sequence (prefix) of conversation turns between the "chosen" and "rejected"
completions and extracts this as the prompt. It then removes this prompt from the respective "chosen" and
"rejected" completions.
Args:
example (`dict[str, list]`):
A dictionary representing a single data entry in the preference dataset. It must contain the keys
`"chosen"` and `"rejected"`, where each value is either conversational or standard (`str`).
Returns:
`dict[str, list]`: A dictionary containing:
- `"prompt"`: The longest common prefix between the "chosen" and "rejected" completions.
- `"chosen"`: The remainder of the "chosen" completion, with the prompt removed.
- `"rejected"`: The remainder of the "rejected" completion, with the prompt removed.
Examples:
```python
>>> 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.'}]}
```
Or, with the `map` method of `datasets.Dataset`:
```python
>>> 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.'}]}
```
"""
# Some dataset add a `"prompt"` column, even though the prompt is implicit and included in the "chosen" and
# "rejected" completions. E.g.:
# {"prompt": "What color is the sky?",
# "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."}]}
# That's why we check if the prompt is also conversational before deciding not to extract it.
if "chosen" not in example or "rejected" not in example: # not a preference example
return example
if "prompt" in example:
# Both conversational or both non-conversational
chosen_conv = is_conversational({"chosen": example["chosen"]})
prompt_conv = is_conversational({"prompt": example["prompt"]})
if (chosen_conv and prompt_conv) or (not chosen_conv and not prompt_conv):
return example
return extract_prompt({"chosen": example["chosen"], "rejected": example["rejected"]})