Upload viz.py

utils/viz.py

@@ -5,7 +5,6 @@ import plotly.graph_objects as go
 import textwrap
 import numpy as np
 import pandas as pd
-from scipy.stats import chi2
 from scipy import stats
 
 
@@ -42,6 +41,8 @@ def create_task_success_heatmap(df, benchmark_name):
     tasks_solved = (pivot_df.sum(axis=0) > 0).astype(int)
     # Total number of tasks (columns)
     total_tasks = len(pivot_df.columns)
+    if 'SWE-bench' in benchmark_name:
+        total_tasks = 50 # TODO - remove hardcoding
     
     # Add the new row to the pivot table
     tasks_solved_df = pd.DataFrame(tasks_solved).T