lm-eval-harness 整体架构 源码精读

01 模块总览

lm-evaluation-harness 的代码组织围绕评测器（evaluator）、任务（task）、模型（model）三大核心展开， 并通过一套注册表（registry）系统将它们松耦合地连接在一起。整个项目的模块关系可以用下图概括：

各模块的核心职责如下：

02 核心抽象类

2.1 LM -- 模型基类

LM（定义于 api/model.py）是所有模型后端必须继承的抽象基类。 它定义了三个抽象方法，分别对应三种不同的推理范式：

class LM(abc.ABC):
    def __init__(self) -> None:
        self._rank = 0
        self._world_size = 1
        self.cache_hook = CacheHook(None)

    @abc.abstractmethod
    def loglikelihood(self, requests) -> List[Tuple[float, bool]]:
        """计算 continuation 在 context 条件下的对数似然。
        返回 (logprob, is_greedy) 元组列表。"""
        pass

    @abc.abstractmethod
    def loglikelihood_rolling(self, requests) -> List[Tuple[float, bool]]:
        """计算完整字符串的对数似然（无截断），用于 perplexity 计算。"""
        pass

    @abc.abstractmethod
    def generate_until(self, requests) -> List[str]:
        """给定 context，贪心生成直到遇到 stopping sequence。"""
        pass

此外，LM 还提供了一个工厂方法 create_from_arg_string()， 用于从命令行参数字符串（如 "pretrained=meta-llama/Llama-2-7b,dtype=float16"）创建模型实例：

@classmethod
def create_from_arg_string(
    cls: Type[T], arg_string: str, additional_config: Optional[dict] = None
) -> T:
    additional_config = {} if additional_config is None else additional_config
    args = utils.simple_parse_args_string(arg_string)
    args2 = {k: v for k, v in additional_config.items() if v is not None}
    return cls(**args, **args2)

rank 和 world_size 属性用于分布式推理场景。默认值为单进程（rank=0、world_size=1）， 在使用多 GPU 时由子类覆盖。

2.2 Task -- 任务基类

Task（定义于 api/task.py）是所有评测任务的抽象基类。 它定义了从数据加载到结果计算的完整生命周期。核心抽象方法包括：

class Task(abc.ABC):
    VERSION = None
    DATASET_PATH: str = None
    DATASET_NAME: str = None
    OUTPUT_TYPE: str = None

    @abc.abstractmethod
    def doc_to_text(self, doc):
        """将一条数据样本转换为模型输入文本（prompt）"""
        pass

    @abc.abstractmethod
    def doc_to_target(self, doc):
        """将一条数据样本转换为目标答案"""
        pass

    @abc.abstractmethod
    def construct_requests(self, doc, ctx, **kwargs):
        """根据文档和上下文构建 Instance 请求"""
        pass

    @abc.abstractmethod
    def process_results(self, doc, results):
        """处理模型输出，返回指标 dict"""
        pass

    @abc.abstractmethod
    def aggregation(self):
        """返回 {metric_name: aggregation_fn} 映射"""
        pass

    @abc.abstractmethod
    def higher_is_better(self):
        """返回 {metric_name: bool} 映射"""
        pass

Task 基类还实现了两个关键的非抽象方法：

build_all_requests() -- 遍历数据集中的每条文档，调用 fewshot_context() 构建 few-shot 上下文， 再调用 construct_requests() 创建 Instance 对象：

