Initial changes.

app.py

@@ -5,6 +5,7 @@ import time
 import uuid
 import os
 import html
+import sys
 
 import pandas as pd
 from huggingface_hub import hf_hub_download
@@ -29,26 +30,52 @@ _pool_path = hf_hub_download(
     token=RATINGS_APP_TOKEN
 )
 _pool_df = pd.read_parquet(_pool_path)
-_pool_group_dfs = {g: gdf for g, gdf in _pool_df.groupby("group")}
+_pool_df["wins"] = 0
+_pool_df["losses"] = 0
+_pool_df["votes"] = 0
+
+WINS_LOC = _pool_df.columns.get_loc("wins")
+LOSSES_LOC = _pool_df.columns.get_loc("losses")
+VOTES_LOC = _pool_df.columns.get_loc("votes")
+
+_md5_to_idx = { md5: idx for idx, md5 in enumerate(_pool_df["md5"]) }
+
 _stats_lock = threading.Lock()
+_pool_lock = threading.Lock()
+
 _stats_last_loaded_at = 0.0
-_stats_by_key: dict[str, tuple[int, int]] = {}
 _explorer_df = pd.DataFrame(columns=["group", "id", "md5", "rating", "sample_url", "image_url", "classifier", "classifier_score", "percentile"])
 
 
 def _reload_stats_if_due(force: bool = False):
-    global _stats_last_loaded_at, _stats_by_key, _explorer_df
+    global _stats_last_loaded_at,_explorer_df
     now = time.time()
+
     if not force and (now - _stats_last_loaded_at) < STATS_RELOAD_S:
         return
+
     with _stats_lock:
         now = time.time()
+
         if not force and (now - _stats_last_loaded_at) < STATS_RELOAD_S:
             return
