loader/__init__.py

# loader/__init__.py

from .models import (
    BaseModel,
    ModelInfo,
    InferenceResult,
    get_model,
)
from .data import (
    list_datasets,
    get_examples,
    get_example_by_id,
)
from .perplexity import (
    score_continuation,
    score_reference_autoregressive_hf,
    score_reference_autoregressive_vllm,
)
from .utils import device_info, set_seed

# --- Optional vLLM hook (only works inside the vLLM container) ---
import traceback as _traceback

_VLLM_IMPORT_ERR = None

try:
    from .models_vllm import VLLMModel
except Exception as exc:  # capture and expose real reason
    VLLMModel = None
    _VLLM_IMPORT_ERR = _traceback.format_exc()

    def get_model_vllm(size: str = "small") -> BaseModel:
        msg = (
            "vLLM is not available in this environment. "
            "Run inside the vllm_meng.sif container to use get_model_vllm()."
        )
        if _VLLM_IMPORT_ERR:
            msg += f"\nUnderlying import error:\n{_VLLM_IMPORT_ERR}"
        raise RuntimeError(msg)  # surface actual import failure

else:

    def get_model_vllm(size: str = "small") -> BaseModel:
        return VLLMModel(size)

# --- Optional vLLM multimodal hook ---
try:
    from .models_mm_vllm import VLLMMultimodalModel, list_mm_model_keys, get_mm_model_info

    def get_model_mm_vllm(key: str = "vl_small") -> BaseModel:
        return VLLMMultimodalModel(key)

except Exception:
    VLLMMultimodalModel = None

    def get_model_mm_vllm(key: str = "vl_small") -> BaseModel:
        raise RuntimeError(
            "Multimodal vLLM is not available in this environment. "
            "Run inside the vLLM container to use get_model_mm_vllm()."
        )

    def list_mm_model_keys() -> list[str]:
        return []

    def get_mm_model_info(_key: str):
        return None


def debug_vllm_status():
    """
    Print basic diagnostics about vLLM availability/import state.
    """
    import sys, os

    print("cwd:", os.getcwd())
    print("sys.path[0:3]:", sys.path[:3])
    try:
        import vllm  # type: ignore

        print("vllm import: OK", getattr(vllm, '__version__', 'unknown'))
    except Exception as exc:  # pragma: no cover - diagnostic helper
        print("vllm import: FAILED", exc)
    if _VLLM_IMPORT_ERR:
        print("Stored VLLM import error trace:\n", _VLLM_IMPORT_ERR)