Move filtering and generation of plots to its own file

src/data_processing.py

@@ -0,0 +1,129 @@
+import pandas as pd
+import plotly.express as px
+
+from config.constants import (
+    CC_BENCHMARKS,
+    LC_BENCHMARKS,
+    NON_RTL_METRICS,
+    RTL_METRICS,
+    S2R_BENCHMARKS,
+    SCATTER_PLOT_X_TICKS,
+    TYPE_COLORS,
+    Y_AXIS_LIMITS,
+)
+from utils import filter_bench, filter_bench_all, filter_RTLRepo, handle_special_cases
+
+
+def filter_leaderboard(task, benchmark, model_type, search_query, max_params, state):
+    """Filter leaderboard data based on user selections."""
+    subset = state.get_current_df().copy()
+
+    # Filter by task specific benchmarks when 'All' benchmarks is selected
+    if task == "Spec-to-RTL":
+        valid_benchmarks = S2R_BENCHMARKS
+        if benchmark == "All":
+            subset = subset[subset["Benchmark"].isin(valid_benchmarks)]
+    elif task == "Code Completion":
+        valid_benchmarks = CC_BENCHMARKS
+        if benchmark == "All":
+            subset = subset[subset["Benchmark"].isin(valid_benchmarks)]
+    elif task == "Line Completion †":
+        valid_benchmarks = LC_BENCHMARKS
+        if benchmark == "All":
+            subset = subset[subset["Benchmark"].isin(valid_benchmarks)]
+
+    if benchmark != "All":
+        subset = state.get_current_df()[state.get_current_df()["Benchmark"] == benchmark]
+
+    if model_type != "All":
+        # without emojis
+        subset = subset[subset["Model Type"] == model_type.split(" ")[0]]
+    if search_query:
+        subset = subset[subset["Model"].str.contains(search_query, case=False, na=False)]
+    max_params = float(max_params)
+    subset = subset[subset["Params"] <= max_params]
+
+    if benchmark == "All":
+        if task == "Spec-to-RTL":
+            return filter_bench_all(subset, state.get_current_agg(), agg_column="Agg S2R")
+        elif task == "Code Completion":
+            return filter_bench_all(subset, state.get_current_agg(), agg_column="Agg MC")
+        elif task == "Line Completion †":
+            return filter_RTLRepo(subset)
+    elif benchmark == "RTL-Repo":
+        return filter_RTLRepo(subset)
+    else:
+        agg_column = None
+        if benchmark == "VerilogEval S2R":
+            agg_column = "Agg VerilogEval S2R"
+        elif benchmark == "VerilogEval MC":
+            agg_column = "Agg VerilogEval MC"
+        elif benchmark == "RTLLM":
+            agg_column = "Agg RTLLM"
+        elif benchmark == "VeriGen":
+            agg_column = "Agg VeriGen"
+
+        return filter_bench(subset, state.get_current_agg(), agg_column)
+
+
+def generate_scatter_plot(benchmark, metric, state):
+    """Generate a scatter plot for the given benchmark and metric."""
+    benchmark, metric = handle_special_cases(benchmark, metric)
+
+    subset = state.get_current_df()[state.get_current_df()["Benchmark"] == benchmark]
+    if benchmark == "RTL-Repo":
+        subset = subset[subset["Metric"].str.contains("EM", case=False, na=False)]
+        detailed_scores = subset.groupby("Model", as_index=False)["Score"].mean()
+        detailed_scores.rename(columns={"Score": "Exact Matching (EM)"}, inplace=True)
+    else:
+        detailed_scores = subset.pivot_table(index="Model", columns="Metric", values="Score").reset_index()
+
+    details = state.get_current_df()[["Model", "Params", "Model Type"]].drop_duplicates("Model")
+    scatter_data = pd.merge(detailed_scores, details, on="Model", how="left").dropna(
+        subset=["Params", metric]
+    )
+
+    scatter_data["x"] = scatter_data["Params"]
+    scatter_data["y"] = scatter_data[metric]
+    scatter_data["size"] = (scatter_data["x"] ** 0.3) * 40
+
+    scatter_data["color"] = scatter_data["Model Type"].map(TYPE_COLORS).fillna("gray")
+
+    y_range = Y_AXIS_LIMITS.get(metric, [0, 80])
+
+    fig = px.scatter(
+        scatter_data,
+        x="x",
+        y="y",
+        log_x=True,
+        size="size",
+        color="Model Type",
+        text="Model",
+        hover_data={metric: ":.2f"},
+        title=f"Params vs. {metric} for {benchmark}",
+        labels={"x": "# Params (Log Scale)", "y": metric},
+        template="plotly_white",
+        height=600,
+        width=1200,
+    )
+
+    fig.update_traces(
+        textposition="top center",
+        textfont_size=10,
+        marker=dict(opacity=0.8, line=dict(width=0.5, color="black")),
+    )
+    fig.update_layout(
+        xaxis=dict(
+            showgrid=True,
+            type="log",
+            tickmode="array",
+            tickvals=SCATTER_PLOT_X_TICKS["tickvals"],
+            ticktext=SCATTER_PLOT_X_TICKS["ticktext"],
+        ),
+        showlegend=False,
+        yaxis=dict(range=y_range),
+        margin=dict(l=50, r=50, t=50, b=50),
+        plot_bgcolor="white",
+    )
+
+    return fig