Repository: oumi-ai/oumi
File: tests/unit/datasets/test_chat_templates.py (Python)

Code before:
        assert len(unique_text_pieces) == 3
        expected_lines = (
            [
                "
system",
                "You are a helpful assistant.
",
                "
user",
                f"{unique_text_pieces[0]}
",
                "
assistant",
                f"{unique_text_pieces[1]}
",
                "
user",
                f"{unique_text_pieces[2]}
",
            ]
            + (["
assistant"] if add_generation_prompt else [])
            + [""]
        )
        expected = "\n".join(expected_lines)
        print(expected)
        print(oumi_result)
        assert oumi_result == expected, debug_tag
        assert hf_result.startswith(
            expected
        ), f"{debug_tag}\n\n{hf_result}\n\n{expected}"

    # With images.
    test_convo_tuple: ConversationTuple = create_test_conversation(3, num_with_images=2)
    for add_generation_prompt in (False, True):
        debug_tag = f"Multi-modal: add_generation_prompt: {add_generation_prompt} "
        oumi_result = oumi_tokenizer.apply_chat_template(
            test_convo_tuple.convo.messages,  # type: ignore
            tokenize=False,
            add_generation_prompt=add_generation_prompt,
        )
        unique_text_pieces = test_convo_tuple.unique_text_pieces
        assert len(unique_text_pieces) == 3
        print(oumi_result)
        expected_lines = (
            [
                "
system",
                "You are a helpful assistant.
",
                "
user",
                f"


{unique_text_pieces[0]}
",
                "
assistant",
                f"


{unique_text_pieces[1]}
",
                "
user",
                f"{unique_text_pieces[2]}
",
            ]
            + (["
assistant"] if add_generation_prompt else [])
            + [""]
        )
        expected = "\n".join(expected_lines)
        assert oumi_result == expected, debug_tag

Review comment (suggestion):
Do you know if HF Qwen model comes with built-in template that supports images? If yes, we could also compute `hf_result` with images and compare with `oumi_result` 

(other models above "Llama and Phi3" come with text-only templates so I was unable to do it for them, but maybe we have better luck with Qwen)

Code after:
        assert len(unique_text_pieces) == 3
        expected_lines = (
            [
                "
system",
                "You are a helpful assistant.
",
                "
user",
                f"{unique_text_pieces[0]}
",
                "
assistant",
                f"{unique_text_pieces[1]}
",
                "
user",
                f"{unique_text_pieces[2]}
",
            ]
            + (["
assistant"] if add_generation_prompt else [])
            + [""]
        )
        expected = "\n".join(expected_lines)
        assert oumi_result == expected, debug_tag
        assert hf_result.startswith(
            expected
        ), f"{debug_tag}\n\n{hf_result}\n\n{expected}"

    # With images.
    test_convo_tuple: ConversationTuple = create_test_conversation(3, num_with_images=2)
    for add_generation_prompt in (False, True):
        debug_tag = f"Multi-modal: add_generation_prompt: {add_generation_prompt} "
        oumi_result = oumi_tokenizer.apply_chat_template(
            test_convo_tuple.convo.messages,  # type: ignore
            tokenize=False,
            add_generation_prompt=add_generation_prompt,
        )
        unique_text_pieces = test_convo_tuple.unique_text_pieces
        assert len(unique_text_pieces) == 3
        expected_lines = (
            [
                "
system",
                "You are a helpful assistant.
",
                "
user",
                f"


{unique_text_pieces[0]}
",
                "
assistant",
                f"


{unique_text_pieces[1]}
",
                "
user",
                f"{unique_text_pieces[2]}
",
            ]
            + (["
assistant"] if add_generation_prompt else [])
            + [""]
        )
        expected = "\n".join(expected_lines)
        assert oumi_result == expected, debug_tag