From 5ccdebf2c0a5e094949f0ca1da4be07ae006c6ff Mon Sep 17 00:00:00 2001 From: "Paul Gauthier (aider)" Date: Thu, 15 Aug 2024 09:50:50 -0700 Subject: [PATCH] refactor: Extract color assignment logic into a separate function --- benchmark/over_time.py | 20 ++++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/benchmark/over_time.py b/benchmark/over_time.py index 565038a8e..f72bac31e 100644 --- a/benchmark/over_time.py +++ b/benchmark/over_time.py @@ -6,6 +6,17 @@ from matplotlib import rc from aider.dump import dump # noqa: 401 +def get_model_color(model): + if "-4o" in model and "gpt-4o-mini" not in model: + return "purple" + elif "gpt-4" in model: + return "red" + elif "gpt-3.5" in model: + return "green" + else: + return "lightblue" + + def plot_over_time(yaml_file): with open(yaml_file, "r") as file: data = yaml.safe_load(file) @@ -49,14 +60,7 @@ def plot_over_time(yaml_file): spine.set_edgecolor("#DDDDDD") spine.set_linewidth(0.5) - colors = [ - ( - "purple" - if "-4o" in model and "gpt-4o-mini" not in model - else "red" if "gpt-4" in model else "green" if "gpt-3.5" in model else "lightblue" - ) - for model in models - ] + colors = [get_model_color(model) for model in models] # Separate data points by color purple_points = [(d, r) for d, r, c in zip(dates, pass_rates, colors) if c == "purple"]