File size: 6,693 Bytes
e840a29
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8062deb
e840a29
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8062deb
 
 
 
 
 
e840a29
8062deb
e840a29
8062deb
e840a29
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
"""Normalized dataset loading and indexing for the read-only Space."""

from __future__ import annotations

import os
from collections import defaultdict
from collections.abc import Iterable, Mapping
from typing import Any


DEFAULT_DATASET_REPO = "montagovian/basedBench"


def _column(table: Any, name: str) -> list[Any]:
    try:
        return list(table[name])
    except (KeyError, TypeError):
        return [row[name] for row in table]


class BenchmarkData:
    """In-memory indexes over the four normalized dataset configs."""

    def __init__(
        self,
        memes: Any,
        predictions: Iterable[Mapping[str, Any]],
        judgments: Iterable[Mapping[str, Any]],
        leaderboard: Iterable[Mapping[str, Any]],
    ) -> None:
        self._memes = memes
        post_ids = [str(value) for value in _column(memes, "post_id")]
        titles = [str(value) for value in _column(memes, "title")]
        subreddits = [str(value) for value in _column(memes, "subreddit")]
        ground_truths = [str(value) for value in _column(memes, "ground_truth")]
        snapshot_ids = [str(value) for value in _column(memes, "snapshot_id")]

        self.post_ids = post_ids
        self._row_index = {post_id: idx for idx, post_id in enumerate(post_ids)}
        self._meta = {
            post_id: {
                "post_id": post_id,
                "title": titles[idx],
                "subreddit": subreddits[idx],
                "ground_truth": ground_truths[idx],
                "snapshot_id": snapshot_ids[idx],
            }
            for idx, post_id in enumerate(post_ids)
        }

        self.predictions_by_post: dict[str, list[dict[str, Any]]] = defaultdict(list)
        self.predictions_by_id: dict[int, dict[str, Any]] = {}
        for source in predictions:
            row = dict(source)
            prediction_id = int(row["prediction_id"])
            post_id = str(row["post_id"])
            self.predictions_by_id[prediction_id] = row
            self.predictions_by_post[post_id].append(row)
        for rows in self.predictions_by_post.values():
            rows.sort(key=lambda row: str(row["model_id"]))

        self.latest_judgments: dict[int, list[dict[str, Any]]] = defaultdict(list)
        self.historical_judgment_counts: dict[int, int] = defaultdict(int)
        for source in judgments:
            row = dict(source)
            prediction_id = int(row["prediction_id"])
            if bool(row.get("is_latest")):
                self.latest_judgments[prediction_id].append(row)
            else:
                self.historical_judgment_counts[prediction_id] += 1
        for rows in self.latest_judgments.values():
            rows.sort(key=lambda row: str(row["judge_model"]))

        self.leaderboard = [dict(row) for row in leaderboard]
        self.leaderboard.sort(
            key=lambda row: (-float(row["accuracy"]), str(row["model_id"]))
        )
        self.models = sorted(
            {
                str(row["model_id"])
                for rows in self.predictions_by_post.values()
                for row in rows
            }
        )

    @property
    def snapshot_id(self) -> str:
        if not self.post_ids:
            return ""
        return str(self._meta[self.post_ids[0]]["snapshot_id"])

    def meme(self, post_id: str) -> dict[str, Any]:
        return self._meta[post_id]

    def image(self, post_id: str) -> Any:
        return self._memes[self._row_index[post_id]]["image"]

    def predictions(self, post_id: str, model_id: str = "all") -> list[dict[str, Any]]:
        rows = self.predictions_by_post.get(post_id, [])
        if model_id == "all":
            return rows
        return [row for row in rows if str(row["model_id"]) == model_id]

    def judgments(self, prediction_id: int) -> list[dict[str, Any]]:
        return self.latest_judgments.get(prediction_id, [])

    def filtered_ids(
        self,
        search: str = "",
        model_id: str = "all",
        outcome: str = "all",
    ) -> list[str]:
        needle = search.strip().casefold()
        matches: list[str] = []
        for post_id in self.post_ids:
            meta = self._meta[post_id]
            if needle and needle not in " ".join(
                (
                    post_id,
                    str(meta["title"]),
                    str(meta["subreddit"]),
                    str(meta["ground_truth"]),
                )
            ).casefold():
                continue

            predictions = self.predictions(post_id, model_id)
            if model_id != "all" and not predictions:
                continue
            verdicts = {
                row.get("consensus_verdict")
                for row in predictions
                if row.get("consensus_verdict") in {"correct", "incorrect"}
            }
            if outcome == "all_correct" and verdicts != {"correct"}:
                continue
            if outcome == "all_incorrect" and verdicts != {"incorrect"}:
                continue
            if outcome == "mixed" and verdicts != {"correct", "incorrect"}:
                continue
            matches.append(post_id)
        return matches

    def leaderboard_rows(self) -> list[list[Any]]:
        return [
            [
                row["model_id"],
                int(row["correct"]),
                int(row["incorrect"]),
                int(row["total"]),
                f"{float(row['accuracy']) * 100:.1f}%",
                (
                    f"{int(row['unanimous_agreements'])}/"
                    f"{int(row['judged_by_multiple'])} "
                    f"({float(row['agreement_rate']) * 100:.1f}%)"
                ),
            ]
            for row in self.leaderboard
        ]


def load_from_hub(repo_id: str | None = None) -> BenchmarkData:
    """Load the published snapshot directly from the Hub, not dataset-server."""
    from datasets import load_dataset

    repo = repo_id or os.getenv("HF_DATASET_REPO", DEFAULT_DATASET_REPO)
    token = os.getenv("HF_TOKEN") or os.getenv("HF_API_KEY")
    kwargs = {"token": token} if token else {}
    try:
        memes = load_dataset(repo, "memes", split="train", **kwargs)
        predictions = load_dataset(repo, "predictions", split="train", **kwargs)
        judgments = load_dataset(repo, "judgments", split="train", **kwargs)
        leaderboard = load_dataset(repo, "leaderboard", split="train", **kwargs)
    except Exception as exc:
        raise RuntimeError(
            f"Unable to load {repo}. For a private dataset, add an HF_TOKEN "
            "with read access to the Space secrets."
        ) from exc
    return BenchmarkData(memes, predictions, judgments, leaderboard)