-        _stats_by_key = load_stats_by_md5(
+
+        stats_by_key = load_stats_by_md5(
             repo_id=POOL_REPO_ID,
             token=RATINGS_APP_TOKEN,
         )
+
+        with _pool_lock:
+            n_missing = 0
+            for md5, stats in stats_by_key.items():
+                if (idx := _md5_to_idx.get(md5)) is not None:
+                    _pool_df.iloc[idx, [WINS_LOC, LOSSES_LOC, VOTES_LOC]] = (*stats, stats[0] + stats[1])
+                else:
+                    n_missing += 1
+
+            if n_missing:
+                print(f"{n_missing} md5s have stats but are not in the pool!", file=sys.stderr)
+
         classifier_scores_path = hf_hub_download(
             repo_id=POOL_REPO_ID,
             filename="classifier_scores.parquet",
@@ -77,62 +104,81 @@ def _reload_stats_if_due(force: bool = False):
 
 _reload_stats_if_due(force=True)
 
-def _pool_fetch_pair(group_name: str) -> tuple:
-    gdf = _pool_group_dfs[group_name]
-    assert len(gdf) >= 2, f"Not enough rows for group: {group_name}"
-    md5_keys = gdf["md5"].astype(str)
-    wins = md5_keys.map(lambda k: _stats_by_key.get(k, (0, 0))[0])
-    losses = md5_keys.map(lambda k: _stats_by_key.get(k, (0, 0))[1])
-
-    def _pick_from_mask(mask: pd.Series):
-        candidate_df = gdf[mask]
-        if len(candidate_df) < 2:
-            return None
-        sample = candidate_df.sample(2, replace=False)
-        return sample.iloc[0], sample.iloc[1]
-    
-    
-    # 1) Repeat the lowest-margin edge participating in a cycle. (To prevent deadlock, stop if all margins are 4+.)
-    #   a) If deadlocked on a cycle with 4+ images and no inner cycles, sample a random missing edge inside the cycle.
-    # 2) Pair images that both have wins only . (One of them will lose/tie. Stop when there is only one left.)
-    # 3) Pair images that both have losses only. (One of them will win/tie. Stop when there is only one left.)
-    # 4) Pair images with only 2 edges.
-    # 5) X% chance, re-sample an existing edge, inversely proportional to existing number of samples.
-    # 6) Y% chance, sample a random missing edge between images already sampled.
-    # 7) Pair an unsampled image with a random sampled image. 
-
-    # 2) Pair images that currently have wins-only records.
-    picked = _pick_from_mask((wins > 0) & (losses == 0))
+def _pick_from_bins(df: pd.DataFrame, field: str) -> tuple[pd.Series, pd.Series, int] | None:
+    if len(df) < 2:
+        return None
+
+    least = df[field].min()
+    if least >= 10:
+        return None # don't push too hard for a total order
+
+    remaining = (df[field] < 10).sum() - 1
+
+    candidates = df[df[field] == least]
+    if len(candidates) > 1:
+        sample = candidates.sample(2, replace=False)
+        return sample.iloc[0], sample.iloc[1], remaining
+
+    first = candidates.iloc[0]
+    while True:
+        least += 1
+        candidates = df[df[field] == least]
+        if candidates.empty:
+            continue
+
+        sample = candidates.sample(1)
+        return first, sample.iloc[0], remaining
+
+def _pick_from(df: pd.DataFrame, weights: pd.Series | None = None) -> tuple[pd.Series, pd.Series, int] | None:
+    if len(df) < 2:
+        return None
+
+    remaining = len(df) - 2
+
+    sample = df.sample(2, weights=weights, replace=False)
+    return sample.iloc[0], sample.iloc[1], remaining
+
+def _pool_fetch_pair(group: str) -> tuple[pd.Series, pd.Series, int, str]:
+    gdf = _pool_df[_pool_df["group"] == group]
+    ranked = gdf[gdf["votes"] > 0]
+
+    # 1) Pair images that have wins-only records.
+    picked = _pick_from_bins(ranked[ranked["losses"] == 0], "wins")
     if picked is not None:
-        return picked[0], picked[1], "wins-only"
+        return *picked, "wins-only"
 
-    # 3) Pair images that currently have losses-only records.
-    picked = _pick_from_mask((wins == 0) & (losses > 0))
+    # 2) Pair images that have losses-only records.
+    picked = _pick_from_bins(ranked[ranked["wins"] == 0], "losses")
     if picked is not None:
-        return picked[0], picked[1], "losses-only"
+        return *picked, "losses-only"
 
-    # 4) Pair images that currently have exactly 2 total edges.
-    vote_totals = wins + losses
-    picked = _pick_from_mask(vote_totals == 2)
+    # 3) Ensure a minimum density of 3.
+    picked = _pick_from(ranked[ranked["votes"] == 2])
     if picked is not None:
-        return picked[0], picked[1], "total_votes=2"
+        return *picked, "sparse"
 
-    # 7) Prefer pairing an unsampled image with a random previously sampled image.
-    unsampled_mask = vote_totals == 0
-    if unsampled_mask.any():
-        unsampled_row = gdf[unsampled_mask].sample(1).iloc[0]
-        sampled_df = gdf[~unsampled_mask]
-        if len(sampled_df) >= 1:
-            sampled_row = sampled_df.sample(1).iloc[0]
-        else:
-            sampled_row = gdf.drop(index=unsampled_row.name).sample(1).iloc[0]
-        return unsampled_row, sampled_row, "unsampled+sampled"
+    # 4) Introduce a new image.
+    if ranked.empty or random.random() < 0.75:
+        unranked = gdf[gdf["votes"] == 0]
 
-    # 8) Safety fall back to low-vote weighted sampling.
-    sample_weights = 1.0 / (vote_totals + 1.0)
-    sample = gdf.sample(2, weights=sample_weights, replace=False)
-    return sample.iloc[0], sample.iloc[1], "low-vote"
+        if ranked.empty: # Very first vote.
+            picked = _pick_from(unranked)
+            if picked is None:
+                raise ValueError("Group is empty.")
 
+            return *picked, "init"
+
+        if not unranked.empty:
+            return (
+                ranked.sample(1).iloc[0],
+                unranked.sample(1).iloc[0],
+                len(unranked) - 1, "new"
+            )
+
+    # 5) Vote-weighted random sampling.
+    picked = _pick_from(ranked, weights=(1.0 / ranked["votes"]))
+    assert picked is not None
+    return *picked, "random"
 
 def _row_image_url(row) -> str:
     sample_url = row.get("sample_url")
@@ -148,39 +194,20 @@ DATASETS: dict[str, dict] = {
         "fetch_pair": _pool_fetch_pair,
         "get_id": lambda row: row["md5"],
         "get_image": _row_image_url,
-        "groups": {g: g for g in sorted(_pool_df["group"].unique())},
+        "groups": sorted(_pool_df["group"].unique()),
     },
 }
 DEFAULT_DATASET = list(DATASETS.keys())[0]
 
 def _select_groups(cfg: dict, rating_pref: str) -> list[str]:
-    groups = list(cfg["groups"].keys())
     if rating_pref == "all":
-        return groups
-    return [g for g in groups if g.endswith(f"_{rating_pref}")]
-
-def _commit_oldest_pending(state: dict):
-    pending = state.setdefault("pending", [])
-    if len(pending) <= 1:
-        return
-    oldest = pending.pop(0)
-    if oldest.get("winner") in ("A", "B"):
-        _apply_local_stats_update(oldest["winner"], oldest["key_a"], oldest["key_b"])
-    threading.Thread(target=VOTE_STORAGE.append_vote_row, args=(oldest.copy(), oldest.get("winner")), daemon=True).start()
-
-
-def _apply_local_stats_update(winner: str, key_a: str, key_b: str):
-    assert winner in ("A", "B")
-    with _stats_lock:
-        wins_a, losses_a = _stats_by_key.get(str(key_a), (0, 0))
-        wins_b, losses_b = _stats_by_key.get(str(key_b), (0, 0))
-        if winner == "A":
-            _stats_by_key[str(key_a)] = (wins_a + 1, losses_a)
-            _stats_by_key[str(key_b)] = (wins_b, losses_b + 1)
-        else:
-            _stats_by_key[str(key_a)] = (wins_a, losses_a + 1)
-            _stats_by_key[str(key_b)] = (wins_b + 1, losses_b)
+        return cfg["groups"]
 
+    return [
+        g
+        for g in cfg["groups"]
+        if g.endswith(f"_{rating_pref}")
+    ]
 
 def _format_rating_post_row(post_id: int, wins: int, losses: int, label: str | None = None) -> str:
     total_votes = wins + losses
@@ -190,8 +217,8 @@ def _format_rating_post_row(post_id: int, wins: int, losses: int, label: str | N
 
 def _render_current(state: dict, submit_status: str = "") -> tuple:
     _reload_stats_if_due()
-    wins_a, losses_a = _stats_by_key.get(str(state["key_a"]), (0, 0))
-    wins_b, losses_b = _stats_by_key.get(str(state["key_b"]), (0, 0))
+    wins_a, losses_a = _pool_df.iloc[_md5_to_idx[state["key_a"]], [WINS_LOC, LOSSES_LOC]]
+    wins_b, losses_b = _pool_df.iloc[_md5_to_idx[state["key_b"]], [WINS_LOC, LOSSES_LOC]]
     title_a = "Image A"
     title_b = "Image B"
     img_a_html = f"<div class=\"rating-card\"><div class=\"rating-card-title\"><strong>{html.escape(title_a)}</strong></div><div class=\"rating-image-frame\"><img src=\"{html.escape(state['url_a'])}\" class=\"rating-image\" loading=\"eager\" referrerpolicy=\"no-referrer\"></div></div>"
@@ -207,7 +234,6 @@ def _render_current(state: dict, submit_status: str = "") -> tuple:
     return img_a_html, img_b_html, link_a, link_b, back_md, group_md, pair_reason_md, status_md, state
 
 
-
 def _normalize_rating_pref(pref: str | None) -> str:
     return pref if pref in ("safe", "all") else "safe"
 
@@ -274,13 +300,12 @@ def new_round(dataset_name: str, rating_pref: str, state: dict) -> tuple:
     cfg = DATASETS[dataset_name]
     groups = _select_groups(cfg, rating_pref)
     assert groups, f"No groups for rating preference: {rating_pref}"
+
     group = random.choice(groups)
-    pair_data = cfg["fetch_pair"](cfg["groups"][group])
-    if len(pair_data) == 3:
-        row_a, row_b, pair_reason = pair_data
-    else:
-        row_a, row_b = pair_data
-        pair_reason = ""
+    row_a, row_b, reason_remaining, pair_reason = cfg["fetch_pair"](group)
+
+    pair_reason = f"{pair_reason} ({reason_remaining})"
+
     state.setdefault("session_id", uuid.uuid4().hex)
     key_a = cfg["get_id"](row_a)
     key_b = cfg["get_id"](row_b)
@@ -295,8 +320,9 @@ def new_round(dataset_name: str, rating_pref: str, state: dict) -> tuple:
 
 def _queue_decision(winner: str | None, state: dict):
     assert state.get("session_id"), "Missing session_id: refusing to record vote"
-    state.setdefault("pending", [])
-    decision = {
+
+    pending = state.setdefault("pending", [])
+    pending.append({
         "winner": winner,
         "key_a": state["key_a"],
         "key_b": state["key_b"],
@@ -309,43 +335,63 @@ def _queue_decision(winner: str | None, state: dict):
         "group": state["group"],
         "pair_reason": state.get("pair_reason", ""),
         "session_id": state["session_id"],
-    }
-    state["pending"].append(decision)
-    state["last_decision"] = decision
-    state["can_go_back"] = True
-    _commit_oldest_pending(state)
+    })
+
+    if len(pending) > 1:
+        VOTE_STORAGE.queue_row(pending.pop(0))
 
 def vote(winner: str | None, state: dict, submit_key: str | None) -> tuple:
-    assert winner in ("A", "B", None)
     if _normalize_submit_key(submit_key) != SUBMIT_KEY:
         return _render_current(state, "Wrong submission key.")
+
     _queue_decision(winner, state)
+
+    with _pool_lock:
+        match winner:
+            case "A":
+                _pool_df.iloc[_md5_to_idx[state["key_a"]], [WINS_LOC, VOTES_LOC]] += 1
+                _pool_df.iloc[_md5_to_idx[state["key_b"]], [LOSSES_LOC, VOTES_LOC]] += 1
+            case "B":
+                _pool_df.iloc[_md5_to_idx[state["key_b"]], [WINS_LOC, VOTES_LOC]] += 1
+                _pool_df.iloc[_md5_to_idx[state["key_a"]], [LOSSES_LOC, VOTES_LOC]] += 1
+            case None:
+                pass
+            case _:
+                raise AssertionError
+
     return new_round(state["dataset"], state["rating_pref"], state)
 
 def go_back(state: dict) -> tuple:
     pending = state.setdefault("pending", [])
-    if not state.get("can_go_back"):
-        return _render_current(state)
-    last = state.get("last_decision")
-    if not last:
-        state["can_go_back"] = False
-        return _render_current(state)
-    if pending and pending[-1] == last:
-        pending.pop()
-    state["can_go_back"] = False
-    state["last_decision"] = None
-    state.update(
-        dataset=last["dataset"],
-        rating_pref=last["rating_pref"],
-        key_a=last["key_a"],
-        key_b=last["key_b"],
-        id_a=last["id_a"],
-        id_b=last["id_b"],
-        url_a=last["url_a"],
-        url_b=last["url_b"],
-        group=last["group"],
-        pair_reason=last.get("pair_reason", ""),
-    )
+
+    if pending:
+        last = pending.pop()
+        state.update(
+            dataset=last["dataset"],
+            rating_pref=last["rating_pref"],
+            key_a=last["key_a"],
+            key_b=last["key_b"],
+            id_a=last["id_a"],
+            id_b=last["id_b"],
+            url_a=last["url_a"],
+            url_b=last["url_b"],
+            group=last["group"],
+            pair_reason=last.get("pair_reason", ""),
+        )
+
+        with _pool_lock:
+            match last["winner"]:
+                case "A":
+                    _pool_df.iloc[_md5_to_idx[state["key_a"]], [WINS_LOC, VOTES_LOC]] -= 1
+                    _pool_df.iloc[_md5_to_idx[state["key_b"]], [LOSSES_LOC, VOTES_LOC]] -= 1
+                case "B":
+                    _pool_df.iloc[_md5_to_idx[state["key_b"]], [WINS_LOC, VOTES_LOC]] -= 1
+                    _pool_df.iloc[_md5_to_idx[state["key_a"]], [LOSSES_LOC, VOTES_LOC]] -= 1
+                case None:
+                    pass
+                case _:
+                    raise AssertionError
+
     return _render_current(state)
 
 # -- UI ---------------------------------------------------------------------

storage.py

@@ -30,79 +30,98 @@ class VoteStorage:
     def __init__(self, mode: str, token: str | None = None):
         assert mode in ("hf", "void"), f"Unsupported storage mode: {mode}"
         self.mode = mode
-        self._token = token
         is_debug_mode = self.mode == "void"
+
         self._flush_every = 3 if is_debug_mode else 50
         self._flush_interval_sec = 15.0 if is_debug_mode else 300.0
-        self._votes_lock = threading.Lock()
+
+        self._shutdown = False
         self._votes_buffer: list[dict] = []
-        self._stop_event = threading.Event()
+
+        self._flush_condition = threading.Condition(threading.Lock())
         self._flush_thread = threading.Thread(target=self._flush_loop, daemon=True)
         self._flush_thread.start()
-        atexit.register(self.close)
 
-    def _hf_token(self) -> str | None:
-        return self._token
+        self.hf_api = HfApi(token=token)
+
+        atexit.register(self.close)
 
     def _empty_votes_df(self) -> pd.DataFrame:
         return pd.DataFrame(columns=VOTE_COLUMNS)
 
     def _upload_votes_batch(self, df: pd.DataFrame, commit_message: str):
-        assert set(VOTE_COLUMNS).issubset(df.columns), "Missing vote columns in upload batch"
         if self.mode == "void":
-            _ = commit_message
             return
+
         ts = int(time.time())
         shard = f"votes_{ts}_{uuid.uuid4().hex}.parquet"
-        api = HfApi(token=self._hf_token())
-        with NamedTemporaryFile(suffix=".parquet", delete=False) as tmp:
-            tmp_path = tmp.name
-        try:
-            df[VOTE_COLUMNS].to_parquet(tmp_path, index=False)
-            api.upload_file(
-                path_or_fileobj=tmp_path,
-                path_in_repo=f"{VOTES_LOG_SUBDIR}/{shard}",
-                repo_id=VOTES_REPO_ID,
-                repo_type=VOTES_REPO_TYPE,
-                commit_message=commit_message,
-            )
-        finally:
-            if os.path.exists(tmp_path):
-                os.remove(tmp_path)
-
-    def _flush_votes(self, force: bool = False):
-        with self._votes_lock:
-            if not self._votes_buffer:
-                return
-            if not force and len(self._votes_buffer) < self._flush_every:
-                return
-            batch = list(self._votes_buffer)
-            self._votes_buffer.clear()
-        incoming = pd.DataFrame(batch)
-        for col in VOTE_COLUMNS:
-            if col not in incoming.columns:
-                incoming[col] = None
-        self._upload_votes_batch(incoming[VOTE_COLUMNS], commit_message=f"append {len(batch)} vote rows")
-
-    def _flush_loop(self):
-        while not self._stop_event.wait(self._flush_interval_sec):
-            self._flush_votes(force=True)
-
-    def close(self):
-        if self._stop_event.is_set():
-            return
-        self._stop_event.set()
-        self._flush_thread.join(timeout=1.0)
-        self._flush_votes(force=True)
 
-    def append_vote_row(self, state: dict, winner: str | None):
+        self.hf_api.upload_file(
+            path_or_fileobj=df.to_parquet(index=False),
+            path_in_repo=f"{VOTES_LOG_SUBDIR}/{shard}",
+            repo_id=VOTES_REPO_ID,
+            repo_type=VOTES_REPO_TYPE,
+            commit_message=commit_message,
+        )
+
+    def _flush_loop(self) -> None:
+        while True:
+            with self._flush_condition:
+                while True:
+                    if self._shutdown:
+                        # Flush last batch of votes.
+                        if self._votes_buffer:
+                            break
+
+                        return
+
+                    # Have enough votes to flush now.
+                    if len(self._votes_buffer) >= self._flush_every:
+                        break
+
+                    # Wait for a notify to flush early or shutdown.
+                    if not self._flush_condition.wait(self._flush_interval_sec):
+                        # Interval elapsed. Flush if there is at least one vote.
+                        if self._votes_buffer:
+                            break
+
+                # Atomically take the batch of votes.
+                batch = self._votes_buffer
+                self._votes_buffer = []
+
+            assert batch
+            batch_df = pd.DataFrame(batch)
+            del batch
+
+            for col in VOTE_COLUMNS:
+                if col not in batch_df.columns:
+                    batch_df[col] = None
+
+            batch_df = batch_df[VOTE_COLUMNS]
+            self._upload_votes_batch(batch_df, commit_message=f"upload {len(batch_df)} vote rows")
+
+    def close(self) -> None:
+        with self._flush_condition:
+            self._shutdown = True
+            self._flush_condition.notify()
+
+        self._flush_thread.join()
+
+    def queue_row(self, state: dict) -> None:
         id_a = int(state["id_a"])
         id_b = int(state["id_b"])
-        winner_md5 = None
-        if winner == "A":
-            winner_md5 = state["key_a"]
-        elif winner == "B":
-            winner_md5 = state["key_b"]
+
+        winner_md5: str | None
+        match state["winner"]:
+            case "A":
+                winner_md5 = state["key_a"]
+            case "B":
+                winner_md5 = state["key_b"]
+            case None:
+                winner_md5 = None
+            case _:
+                raise AssertionError
+
         vote_row = {
             "vote_id": uuid.uuid4().hex,
             "timestamp": datetime.now(timezone.utc).isoformat(timespec="seconds"),
@@ -115,6 +134,9 @@ class VoteStorage:
             "group": state["group"],
             "session_id": state["session_id"],
         }
-        with self._votes_lock:
+
+        with self._flush_condition:
             self._votes_buffer.append(vote_row)
-        self._flush_votes()
+
+            if len(self._votes_buffer) == self._flush_every:
+                self._flush_condition.notify()