Repository: oumi-ai/oumi
File: src/oumi/core/configs/internal/supported_models.py (Python)

Code before:
from typing import NamedTuple, Optional

import transformers

from oumi.core.configs import ModelParams
from oumi.core.configs.internal.internal_model_config import (
    InternalFeatureFirstDimAction,
    InternalFeatureSpec,
    InternalModelConfig,
    InternalVisualModelConfig,
)
from oumi.core.registry import REGISTRY, RegistryType
from oumi.utils.logging import logger


@functools.cache
def find_model_hf_config(model_name: str, *, trust_remote_code: bool):
    """Finds HF model config by model name."""
    try:
        hf_config, unused_kwargs = transformers.AutoConfig.from_pretrained(
            model_name,
            trust_remote_code=trust_remote_code,
            return_unused_kwargs=True,
        )
    except OSError:
        # If the model is not found, it's likely a .gguf model. Instead of failing, we
        # return None, indicating that no HF config was found for this model.
        return None

    if unused_kwargs:
        logger.warning(
            f"Unused kwargs found in '{model_name}' config: {unused_kwargs}."
        )
    return hf_config


class _ModelTypeInfo(NamedTuple):
    model_type: str
    model_class: type
    config: InternalModelConfig
    tested: bool = False


def _create_default_vlm_config(
    pixel_values_variable_shape: bool = False,
) -> InternalModelConfig:
    config = InternalModelConfig()
    config.chat_template = "llava"
    config.model_input_features.update(
        {
            "pixel_values": InternalFeatureSpec(

Review comment (performance):
it can also be an `oumi` custom model (but they aren't normally reaching this code path). 

Can we log some warning here with model name?  The function is cached  so the logging should happen only once.

Code after:
from typing import NamedTuple, Optional

import transformers

from oumi.core.configs import ModelParams
from oumi.core.configs.internal.internal_model_config import (
    InternalFeatureFirstDimAction,
    InternalFeatureSpec,
    InternalModelConfig,
    InternalVisualModelConfig,
)
from oumi.core.registry import REGISTRY, RegistryType
from oumi.utils.logging import logger


@functools.cache
def find_model_hf_config(model_name: str, *, trust_remote_code: bool):
    """Finds HF model config by model name."""
    hf_config, unused_kwargs = transformers.AutoConfig.from_pretrained(
        model_name,
        trust_remote_code=trust_remote_code,
        return_unused_kwargs=True,
    )
    if unused_kwargs:
        logger.warning(
            f"Unused kwargs found in '{model_name}' config: {unused_kwargs}."
        )
    return hf_config


class _ModelTypeInfo(NamedTuple):
    model_type: str
    model_class: type
    config: InternalModelConfig
    tested: bool = False


def _create_default_vlm_config(
    pixel_values_variable_shape: bool = False,
) -> InternalModelConfig:
    config = InternalModelConfig()
    config.chat_template = "llava"
    config.model_input_features.update(
        {
            "pixel_values": InternalFeatureSpec(
                name="pixel_values",
                required=True,
                variable_shape=pixel_values_variable_shape,
                first_dim_action=InternalFeatureFirstDimAction.DROP_IF_DUMMY,
            )
        }