New eval/plot pipeline

.gitignore

@@ -0,0 +1,183 @@
+.venv
+lm-outputs*
+lmo
+*.png
+.vscode/launch.json
+scratch
+
+# Created by https://www.toptal.com/developers/gitignore/api/python
+# Edit at https://www.toptal.com/developers/gitignore?templates=python
+
+### Python ###
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+### Python Patch ###
+# Poetry local configuration file - https://python-poetry.org/docs/configuration/#local-configuration
+poetry.toml
+
+# ruff
+.ruff_cache/
+
+# LSP config files
+pyrightconfig.json
+
+# End of https://www.toptal.com/developers/gitignore/api/python

.pre-commit-config.yaml

@@ -0,0 +1,13 @@
+repos:
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v2.3.0
+    hooks:
+      - id: check-yaml
+      - id: check-toml
+      - id: end-of-file-fixer
+      - id: trailing-whitespace
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.5.5
+    hooks:
+      - id: ruff
+      - id: ruff-format

pyproject.toml

@@ -0,0 +1,19 @@
+[build-system]
+requires = ["setuptools>=68", "setuptools_scm[toml]>=8"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "rbeval"
+requires-python = ">=3.8"
+dynamic = ["version"]
+dependencies = [
+    "seaborn>=0.13.2"
+]
+
+[project.optional-dependencies]
+eval = [
+    "lm-eval[vllm]==0.4.3"
+]
+
+# Enables the usage of setuptools_scm
+[tool.setuptools_scm]

src/rbeval/eval.py

@@ -0,0 +1,104 @@
+import subprocess
+import argparse
+from typing import Optional
+import torch
+import warnings
+import os
+from pathlib import Path
+
+from rbeval.eval_spec import EvalSpec, rand_uid
+
+
+def run_lm_eval(
+    lm_eval_path: Optional[str],
+    model_args: str,
+    tasks: str,
+    num_fewshot: int,
+    output_path: str,
+):
+    lm_eval_path = lm_eval_path or "lm_eval"
+    cmd = [lm_eval_path]
+    cmd += ["--model_args", model_args]
+    cmd += ["--tasks", tasks]
+    cmd += ["--num_fewshot", num_fewshot]
+    cmd += ["--output_path", output_path]
+    cmd += ["--log_samples"]
+    cmd += ["--cache_requests", "true"]
+    cmd += ["--cache_requests", "true"]
+    print(f"Running: {' '.join(cmd)}")
+    subprocess.check_call(cmd, env=os.environ)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Run eval for a given model")
+
+    parser.add_argument("output", type=str, help="output directory")
+    parser.add_argument("model", type=str, help="model path")
+    parser.add_argument("group", type=str, default=None)
+
+    parser.add_argument("--lm_eval_path", type=str)
+    parser.add_argument("--req_cache_path", type=str, default="/tmp/lm_eval_cache")
+    parser.add_argument("--tasks", type=str, default="mmlu")
+    parser.add_argument("--min_fewshot", type=int, default=0)
+    parser.add_argument("--max_fewshot", type=int, default=0)
+    parser.add_argument("-r", "--reformat", type=str)
+
+    args = parser.parse_args()
+    model: str = args.model
+    output_path: Path = Path(args.output)
+    group: Optional[str] = args.group
+    reformat: Optional[str] = args.reformat
+    lm_eval_path: Optional[str] = args.lm_eval_path
+    req_cache_path = Path(args.req_cache_path)
+    tasks: str = args.tasks
+    min_fewshot: int = args.min_fewshot
+    max_fewshot: int = args.max_fewshot
+    max_fewshot = max(min_fewshot, max_fewshot)
+
+    if not output_path.exists():
+        warnings.warn(f"Output path {output_path} does not exist, creating it")
+        output_path.mkdir(parents=True)
+    if not req_cache_path.exists():
+        warnings.warn(f"Request cache path, {str(req_cache_path)} does not exist")
+
+    os.environ["LM_HARNESS_CACHE_PATH"] = args.req_cache_path
+
+    n_gpu = torch.cuda.device_count()
+    os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(map(str, range(n_gpu)))
+    model_args = f"pretrained={model},dtype=auto,gpu_memory_utilization=0.7,tensor_parallel_size=1,data_parallel_size={n_gpu},max_model_len=4096"
+    fewshot = list(range(min_fewshot, max_fewshot + 1))
+
+    for num_fewshot in fewshot:
+        cfg = EvalSpec(
+            uid=rand_uid(),
+            model=model,
+            model_name=model.split("/")[-1],
+            group=group,
+            model_args=model_args,
+            fewshot=num_fewshot,
+            tasks=tasks,
+        )
+        spec_name = cfg.name()
+        lm_eval_output_path = output_path / spec_name
+        if reformat:
+            ref_path = Path(reformat)
+            assert ref_path.exists()
+            ref_files = list(ref_path.glob("**/*.json*"))
+            print(f"Found dir to reformat with {len(ref_files)} files")
+            assert len(ref_files) > 0
+            lm_eval_output_path.mkdir(parents=True, exist_ok=False)
+            for file in ref_files:
+                file.rename(lm_eval_output_path / file.name)
+        else:
+            run_lm_eval(
+                lm_eval_path, model_args, tasks, num_fewshot, str(lm_eval_output_path)
+            )
+
+        # Succeeded, write config
+        cfg_path = output_path / f"{spec_name}.json"
+        with open(cfg_path, "w") as f:
+            f.write(cfg.json())
+
+
+if __name__ == "__main__":
+    main()

src/rbeval/eval_spec.py

@@ -0,0 +1,29 @@
+from dataclasses import dataclass, asdict
+import json
+import uuid
+
+
+def rand_uid():
+    return uuid.uuid4().hex[:8]
+
+
+@dataclass(frozen=True)
+class EvalSpec:
+    uid: str
+    model: str
+    model_name: str
+    group: str
+    model_args: str
+    fewshot: int
+    tasks: str
+
+    def json(self) -> str:
+        return json.dumps(asdict(self))
+
+    def name(self) -> str:
+        return (
+            f"{self.group}_{self.model_name}_fs{self.fewshot}_{self.tasks}_{self.uid}"
+        )
+
+    def pretty_name(self) -> str:
+        return f"{self.model_name} fs{self.fewshot}"

src/rbeval/plot/__main__.py

@@ -0,0 +1,76 @@
+import argparse
+from pathlib import Path
+import json
+from typing import Dict, Optional
+import numpy as np
+import re
+from rbeval.eval_spec import EvalSpec
+from rbeval.plot.data import Eval, EvalGroup, ModelEval
+from rbeval.plot.score_cdf import score_cdf
+from tqdm import tqdm
+
+plot_fns = [score_cdf]
+
+
+def get_samples(inp: Path, name_filter: str) -> Dict[str, EvalGroup]:
+    groups: Dict[str, EvalGroup] = {}
+
+    for spec_file in (pbar := tqdm(list(inp.glob("*.json")), desc="Reading specs")):
+        pbar.set_description(f"Reading spec {spec_file.stem}")
+        with open(spec_file) as f:
+            spec = EvalSpec(**json.load(f))
+
+        if name_filter:
+            if re.match(name_filter, spec.model_name) is None:
+                print(f"Skipping spec {spec_file.stem}")
+                continue
+
+        group = groups.setdefault(spec.group, EvalGroup(group=spec.group))
+        model_eval = ModelEval(model_spec=spec)
+        group.model_evals.append(model_eval)
+        for samples_file in (spec_file.parent / spec_file.stem).glob(
+            "**/samples_*.json*"
+        ):
+            with open(samples_file, "r") as f:
+                if samples_file.suffix == ".jsonl":
+                    docs = [json.loads(s) for s in f.readlines()]
+                else:
+                    assert samples_file.suffix == ".json"
+                    docs = json.load(f)
+
+            cor_logprobs = []
+            inc_logprobs = []
+            for doc in docs:
+                target = doc["target"]
+                probs = [float(a[0][0]) for a in doc["resps"]]
+                cor_logprobs.append(probs.pop(target))
+                inc_logprobs.append(probs)
+            model_eval.evals.append(
+                Eval(
+                    name=samples_file.stem,
+                    cor_logprobs=np.array(cor_logprobs),
+                    inc_logprobs=np.array(inc_logprobs),
+                )
+            )
+
+    return groups
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Generate performance curves")
+    parser.add_argument("eval_dir", type=str)
+    parser.add_argument("figure_dir", type=str)
+    parser.add_argument("-n", "--name", type=str)
+    args, rest = parser.parse_known_args()
+
+    name_filter: Optional[str] = args.name
+    eval_dir = Path(args.eval_dir)
+    figure_dir = Path(args.figure_dir)
+    samples = get_samples(eval_dir, name_filter)
+
+    for fn in plot_fns:
+        fn(samples, figure_dir, rest)
+
+
+if __name__ == "__main__":
+    main()

src/rbeval/plot/data.py

@@ -0,0 +1,35 @@
+from dataclasses import dataclass, field
+from typing import List
+
+import numpy as np
+
+from rbeval.eval_spec import EvalSpec
+
+
+@dataclass
+class Eval:
+    name: str
+    cor_logprobs: np.ndarray
+    """shape [n] array of correct logprobs"""
+    inc_logprobs: np.ndarray
+    """shape [n, k] array of incorrect logprobs"""
+
+
+@dataclass
+class ModelEval:
+    """The evaluations for a given model"""
+
+    model_spec: EvalSpec
+    evals: List[Eval] = field(default_factory=list)
+
+    @property
+    def model_name(self) -> str:
+        return self.model_spec.model_name
+
+
+@dataclass
+class EvalGroup:
+    """A group of model evals"""
+
+    group: str
+    model_evals: List[ModelEval] = field(default_factory=list)

src/rbeval/plot/score_cdf.py

@@ -0,0 +1,91 @@
+from pathlib import Path
+from typing import Dict, List
+
+from matplotlib import pyplot as plt
+from rbeval.plot.data import EvalGroup, ModelEval
+from matplotlib import colormaps
+from matplotlib.axes import Axes
+import numpy as np
+
+
+def score_cdf(samples: Dict[str, EvalGroup], figure_dir: Path, args: List[str]):
+    fig, axs = plt.subplots(
+        1, len(samples), figsize=(5 * len(samples), 5), dpi=320, sharey=True
+    )
+
+    for ax, (group_name, group) in zip(axs, samples.items()):
+        group: EvalGroup
+
+        model_names = set(m.model_spec.model_name for m in group.model_evals)
+        max_fewshot = {}
+        for m in group.model_evals:
+            max_fewshot[m.model_name] = max(
+                max_fewshot.get(m.model_name, 0), m.model_spec.fewshot
+            )
+
+        scales = ["Purples", "Greens", "Oranges", "Reds"]
+        model_cmaps = {}
+        for i, (scale, n) in enumerate(zip(scales, model_names)):
+            mfs = max_fewshot[n]
+            if mfs > 0:
+                model_cmaps[n] = colormaps[scale](
+                    np.linspace(0.4, 1, max_fewshot[n] + 1)
+                )
+            else:
+                model_cmaps[n] = colormaps[scale]([1.0])
+
+        for model_eval in group.model_evals:
+            spec = model_eval.model_spec
+            color = model_cmaps[spec.model_name][spec.fewshot]
+            plot_samples(
+                ax, model_eval, model_eval.model_spec.pretty_name(), color=color
+            )
+
+        label_ax(ax, title=group_name)
+
+        handles, labels = ax.get_legend_handles_labels()
+        # sort both labels and handles by labels
+        labels, handles = zip(*sorted(zip(labels, handles), key=lambda t: t[0]))
+        ax.legend(handles, labels)
+
+    fig.savefig(figure_dir / "score_cdf.png")
+
+
+def label_ax(ax, title=True, y=True, x=True):
+    if x:
+        ax.set_xlabel("Model output probability")
+    if y:
+        ax.set_ylabel("Percent of samples with correct model output prob > p")
+    if title:
+        ax.set_title("Performance curve for mmlu")
+
+
+def get_base_logits(probs):
+    logits = np.zeros(len(probs))
+    logits[0] = 1
+    rest = 1 - np.sum(probs)
+    assert (rest >= 0) and (rest <= 1)
+
+
+def plot_samples(ax: Axes, meval: ModelEval, name: str, norm_by_stat=True, color=None):
+    bulk = np.concatenate([np.exp(e.cor_logprobs) for e in meval.evals])
+    num_cats = len(meval.evals)
+    weights = []
+    if norm_by_stat:
+        for e in meval.evals:
+            n = len(e.cor_logprobs)
+            # Each eval gets a total weight of 1/num_cats
+            # So each sample should have a weight of 1/num_cats/n
+            weights.append(np.ones(n) / (num_cats * n))
+        weights = np.concatenate(weights)
+    else:
+        weights = np.ones_like(bulk) / len(bulk)
+
+    sort_perm = bulk.argsort()
+    bulk = bulk[sort_perm]
+    weights = weights[sort_perm]
+    cdf_p = 1 - np.cumsum(weights)
+
+    ax.plot(bulk, cdf_p, label=name, color=color, alpha=0.8)
+    ax.set_xlim(0, 1)
+    ax.set_ylim(0, 1)