mirror of
https://github.com/EleutherAI/lm-evaluation-harness.git
synced 2026-10-11 21:59:53 +00:00
Tai An
d1193e39a2
The `__main__` block passed `model=` and `device=` to `run_model_for_task_caching()`, whose signature is `(tasks, cache_requests)` -- both are read from module constants inside the function, so the call raised `TypeError: run_model_for_task_caching() got an unexpected keyword argument 'model'` before anything ran. With that fixed the script then failed with `NameError: name 'request_caching_arg_to_dict' is not defined`, because that helper was defined *below* the `__main__` block and so was not yet bound when the block executed. Drop the two extra keyword arguments and move the helper above the `__main__` block, next to its only caller.
104 lines
2.5 KiB
Python
104 lines
2.5 KiB
Python
"""
|
|
Usage:
|
|
python requests_caching.py --tasks=comma,separated,list,of,tasks --cache_requests=<true|refresh|delete]>
|
|
"""
|
|
|
|
import argparse
|
|
import logging
|
|
import os
|
|
|
|
import torch
|
|
from transformers import (
|
|
pipeline as trans_pipeline,
|
|
)
|
|
|
|
from lm_eval import simple_evaluate
|
|
|
|
|
|
eval_logger = logging.getLogger(__name__)
|
|
|
|
|
|
MODULE_DIR = os.path.dirname(os.path.realpath(__file__))
|
|
|
|
# Used to specify alternate cache path, useful if run in a docker container
|
|
# NOTE raw datasets will break if you try to transfer the cache from your host to a docker image
|
|
LM_HARNESS_CACHE_PATH = os.getenv("LM_HARNESS_CACHE_PATH")
|
|
|
|
|
|
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
|
|
|
|
MODEL = "EleutherAI/pythia-70m"
|
|
|
|
TASK = "text-generation"
|
|
|
|
|
|
def run_model_for_task_caching(tasks: list[str], cache_requests: str):
|
|
eval_logger.info(f"Loading HF model: {MODEL}")
|
|
|
|
trans_pipe = trans_pipeline(
|
|
task=TASK, model=MODEL, device=DEVICE, trust_remote_code=True
|
|
)
|
|
|
|
model = trans_pipe.model
|
|
tokenizer = trans_pipe.tokenizer
|
|
|
|
eval_logger.info(
|
|
f"Running simple_evaluate to cache request objects for tasks: {tasks}"
|
|
)
|
|
|
|
cache_args = request_caching_arg_to_dict(cache_requests=cache_requests)
|
|
|
|
eval_logger.info(
|
|
f"The following operations will be performed on the cache: {cache_requests}"
|
|
)
|
|
|
|
eval_data = simple_evaluate(
|
|
model="hf-auto",
|
|
model_args={
|
|
"pretrained": model,
|
|
"tokenizer": tokenizer,
|
|
},
|
|
limit=1,
|
|
device=DEVICE,
|
|
tasks=tasks,
|
|
write_out=True,
|
|
**cache_args,
|
|
)
|
|
|
|
return eval_data
|
|
|
|
|
|
def request_caching_arg_to_dict(cache_requests: str) -> dict:
|
|
request_caching_args = {
|
|
"cache_requests": cache_requests in {"true", "refresh"},
|
|
"rewrite_requests_cache": cache_requests == "refresh",
|
|
"delete_requests_cache": cache_requests == "delete",
|
|
}
|
|
|
|
return request_caching_args
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument(
|
|
"--tasks",
|
|
"-t",
|
|
default=None,
|
|
metavar="task1,task2",
|
|
)
|
|
parser.add_argument(
|
|
"--cache_requests",
|
|
type=str,
|
|
default=None,
|
|
choices=["true", "refresh", "delete"],
|
|
help="Speed up evaluation by caching the building of dataset requests. `None` if not caching.",
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
tasks = args.tasks.split(",")
|
|
|
|
eval_data = run_model_for_task_caching(
|
|
tasks=tasks, cache_requests=args.cache_requests
|
|
)
|