|
import evaluate |
|
import lm_eval |
|
from typing import Union, List, Optional |
|
from dmx.compressor.dmx import config_rules, DmxModel |
|
import datasets |
|
import torch |
|
|
|
_DESCRIPTION = """ |
|
Evaluation function using lm-eval with d-Matrix integration. |
|
This function allows for the evaluation of language models across various tasks, |
|
with the option to use d-Matrix compressed models. |
|
""" |
|
|
|
_KWARGS_DESCRIPTION = """ |
|
Args: |
|
model (str): The name or path of the model to evaluate. |
|
tasks (Union[str, List[str]]): The task or list of tasks to evaluate on. |
|
dmx_config (Optional[str]): Configuration string for d-Matrix transformations, defaults to None. |
|
num_fewshot (Optional[int]): Number of examples in few-shot context, defaults to None. |
|
batch_size (Optional[Union[int, str]]): Batch size for model, defaults to None. |
|
max_batch_size (Optional[int]): Maximum batch size to try with automatic batch size detection, defaults to None. |
|
limit (Optional[Union[int, float]]): Limit the number of examples per task, defaults to None. |
|
device (Optional[str]): Device to run on. If None, defaults to 'cuda' if available, otherwise 'cpu'. |
|
revision (str): Model revision to use, defaults to 'main'. |
|
trust_remote_code (bool): Whether to trust remote code, defaults to False. |
|
log_samples (bool): If True, logs all model outputs and documents, defaults to True. |
|
verbosity (str): Logging verbosity level, defaults to 'INFO'. |
|
**kwargs: Additional keyword arguments to pass to `lm_eval.evaluate`. |
|
|
|
Returns: |
|
dict: A dictionary containing the evaluation results. |
|
""" |
|
|
|
@evaluate.utils.file_utils.add_start_docstrings(_DESCRIPTION, _KWARGS_DESCRIPTION) |
|
class DmxMetric(evaluate.Metric): |
|
def _info(self): |
|
return evaluate.MetricInfo( |
|
module_type="metric", |
|
description=_DESCRIPTION, |
|
citation="", |
|
inputs_description=_KWARGS_DESCRIPTION, |
|
features=datasets.Features( |
|
{ |
|
"references": datasets.Value("string"), |
|
} |
|
), |
|
reference_urls=["https://github.com/EleutherAI/lm-evaluation-harness"], |
|
) |
|
|
|
def _compute( |
|
self, |
|
model: str, |
|
tasks: Union[str, List[str]], |
|
dmx_config: Optional[str] = None, |
|
num_fewshot: Optional[int] = None, |
|
batch_size: Optional[Union[int, str]] = None, |
|
max_batch_size: Optional[int] = None, |
|
limit: Optional[Union[int, float]] = None, |
|
device: Optional[str] = None, |
|
revision: str = "main", |
|
trust_remote_code: bool = False, |
|
log_samples: bool = True, |
|
verbosity: str = "INFO", |
|
**kwargs |
|
): |
|
""" |
|
Evaluate a model on multiple tasks and metrics using lm-eval with optional d-Matrix integration. |
|
""" |
|
if device is None: |
|
device = "cuda" if torch.cuda.is_available() else "cpu" |
|
|
|
model_args = f"pretrained={model},revision={revision},trust_remote_code={str(trust_remote_code)},device={device}" |
|
|
|
lm = lm_eval.api.registry.get_model("hf").create_from_arg_string( |
|
model_args, |
|
{ |
|
"batch_size": batch_size, |
|
"max_batch_size": max_batch_size, |
|
} |
|
) |
|
|
|
if dmx_config: |
|
lm._model = DmxModel.from_torch(lm._model) |
|
lm._model.transform(lm._model.dmx_config, *eval(f"config_rules.{dmx_config}")) |
|
|
|
task_dict = lm_eval.tasks.get_task_dict(tasks if isinstance(tasks, list) else [tasks]) |
|
|
|
for task in task_dict.values(): |
|
if num_fewshot is not None: |
|
task.set_config(key="num_fewshot", value=num_fewshot) |
|
|
|
eval_params = { |
|
'lm': lm, |
|
'task_dict': task_dict, |
|
'limit': limit, |
|
'log_samples': log_samples, |
|
'verbosity': verbosity, |
|
**kwargs |
|
} |
|
|
|
results = lm_eval.evaluate(**eval_params) |
|
return results.get('results', {}) |
