Merge pull request #2 from macrocosm-os/update-api

.gitignore

@@ -0,0 +1,3 @@
+.venv
+__pycache__/
+.env

README.md

@@ -1,6 +1,6 @@
 ---
-title: Sn9
-emoji: 👁
+title: Finetuning subnet
+emoji: :em
 colorFrom: purple
 colorTo: pink
 sdk: gradio

app.py

@@ -2,25 +2,26 @@
 
 import os
 import datetime
+from typing import Dict
 import gradio as gr
 
 from dotenv import load_dotenv
 from huggingface_hub import HfApi
 from apscheduler.schedulers.background import BackgroundScheduler
 
+import competitions
 import utils
 
 FONT = (
     """<link href="https://fonts.cdnfonts.com/css/jmh-typewriter" rel="stylesheet">"""
 )
-TITLE = """<h1 align="center" id="space-title" class="typewriter">Subnet 9 Leaderboard</h1>"""
-HEADER = """<h2 align="center" class="typewriter"><a href="https://github.com/macrocosm-os/pretraining" target="_blank">Subnet 9</a> is a <a href="https://bittensor.com/" target="_blank">Bittensor</a> subnet that rewards miners for producing pretrained Foundation-Models on the <a href="https://huggingface.co/datasets/tiiuae/falcon-refinedweb" target="_blank">Falcon Refined Web dataset</a>. It acts like a continuous benchmark whereby miners are rewarded for attaining the best losses on randomly sampled pages of Falcon.<br/>The models with the best head-to-head loss on the evaluation data receive a steady emission of TAO.</h3>"""
-
+TITLE = """<h1 align="center" id="space-title" class="typewriter">Finetuning Subnet Leaderboard</h1>"""
+HEADER = """<h2 align="center" class="typewriter"><a href="https://github.com/macrocosm-os/finetuning" target="_blank">Finetuning</a> is a <a href="https://bittensor.com/" target="_blank">Bittensor</a> subnet that rewards miners for producing finetuned models in defined competitions. The model with the best head-to-head score in each competition receive a steady emission of TAO.</h3>"""
+# TODO: Update links once subnet is regged.
 EVALUATION_DETAILS = """<ul><li><b>Name:</b> the 🤗 Hugging Face model name (click to go to the model card)</li><li><b>Rewards / Day:</b> the expected rewards per day based on current ranking.</li><li><b>Last Average Loss:</b> the last loss value on the evaluation data for the model as calculated by a validator (lower is better)</li><li><b>UID:</b> the Bittensor UID of the miner</li><li><b>Block:</b> the Bittensor block that the model was submitted in</li></ul><br/>More stats on <a href="https://taostats.io/subnets/netuid-9/" target="_blank">taostats</a>."""
 EVALUATION_HEADER = """<h3 align="center">Shows the latest internal evaluation statistics as calculated by the Opentensor validator</h3>"""
 
-
-HF_REPO_ID = "macrocosm-os/pretraining-leaderboard"
+HF_REPO_ID = "macrocosm-os/finetuning-leaderboard"
 SECONDS_PER_BLOCK = 12
 
 load_dotenv()
@@ -68,6 +69,7 @@ def main():
         # TODO: Re-enable once ""SubtensorModule.BlocksSinceEpoch" not found" issue is resolved.
         # gr.HTML(value=get_next_update_div(current_block, next_epoch_block))
 
