scripts/run_panel.py

"""Drive aligne's stated-preference panel against a vLLM-served Gemma-3-27B
(base + functional-wellbeing LoRA), on two concept sets:
  1. aligne's bundled 155-item concept dataset (food / animals / people / ideas)
  2. our custom 155-item neutral-emoji set (data/emoji_concepts.json), with
     the 3 maze tiles 🧾 📇 📐 at indices 0..2.

Runs 4 panels total (base vs FT × concepts vs emoji), writing all aligne
artefacts (`panel.json`, `mu.json`, `edges.jsonl`) under
`logs/<ts>_<run-name>/aligne/`.

Designed to be invoked under `setsid` on the GPU pod so it survives ssh drop.

Usage:
    uv run python scripts/run_panel.py \
        --target-url http://localhost:8000/v1 \
        --base-model google/gemma-3-27b-it \
        --ft-model functional-wellbeing \
        --runs all

    # or one at a time:
    uv run python scripts/run_panel.py ... --runs base_concepts
"""

from __future__ import annotations

import argparse
import asyncio
import json
import os
import sys
from dataclasses import asdict
from pathlib import Path

# Add scripts/ to path so we can import _logging when invoked directly.
sys.path.insert(0, str(Path(__file__).resolve().parent))
from _logging import make_run_dir  # noqa: E402

from aligne.client import ChatClient, Endpoint  # noqa: E402
from aligne.metrics.preferences import PanelConfig, run_panel  # noqa: E402


REPO = Path(__file__).resolve().parent.parent
DEFAULT_EMOJI_PATH = REPO / "data" / "emoji_concepts.json"

ALL_RUNS = ("base_concepts", "ft_concepts", "base_emoji", "ft_emoji")


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--target-url", required=True, help="OpenAI-compatible base URL.")
    p.add_argument("--base-model", required=True, help="vLLM served-model-name for the base Gemma.")
    p.add_argument("--ft-model", required=True, help="vLLM served-model-name for the FT (LoRA) Gemma.")
    p.add_argument("--api-key", default=os.environ.get("ALIGNE_TARGET_KEY", "dummy"))
    p.add_argument("--emoji-path", default=str(DEFAULT_EMOJI_PATH), help="Path to emoji JSON list.")
    p.add_argument("--n-concepts", type=int, default=155)
    p.add_argument("--seed", type=int, default=0)
    p.add_argument(
        "--runs",
        nargs="+",
        default=["all"],
        choices=["all", *ALL_RUNS],
        help="Which of the 4 runs to execute.",
    )
    return p.parse_args()


def make_cfg(emoji_path: str | None, n_concepts: int, seed: int) -> PanelConfig:
    """PanelConfig with our preferred defaults (155 items, seed 0)."""
    return PanelConfig(
        n_concepts=n_concepts,
        seed=seed,
        concepts_path=Path(emoji_path) if emoji_path else None,
    )


async def one_run(
    name: str,
    target_url: str,
    model_name: str,
    api_key: str,
    cfg: PanelConfig,
) -> Path:
    run_dir = make_run_dir(
        name,
        config={
            "name": name,
            "target_url": target_url,
            "model_name": model_name,
            "panel_cfg": asdict(cfg),
        },
    )
    out_dir = run_dir / "aligne"
    out_dir.mkdir(parents=True, exist_ok=True)
    endpoint = Endpoint(base_url=target_url, model=model_name, api_key=api_key)
    # On-disk request cache per (target_url, model) so reruns are idempotent.
    cache_path = run_dir / "client_cache.sqlite"
    client = ChatClient(endpoint=endpoint, concurrency=32, cache_path=cache_path)
    print(f"[{name}] starting panel against {target_url} model={model_name}", flush=True)
    panel = await run_panel(client, cfg, out_dir)
    # `panel` is the same dict written to panel.json by run_panel; persist a copy
    # at the run-dir root for convenience.
    (run_dir / "panel_summary.json").write_text(json.dumps(panel, indent=2, default=str))
    print(
        f"[{name}] decisiveness={panel.get('decisiveness')!r} "
        f"unidim_r2={panel.get('unidim_r2')!r} transitivity_triad={panel.get('transitivity_triad')!r}",
        flush=True,
    )
    return run_dir


async def main() -> None:
    args = parse_args()
    runs = ALL_RUNS if "all" in args.runs else tuple(args.runs)

    plans: dict[str, tuple[str, str | None]] = {
        # (model_name, concepts_path)  -- concepts_path None => use aligne's bundled 155 concepts
        "base_concepts": (args.base_model, None),
        "ft_concepts": (args.ft_model, None),
        "base_emoji": (args.base_model, args.emoji_path),
        "ft_emoji": (args.ft_model, args.emoji_path),
    }

    for run_name in runs:
        model_name, concepts_path = plans[run_name]
        cfg = make_cfg(concepts_path, args.n_concepts, args.seed)
        await one_run(
            run_name,
            args.target_url,
            model_name,
            args.api_key,
            cfg,
        )

    print("all requested runs complete", flush=True)


if __name__ == "__main__":
    asyncio.run(main())