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

Code before:
        model_name="openai-community/gpt2",
        test_image=True,
        image_placeholder="
",
    ),
    ChatTemplateTestSpec(
        chat_template_name="llava",
        model_name="llava-hf/llava-1.5-7b-hf",
        test_image=True,
        image_placeholder="<image>",
    ),
    ChatTemplateTestSpec(
        chat_template_name="phi3-instruct",
        model_name="microsoft/Phi-3-vision-128k-instruct",
        test_image=True,
        trust_remote_code=True,
        image_placeholder="
",
    ),
    ChatTemplateTestSpec(
        chat_template_name="qwen2-vl-instruct",
        model_name="Qwen/Qwen2-VL-2B-Instruct",
        test_image=True,
        image_placeholder="


",
    ),
    ChatTemplateTestSpec(
        chat_template_name="qwen2.5-vl-instruct",
        model_name="Qwen/Qwen2.5-VL-3B-Instruct",
        test_image=True,
        image_placeholder="


",
    ),
    ChatTemplateTestSpec(
        chat_template_name="zephyr",
        model_name="openai-community/gpt2",
    ),
]


def _generate_all_test_specs() -> list[ChatTemplateTestSpec]:
    result = copy.deepcopy(_ALL_CHAT_TEMPLATE_TESTS)

    # Backfill with templates for which there is no explicit test defined yet.
    known_template_names = {t.chat_template_name for t in result}
    chat_template_dir = get_oumi_root_directory() / "datasets" / "chat_templates"
    for f in chat_template_dir.glob("*.jinja"):
        template_name = f.stem
        if template_name in known_template_names:
            continue
        logger.warning(
            f"No explicit chat template test is configured for '{f}' yet! "
            "Consider adding a new entry to _ALL_CHAT_TEMPLATE_TESTS."
        )
        result.append(

Review comment (suggestion):
this model is unusable during testing because it requires building `transformers` from source 

Let's plug `Qwen/Qwen2-VL-2B-Instruct` temporarily, and add a comment about this

Code after:
        model_name="openai-community/gpt2",
        test_image=True,
        image_placeholder="
",
    ),
    ChatTemplateTestSpec(
        chat_template_name="llava",
        model_name="llava-hf/llava-1.5-7b-hf",
        test_image=True,
        image_placeholder="<image>",
    ),
    ChatTemplateTestSpec(
        chat_template_name="phi3-instruct",
        model_name="microsoft/Phi-3-vision-128k-instruct",
        test_image=True,
        trust_remote_code=True,
        image_placeholder="
",
    ),
    ChatTemplateTestSpec(
        chat_template_name="qwen2-vl-instruct",
        model_name="Qwen/Qwen2-VL-2B-Instruct",
        test_image=True,
        image_placeholder="


",
    ),
    # TODO: To be added back once we update Oumi to the latest (>=4.49.0) version of
    # transformers, which is required for this model. The test passes locally with
    # 4.49.0.dev0
    # ChatTemplateTestSpec(
    #     chat_template_name="qwen2.5-vl-instruct",
    #     model_name="Qwen/Qwen2.5-VL-3B-Instruct",
    #     test_image=True,
    #     image_placeholder="


",
    # ),
    ChatTemplateTestSpec(
        chat_template_name="zephyr",
        model_name="openai-community/gpt2",
    ),
]


def _generate_all_test_specs() -> list[ChatTemplateTestSpec]:
    result = copy.deepcopy(_ALL_CHAT_TEMPLATE_TESTS)

    # Backfill with templates for which there is no explicit test defined yet.
    known_template_names = {t.chat_template_name for t in result}
    chat_template_dir = get_oumi_root_directory() / "datasets" / "chat_templates"
    for f in chat_template_dir.glob("*.jinja"):
        template_name = f.stem
        if template_name in known_template_names:
            continue
        logger.warning(
            f"No explicit chat template test is configured for '{f}' yet! "