+        # TODO: Figure out the best approach to showing the per competition rewards.
         gr.Label(
             value={
                 f"{c.namespace}/{c.name} ({c.commit[0:8]}) · (τ{round(c.emission, 2):,})": c.incentive
@@ -85,28 +87,34 @@ def main():
         with gr.Accordion("Evaluation Stats"):
             gr.HTML(EVALUATION_HEADER)
             show_stale = gr.Checkbox(label="Show Stale", interactive=True)
-            leaderboard_table = gr.components.Dataframe(
-                value=utils.leaderboard_data(model_data, scores, show_stale.value),
-                headers=["Name", "Win Rate", "Average Loss", "Weight", "UID", "Block"],
-                datatype=["markdown", "number", "number", "number", "number", "number"],
-                elem_id="leaderboard-table",
-                interactive=False,
-                visible=True,
-            )
+            competition_leaderboards = []
+            # TODO: Dynamically generate per-competition leaderboards based on model_data.
+            competition_details = competitions.COMPETITION_DETAILS[1]
+            with gr.Accordion(f"{competition_details.name} competition"):
+                gr.HTML(competition_details.html_description)
+                competition_leaderboards.append(gr.components.Dataframe(
+                    value=utils.leaderboard_data(model_data, scores, show_stale.value),
+                    headers=["Name", "Win Rate", "Average Loss", "Weight", "UID", "Block"],
+                    datatype=["markdown", "number", "number", "number", "number", "number"],
+                    elem_id="leaderboard-table",
+                    interactive=False,
+                    visible=True,
+                ))
             gr.HTML(EVALUATION_DETAILS)
             show_stale.change(
                 lambda stale: utils.leaderboard_data(model_data, scores, stale),
                 inputs=[show_stale],
-                outputs=leaderboard_table,
+                outputs=competition_leaderboards,
             )
 
+            # TODO: Make this a multi-competition line plot
             gr.LinePlot(
                 utils.get_losses_over_time(vali_runs),
                 x="timestamp",
                 x_title="Date",
-                y="best_loss",
+                y="SN9_MODEL",
                 y_title="Average Loss",
-                tooltip="best_loss",
+                tooltip="SN9_MODEL",
                 interactive=True,
                 visible=True,
                 width=1024,

competitions.py

@@ -0,0 +1,20 @@
+from dataclasses import dataclass
+from typing import Dict
+
+
+@dataclass(frozen=True)
+class CompetitionDetails:
+    # The display name of the competition.
+    name: str
+
+    # The HTML description of the competition.
+    html_description: str
+
+
+# A map of competition IDs to HTML descriptions.
+COMPETITION_DETAILS: Dict[int, CompetitionDetails] = {
+    1: CompetitionDetails(
+        name="SN9_MODEL",
+        html_description="""<b>Competition ID 1</b><br/>Produce the best fine-tuned model from a Subnet 9 pretrained model. Models are evaluated using synthetic prompt/response data from Subnet 18.""",
+    )
+}

requirements.txt

@@ -1,6 +1,7 @@
 bittensor
 requests
-wandb
+wandb==0.17.1
+numpy==1.26.4
 python-dotenv
 APScheduler
 huggingface-hub

utils.py

@@ -1,25 +1,28 @@
-import os
-import math
-import time
-import json
-import wandb
-import pickle
-import datetime
 import argparse
+import datetime
 import functools
+import json
+import math
+import os
+import time
 import traceback
+from collections import defaultdict
+from dataclasses import dataclass
+from email.policy import default
+from typing import Any, Dict, List, Optional, Tuple
 
-import pandas as pd
-import numpy as np
 import bittensor as bt
-
-from dotenv import load_dotenv
-from dataclasses import dataclass
-from typing import Dict, List, Any, Optional, Tuple
+import numpy as np
+import pandas as pd
+import wandb
 from bittensor.extrinsics.serving import get_metadata
+from dotenv import load_dotenv
+from wandb.apis.public.history import HistoryScan
 
+import competitions
 
-NETUID = 9
+# TODO: Update once registered
+NETUID = 179
 DELAY_SECS = 3
 RETRIES = 3
 
@@ -27,18 +30,22 @@ load_dotenv()
 
 WANDB_TOKEN = os.environ.get("WANDB_API_KEY", None)
 SUBTENSOR_ENDPOINT = os.environ.get("SUBTENSOR_ENDPOINT", None)
-VALIDATOR_WANDB_PROJECT = "opentensor-dev/pretraining-subnet"
-BENCHMARK_WANDB_PROJECT = "pretraining-benchmark-data"
+VALIDATOR_WANDB_PROJECT = "rusticluftig/finetuning"
+BENCHMARK_WANDB_PROJECT = ""
 BENCHMARK_FLAG = os.environ.get("BENCHMARK_FLAG", None)
 
-@dataclass
+
+@dataclass(frozen=True)
 class ModelData:
     uid: int
     hotkey: str
+    competition_id: int
     namespace: str
     name: str
     commit: str
-    hash: str
+
+    # Hash of (hash(model) + hotkey)
+    secure_hash: str
     block: int
     incentive: float
     emission: float
@@ -60,8 +67,9 @@ class ModelData:
             hotkey=hotkey,
             namespace=tokens[0],
             name=tokens[1],
-            commit=tokens[2] if tokens[2] != "None" else None,
-            hash=tokens[3] if tokens[3] != "None" else None,
+            commit=tokens[2],
+            secure_hash=tokens[3],
+            competition_id=int(tokens[4]),
             block=block,
             incentive=incentive,
             emission=emission,
@@ -69,6 +77,7 @@ class ModelData:
 
 
 def run_with_retries(func, *args, **kwargs):
+    """Runs a provided function with retries in the event of a failure."""
     for i in range(0, RETRIES):
         try:
             return func(*args, **kwargs)
@@ -81,12 +90,18 @@ def run_with_retries(func, *args, **kwargs):
 
 
 def get_subtensor_and_metagraph() -> Tuple[bt.subtensor, bt.metagraph]:
+    """Returns a subtensor and metagraph for the finetuning subnet."""
 
     def _internal() -> Tuple[bt.subtensor, bt.metagraph]:
         if SUBTENSOR_ENDPOINT:
             parser = argparse.ArgumentParser()
             bt.subtensor.add_args(parser)
-            subtensor = bt.subtensor(config=bt.config(parser=parser, args=["--subtensor.chain_endpoint", SUBTENSOR_ENDPOINT]))
+            subtensor = bt.subtensor(
+                config=bt.config(
+                    parser=parser,
+                    args=["--subtensor.chain_endpoint", SUBTENSOR_ENDPOINT],
+                )
+            )
         else:
             subtensor = bt.subtensor("finney")
 
@@ -138,9 +153,10 @@ def get_subnet_data(
 
 def get_wandb_runs(project: str, filters: Dict[str, Any]) -> List:
     """Get the latest runs from Wandb, retrying infinitely until we get them.
-    
+
     Returns:
-        List: List of runs matching the provided filters, newest run (by creation time) first."""
+        List: List of runs matching the provided filters, newest run (by creation time) first.
+    """
     while True:
         api = wandb.Api(api_key=WANDB_TOKEN)
         runs = list(
@@ -162,7 +178,7 @@ def get_scores(
     wandb_runs: List,
 ) -> Dict[int, Dict[str, Optional[float]]]:
     """Returns the most recent scores for the provided UIDs.
-    
+
     Args:
         uids (List[int]): List of UIDs to get scores for.
         wandb_runs (List): List of validator runs from Wandb. Requires the runs are provided in descending order.
@@ -194,6 +210,7 @@ def get_scores(
                     "win_rate": uid_data.get("win_rate", None),
                     "win_total": uid_data.get("win_total", None),
                     "weight": uid_data.get("weight", None),
+                    "competition_id": uid_data.get("competition_id", None),
                     "fresh": is_fresh,
                 }
         if len(result) == len(uids):
@@ -223,25 +240,48 @@ def get_validator_weights(
 def get_losses_over_time(wandb_runs: List) -> pd.DataFrame:
     """Returns a dataframe of the best average model loss over time."""
     timestamps = []
-    best_losses = []
-
+    datapoints_per_comp_id = {id: [] for id in competitions.COMPETITION_DETAILS}
+    
     for run in wandb_runs:
-        if "original_format_json" not in run.summary:
-            continue
-        data = json.loads(run.summary["original_format_json"])
-        all_uid_data = data["uid_data"]
-        timestamp = datetime.datetime.fromtimestamp(data["timestamp"])
-        best_loss = math.inf
-        for _, uid_data in all_uid_data.items():
-            loss = uid_data.get("average_loss", math.inf)
-            # Filter out the numbers from the exploit and when validators lost the best model.
-            if loss < best_loss and (loss > 2.5 or timestamp > datetime.datetime(2024,2,12)) and (loss < 5 or timestamp > datetime.datetime(2024,3,27)):
-                best_loss = uid_data["average_loss"]
-        if best_loss != math.inf:
-            timestamps.append(timestamp)
-            best_losses.append(best_loss)
-
-    return pd.DataFrame({"timestamp": timestamps, "best_loss": best_losses})
+        # For each run, check the 10 most recent steps.
+        best_loss_per_competition_id = defaultdict(lambda: math.inf)
+        should_add_datapoint = False
+        min_step = max(0, run.lastHistoryStep - 10)
+        history_scan = HistoryScan(
+            run.client, run, min_step, run.lastHistoryStep, page_size=10
+        )
+        max_timestamp = None
+        for step in history_scan:
+            if "original_format_json" not in step:
+                continue
+            data = json.loads(step["original_format_json"])
+            all_uid_data = data["uid_data"]
+            timestamp = datetime.datetime.fromtimestamp(data["timestamp"])
+            if max_timestamp is None:
+                max_timestamp = timestamp
+            max_timestamp = max(max_timestamp, timestamp)
+            
+            for _, uid_data in all_uid_data.items():
+                loss = uid_data.get("average_loss", math.inf)
+                competition_id = uid_data.get("competition_id", None)
+                if not competition_id:
+                    continue
+                
+                if loss < best_loss_per_competition_id[competition_id]:
+                    best_loss_per_competition_id[competition_id] = uid_data["average_loss"]
+                    should_add_datapoint = True
+        # Now that we've processed the run's most recent steps, check if we should add a datapoint.
+        if should_add_datapoint:
+            timestamps.append(max_timestamp)
+            # Iterate through all possible competitions and add the best loss for each.
+            # Set None for any that aren't active during this run.
+            for id, losses in datapoints_per_comp_id.items():
+                losses.append(best_loss_per_competition_id.get(id, None))
+
+    # Create a dictionary of competitions to lists of losses.
+    output_columns = {competitions.COMPETITION_DETAILS[id].name: losses for id, losses in datapoints_per_comp_id.items()}
+    
+    return pd.DataFrame({"timestamp": timestamps, **output_columns})
 
 
 def next_epoch(subtensor: bt.subtensor, block: int) -> int:
@@ -298,53 +338,67 @@ def get_benchmarks() -> Tuple[pd.DataFrame, datetime.datetime]:
         if artifacts:
             table = artifacts[-1].get("benchmarks")
             if table:
-                return table.get_dataframe(), datetime.datetime.strptime(run.metadata["startedAt"], "%Y-%m-%dT%H:%M:%S.%f")
+                return table.get_dataframe(), datetime.datetime.strptime(
+                    run.metadata["startedAt"], "%Y-%m-%dT%H:%M:%S.%f"
+                )
     bt.logging.error("Failed to get benchmarks from Wandb.")
     return None, None
 
 
-def make_validator_dataframe(validator_df: pd.DataFrame, model_data: ModelData) -> pd.DataFrame:
+def make_validator_dataframe(
+    validator_df: pd.DataFrame, model_data: ModelData
+) -> pd.DataFrame:
 
     values = [
-            [uid, int(validator_df[uid][1]), round(validator_df[uid][0], 4)]
-            + [
-                validator_df[uid][-1].get(c.uid)
-                for c in model_data
-                if c.incentive
-            ]
-            for uid, _ in sorted(
-                zip(
-                    validator_df.keys(),
-                    [validator_df[x][1] for x in validator_df.keys()],
-                ),
-                key=lambda x: x[1],
-                reverse=True,
-            )
-        ]
-    dtypes = {"UID":int, "Stake (τ)":float, "V-Trust":float}
-    dtypes.update({
+        [uid, int(validator_df[uid][1]), round(validator_df[uid][0], 4)]
+        + [validator_df[uid][-1].get(c.uid) for c in model_data if c.incentive]
+        for uid, _ in sorted(
+            zip(
+                validator_df.keys(),
+                [validator_df[x][1] for x in validator_df.keys()],
+            ),
+            key=lambda x: x[1],
+            reverse=True,
+        )
+    ]
+    dtypes = {"UID": int, "Stake (τ)": float, "V-Trust": float}
+    dtypes.update(
+        {
             f"{c.namespace}/{c.name} ({c.commit[0:8]})": float
             for c in model_data
             if c.incentive
-        })
+        }
+    )
     return pd.DataFrame(values, columns=dtypes.keys()).astype(dtypes)
 
+
 def make_metagraph_dataframe(metagraph: bt.metagraph, weights=False) -> pd.DataFrame:
 
-    cols = ['stake','emission','trust','validator_trust','dividends','incentive','R', 'consensus','validator_permit']
+    cols = [
+        "stake",
+        "emission",
+        "trust",
+        "validator_trust",
+        "dividends",
+        "incentive",
+        "R",
+        "consensus",
+        "validator_permit",
+    ]
 
     frame = pd.DataFrame({k: getattr(metagraph, k) for k in cols})
-    frame['block'] = metagraph.block.item()
-    frame['netuid'] = NETUID
-    frame['uid'] = range(len(frame))
-    frame['hotkey'] = [axon.hotkey for axon in metagraph.axons]
-    frame['coldkey'] = [axon.coldkey for axon in metagraph.axons]
+    frame["block"] = metagraph.block.item()
+    frame["netuid"] = NETUID
+    frame["uid"] = range(len(frame))
+    frame["hotkey"] = [axon.hotkey for axon in metagraph.axons]
+    frame["coldkey"] = [axon.coldkey for axon in metagraph.axons]
     if weights and metagraph.W is not None:
         # convert NxN tensor to a list of lists so it fits into the dataframe
-        frame['weights'] = [w.tolist() for w in metagraph.W]
+        frame["weights"] = [w.tolist() for w in metagraph.W]
 
     return frame
 
+
 def load_state_vars() -> dict[Any]:
     while True:
         try:
@@ -355,8 +409,12 @@ def load_state_vars() -> dict[Any]:
             model_data: List[ModelData] = get_subnet_data(subtensor, metagraph)
             model_data.sort(key=lambda x: x.incentive, reverse=True)
 
-            bt.logging.success(f'Loaded {len(model_data)} models')
-            vali_runs = get_wandb_runs(project=VALIDATOR_WANDB_PROJECT, filters={"config.type": "validator", "config.uid": 238})
+            bt.logging.success(f"Loaded {len(model_data)} models")
+            vali_runs = get_wandb_runs(
+                project=VALIDATOR_WANDB_PROJECT,
+                # TODO: Update to point to the OTF vali on finetuning
+                filters={"config.type": "validator", "config.uid": 0},
+            )
 
             scores = get_scores([x.uid for x in model_data], vali_runs)
 
@@ -385,40 +443,98 @@ def load_state_vars() -> dict[Any]:
             time.sleep(30)
 
     return {
-        'metagraph': metagraph,
+        "metagraph": metagraph,
         "model_data": model_data,
         "vali_runs": vali_runs,
         "scores": scores,
         "validator_df": validator_df,
         "benchmarks": benchmarks,
-        "benchmark_timestamp": benchmark_timestamp
+        "benchmark_timestamp": benchmark_timestamp,
     }
 
-def test_load_state_vars():
 
+def test_load_state_vars():
+    # TODO: Change to finetuning data.
     subtensor = bt.subtensor("finney")
     metagraph = subtensor.metagraph(NETUID, lite=True)
     model_data = [
-        ModelData(uid=253, hotkey='5DjoPAgZ54Zf6NsuiVYh8RjonnWWWREE2iXBNzM2VDBMQDPm', namespace='jw-hf-test', name='jw2', commit='aad131f6b02219964e6dcf749c2a23e75a7ceca8', hash='L1ImYzWJwV+9KSnZ2TYW0Iy2KMcVjJVTd30YJoRkpbw=', block=3131103, incentive=1.0, emission=209.06051635742188),
-        ModelData(uid=1, hotkey='5CccVtjk4yamCao6QYgEg7jc8vktdj16RbLKNUftHfEsjuJS', namespace='borggAI', name='bittensor-subnet9-models', commit='d373864bc6c972872edb8db95eed570958054bac', hash='+drdTIKYEGYClW2FFVVID6A2Dh//4rLmExRFCJsH6Y4=', block=2081837, incentive=0.0, emission=0.0),
-        ModelData(uid=2, hotkey='5HYwoXaczs3jAptbb5mk4aUCkgZqeNcNzJKxSec97GwasfLy', namespace='jungiebeen', name='pretrain1', commit='4c0c6bfd0f92e243d6c8a82209142e7204c852c3', hash='ld/agc0XIWICom/Cpj0fkQLcMogMNj/F65MJogK5RLY=', block=2467482, incentive=0.0, emission=0.0),
-        ModelData(uid=3, hotkey='5Dnb6edh9yTeEp5aasRPZVPRAkxvQ6qnERVcXw22awMZ5rxm', namespace='jungiebeen', name='pretrain2', commit='e827b7281c92224adb11124489cc45356553a87a', hash='ld/agc0XIWICom/Cpj0fkQLcMogMNj/F65MJogK5RLY=', block=2467497, incentive=0.0, emission=0.0),
-        ModelData(uid=4, hotkey='5FRfca8NbnH424WaX43PMhKBnbLA1bZpRRoXXiVs6HgsxN4K', namespace='ZainAli60', name='mine_modeles', commit='8a4ed4ad1f1fb58d424fd22e8e9874b87d32917c', hash='tVcbZAFoNIOF+Ntxq31OQ2NrLXf5iFCmmPUJlpkMYYo=', block=2508509, incentive=0.0, emission=0.0)
+        ModelData(
+            uid=253,
+            hotkey="5DjoPAgZ54Zf6NsuiVYh8RjonnWWWREE2iXBNzM2VDBMQDPm",
+            namespace="jw-hf-test",
+            name="jw2",
+            commit="aad131f6b02219964e6dcf749c2a23e75a7ceca8",
+            secure_hash="L1ImYzWJwV+9KSnZ2TYW0Iy2KMcVjJVTd30YJoRkpbw=",
+            block=3131103,
+            incentive=1.0,
+            emission=209.06051635742188,
+        ),
+        ModelData(
+            uid=1,
+            hotkey="5CccVtjk4yamCao6QYgEg7jc8vktdj16RbLKNUftHfEsjuJS",
+            namespace="borggAI",
+            name="bittensor-subnet9-models",
+            commit="d373864bc6c972872edb8db95eed570958054bac",
+            secure_hash="+drdTIKYEGYClW2FFVVID6A2Dh//4rLmExRFCJsH6Y4=",
+            block=2081837,
+            incentive=0.0,
+            emission=0.0,
+        ),
+        ModelData(
+            uid=2,
+            hotkey="5HYwoXaczs3jAptbb5mk4aUCkgZqeNcNzJKxSec97GwasfLy",
+            namespace="jungiebeen",
+            name="pretrain1",
+            commit="4c0c6bfd0f92e243d6c8a82209142e7204c852c3",
+            secure_hash="ld/agc0XIWICom/Cpj0fkQLcMogMNj/F65MJogK5RLY=",
+            block=2467482,
+            incentive=0.0,
+            emission=0.0,
+        ),
+        ModelData(
+            uid=3,
+            hotkey="5Dnb6edh9yTeEp5aasRPZVPRAkxvQ6qnERVcXw22awMZ5rxm",
+            namespace="jungiebeen",
+            name="pretrain2",
+            commit="e827b7281c92224adb11124489cc45356553a87a",
+            secure_hash="ld/agc0XIWICom/Cpj0fkQLcMogMNj/F65MJogK5RLY=",
+            block=2467497,
+            incentive=0.0,
+            emission=0.0,
+        ),
+        ModelData(
+            uid=4,
+            hotkey="5FRfca8NbnH424WaX43PMhKBnbLA1bZpRRoXXiVs6HgsxN4K",
+            namespace="ZainAli60",
+            name="mine_modeles",
+            commit="8a4ed4ad1f1fb58d424fd22e8e9874b87d32917c",
+            secure_hash="tVcbZAFoNIOF+Ntxq31OQ2NrLXf5iFCmmPUJlpkMYYo=",
+            block=2508509,
+            incentive=0.0,
+            emission=0.0,
+        ),
     ]
-    vali_runs = get_wandb_runs(project=VALIDATOR_WANDB_PROJECT, filters={"config.type": "validator", "config.uid": 238})
+    vali_runs = get_wandb_runs(
+        project=VALIDATOR_WANDB_PROJECT,
+        filters={"config.type": "validator", "config.uid": 238},
+    )
 
     scores = get_scores([x.uid for x in model_data], vali_runs)
 
     validator_df = {
         28: (1.0, 33273.4453125, {253: 1.0}),
-        49: (0.9127794504165649,
-        10401.677734375,
-        {7: 0.0867,
-        217: 0.0001,
-        219: 0.0001,
-        241: 0.0001,
-        248: 0.0001,
-        253: 0.9128}),
+        49: (
+            0.9127794504165649,
+            10401.677734375,
+            {
+                7: 0.0867,
+                217: 0.0001,
+                219: 0.0001,
+                241: 0.0001,
+                248: 0.0001,
+                253: 0.9128,
+            },
+        ),
         78: (1.0, 26730.37109375, {253: 1.0}),
         116: (1.0, 629248.4375, {253: 1.0}),
         150: (1.0, 272634.53125, {253: 1.0}),
@@ -438,11 +554,11 @@ def test_load_state_vars():
         249: (1.0, 478127.3125, {253: 1.0}),
         252: (1.0, 442395.03125, {253: 1.0}),
         254: (1.0, 46845.2109375, {253: 1.0}),
-        255: (1.0, 28977.56640625, {253: 1.0})
+        255: (1.0, 28977.56640625, {253: 1.0}),
     }
 
     return {
-        'metagraph': metagraph,
+        "metagraph": metagraph,
         "model_data": model_data,
         "vali_runs": vali_runs,
         "scores": scores,