Refactor barplot df generation code

app.py

@@ -9,27 +9,31 @@ DATASET_REPO = "uclanecl/NECL_GPUs"
 USER_BLACKLIST = ["gdm", "necl", "ncel", "root"]
 
 
-def split_multi_users(dfs):
-    df = dfs.copy()
-    df["users"] = df["users"].apply(lambda x: x.split(","))
+def get_plot_df(df):
+    df["users"] = df["users"].str.split(",")
+
+    # Clean users list with "Free" as the empty placeholder value
     df["users"] = df["users"].apply(
-        lambda user_list: (
-            [u for u in user_list if u.strip() and u.strip() not in USER_BLACKLIST] or ["Free"]
-        )
+        lambda user_list: [u for u in user_list if u not in USER_BLACKLIST] or ["Free"]
     )
-    df["count"] = 1
+
+    # Construct a new DataFrame for the barplot
     new_df = []
-    for row in df.to_dict(orient="records"):
+    for row in df.iterrows():
+        users_list = row["users"]
         gpu_users_num = len(row["users"])
-        for username in row["users"]:
+
+        for username in users_list:
             new_row = row.copy()
+            new_row.drop(labels="users")
+
             new_row["count"] = 1 / gpu_users_num
             new_row["username"] = username
+
             new_df.append(new_row)
+
     df = pd.DataFrame(new_df)
 
-    # Drop duplicate user columns
-    df = df.drop(columns=["users"])
     return df
 
 
@@ -57,16 +61,16 @@ def plot():
         if not df_list:
             return None, pd.DataFrame()
 
-        dfs = pd.concat(df_list, ignore_index=True)
-        dfs = dfs.fillna({"fan_speed": "N/A", "users": ""})
-        dfs = split_multi_users(dfs)
+        df = pd.concat(df_list, ignore_index=True)
+        df = df.fillna({"fan_speed": "N/A", "users": ""})
+        df = get_plot_df(df)
 
-        last_update_time = max(dfs["query_time"])
+        last_update_time = max(df["query_time"])
         last_update_time = datetime.fromisoformat(last_update_time).strftime("%Y-%m-%d %H:%M:%S")
 
         # Create the bar plot
         fig = px.bar(
-            dfs,
+            df,
             x="count",
             y="server",
             color="username",
@@ -74,27 +78,27 @@ def plot():
             color_discrete_map={
                 "Free": "black",
             },
-            text=dfs["index"].astype(str) + "<br>" + dfs["username"].astype(str),
+            text=df["index"].astype(str) + "<br>" + df["username"].astype(str),
         )
         # fig.update_layout(
         #     yaxis={"categoryorder": "array", "categoryarray": dfs["server"].unique()[::-1]},
         # )
         fig.update_traces(textposition="inside", insidetextanchor="middle")
 
-        return fig, dfs
+        return fig, df
 
     except Exception as e:
         print(f"Global error in plot: {e}")
         return None, pd.DataFrame()
 
 
-def plot_figs():
-    fig, dfn = plot()
-    return fig, dfn
+def generate_interface():
+    fig, df = plot()
+    return fig, df
 
 
 demo = gr.Interface(
-    fn=plot_figs,
+    fn=generate_interface,
     inputs=[],
     outputs=[
         gr.Plot(label="GPU Status", elem_classes="plotcss"),