def _load_models()

in yourbench/utils/inference/inference_core.py [0:0]

• 25 lines of code
• 7 McCabe index (conditional complexity)

def _load_models(base_config: Dict[str, Any], step_name: str) -> List[Model]:
    """
    Load only the models assigned to this step from the config's 'model_list' and 'model_roles'.
    If no model role is defined for the step, use the first model from model_list.
    """
    all_configured_models = base_config.get("model_list", [])
    role_models = base_config.get("model_roles", {}).get(step_name, [])

    # If no role models are defined for this step, use the first model from model_list
    if not role_models and all_configured_models:
        first_model_config = all_configured_models[0]
        logger.info(
            "No models defined in model_roles for step '{}'. Using the first model from model_list: {}",
            step_name,
            first_model_config["model_name"],
        )
        return [
            Model(**{**first_model_config, "encoding_name": first_model_config.get("encoding_name", "cl100k_base")})
        ]

    # Filter out only those with a matching 'model_name'
    matched = []
    for m_config in all_configured_models:
        if m_config["model_name"] in role_models:
            model_instance = Model(**{**m_config, "encoding_name": m_config.get("encoding_name", "cl100k_base")})
            matched.append(model_instance)

    logger.info(
        "Found {} models in config for step '{}': {}",
        len(matched),
        step_name,
        [m.model_name for m in matched],
    )
    return matched