[BUG]`LocalChatCompletion` batch_size always 1

[jikunshang]

https://github.com/EleutherAI/lm-evaluation-harness/blob/main/lm_eval/models/openai_completions.py#L169-L173

seems openai server batch API support chat/completion endpoint https://developers.openai.com/api/docs/guides/batch/

[inf3rnus]

+1 on this, would be awesome to have

[inf3rnus]

Also, this is a feature request, not a bug

[inf3rnus]

@jikunshang This should help in the meantime, runs a worker pool and recomputes the performance for each shard of a given test task

import os
from copy import deepcopy
from concurrent.futures import ProcessPoolExecutor, as_completed

from lm_eval import simple_evaluate

os.environ.setdefault("OPENAI_API_KEY", "REPLACE_ME")


def chunk_indices(n: int, k: int) -> list[list[int]]:
    base, rem = divmod(n, k)
    chunks = []
    start = 0
    for i in range(k):
        size = base + (1 if i < rem else 0)
        end = start + size
        if start < end:
            chunks.append(list(range(start, end)))
        start = end
    return chunks


def to_picklable(obj):
    """Recursively convert result objects into plain picklable data."""
    if obj is None or isinstance(obj, (str, int, float, bool)):
        return obj

    if isinstance(obj, dict):
        out = {}
        for k, v in obj.items():
            # keep keys picklable too
            if callable(v):
                continue
            out[to_picklable(k)] = to_picklable(v)
        return out

    if isinstance(obj, (list, tuple, set)):
        return [to_picklable(x) for x in obj if not callable(x)]

    # Drop callables explicitly
    if callable(obj):
        return None

    # Try common object patterns
    if hasattr(obj, "__dict__"):
        return to_picklable(vars(obj))

    # Last resort: stringify unknown objects
    return repr(obj)


def run_worker(task: str, sample_ids: list[int]) -> dict:
    result = simple_evaluate(
        model="openai-chat-completions",
        model_args={
            "base_url": "http://localhost:3000/api/agi",
            "model": "agi",
        },
        tasks=[task],
        samples={task: sample_ids},
        apply_chat_template=True,
        write_out=False,
        num_fewshot=0,
        batch_size=1,
        cache_requests=False,
        log_samples=True,
    )

    return to_picklable(result)


def merge_results(results: list[dict], task_name: str) -> dict:
    if not results:
        raise ValueError("results must not be empty")

    merged = deepcopy(results[0])

    merged.setdefault("samples", {})
    merged["samples"][task_name] = []

    for r in results:
        merged["samples"][task_name].extend(r.get("samples", {}).get(task_name, []))

    merged["samples"][task_name].sort(key=lambda x: x.get("doc_id", -1))

    merged.setdefault("results", {})
    merged["results"].setdefault(task_name, {})

    # Find the aggregate metric key from the original result, e.g. "exact_match,custom-extract"
    metric_key = next(
        (
            k
            for k in merged["results"][task_name]
            if k not in {"alias"}
            and not k.endswith("_stderr,custom-extract")
            and "stderr" not in k
        ),
        None,
    )

    if metric_key:
        sample_metric_key = metric_key.split(",", 1)[0]  # "exact_match"
        metric_values = [
            s[sample_metric_key]
            for s in merged["samples"][task_name]
            if sample_metric_key in s
        ]
        if metric_values:
            merged["results"][task_name][metric_key] = sum(metric_values) / len(
                metric_values
            )

    merged.setdefault("n-samples", {})
    merged["n-samples"].setdefault(task_name, {})
    merged["n-samples"][task_name]["effective"] = len(merged["samples"][task_name])

    return merged


if __name__ == "__main__":
    TASK = "mmlu_pro_biology"
    TOTAL_SAMPLES = 8
    NUM_WORKERS = 8

    shards = chunk_indices(TOTAL_SAMPLES, NUM_WORKERS)
    worker_outputs = []

    with ProcessPoolExecutor(max_workers=NUM_WORKERS) as executor:
        futures = [executor.submit(run_worker, TASK, shard) for shard in shards]

        for future in as_completed(futures):
            worker_outputs.append(future.result())

    final = merge_results(worker_outputs, TASK)

    print("Combined sample count:", len(final["samples"][TASK]))

    first = final["samples"][TASK][0]
    print("First doc_id:", first.get("doc_id"))
    print("Target:", first.get("target"))
    print("Response:", first.get("resps", [None])[0])

    print("\nCombined metrics:")
    print(final.get("results", {}).get(TASK, {}))