benchmark.py

import wandb
import json
import os

if os.path.exists("results.json"):
    with open('results.json', 'r') as file:
        results = json.load(file)
else:
    results = {}

api = wandb.Api()
run = api.run('kf_ps/convergence/r6wp7fa8')

df = run.history()

for i in df.index:
    results["distributed/optimized-gpt2-250m"]["baseline"][(str(df.loc[i, "step"]))] = [df.loc[i, "train_loss"]]

with open("results.json", "w") as outfile:
    json.dump(results, outfile, indent = 4)