"""Autoregressive baselines on the same lm-eval tasks, zero-shot, one task at a time."""
import os, json, lm_eval
tasks = os.environ.get("EV_TASKS", "hellaswag,arc_easy,arc_challenge,piqa,winogrande,openbookqa,sciq,lambada_openai,mmlu").split(",")
name = os.environ["EV_MODEL"]
out = os.environ.get("EV_OUT", "/tmp/bench_ar.json")
res = {}
for t in tasks:
    try:
        r = lm_eval.simple_evaluate(model="hf", model_args="pretrained=%s,dtype=bfloat16" % name, tasks=[t],
                                    num_fewshot=0, batch_size="auto", bootstrap_iters=0)
        res.update(r["results"])
        print("RESULT", t, json.dumps(r["results"].get(t, {})), flush=True)
    except Exception as e:
        print("TASKFAIL", t, str(e)[:200], flush=True)
    json.dump(res, open(out, "w"))
print("BENCH DONE", flush=True)