scripts/recsys_benchmark.py

#!/usr/bin/env python
"""Run the OTS recommender export and diversity benchmark."""
from __future__ import annotations

import json
from pathlib import Path

import click

from backend.core import Backend
from backend.recsys_benchmark import run_recsys_benchmark


@click.command(help=__doc__)
@click.option("--db", type=click.Path(path_type=Path), default=Path("data/store_slim.db"))
@click.option("--out", type=click.Path(path_type=Path), default=Path("data/recsys/benchmark.json"))
@click.option("--candidate-cache", type=click.Path(path_type=Path), default=Path("data/recsys/candidates.json"))
@click.option("--recbole-dir", type=click.Path(path_type=Path), default=Path("data/recsys/recbole"))
@click.option("--model", type=click.Choice(["ease", "svd", "mf", "bpr", "lightgcn"]), default="svd")
@click.option("--limit", type=int, default=20)
@click.option("--candidate-limit", type=int, default=200)
@click.option("--journey-steps", type=int, default=8)
@click.option("--skip-graph", is_flag=True, help="Skip current graph recommender metrics and journey.")
@click.option(
    "--mode",
    "modes",
    multiple=True,
    type=click.Choice(["graph", "ots", "hybrid"]),
    help="Serving mode to measure; repeat for multiple modes. Defaults to all modes.",
)
@click.option("--user", "users", multiple=True, help="User id to benchmark; repeat for multiple users.")
def main(
    db: Path,
    out: Path,
    candidate_cache: Path,
    recbole_dir: Path,
    model: str,
    limit: int,
    candidate_limit: int,
    journey_steps: int,
    skip_graph: bool,
    modes: tuple[str, ...],
    users: tuple[str, ...],
) -> None:
    mode_list = list(modes) if modes else (["ots"] if skip_graph else ["graph", "ots", "hybrid"])
    if skip_graph and any(mode in {"graph", "hybrid"} for mode in mode_list):
        raise click.UsageError("--skip-graph can only be used with --mode ots, because hybrid calls graph mode.")
    backend = Backend(db)
    user_ids = list(users) if users else ["vocal-gnu-15", "neat-bobcat-85"]
    report = run_recsys_benchmark(
        backend,
        user_ids=user_ids,
        model_name=model,
        limit=limit,
        candidate_limit=candidate_limit,
        recbole_dir=recbole_dir,
        candidate_cache_path=candidate_cache,
        include_graph=not skip_graph,
        journey_steps=journey_steps,
        modes=mode_list,
    )
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(json.dumps(report, indent=2, sort_keys=True), encoding="utf-8")
    click.echo(json.dumps({"out": str(out), "candidate_cache": str(candidate_cache)}, indent=2))


if __name__ == "__main__":
    main()