# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0.
#!/usr/bin/env python3
"""
MTEB evaluation for nvidia/llama-nv-embed-reasoning-3b on Bright(v1.1) benchmark.

Usage:
  # pip install mteb==2.8.1
  python eval_bright.py --model_name nvidia/llama-nv-embed-reasoning-3b --benchmark "BRIGHT(v1.1)" --task-list BrightBiologyRetrieval
"""

import argparse
import os
from pathlib import Path

import mteb
from mteb_llama_nv_embed_reasoning_3b import LLAMA_NV_EMBED_REASONING_3B_META


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Run MTEB evaluation for nvidia/llama-nv-embed-reasoning-3b (text embedding)."
    )
    parser.add_argument(
        "--model_name",
        type=str,
        default="nvidia/llama-nv-embed-reasoning-3b",
        help="HuggingFace model id (used for result filename).",
    )
    parser.add_argument(
        "--batch_size",
        type=int,
        default=2,
        help="Encode batch size.",
    )
    parser.add_argument(
        "--results_folder",
        type=str,
        default="results_json",
        help="Folder to save JSON results.",
    )
    parser.add_argument(
        "--predictions_folder",
        type=str,
        default=None,
        help="Optional folder to save predictions.",
    )
    parser.add_argument(
        "--benchmark",
        type=str,
        default="BRIGHT(v1.1)",
        help="Benchmark name (e.g. BRIGHT(v1.1), MTEB(eng, v2), BEIR).",
    )
    parser.add_argument(
        "--task-list",
        dest="task_list",
        type=str,
        nargs="*",
        default=None,
        help="Optional list of task class names. If not set, all tasks in the benchmark are run.",
    )
    parser.add_argument(
        "--overwrite_strategy",
        type=str,
        default="always",
        choices=["always", "never", "only-missing", "only-cache"],
        help="Overwrite strategy for cached results.",
    )
    parser.add_argument(
        "--hf-token",
        dest="hf_token",
        type=str,
        default=None,
        help="Hugging Face token (or set HF_TOKEN / HUGGING_FACE_HUB_TOKEN). ",
    )
    args = parser.parse_args()

    
    hf_token = args.hf_token or os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
    if hf_token:
        os.environ["HUGGING_FACE_HUB_TOKEN"] = hf_token

    model = LLAMA_NV_EMBED_REASONING_3B_META
    all_tasks = mteb.get_benchmark(args.benchmark).tasks
    all_tasks_names = " ".join([task.__class__.__name__ for task in all_tasks])
    print(f"Available tasks in benchmark {args.benchmark}: {all_tasks_names}")

    if args.task_list:
        print(f"Running evaluation on specified tasks: {args.task_list}")
        requested = set(args.task_list)
        tasks = [t for t in all_tasks if t.__class__.__name__ in requested]
        found = {t.__class__.__name__ for t in tasks}
        missing = requested - found
        if missing:
            print(f"Warning: tasks not found (skipped): {missing}")
    else:
        print("Running evaluation on all tasks in the benchmark.")
        tasks = all_tasks

    tasks_names = " ".join([task.__class__.__name__ for task in tasks])
    print(f"Evaluating tasks: {tasks_names}")

    results = mteb.evaluate(
        model=model,
        tasks=tasks,
        encode_kwargs={"batch_size": args.batch_size},
        prediction_folder=args.predictions_folder,
        overwrite_strategy=args.overwrite_strategy,
    )

    print(results)

    out_dir = Path(args.results_folder)
    out_dir.mkdir(parents=True, exist_ok=True)
    model_slug = args.model_name.replace("/", "_")
    out_path = out_dir / f"{model_slug}-{tasks_names.replace(' ', '-')}.json"
    out_path.write_text(results.model_dump_json(indent=2), encoding="utf-8")
    print(f"Results saved to {out_path}")


if __name__ == "__main__":
    main()