Switch to generation time

app.py

@@ -10,11 +10,13 @@ import plotly.graph_objects as go
 import wandb
 from substrateinterface import Keypair
 from wandb.apis.public import Run
+import numpy as np
 
 WANDB_RUN_PATH = os.environ["WANDB_RUN_PATH"]
 SOURCE_VALIDATOR_UID = int(os.environ["SOURCE_VALIDATOR_UID"])
+BASELINE_AVERAGE = float(os.environ["BASELINE_AVERAGE"])
+
 REFRESH_RATE = 60 * 30  # 30 minutes
-BASELINE = 0.0
 GRAPH_HISTORY_DAYS = 30
 MAX_GRAPH_ENTRIES = 10
 
@@ -28,6 +30,8 @@ class LeaderboardEntry:
     uid: int
     model: str
     score: float
+    model_average: float
+    similarity: float
     hotkey: str
     previous_day_winner: bool
     rank: int
@@ -36,9 +40,11 @@ class LeaderboardEntry:
 @dataclass
 class GraphEntry:
     dates: list[datetime]
+    generation_times: list[float]
+    similarities: list[float]
     scores: list[float]
     models: list[str]
-    max_score: float
+    best_time: float
 
 
 def is_valid_run(run: Run):
@@ -60,6 +66,13 @@ def is_valid_run(run: Run):
         return False
 
 
+def calculate_score(model_average: float, similarity: float) -> float:
+    return max(
+        0.0,
+        BASELINE_AVERAGE - model_average
+    ) * similarity
+
+
 def get_graph_entries(runs: list[Run]) -> dict[int, GraphEntry]:
     entries: dict[int, GraphEntry] = {}
 
@@ -70,22 +83,30 @@ def get_graph_entries(runs: list[Run]) -> dict[int, GraphEntry]:
             if key.startswith("_"):
                 continue
 
+            if "score" in value:
+                continue
+
             uid = int(key)
-            score = value["score"]
+
+            generation_time = value["generation_time"]
+            similarity = min(1, value["similarity"])
+            score = calculate_score(generation_time, similarity)
             model = value["model"]
 
             if uid not in entries:
-                entries[uid] = GraphEntry([date], [score], [model], score)
+                entries[uid] = GraphEntry([date], [generation_time], [similarity], [score], [model], generation_time)
             else:
-                if score > entries[uid].max_score:
-                    entries[uid].max_score = score
+                if generation_time < entries[uid].best_time:
+                    entries[uid].best_time = generation_time
 
                 data = entries[uid]
                 data.dates.append(date)
-                data.scores.append(data.max_score)
+                data.generation_times.append(data.best_time)
+                data.similarities.append(similarity)
+                data.scores.append(score)
                 data.models.append(model)
 
-    return dict(sorted(entries.items(), key=lambda entry: entry[1].max_score, reverse=True)[:MAX_GRAPH_ENTRIES])
+    return dict(sorted(entries.items(), key=lambda entry: entry[1].best_time)[:MAX_GRAPH_ENTRIES])
 
 
 def create_graph(runs: list[Run]) -> go.Figure:
@@ -95,14 +116,16 @@ def create_graph(runs: list[Run]) -> go.Figure:
     for uid, data in entries.items():
         fig.add_trace(go.Scatter(
             x=data.dates,
-            y=data.scores,
-            customdata=data.models,
+            y=data.generation_times,
+            customdata=np.stack((data.similarities, data.scores, data.models), axis=-1),
             mode="lines+markers",
             name=uid,
             hovertemplate=(
                     "<b>Date:</b> %{x|%Y-%m-%d}<br>" +
-                    "<b>Score:</b> %{y}<br>" +
-                    "<b>Model:</b> %{customdata}<br>"
+                    "<b>Generation Time:</b> %{y}s<br>" +
+                    "<b>Similarity:</b> %{customdata[0]}<br>" +
+                    "<b>Score:</b> %{customdata[1]}<br>" +
+                    "<b>Model:</b> %{customdata[2]}<br>"
             ),
         ))
 
@@ -110,7 +133,7 @@ def create_graph(runs: list[Run]) -> go.Figure:
 
     fig.add_trace(go.Scatter(
         x=date_range,
-        y=[BASELINE] * len(date_range),
+        y=[BASELINE_AVERAGE] * len(date_range),
         line=dict(color="#ff0000", width=3),
         mode="lines",
         name="Baseline",
@@ -119,8 +142,8 @@ def create_graph(runs: list[Run]) -> go.Figure:
     background_color = gr.themes.default.colors.slate.c800
 
     fig.update_layout(
-        title="Score Improvements",
-        yaxis_title="Score",
+        title="Generation Time Improvements",
+        yaxis_title="Generation Time (s)",
         plot_bgcolor=background_color,
         paper_bgcolor=background_color,
         template="plotly_dark"
@@ -142,12 +165,17 @@ def create_leaderboard(runs: list[Run]) -> list[tuple]:
 
             try:
                 uid = int(key)
+                generation_time = value.get("generation_time", 0)
+                similarity = min(1, value.get("similarity", 0))
+                score = value.get("score", calculate_score(generation_time, similarity))
 
                 entries[uid] = LeaderboardEntry(
                     uid=uid,
                     rank=value["rank"],
                     model=value["model"],
-                    score=value["score"],
+                    score=score,
+                    model_average=generation_time,
+                    similarity=similarity,
                     hotkey=value["hotkey"],
                     previous_day_winner=value["multiday_winner"],
                 )
@@ -158,7 +186,8 @@ def create_leaderboard(runs: list[Run]) -> list[tuple]:
             break
 
     leaderboard: list[tuple] = [
-        (entry.rank + 1, entry.uid, entry.model, entry.score, entry.hotkey, entry.previous_day_winner)
+        (entry.rank + 1, entry.uid, entry.model, entry.score, f"{entry.model_average:.3f}s", f"{entry.similarity:.3f}",
+         entry.hotkey, entry.previous_day_winner)
         for entry in sorted(entries.values(), key=lambda entry: (entry.score, entry.rank), reverse=True)
     ]
 
@@ -227,8 +256,9 @@ def refresh():
 
             leaderboard = gr.components.Dataframe(
                 create_leaderboard(runs[dropdown.value]),
-                headers=["Rank", "Uid", "Model", "Score", "Hotkey", "Previous day winner"],
-                datatype=["number", "number", "markdown", "number", "markdown", "bool"],
+                headers=["Rank", "Uid", "Model", "Score", "Generation Time", "Similarity", "Hotkey",
+                         "Previous day winner"],
+                datatype=["number", "number", "markdown", "number", "markdown", "markdown", "markdown", "bool"],
                 elem_id="leaderboard-table",
                 interactive=False,
                 visible=True,