def build_all_requests(self, limit=None, rank=None, world_size=None) -> None:
    """Build a set of Instances for a task, and store them in task.instances"""
    if self.has_test_docs():
        docs = self.test_docs()
    elif self.has_validation_docs():
        docs = self.validation_docs()

    instances = []
    for doc_id, doc in utils.create_iterator(
        enumerate(docs), rank, world_size, limit
    ):
        fewshot_ctx = self.fewshot_context(
            doc,
            0 if self.config.num_fewshot is None else self.config.num_fewshot,
        )
        # 添加 prefix / suffix / chatml_prefix / chatml_suffix
        fewshot_ctx["text"] = self.config.prefix + fewshot_ctx["text"] + self.config.suffix
        fewshot_ctx["text"] = self.config.chatml_prefix + fewshot_ctx["text"] + self.config.chatml_suffix

        inst = self.construct_requests(
            doc=doc, ctx=fewshot_ctx,
            metadata=(self.config["task"], doc_id, self.config.repeats),
        )
        instances.extend(inst if isinstance(inst, list) else [inst])

    self._instances = instances

fewshot_context() -- 从训练集中随机采样 num_fewshot 条样本， 拼接成 "description + few-shot examples + current question" 格式的 prompt：

def fewshot_context(self, doc, num_fewshot, rnd=random.Random(1234), description=None):
    if num_fewshot == 0:
        labeled_examples = {"text": ""}
    else:
        fewshotex = self.fewshot_examples(k=num_fewshot, rnd=rnd)
        labeled_examples = "\n\n".join(
            [self.doc_to_text(doc) + self.doc_to_target(doc) for doc in fewshotex]
        ) + "\n\n"
        labeled_examples = {"text": labeled_examples}

    example = self.doc_to_text(doc)
    labeled_examples["text"] = description + labeled_examples["text"] + example
    return labeled_examples

# 假设 GSM8K 数据集中有以下样本：

# 当前待评测的文档
doc = {"question": "Janet's ducks lay 16 eggs per day...", "answer": "#### 18"}

# fewshot_context(doc, num_fewshot=2) 的执行过程：

# 1. 从训练集随机采样 2 条示例
fewshot_examples = [
    {"question": "Natalia sold clips...", "answer": "#### 72"},
    {"question": "Weng earns $12/hour...", "answer": "#### 10"},
]

# 2. 对每条示例调用 doc_to_text() + doc_to_target()，再用 "\n\n" 拼接
labeled_examples = (
    "Natalia sold clips...\n#### 72\n\n"
    "Weng earns $12/hour...\n#### 10\n\n"
)

# 3. 拼接当前文档的 doc_to_text()
example = "Janet's ducks lay 16 eggs per day..."
final_prompt = labeled_examples + example

# 最终 prompt（发给模型的完整输入）：
# "Natalia sold clips...\n#### 72\n\n"
# "Weng earns $12/hour...\n#### 10\n\n"
# "Janet's ducks lay 16 eggs per day..."

apply_filters() -- 在模型推理完成后，对每个 Instance 的输出应用过滤器管道（默认为 take_first，即取第一个结果）。

2.3 ConfigurableTask -- YAML 驱动的任务子类

ConfigurableTask 继承自 Task，是绝大多数评测任务的实际类型。 它通过 TaskConfig dataclass 接收 YAML 配置，自动实现父类的所有抽象方法：

class ConfigurableTask(Task):
    VERSION = "Yaml"
    OUTPUT_TYPE = None
    CONFIG = None

    def __init__(self, data_dir=None, cache_dir=None, download_mode=None,
                 config: dict = None, model_name: Optional[str] = None) -> None:
        self._config = self.CONFIG  # 使用预配置（来自动态子类）
        if self.config is None:
            self._config = TaskConfig(**config)
        else:
            if config is not None:
                self._config.__dict__.update(config)  # 允许运行时覆盖

        if self.config.output_type is not None:
            assert self.config.output_type in ALL_OUTPUT_TYPES
            self.OUTPUT_TYPE = self.config.output_type
        ...

TaskConfig 是一个包含 30+ 字段的 dataclass，核心字段包括：

