vivid_eval.py

"""
VIVID Evaluation Framework — CLI Entry Point
============================================
Run any VIVID benchmark task from the command line.

Commands:
  generate       Generate idiom explanations (vLLM or API)
  judge          Score explanations with LLM-as-a-Judge (GPT-4.1)
  discriminate   Topic / pattern classification via lm_eval
  full-pipeline  generate → judge in one command

Run `python vivid_eval.py <command> --help` for per-command options.
"""

import argparse
from datetime import datetime

from vivid.constants import DEFAULT_DATASET, DEFAULT_RESULTS_DIR
from vivid.generate import run_generate
from vivid.judge import run_judge
from vivid.discriminate import run_discriminate
from vivid.pipeline import run_full_pipeline


# ── Shared argument group ─────────────────────────────────────────────────────

def _shared_parser():
    p = argparse.ArgumentParser(add_help=False)
    p.add_argument("--results-dir", default=DEFAULT_RESULTS_DIR,
                   help="Directory to save all results  (default: ./results)")
    p.add_argument("--api-key", default=None,
                   help="OpenAI API key  (or set OPENAI_API_KEY env var)")
    p.add_argument("--gemini-api-key", default=None,
                   help="Google Gemini API key  (or set GEMINI_API_KEY env var)")
    p.add_argument("--api-delay", type=float, default=1.0,
                   help="Seconds between API calls for rate-limiting  (default: 1.0)")
    return p


# ── vLLM argument group (reused by generate, discriminate, full-pipeline) ─────

def _add_vllm_args(p):
    p.add_argument("--tensor-parallel", type=int, default=1,
                   help="Tensor parallel size for vLLM  (default: 1)")
    p.add_argument("--gpu-memory-util", type=float, default=0.85,
                   help="GPU memory utilisation for vLLM  (default: 0.85)")
    p.add_argument("--dtype", default="auto",
                   help="Model dtype: auto | float16 | bfloat16  (default: auto)")
    p.add_argument("--max-model-len", type=int, default=None,
                   help="Override context length for vLLM  (optional)")


# ── Parser ────────────────────────────────────────────────────────────────────

def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="vivid_eval",
        description="VIVID Benchmark Evaluation Framework",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples
--------
# Open-source model — zero-shot generation
python vivid_eval.py generate --model Qwen/Qwen3-14B --prompt zero-shot

# OpenAI model — few-shot generation
python vivid_eval.py generate --model gpt-4o --prompt few-shot --api-key sk-...

# Gemini model — zero-shot generation  (uses google-genai library)
python vivid_eval.py generate --model gemini-2.5-flash --prompt zero-shot --gemini-api-key AIza...

# Score with OpenAI judge
python vivid_eval.py judge \\
    --input  results/generate_Qwen_Qwen3-14B_zero-shot_<ts>.csv \\
    --model-col Qwen_Qwen3-14B_zero_shot_explanation \\
    --api-key sk-...

# Score with Gemini judge  (uses google-genai library)
python vivid_eval.py judge \\
    --input  results/generate_Qwen_Qwen3-14B_zero-shot_<ts>.csv \\
    --model-col Qwen_Qwen3-14B_zero_shot_explanation \\
    --judge-model gemini-2.5-flash \\
    --gemini-api-key AIza...

# Discriminative evaluation — open-source (vLLM backend)
python vivid_eval.py discriminate --model Qwen/Qwen3-14B --task both

# Discriminative evaluation — Gemini  (uses OpenAI-compat endpoint via lm_eval)
python vivid_eval.py discriminate --model gemini-2.5-flash --task topic --api-key AIza...

