Ready for spaces

README.md

@@ -0,0 +1,10 @@
+---
+title: rbeval
+emoji: 💩
+colorFrom: yellow
+colorTo: orange
+sdk: streamlit
+sdk_version: 1.25.0
+app_file: app.py
+pinned: false
+---

app.py

@@ -0,0 +1,3 @@
+import rbeval.dash
+
+rbeval.dash.main()

pyproject.toml

@@ -7,6 +7,12 @@ name = "rbeval"
 requires-python = ">=3.8"
 dynamic = ["version"]
 dependencies = [
+    "pandas>=2.2.2",
+    "matplotlib>=3.9.1",
+    "huggingface-hub>=0.24.2",
+    "tqdm>=4.66.4",
+    "numpy>=1.26.4",
+    "dacite>=1.8.1",
     "seaborn>=0.13.2"
 ]
 

requirements.txt

@@ -0,0 +1,8 @@
+pandas>=2.2.2
+matplotlib>=3.9.1
+huggingface-hub>=0.24.2
+tqdm>=4.66.4
+numpy>=1.26.4
+dacite>=1.8.1
+seaborn>=0.13.1
+.

src/rbeval/dash.py

@@ -0,0 +1,81 @@
+from pathlib import Path
+from typing import List, Optional
+import streamlit as st
+import argparse
+
+from rbeval.plot.data import EvalGroup, get_samples
+from rbeval.plot.score_cdf_altair import (
+    plot_with_data,
+    get_plot_data,
+    plot_cfgs,
+)
+from rbeval.plot import model_comp
+from huggingface_hub import snapshot_download
+
+
+@st.cache_resource
+def cached_samples(dir: Path, name_filter: Optional[str]) -> List[EvalGroup]:
+    if not dir.exists():
+        dir = Path(snapshot_download("mli-will/rbeval"))
+    samples = get_samples(dir, name_filter)
+    return samples
+
+
+@st.cache_data
+def cached_score_cdf(dir, name_filter):
+    samples = cached_samples(dir, name_filter)
+    cfgs = plot_cfgs()
+    data = [get_plot_data(cfg, samples) for cfg in cfgs]
+    return data, cfgs
+
+
+@st.cache_data
+def cache_compare(dir, name_filter, base_name, compare_name):
+    samples = cached_samples(dir, name_filter)
+    grouped, base_name, comp_name = model_comp.get_scores(
+        samples, base_name + "$", compare_name + "$"
+    )
+    grouped_dict = {k: [vi.to_dict() for vi in v] for k, v in grouped.items()}
+    return grouped_dict, base_name, comp_name
+
+
+def main():
+    parser = argparse.ArgumentParser(description="rbeval dashboard")
+    parser.add_argument("eval_dir", type=str)
+    args, rest = parser.parse_known_args()
+    eval_dir = Path(args.eval_dir)
+    # Show all the models
+    score_cdf_data, cfgs = cached_score_cdf(eval_dir, None)
+    for data, cfg in zip(score_cdf_data, cfgs):
+        figs = plot_with_data(cfg, data)
+        with st.expander(cfg.name):
+            for fig in figs:
+                st.altair_chart(fig.chart)
+
+    model_names = set(
+        [
+            m.model_name
+            for group in cached_samples(eval_dir, None)
+            for m in group.model_evals
+        ]
+    )
+    base_model = st.selectbox("Base model", model_names)
+    compare_model = st.selectbox("Compare model", model_names)
+    st.write(f"Comparing {base_model} with {compare_model}")
+    if base_model and compare_model:
+        if base_model == compare_model:
+            st.write("Base and compare models are the same")
+            return
+        grouped, base_name, comp_name = cache_compare(
+            eval_dir, None, base_model, compare_model
+        )
+        grouped = {
+            k: [model_comp.Scores.from_dict(vi) for vi in v] for k, v in grouped.items()
+        }
+        for fig in model_comp.get_figures(grouped, base_name, comp_name):
+            st.write(fig.name)
+            st.altair_chart(fig.chart)
+
+
+if __name__ == "__main__":
+    main()

src/rbeval/plot/data.py

@@ -7,6 +7,7 @@ from typing import Dict, List, Optional
 from collections import defaultdict
 import altair as alt
 
+from dacite import from_dict
 import numpy as np
 from tqdm import tqdm
 
@@ -26,18 +27,21 @@ def get_samples(inp: Path, name_filter: Optional[str]) -> List["EvalGroup"]:
                 print(f"Skipping spec {spec_file.stem}")
                 continue
 