@dataclass
class TaskConfig(dict):
    task: str = None                          # 任务唯一标识
    group: Union[str, list] = None            # 所属分组
    dataset_path: str = None                  # HF 数据集路径
    dataset_name: str = None                  # 数据集子集名称
    dataset_kwargs: dict = None               # 数据集加载参数
    test_split: str = None                    # 测试集 split 名
    doc_to_text: Union[Callable, str] = None  # prompt 模板 / Jinja2 / 函数
    doc_to_target: Union[Callable, str] = None # 答案模板
    process_results: Union[Callable, str] = None # 自定义评分函数
    output_type: str = "generate_until"       # 输出类型
    generation_kwargs: dict = None            # 生成参数
    metric_list: list = None                  # 自定义指标列表
    num_fewshot: int = None                   # few-shot 数量
    filter_list: Union[str, list] = None      # 输出过滤器
    custom_infer: Callable = None             # 自定义推理函数
    repeats: int = 1                          # 每条样本推理次数

ConfigurableTask 的 construct_requests() 方法根据 OUTPUT_TYPE 分支构建不同类型的 Instance：

def construct_requests(self, doc: dict, ctx: str, **kwargs
) -> Union[List[Instance], Instance]:

    if self.OUTPUT_TYPE == "loglikelihood":
        arguments = (ctx, self.doc_to_target(doc), self.config.generation_kwargs)

    elif self.OUTPUT_TYPE == "loglikelihood_rolling":
        arguments = (self.doc_to_target(doc),)

    elif self.OUTPUT_TYPE == "multiple_choice":
        choices = self.doc_to_choice(doc)
        # 为每个选项创建一个 loglikelihood Instance
        arguments = [(ctx, f"{target_delimiter}{cont}",
                       self.config.generation_kwargs) for cont in choices]
        request_list = [
            Instance(request_type="loglikelihood", doc=doc,
                     arguments=arg, idx=i, **kwargs)
            for i, arg in enumerate(arguments)
        ]
        return request_list

    elif self.OUTPUT_TYPE == "generate_until":
        arguments = (ctx, self.config.generation_kwargs)

    return Instance(
        request_type=self.OUTPUT_TYPE, doc=doc,
        arguments=arguments, idx=0, **kwargs
    )

2.4 Instance -- 评测实例

Instance（定义于 api/instance.py）是一个轻量级数据类， 封装了单条评测请求的所有信息：

@dataclass
class Instance:
    request_type: Literal["loglikelihood", "loglikelihood_rolling", "generate_until"]
    doc: dict                                    # 原始文档
    arguments: tuple                             # 传给 LM 方法的参数
    idx: int                                     # 在同一 doc 中的序号
    metadata: Tuple[str, int, int] = ...         # (task_name, doc_id, repeats)
    resps: list = field(default_factory=list)     # 模型原始输出
    filtered_resps: dict = field(default_factory=dict)  # 过滤后的输出
    think_resps: list = field(default_factory=list)     # thinking tokens
    raw_resps: list = field(default_factory=list)       # 完整原始响应

    # 自动从 metadata 解包
    task_name: str = None
    doc_id: str = None
    repeats: str = None

    def __post_init__(self) -> None:
        self.task_name, self.doc_id, self.repeats = self.metadata

    @property
    def args(self):
        return (
            self.arguments if isinstance(self.arguments, tuple)
            else (self.arguments,)
        )

03 注册表机制

lm-eval-harness 采用全局注册表 + 装饰器的模式来管理模型、任务、指标等扩展点。 所有注册表定义于 api/registry.py。

3.1 MODEL_REGISTRY -- 模型注册表

MODEL_REGISTRY 是一个 dict，将字符串名称映射到 LM 子类。 通过 @register_model 装饰器注册：

MODEL_REGISTRY = {}

def register_model(*names):
    def decorate(cls):
        for name in names:
            assert issubclass(cls, LM), \
                f"Model '{name}' ({cls.__name__}) must extend LM class"
            assert name not in MODEL_REGISTRY, \
                f"Model named '{name}' conflicts with existing model!"
            MODEL_REGISTRY[name] = cls
        return cls
    return decorate

