Add baseline

app.py

@@ -9,15 +9,46 @@ with open('results.json', 'r') as file:
 models = [key for key in results.keys()]
 demo = gr.Blocks()
 
-df = pd.DataFrame.from_dict(results[models[0]], orient = "index").reset_index()
+
+from random import randint, random
+
+food_rating_data = pd.DataFrame(
+    {
+        "cuisine": [["Italian", "Mexican", "Chinese"][i % 3] for i in range(100)],
+        "rating": [random() * 4 + 0.5 * (i % 3) for i in range(100)],
+        "price": [randint(10, 50) + 4 * (i % 3) for i in range(100)],
+        "wait": [random() for i in range(100)],
+    }
+)
+
+df = pd.DataFrame.from_dict(results[models[0]]["main-net"], orient = "index").reset_index()
 df.columns = ["Step", "Loss"]
 df["Step"] = pd.to_numeric(df["Step"])
+df["Test"] = "Main-net"
+
+if "baseline" in results[models[0]]:
+    df_baseline = pd.DataFrame.from_dict(results[models[0]]["baseline"], orient = "index").reset_index()
+    df_baseline.columns = ["Step", "Loss"]
+    df_baseline["Step"] = pd.to_numeric(df_baseline["Step"])
+    df_baseline["Test"] = "Baseline"
+
+    df = pd.concat([df, df_baseline])
 
 def return_results(model_name):
     print(model_name)
-    df = pd.DataFrame.from_dict(results[model_name], orient = "index").reset_index()
+    df = pd.DataFrame.from_dict(results[model_name]["main-net"], orient = "index").reset_index()
     df.columns = ["Step", "Loss"]
     df["Step"] = pd.to_numeric(df["Step"])
+    df["Test"] = "Main-net"
+
+    if "baseline" in results[model_name]:
+        df_baseline = pd.DataFrame.from_dict(results[model_name]["baseline"], orient = "index").reset_index()
+        df_baseline.columns = ["Step", "Loss"]
+        df_baseline["Step"] = pd.to_numeric(df_baseline["Step"])
+        df_baseline["Test"] = "Baseline"
+
+        df = pd.concat([df, df_baseline])
+
     return df
 
 with demo:
@@ -27,7 +58,7 @@ with demo:
         dropdown_1 = gr.Dropdown(choices = models, value = models[0])
         button_1 = gr.Button("Submit")
     with gr.Row():
-        chart = gr.LinePlot(df, "Step", "Loss")
+        chart = gr.LinePlot(df, "Step", "Loss", color="Test", x_lim = (0, 2000))
     
     button_1.click(return_results, dropdown_1, chart)
 

evaluate.py

@@ -2,17 +2,21 @@ import json
 import random
 
 import torch
+import os
 from distributed_training.data.dataset import DataLoader
 from huggingface_hub import list_repo_refs
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 device = "cuda"
-test_indices_length = 10
+test_indices_length = 1000
 
 models = ["distributed/optimized-gpt2-250m", "distributed/optimized-gpt2-250m-v0.1.1", "distributed/gpt2-94m"]
 
-with open('results.json', 'r') as file:
-    results = json.load(file)
+if os.path.exists("results.json"):
+    with open('results.json', 'r') as file:
+        results = json.load(file)
+else:
+    results = {}
 
 for model_name in models:
 
@@ -24,7 +28,7 @@ for model_name in models:
     refs = list_repo_refs(model_name, repo_type="model")
     global_epoch = max([int(tag.name) for tag in refs.tags]) if refs.tags else None
 
-    for epoch in range(0, global_epoch):
+    for epoch in range(0,global_epoch, 5):
 
         if str(epoch) in results[model_name].keys():
             continue