#!/usr/bin/env python3
"""
Runs lm-evaluation-harness against Ivme-Conversate-v2 using the ivme_lm.py adapter.

Example:
    python run_eval.py \
        --checkpoint /path/to/ckpt_final.pt \
        --tokenizer /path/to/tokenizer.json \
        --model_code_dir /path/to/model_folder \
        --tasks wikitext,arc_easy,blimp \
        --device cuda:0 \
        --batch_size 16

Notes:
- `--model_code_dir` should point at the folder containing the model/ package
  (config.py, transformer.py, etc.) -- i.e. what snapshot_download gave you,
  or wherever you cloned the repo.
- `blimp` here means the actual harness BLiMP group task, which covers the
  real 67 paradigms with their real, correct config names on nyu-mll/blimp.
  This replaces the paradigm list in the old custom script, which had several
  fabricated/misspelled task names.
- generate_until (free-form generation tasks) is not implemented in the
  adapter since this is a non-instruction-tuned base model -- stick to
  loglikelihood-based tasks (wikitext, arc_easy, blimp, hellaswag, piqa, etc.)
"""
import argparse

import ivme_lm  # noqa: F401  (registers the "ivme" model with lm_eval on import)
import lm_eval
from lm_eval.utils import make_table


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint", type=str, required=True)
    parser.add_argument("--tokenizer", type=str, required=True)
    parser.add_argument("--model_code_dir", type=str, default="")
    parser.add_argument("--tasks", type=str, default="wikitext,arc_easy,blimp")
    parser.add_argument("--device", type=str, default="cuda:0")
    parser.add_argument("--batch_size", type=int, default=16)
    parser.add_argument("--limit", type=float, default=None,
                         help="Optional: cap number of docs per task, for a quick sanity run first.")
    args = parser.parse_args()

    model_args = (
        f"checkpoint={args.checkpoint},"
        f"tokenizer={args.tokenizer},"
        f"model_code_dir={args.model_code_dir},"
        f"device={args.device},"
        f"batch_size={args.batch_size}"
    )

    results = lm_eval.simple_evaluate(
        model="ivme",
        model_args=model_args,
        tasks=args.tasks.split(","),
        limit=args.limit,
    )

    print(make_table(results))
    if "groups" in results and results["groups"]:
        print(make_table(results, "groups"))


if __name__ == "__main__":
    main()