def get_model(model_name):
    try:
        return MODEL_REGISTRY[model_name]
    except KeyError:
        raise ValueError(
            f"Attempted to load model '{model_name}', but no model for this name found! "
            f"Supported model names: {', '.join(MODEL_REGISTRY.keys())}"
        )

使用示例：

@register_model("vllm_norm")
class VLLM(LM):
    def loglikelihood(self, requests): ...
    def loglikelihood_rolling(self, requests): ...
    def generate_until(self, requests): ...

3.2 TASK_REGISTRY 与 GROUP_REGISTRY -- 任务注册表

TASK_REGISTRY = {}
GROUP_REGISTRY = {}
ALL_TASKS = set()
func2task_index = {}

def register_task(name):
    def decorate(fn):
        assert name not in TASK_REGISTRY, \
            f"task named '{name}' conflicts with existing registered task!"
        TASK_REGISTRY[name] = fn
        ALL_TASKS.add(name)
        func2task_index[fn.__name__] = name
        return fn
    return decorate

def register_group(name):
    def decorate(fn):
        func_name = func2task_index[fn.__name__]
        if name in GROUP_REGISTRY:
            GROUP_REGISTRY[name].append(func_name)
        else:
            GROUP_REGISTRY[name] = [func_name]
            ALL_TASKS.add(name)
        return fn
    return decorate

具体示例

假设 tasks/ 目录下有以下 3 个 YAML 文件：

# ① gsm8k_0shot_cot.yaml — 纯任务，无分组
task: gsm8k_0shot_cot
dataset_path: parquet
output_type: generate_until

# ② charxiv_descriptive.yaml — 任务 + 隶属 "charxiv" 分组
task: "charxiv_descriptive"
group: "charxiv"
dataset_path: json
output_type: generate_until

# ③ charxiv_reasoning.yaml — 同样隶属 "charxiv" 分组
task: "charxiv_reasoning"
group: "charxiv"
dataset_path: json
output_type: generate_until

框架扫描后的注册流程（include_task_folder 两遍扫描）：

# —— 第一遍：register_task（task 字段为 str 的 YAML） ——

# ① gsm8k_0shot_cot.yaml
SubClass = type("gsm8k_0shot_cotConfigurableTask", (ConfigurableTask,), {"CONFIG": TaskConfig(...)})
register_task("gsm8k_0shot_cot")(SubClass)
# → TASK_REGISTRY["gsm8k_0shot_cot"] = SubClass
# → ALL_TASKS = {"gsm8k_0shot_cot"}

# ② charxiv_descriptive.yaml
SubClass2 = type("charxiv_descriptiveConfigurableTask", (ConfigurableTask,), {"CONFIG": TaskConfig(...)})
register_task("charxiv_descriptive")(SubClass2)
register_group("charxiv")(SubClass2)
# → TASK_REGISTRY["charxiv_descriptive"] = SubClass2
# → GROUP_REGISTRY["charxiv"] = ["charxiv_descriptive"]
# → ALL_TASKS = {"gsm8k_0shot_cot", "charxiv_descriptive", "charxiv"}

# ③ charxiv_reasoning.yaml（同理）
register_task("charxiv_reasoning")(SubClass3)
register_group("charxiv")(SubClass3)
# → GROUP_REGISTRY["charxiv"] = ["charxiv_descriptive", "charxiv_reasoning"]  ← append

最终注册表状态：

3.3 DEFAULT_METRIC_REGISTRY -- 默认指标映射

当 YAML 配置没有显式指定 metric_list 时，框架会根据 output_type 从此映射中取默认指标：

DEFAULT_METRIC_REGISTRY = {
    "loglikelihood":          ["perplexity", "acc"],
    "loglikelihood_rolling":  ["word_perplexity", "byte_perplexity", "bits_per_byte"],
    "multiple_choice":        ["acc", "acc_norm"],
    "generate_until":         ["exact_match", "fc_accuracy"],
}

3.4 METRIC_REGISTRY 与 AGGREGATION_REGISTRY

