wehub-resource-sync
162 lines
4.7 KiB
Python
162 lines
4.7 KiB
Python
import argparse
|
|
import contextlib
|
|
import json
|
|
import os
|
|
from tempfile import NamedTemporaryFile
|
|
|
|
from gemhall.adapters.mmlu import SUBJECTS_ALL
|
|
from gemhall.adapters.mmlu import export_temp_csv as mmlu_export
|
|
from gemhall.eval import evaluate
|
|
|
|
|
|
def add_common_args(ap: argparse.ArgumentParser) -> None:
|
|
ap.add_argument(
|
|
"--thresholds",
|
|
nargs="+",
|
|
type=float,
|
|
required=True,
|
|
help="Confidence thresholds, e.g., 0.5 0.75 0.9",
|
|
)
|
|
ap.add_argument("--out", default="outputs", help="Output directory")
|
|
ap.add_argument(
|
|
"--model",
|
|
default="gemini-2.5-flash",
|
|
choices=["gemini-2.0-flash", "gemini-2.5-flash", "gemini-2.5-flash-lite"],
|
|
help="Gemini model id",
|
|
)
|
|
ap.add_argument(
|
|
"--temperature", type=float, default=0.0, help="Sampling temperature"
|
|
)
|
|
ap.add_argument(
|
|
"--thinking-budget",
|
|
type=int,
|
|
default=0,
|
|
dest="thinking_budget",
|
|
help="Thinking budget (0 disables) - not supported at the moment",
|
|
)
|
|
ap.add_argument("--seed", type=int, default=1234, help="Sampling seed")
|
|
ap.add_argument(
|
|
"--judge",
|
|
choices=["exact", "llm"],
|
|
default="exact",
|
|
help="Validity judge to use",
|
|
)
|
|
ap.add_argument(
|
|
"--async",
|
|
action="store_true",
|
|
dest="use_async",
|
|
help="Use async client for parallelism",
|
|
)
|
|
ap.add_argument(
|
|
"--concurrency",
|
|
type=int,
|
|
default=8,
|
|
help="Max concurrent requests in async mode",
|
|
)
|
|
ap.add_argument("--progress", action="store_true", help="Show a progress bar")
|
|
ap.add_argument(
|
|
"--rpm-limit",
|
|
type=int,
|
|
default=None,
|
|
help="Optional requests-per-minute cap (client-side)",
|
|
)
|
|
ap.add_argument(
|
|
"--max-retries",
|
|
type=int,
|
|
default=6,
|
|
help="Max retries on 429/RESOURCE_EXHAUSTED",
|
|
)
|
|
|
|
|
|
def main() -> None:
|
|
p = argparse.ArgumentParser(
|
|
prog="gemhall", description="Gemini confidence-targeted hallucination evaluator"
|
|
)
|
|
sub = p.add_subparsers(dest="cmd", required=True)
|
|
|
|
# CSV mode
|
|
run = sub.add_parser("run", help="Run evaluation from CSV")
|
|
run.add_argument(
|
|
"--data", required=True, help="Path to CSV with id,question,gold,unknown_ok"
|
|
)
|
|
add_common_args(run)
|
|
|
|
# MMLU mode (direct from HF)
|
|
mmlu = sub.add_parser(
|
|
"mmlu", help="Run evaluation on MMLU (direct from Hugging Face Datasets)"
|
|
)
|
|
mmlu.add_argument("--split", default="test", help="Dataset split (e.g., test, dev)")
|
|
mmlu.add_argument(
|
|
"--subjects", nargs="+", default=["all"], help="Subjects or 'all'"
|
|
)
|
|
mmlu.add_argument(
|
|
"--limit",
|
|
type=int,
|
|
default=None,
|
|
help="Randomly sample N items after filtering subjects",
|
|
)
|
|
mmlu.add_argument(
|
|
"--idk-frac",
|
|
type=float,
|
|
default=0.0,
|
|
help="Fraction [0..1] of sampled items to convert to IDK-only",
|
|
)
|
|
add_common_args(mmlu)
|
|
|
|
args = p.parse_args()
|
|
|
|
if args.cmd == "mmlu":
|
|
with NamedTemporaryFile("w+", suffix=".csv", delete=False) as tf:
|
|
tmp_csv = tf.name
|
|
subjects = SUBJECTS_ALL if args.subjects == ["all"] else args.subjects
|
|
csv_path = mmlu_export(
|
|
out_csv=tmp_csv,
|
|
split=args.split,
|
|
subjects=subjects,
|
|
limit=args.limit,
|
|
seed=args.seed,
|
|
idk_frac=args.idk_frac,
|
|
)
|
|
res = evaluate(
|
|
data_csv=csv_path,
|
|
thresholds=args.thresholds,
|
|
out_dir=args.out,
|
|
model=args.model,
|
|
temperature=args.temperature,
|
|
thinking_budget=0,
|
|
seed=args.seed,
|
|
judge=args.judge,
|
|
use_async=args.use_async,
|
|
concurrency=args.concurrency,
|
|
show_progress=args.progress,
|
|
rpm_limit=args.rpm_limit,
|
|
max_retries=args.max_retries,
|
|
)
|
|
print(json.dumps(res, indent=2))
|
|
with contextlib.suppress(Exception):
|
|
os.unlink(csv_path)
|
|
return
|
|
|
|
if args.cmd == "run":
|
|
res = evaluate(
|
|
data_csv=args.data,
|
|
thresholds=args.thresholds,
|
|
out_dir=args.out,
|
|
model=args.model,
|
|
temperature=args.temperature,
|
|
thinking_budget=0,
|
|
seed=args.seed,
|
|
judge=args.judge,
|
|
use_async=args.use_async,
|
|
concurrency=args.concurrency,
|
|
show_progress=args.progress,
|
|
rpm_limit=args.rpm_limit,
|
|
max_retries=args.max_retries,
|
|
)
|
|
print(json.dumps(res, indent=2))
|
|
return
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|