src/radial/radial.py

import plotly.graph_objects as go
import random
from src.leaderboard.build_leaderboard import build_leadearboard_df
import numpy as np
import itertools as it


def create_plot(selected_models):
    models = build_leadearboard_df()
    metrics = ["musicmc", "lawmc", "moviesmc", "booksmc", "mmluproru"]
    MIN_COLOUR_DISTANCE_BETWEEN_MODELS = 100
    seed = 42
    
    def generate_colours(min_distance, seed):
        colour_mapping = {}
        all_models = selected_models
        
        for i in it.count():
            min_colour_distance = min_distance - i
            retries_left = 10 * len(all_models)
            
            for model_id in all_models:
                random.seed(hash(model_id) + i + seed)
                r, g, b = 0, 0, 0
                too_bright, similar_to_other_model = True, True
                
                while (too_bright or similar_to_other_model) and retries_left > 0:
                    r, g, b = tuple(random.randint(0, 255) for _ in range(3))
                    too_bright = np.min([r, g, b]) > 200
                    similar_to_other_model = any(
                        np.abs(np.array(colour) - np.array([r, g, b])).sum() < min_colour_distance
                        for colour in colour_mapping.values()
                    )
                    retries_left -= 1
                
                colour_mapping[model_id] = (r, g, b)
            if len(colour_mapping) == len(all_models):
                break
        
        return colour_mapping

    colour_mapping = generate_colours(MIN_COLOUR_DISTANCE_BETWEEN_MODELS, seed)
    fig = go.Figure()
    for _, model_data in models.iterrows():
        model_name = model_data["model"]
        if not model_name in selected_models:
            continue
        values = [model_data[metric] for metric in metrics]
        color = f'rgb{colour_mapping[model_name]}'

        fig.add_trace(go.Scatterpolar(
            r=values,
            theta=metrics,
            name=model_name,
            fill='toself',
            fillcolor=f'rgba{colour_mapping[model_name] + (0.6,)}',
            line=dict(color=color)
        ))
    fig.update_layout(
        polar=dict(radialaxis=dict(visible=True)),
        showlegend=True,
        title='Models metrics',
        template="plotly_dark",
    )
    return fig