METRIC_REGISTRY 存放指标计算函数，AGGREGATION_REGISTRY 存放聚合函数（如 mean）。 通过 @register_metric 和 @register_aggregation 装饰器注册：

METRIC_REGISTRY = {}
AGGREGATION_REGISTRY = {}
METRIC_AGGREGATION_REGISTRY = {}
HIGHER_IS_BETTER_REGISTRY = {}

def register_metric(**args):
    def decorate(fn):
        assert "metric" in args
        name = args["metric"]
        METRIC_REGISTRY[name] = fn
        if "aggregation" in args:
            METRIC_AGGREGATION_REGISTRY[name] = AGGREGATION_REGISTRY[args["aggregation"]]
        if "higher_is_better" in args:
            HIGHER_IS_BETTER_REGISTRY[name] = args["higher_is_better"]
        return fn
    return decorate

def register_aggregation(name):
    def decorate(fn):
        AGGREGATION_REGISTRY[name] = fn
        return fn
    return decorate

3.5 任务注册流程

从文件系统中的 YAML 到可用的评测任务，完整的注册流程如下：

# tasks/__init__.py

def include_path(task_dir):
    """先注册所有单任务，再注册所有分组"""
    include_task_folder(task_dir)                      # 第一遍：注册单任务
    include_task_folder(task_dir, register_task=False)  # 第二遍：注册分组

def include_task_folder(task_dir, register_task=True):
    """遍历 task_dir 下所有 .yaml 文件"""
    for root, subdirs, file_list in os.walk(task_dir):
        for f in file_list:
            if f.endswith(".yaml"):
                config = utils.load_yaml_config(os.path.join(root, f))
                if "task" not in config:
                    continue
                if register_task:
                    if type(config["task"]) == str:
                        register_configurable_task(config)  # 注册单任务
                else:
                    if type(config["task"]) == list:
                        register_configurable_group(config) # 注册分组

def register_configurable_task(config):
    """动态创建 ConfigurableTask 子类并注册"""
    SubClass = type(
        config["task"] + "ConfigurableTask",
        (ConfigurableTask,),
        {"CONFIG": TaskConfig(**config)},
    )
    task_name = config["task"]
    register_task(task_name)(SubClass)

    if "group" in config:
        for group in (config["group"] if isinstance(config["group"], list)
                      else [config["group"]]):
            register_group(group)(SubClass)

为什么需要两遍扫描？

以 tasks/mvbench/ 目录为例，目录下有 两种 YAML：

tasks/mvbench/
├── mvbench.yaml                # ← 分组 YAML（task 是 list）
├── mvbench_action_count.yaml   # ← 单任务 YAML（task 是 str）
├── mvbench_moving_count.yaml
├── mvbench_action_sequence.yaml
└── ...（共 19 个单任务 YAML）

单任务 YAML（task 字段为 str）—— 定义一个具体的评测任务：

# mvbench_action_count.yaml
include: _default_template_yaml
task: mvbench_action_count          # ← str，表示这是一个单任务
dataset_path: json
test_split: train

分组 YAML（task 字段为 list）—— 仅声明一个分组名和它包含哪些任务：

# mvbench.yaml
group: mvbench                      # ← 分组名
task:                               # ← list，表示这是一个分组定义
  - mvbench_action_count
  - mvbench_moving_count
  - mvbench_action_sequence
  - ...（共 19 个子任务）

分组 YAML 依赖子任务已经在 TASK_REGISTRY 中注册好，否则 register_configurable_group() 查不到它们。因此 必须先跑一遍注册单任务，再跑一遍注册分组：

def include_path(task_dir):
    include_task_folder(task_dir)                      # 第一遍
    include_task_folder(task_dir, register_task=False)  # 第二遍

