test_models.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Multi-GPU LLM benchmark runner for guardrails tasks.

Loads .env from project root if present (e.g. HF_TOKEN for gated models like Gemma).

Install:
  pip install -U "torch" "transformers>=4.45" "accelerate>=0.33" "safetensors"
  pip install vllm   # optional, for --backend vllm

Recommended: serve with vllm, then benchmark:
  vllm serve Qwen/Qwen2.5-32B-Instruct --tensor-parallel-size 6
  python test_models.py --backend vllm_serve --model Qwen/Qwen2.5-32B-Instruct

Usage:
  python test_models.py --list-models
  python test_models.py --list-presets
  python test_models.py --backend vllm --model-key qwen25_32b_instruct
  python test_models.py --preset around_30b_multilingual --subprocess
"""


def _load_dotenv():
    """Load .env from project root so HF_TOKEN etc. are available for gated models."""
    import os as _os
    from pathlib import Path
    env_path = Path(__file__).resolve().parent / ".env"
    if not env_path.is_file():
        return
    for line in env_path.read_text().splitlines():
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        if "=" in line:
            key, _, value = line.partition("=")
            key = key.strip()
            value = value.strip().strip('"').strip("'")
            if key and _os.environ.get(key) is None:
                _os.environ[key] = value


import os

_load_dotenv()

from bench.cli import main

if __name__ == "__main__":
    main(script_path=os.path.abspath(__file__))