-        group = groups.setdefault(spec.group, EvalGroup(name=spec.group))
-        model_eval = ModelEval(eval_spec=spec)
-        group.model_evals.append(model_eval)
-        for samples_file in (spec_file.parent / spec_file.stem).glob(
-            "**/samples_*.json*"
-        ):
-            cache_file = samples_file.with_suffix(".npy")
-            if samples_file.with_suffix(".npy").exists():
-                model_eval.evals.append(
-                    Eval(**np.load(str(cache_file), allow_pickle=True).item())
-                )
-            else:
+        group_cache_file = Path(
+            spec_file.with_stem(spec_file.stem + "_group_cache")
+        ).with_suffix(".npy")
+        if group_cache_file.exists():
+            res_dict = np.load(str(group_cache_file), allow_pickle=True).item()
+            group = from_dict(data_class=EvalGroup, data=res_dict)
+            groups[group.name] = group
+            continue
+        else:
+            group = groups.setdefault(spec.group, EvalGroup(name=spec.group))
+            model_eval = ModelEval(eval_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()]
@@ -57,8 +61,8 @@ def get_samples(inp: Path, name_filter: Optional[str]) -> List["EvalGroup"]:
                     cor_logprobs=np.array(cor_logprobs),
                     inc_logprobs=np.array(inc_logprobs),
                 )
-                np.save(str(cache_file), asdict(eval))  # type: ignore
                 model_eval.evals.append(eval)
+            np.save(str(group_cache_file), asdict(group))  # type: ignore
 
     return list(groups.values())
 

src/rbeval/plot/model_comp.py

@@ -2,7 +2,7 @@ import argparse
 import altair as alt
 import pandas as pd
 from collections import defaultdict
-from dataclasses import dataclass, field
+from dataclasses import asdict, dataclass, field
 import itertools
 from typing import Dict, List, Optional
 import warnings
@@ -20,21 +20,18 @@ class Scores:
     cor_minus_inc_samples: List[np.ndarray] = field(default_factory=list)
     cor_samples: List[np.ndarray] = field(default_factory=list)
 
+    def to_dict(self):
+        return asdict(self)
 
-def model_comparer(samples: List[EvalGroup], rem_args: List[str]) -> List[Figure]:
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--base", type=str)
-    parser.add_argument("--compare", type=str)
-    args = parser.parse_args(rem_args)
-    base_name_filt: Optional[str] = args.base
-    comp_name_filt: Optional[str] = args.compare
+    @classmethod
+    def from_dict(cls, d: dict):
+        d["spec"] = EvalSpec(**d["spec"])
+        return cls(**d)
 
-    if base_name_filt is None or comp_name_filt is None:
-        warnings.warn(
-            "Skipping model comparison plot, need to specify base and compare"
-        )
-        return []
 