# ═══════════════════════════════════════════════════════════
# 第一遍 register_task=True：只处理 type(config["task"]) == str
# ═══════════════════════════════════════════════════════════
#   扫到 mvbench_action_count.yaml  → task 是 str → 注册
#   扫到 mvbench_moving_count.yaml  → task 是 str → 注册
#   ...
#   扫到 mvbench.yaml               → task 是 list → 跳过！
#
# 第一遍结束后：
#   TASK_REGISTRY = {
#       "mvbench_action_count": SubClass1,
#       "mvbench_moving_count": SubClass2,
#       ...（共 19 个）
#   }
#   GROUP_REGISTRY = {}  ← 还是空的
#
# ═══════════════════════════════════════════════════════════
# 第二遍 register_task=False：只处理 type(config["task"]) == list
# ═══════════════════════════════════════════════════════════
#   扫到 mvbench_action_count.yaml  → task 是 str → 跳过！
#   ...
#   扫到 mvbench.yaml               → task 是 list → 注册分组
#       → GROUP_REGISTRY["mvbench"] = [
#             "mvbench_action_count",
#             "mvbench_moving_count",
#             ...（共 19 个引用）
#         ]
#       → ALL_TASKS.add("mvbench")
#
# 最终效果：
#   --tasks mvbench          → 展开为 19 个子任务，逐个评测
#   --tasks mvbench_action_count  → 直接评测单个任务

04 评测主循环

评测的核心逻辑位于 evaluator.py 的 evaluate() 函数。 simple_evaluate() 是面向用户的高层接口，负责模型实例化和任务解析，最终委托 evaluate() 执行。

4.1 simple_evaluate() -- 高层入口

def simple_evaluate(
    model, model_args=None, tasks=[], num_fewshot=None,
    batch_size=None, device=None, limit=None, gen_kwargs=None, ...
):
    # 1. 实例化模型
    if isinstance(model, str):
        lm = lm_eval.api.registry.get_model(model).create_from_arg_string(
            model_args,
            {"batch_size": batch_size, "max_batch_size": max_batch_size, "device": device},
        )
    else:
        lm = model  # 直接传入 LM 对象

    # 2. 可选：启用缓存
    if use_cache is not None:
        lm = lm_eval.api.model.CachingLM(lm, use_cache + "_rank" + str(lm.rank) + ".db")

    # 3. 解析任务
    task_dict = lm_eval.tasks.get_task_dict(tasks)

    # 4. 覆盖 num_fewshot / gen_kwargs
    for task_name, task_obj in task_dict.items():
        config = task_obj._config
        if config["output_type"] == "generate_until" and gen_kwargs is not None:
            config["generation_kwargs"].update(gen_kwargs)
        if num_fewshot is not None:
            task_obj._config["num_fewshot"] = num_fewshot

    # 5. 调用 evaluate()
    results = evaluate(lm=lm, task_dict=task_dict, limit=limit, ...)
    return results

4.2 evaluate() -- 核心循环

evaluate() 的执行流程可分为五个阶段：

阶段一：构建请求

requests = collections.defaultdict(list)

for task_name, task in task_dict.items():
    if type(task) == tuple:
        group_name, task = task

    # 构建所有 Instance（内部调用 fewshot_context + construct_requests）
    task.build_all_requests(limit=limit, rank=lm.rank, world_size=lm.world_size)

    # 按 request_type 分组
    for instance in task.instances:
        reqtype = instance.request_type
        requests[reqtype].append(instance)

阶段二：模型推理

for reqtype, reqs in requests.items():
    # 根据 repeats 展开请求
    cloned_reqs = [req for req in reqs for _ in range(req.repeats)]

    # 调用模型对应的方法
    resps = getattr(lm, reqtype)(cloned_reqs)

    # 将响应写回 Instance
    for x, req in zip(resps, cloned_reqs):
        if isinstance(x, str):
            # 处理 thinking tokens：过滤 <think>...</think>
            filtered_think_resp = x.split(THINKING_END_TOKEN)[1] \
                if THINKING_END_TOKEN in x else x
            req.resps.append(filtered_think_resp)
            req.think_resps.append(think_resp)
            req.raw_resps.append(x)
        else:
            req.resps.append(x)