# Full generate + judge pipeline
python vivid_eval.py full-pipeline \\
    --model Qwen/Qwen3-14B --prompt zero-shot --api-key sk-...
        """,
    )

    sub = parser.add_subparsers(dest="command", required=True)
    shared = _shared_parser()

    # ── generate ─────────────────────────────────────────────────────────────
    p_gen = sub.add_parser(
        "generate", parents=[shared],
        help="Generate idiom/proverb explanations with any model",
    )
    p_gen.add_argument("--model", required=True,
                       help="HuggingFace model ID or API model name")
    p_gen.add_argument("--prompt", choices=["zero-shot", "few-shot"], default="zero-shot",
                       help="Prompting strategy  (default: zero-shot)")
    p_gen.add_argument("--dataset", default=DEFAULT_DATASET,
                       help="Path to VIVID_Dataset.csv")
    p_gen.add_argument("--batch-size", type=int, default=64,
                       help="vLLM inference batch size  (default: 64)")
    p_gen.add_argument("--max-tokens", type=int, default=150,
                       help="Max tokens to generate  (default: 150)")
    p_gen.add_argument("--temperature", type=float, default=0.7,
                       help="Sampling temperature  (default: 0.7)")
    _add_vllm_args(p_gen)

    # ── judge ─────────────────────────────────────────────────────────────────
    p_jud = sub.add_parser(
        "judge", parents=[shared],
        help="Score explanations with LLM-as-a-Judge (GPT-4.1)",
    )
    p_jud.add_argument("--input", required=True,
                       help="CSV produced by the `generate` command")
    p_jud.add_argument("--model-col", required=True,
                       help="Column name that contains the LLM explanations")
    p_jud.add_argument("--judge-model", default="gpt-4.1",
                       help="Judge model to use  (default: gpt-4.1)")

    # ── discriminate ──────────────────────────────────────────────────────────
    p_disc = sub.add_parser(
        "discriminate", parents=[shared],
        help="Run topic / pattern classification via lm_eval",
    )
    p_disc.add_argument("--model", required=True,
                        help="HuggingFace model ID or API model name")
    p_disc.add_argument("--task", choices=["topic", "pattern", "both"], default="both",
                        help="Which classification task to run  (default: both)")
    p_disc.add_argument("--dataset", default=DEFAULT_DATASET,
                        help="Path to VIVID_Dataset.csv  (used to build eval data)")
    p_disc.add_argument("--rebuild-eval-data", action="store_true",
                        help="Force rebuild of topic.json and pattern.json from dataset")
    p_disc.add_argument("--num-fewshot", type=int, default=3,
                        help="Few-shot examples for lm_eval  (default: 3)")
    p_disc.add_argument("--batch-size", type=int, default=32,
                        help="lm_eval batch size for open-source models  (default: 32)")
    _add_vllm_args(p_disc)

    # ── full-pipeline ─────────────────────────────────────────────────────────
    p_full = sub.add_parser(
        "full-pipeline", parents=[shared],
        help="Run generate → judge in one command",
    )
    p_full.add_argument("--model", required=True,
                        help="HuggingFace model ID or API model name")
    p_full.add_argument("--prompt", choices=["zero-shot", "few-shot"], default="zero-shot",
                        help="Prompting strategy  (default: zero-shot)")
    p_full.add_argument("--dataset", default=DEFAULT_DATASET,
                        help="Path to VIVID_Dataset.csv")
    p_full.add_argument("--batch-size", type=int, default=64,
                        help="vLLM inference batch size  (default: 64)")
    p_full.add_argument("--max-tokens", type=int, default=150,
                        help="Max tokens to generate  (default: 150)")
    p_full.add_argument("--temperature", type=float, default=0.7,
                        help="Sampling temperature  (default: 0.7)")
    p_full.add_argument("--judge-model", default="gpt-4.1",
                        help="Judge model  (default: gpt-4.1)")
    _add_vllm_args(p_full)

    return parser


# ── Entry point ───────────────────────────────────────────────────────────────

def main():
    parser = build_parser()
    args = parser.parse_args()

    print(f"\n{'─' * 60}")
    print(f"  VIVID Evaluation Framework  |  {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"{'─' * 60}")

    dispatch = {
        "generate"      : run_generate,
        "judge"         : run_judge,
        "discriminate"  : run_discriminate,
        "full-pipeline" : run_full_pipeline,
    }
    dispatch[args.command](args)


if __name__ == "__main__":
    main()