+def get_scores(
+    samples: List[EvalGroup], base_name_filt: str, comp_name_filt: str
+) -> tuple[Dict[str, List[Scores]], str, str]:
     bases: List[ModelEval] = list(
         itertools.chain.from_iterable(
             g.collect_with_name(base_name_filt) for g in samples
@@ -107,6 +104,10 @@ def model_comparer(samples: List[EvalGroup], rem_args: List[str]) -> List[Figure
         for title, scores in scores_by_mask.items():
             grouped[title].append(scores)
 
+    return grouped, base_name, comp_name
+
+
+def get_figures(grouped: Dict[str, List[Scores]], base_name, comp_name) -> List[Figure]:
     cmp_name = f"{base_name} to {comp_name}"
     return [
         Figure(name=f"{cmp_name} prob diff perf curves", chart=plot_diff_cdf(grouped)),
@@ -115,6 +116,24 @@ def model_comparer(samples: List[EvalGroup], rem_args: List[str]) -> List[Figure
     ]
 
 
+def model_comparer(samples: List[EvalGroup], rem_args: List[str]) -> List[Figure]:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--base", type=str)
+    parser.add_argument("--compare", type=str)
+    args = parser.parse_args(rem_args)
+    base_name_filt: Optional[str] = args.base
+    comp_name_filt: Optional[str] = args.compare
+
+    if base_name_filt is None or comp_name_filt is None:
+        warnings.warn(
+            "Skipping model comparison plot, need to specify base and compare"
+        )
+        return []
+
+    grouped, base_name, comp_name = get_scores(samples, base_name_filt, comp_name_filt)
+    return get_figures(grouped, base_name, comp_name)
+
+
 def plot_diff_cdf(grouped: Dict[str, List[Scores]]) -> alt.HConcatChart:
     charts = []
     for title, score_list in grouped.items():

src/rbeval/plot/score_cdf_altair.py

@@ -1,3 +1,4 @@
+from dataclasses import dataclass, field
 from typing import List
 
 from rbeval.plot.data import Eval, EvalGroup, Figure
@@ -9,25 +10,31 @@ import pandas as pd
 from rbeval.plot.utils import CdfData, renormed
 
 
+@dataclass
+class PlotData:
+    renorm: List[pd.DataFrame] = field(default_factory=list)
+    norenorm: List[pd.DataFrame] = field(default_factory=list)
+
+
+def plot_cfgs():
+    return [CorrectProbCdfPlot(), CorrIncorrDiffConfig()]
+
+
 def score_cdf(samples: List[EvalGroup], args: List[str]) -> List[Figure]:
     return [
-        Figure(
-            name="Correct Prob Perf Curve",
-            chart=plot_with_config(CorrectProbCdfPlot(), samples),
-        ),
-        Figure(
-            name="Corr-Incorr Gap Perf Curve",
-            chart=plot_with_config(CorrIncorrDiffConfig(), samples),
-        ),
+        a
+        for cfg in plot_cfgs()
+        for a in plot_with_data(cfg, get_plot_data(cfg, samples))
     ]
 
 
-def plot_with_config(
+def get_plot_data(
     cfg: "CdfPlotConfig",
     samples: List[EvalGroup],
-) -> alt.ConcatChart:
-    group_dfs = []
+) -> PlotData:
+    data = PlotData()
     for renorm in [True, False]:
+        gfs = data.renorm if renorm else data.norenorm
         for group in samples:
             dfs = []
             for m in group.model_evals:
@@ -38,43 +45,52 @@ def plot_with_config(
                         "x": cdf.scores,
                         "y": cdf.cdf_p,
                         "label": m.model_name,
+                        "group": group.name,
                         "renorm": renorm,
                         "fewshot": spec.fewshot,
                     }
                 )
                 dfs.append(df)
-            group_dfs.append(pd.concat(dfs))
-
-    selection = alt.selection_point(fields=["label"], bind="legend")
-    charts = []
-    for group, df in zip(samples, group_dfs):
-        chart = (
-            alt.Chart(df)
-            .mark_line()
-            .encode(
-                x=alt.X("x:Q", title=cfg.xlabel),
-                y=alt.Y("y:Q", title=cfg.ylabel),
-                color=alt.Color("label:N", legend=alt.Legend(symbolOpacity=1.0)),
-                opacity=alt.condition(
-                    selection, alt.Opacity("fewshot:O"), alt.value(0.1)
-                ),
-            )
-            .properties(title=cfg.title(group.name, renorm))
-            .resolve_legend(color="independent")
-        )
+            gfs.append(pd.concat(dfs))
+    return data
 
-        charts.append(chart)
 
-    final_chart = (
-        alt.concat(*charts, columns=len(samples)).add_params(selection).interactive()
-    )
-    return final_chart
+def plot_with_data(
+    cfg: "CdfPlotConfig",
+    data: PlotData,
+) -> List[Figure]:
+    figures = []
+    for renorm, group_dfs in zip([True, False], [data.renorm, data.norenorm]):
+        for df in group_dfs:
+            group_name = df["group"].iloc[0]
+            selection = alt.selection_point(fields=["label"], bind="legend")
+            chart = (
+                alt.Chart(df)
+                .mark_line()
+                .encode(
+                    x=alt.X("x:Q", title=cfg.xlabel),
+                    y=alt.Y("y:Q", title=cfg.ylabel),
+                    color=alt.Color("label:N", legend=alt.Legend(symbolOpacity=1.0)),
+                    opacity=alt.condition(
+                        selection, alt.Opacity("fewshot:O"), alt.value(0.1)
+                    ),
+                )
+                .properties(title=cfg.title(group_name, renorm), width=800, height=400)
+                .resolve_legend(color="independent")
+                .resolve_axis(y="independent", x="independent")
+                .add_params(selection)
+                .interactive()
+            )
+            figures.append(Figure(name=f"{group_name} {cfg.name}", chart=chart))
+
+    return figures
 
 
 class CdfPlotConfig(ABC):
     plot_type: str
     xlabel: str
     ylabel: str
+    name: str = ""
 
     @abstractmethod
     def get_cdf(self, evals: List[Eval], prob_renorm: bool) -> "CdfData":
@@ -92,6 +108,8 @@ class CdfPlotConfig(ABC):
 
 
 class CorrectProbCdfPlot(CdfPlotConfig):
+    name = "Correct Prob Perf Curve"
+
     def __init__(self):
         self.plot_type = "corr perf plot"
         self.xlabel = "Correct answer probability"
@@ -112,6 +130,8 @@ class CorrectProbCdfPlot(CdfPlotConfig):
 
 
 class CorrIncorrDiffConfig(CdfPlotConfig):
+    name = "Corr-Incorr Gap Perf Curve"
+
     def __init__(self):
         self.plot_type = "corr-max(incor) perf plot"
         self.xlabel = "corr prob - max(incor prob)"

src/rbeval/plot/utils.py

@@ -59,7 +59,7 @@ class CdfData:
 
     @classmethod
     def from_weights(
-        cls, weights: np.ndarray, scores: np.ndarray, max_p=1000
+        cls, weights: np.ndarray, scores: np.ndarray, max_p=600
     ) -> "CdfData":
         sort_perm = scores.argsort()
         base_weights = weights[sort_perm]