# 模型原始输出（raw_resps）
raw_output = "<think>\n让我分析一下这道题...\n先计算16-3=13\n再计算13-5=8\n最后8+10=18\n</think>\n#### 18"

# 框架检测到 </think> 标记，自动拆分为三个字段：

# 1. raw_resps —— 完整的原始输出，原封不动保存
instance.raw_resps = [
    "<think>\n让我分析一下这道题...\n先计算16-3=13\n再计算13-5=8\n最后8+10=18\n</think>\n#### 18"
]

# 2. think_resps —— 仅 thinking 部分（<think> 和 </think> 之间的内容）
instance.think_resps = [
    "让我分析一下这道题...\n先计算16-3=13\n再计算13-5=8\n最后8+10=18"
]

# 3. resps —— 去掉 thinking 后的最终回答（用于后续评分）
instance.resps = ["#### 18"]

阶段三：应用过滤器

for task_name, task in task_dict.items():
    if type(task) == tuple:
        group, task = task
    task.apply_filters()  # 应用过滤管道，结果写入 instance.filtered_resps

阶段四：计算指标

vals = collections.defaultdict(list)

for task_name, task in task_dict.items():
    for key in task.instances[0].filtered_resps.keys():
        doc_iterator = itertools.islice(
            enumerate(task.test_docs()), lm.rank, limit, lm.world_size
        )
        for doc_id, doc in doc_iterator:
            requests = [x for x in task.instances if x.doc_id == doc_id]
            requests.sort(key=lambda x: x.idx)

            # 调用 task.process_results() 计算每条样本的指标
            metrics = task.process_results(
                doc, [req.filtered_resps[key] for req in requests]
            )
            for metric, value in metrics.items():
                vals[(task_name, key, metric)].append(value)

# 假设任务 gsm8k 有 3 条样本，process_results 逐条计算：

# 样本 1：模型输出 "42"，正确答案 "42"
process_results(doc1, ["42"]) → {"exact_match": 1.0}

# 样本 2：模型输出 "15"，正确答案 "18"
process_results(doc2, ["15"]) → {"exact_match": 0.0}

# 样本 3：模型输出 "1024"，正确答案 "1024"
process_results(doc3, ["1024"]) → {"exact_match": 1.0}

# vals 字典累积结果：
vals = {
    ("gsm8k", "get-answer", "exact_match"): [1.0, 0.0, 1.0]
}

# 阶段五聚合：mean([1.0, 0.0, 1.0]) = 0.667

阶段五：聚合与返回

results = collections.defaultdict(dict)

for (task_name, key, metric), items in vals.items():
    task = task_dict[task_name]
    # 调用聚合函数（如 mean）
    agg_fn = task.aggregation()[metric]
    results[task_name][metric + "," + key] = agg_fn(items)

    # 计算 bootstrap 标准误
    stderr = lm_eval.api.metrics.stderr_for_metric(
        metric=agg_fn, bootstrap_iters=bootstrap_iters,
    )
    if stderr is not None and len(items) > 1:
        results[task_name][metric + "_stderr," + key] = stderr(items)

return {
    "results": dict(results_agg.items()),
    "groups": dict(groups_agg.items()),
    "configs": dict(sorted(configs.items())),
    "versions": dict(sorted(versions.items())),
    "n-shot": dict(sorted(num_fewshot.items())),
    "samples": dict(samples),
}

05 数据流

下图展示了从 YAML 配置到最终评测分数的完整数据流：

更细致地看，一条文档的完整生命周期如下：

分布式场景下的数据分片：build_all_requests() 通过 utils.create_iterator(enumerate(docs), rank, world_size, limit) 将数据集按 rank 进行切片，每个 rank 只处理自己负责的文档子集。 推理完成后，evaluate() 使用 torch.distributed.all_gather_object() 收集所有 rank 的结果进行聚合。

