Repository: oumi-ai/oumi
File: scripts/datasets/save_conversations.py (Python)

Code before:

    dataset_kwargs = copy.deepcopy(dataset_kwargs) if dataset_kwargs is not None else {}
    if processor_name:
        dataset_kwargs["processor_name"] = processor_name

    dataset = dataset_class(
        dataset_name=dataset_name,
        dataset_path=dataset_path,
        split=dataset_split,
        subset=dataset_subset,
        tokenizer=tokenizer,
        trust_remote_code=trust_remote_code,
        **dataset_kwargs,
    )
    if not isinstance(dataset, BaseSftDataset):
        raise ValueError(
            f"Dataset '{dataset_name}' is not a subclass of BaseSftDataset. "
            f"Actual type: {type(dataset)}"
        )
    return dataset


def _split_message_list_on_user(messages: list[Message]) -> list[list[Message]]:
    turn_start_indices: list[int] = []
    for idx, m in enumerate(messages):
        if m.role == Role.USER:
            turn_start_indices.append(idx)

    num_turns = len(turn_start_indices)
    if num_turns == 0:
        return [messages]

    # SYSTEM messages
    lead_slice = messages[0 : turn_start_indices[0]]
    result = []

    image_content_items: list[ContentItem] = []

    for idx, start_index in enumerate(turn_start_indices):
        turn_messages = copy.copy(lead_slice)

        user_message = messages[start_index]
        if idx == 0:
            image_content_items = user_message.image_content_items
        elif len(image_content_items) > 0 and not user_message.contains_images():
            # Copy images from the first turn
            user_message = Message(
                role=Role.USER,
                content=(image_content_items + user_message.text_content_items),
            )
        turn_messages.append(user_message)

Review comment (suggestion):
Food for thought here: should you add a check for an empty input? I think `_split_message_list_on_user([])` should return `[]` not `[ [] ]`

Code after:

    dataset_kwargs = copy.deepcopy(dataset_kwargs) if dataset_kwargs is not None else {}
    if processor_name:
        dataset_kwargs["processor_name"] = processor_name

    dataset = dataset_class(
        dataset_name=dataset_name,
        dataset_path=dataset_path,
        split=dataset_split,
        subset=dataset_subset,
        tokenizer=tokenizer,
        trust_remote_code=trust_remote_code,
        **dataset_kwargs,
    )
    if not isinstance(dataset, BaseSftDataset):
        raise ValueError(
            f"Dataset '{dataset_name}' is not a subclass of BaseSftDataset. "
            f"Actual type: {type(dataset)}"
        )
    return dataset


def _split_message_list_on_user(messages: list[Message]) -> list[list[Message]]:
    if len(messages) == 0:
        return []
    turn_start_indices: list[int] = []
    for idx, m in enumerate(messages):
        if m.role == Role.USER:
            turn_start_indices.append(idx)

    num_turns = len(turn_start_indices)
    if num_turns == 0:
        return [messages]

    # SYSTEM messages
    lead_slice = messages[0 : turn_start_indices[0]]
    result = []

    image_content_items: list[ContentItem] = []

    for idx, start_index in enumerate(turn_start_indices):
        turn_messages = copy.copy(lead_slice)

        user_message = messages[start_index]
        if idx == 0:
            image_content_items = user_message.image_content_items
        elif len(image_content_items) > 0 and not user_message.contains_images():
            # Copy images from the first turn
            user_message = Message(
                role=Role.USER,
                content=(image_content_items + user_message.text_content_items),