Function prepare_dialogue

chat/dialogues.py:205–229

Format example to single- or multi-turn dialogue.

(example, dialogue_template, is_train=True)

def prepare_dialogue(example, dialogue_template, is_train=True):
    """Format example to single- or multi-turn dialogue."""
    # TODO: make this simpler by just ensuring every dataset has a messages column
    if "messages" in example.keys() and example["messages"] is not None:
        dialogue_template.messages = example["messages"]
    elif all(k in example.keys() for k in ("prompt", "completion")):
        # Construct single-turn dialogue from prompt and completion
        dialogue_template.messages = [
            {"role": "user", "content": example["prompt"]},
            {"role": "assistant", "content": example["completion"]},
        ]
    elif "prompt" in example.keys():
        # Construct single-turn dialogue from prompt (inference only)
        dialogue_template.messages = [
            {"role": "user", "content": example["prompt"]},
        ]
    else:
        raise ValueError(
            f"Could not format example as dialogue! Require either `messages` or `[prompt, completion]` or `[prompt]` keys but found {list(example.keys())}"
        )
    if is_train:
        example["text"] = dialogue_template.get_training_prompt()
    else:
        example["text"] = dialogue_template.get_inference_prompt()
    return example


def mask_user_labels(tokenizer, dialogue_template, labels):