06 扩展点

6.1 添加新模型

要接入一个新的推理后端，只需三步：

1. 创建一个继承 LM 的子类
2. 实现三个抽象方法：loglikelihood()、loglikelihood_rolling()、generate_until()
3. 使用 @register_model("your_model_name") 装饰器注册

from lm_eval.api.model import LM
from lm_eval.api.registry import register_model

@register_model("my_custom_model")
class MyCustomModel(LM):
    def __init__(self, pretrained, batch_size=1, **kwargs):
        super().__init__()
        self.model = load_my_model(pretrained)
        self.batch_size = int(batch_size)

    def loglikelihood(self, requests):
        # 实现条件对数似然计算
        results = []
        for req in requests:
            context, continuation = req.args[0], req.args[1]
            logprob, is_greedy = self._compute_ll(context, continuation)
            results.append((logprob, is_greedy))
        return results

    def loglikelihood_rolling(self, requests):
        # 实现无条件对数似然计算
        ...

    def generate_until(self, requests):
        # 实现文本生成
        results = []
        for req in requests:
            context = req.args[0]
            gen_kwargs = req.args[1]
            output = self.model.generate(context, **gen_kwargs)
            results.append(output)
        return results

注册后即可通过命令行使用：

lm_eval --model my_custom_model --model_args pretrained=/path/to/model --tasks mmlu

6.2 添加新任务

最常用的方式是编写一个 YAML 配置文件，并可选地提供一个 utils.py 辅助脚本。 将它们放入 lm_eval/tasks/ 目录下的适当子文件夹中即可。

# tasks/my_benchmark/my_task.yaml
task: my_task
group: my_benchmark
dataset_path: /data/my_dataset
dataset_kwargs:
  data_dir: /data/my_dataset
test_split: test
output_type: generate_until
doc_to_text: !function utils.doc_to_text
doc_to_target: !function utils.doc_to_target
process_results: !function utils.process_results
generation_kwargs:
  max_gen_toks: 2048
  temperature: 0.0
metric_list:
  - metric: !function utils.my_custom_metric
    aggregation: mean
    higher_is_better: true

# tasks/my_benchmark/utils.py
def doc_to_text(doc):
    return f"Question: {doc['question']}\nAnswer:"

def doc_to_target(doc):
    return doc["answer"]

def process_results(doc, results):
    pred = results[0].strip()
    gold = doc["answer"]
    return {"accuracy": 1.0 if pred == gold else 0.0}

def my_custom_metric(items):
    return sum(items) / len(items)

6.3 添加新指标

自定义指标可以通过 @register_metric 和 @register_aggregation 装饰器注册到全局注册表中：

from lm_eval.api.registry import register_metric, register_aggregation

@register_aggregation("geometric_mean")
def geometric_mean(items):
    import math
    return math.exp(sum(math.log(x) for x in items if x > 0) / len(items))

@register_metric(metric="my_f1", aggregation="mean", higher_is_better=True)
def my_f1(references, predictions, **kwargs):
    # 计算 F1 分数
    ...

注册后，在 YAML 的 metric_list 中引用即可：

metric_list:
  - metric: my_f1
    aggregation: mean
    higher_is_better: true

6.4 添加新 output_type

如果需要支持全新的推理模式（例如多轮对话），需要修改以下几处：

1. 在 api/task.py 的 ALL_OUTPUT_TYPES 列表中添加新类型：

   ALL_OUTPUT_TYPES = [
       "loglikelihood",
       "multiple_choice",
       "loglikelihood_rolling",
       "generate_until",
       "generate_until_multi_round",  # 新增
   ]

2. 在 ConfigurableTask.construct_requests() 中添加对应的分支逻辑
3. 在 LM 基类中添加对应的抽象方法（或在子类中实现同名方法）
4. 在 DEFAULT_METRIC_REGISTRY 中添加该 output_type 的默认指标