Function visualize

opencompass/summarizers/needlebench.py:256–365

(df_raw, save_path: str,model_name: str ,dataset_type:str)

    return model_scores

def visualize(df_raw, save_path: str,model_name: str ,dataset_type:str):
    df = df_raw.copy()
    if df.empty:
        return -1
    df['Context Length'] = df['dataset'].apply(
        lambda x: int(x.split('Length')[1].split('Depth')[0]))
    df['Document Depth'] = df['dataset'].apply(
        lambda x: float(x.split('Depth')[1].split('_')[0]))

    model_columns = [
        col for col in df.columns
        if col not in ['Context Length', 'Document Depth']
    ]

    for model_name in model_columns[1:]:
        model_df = df[['Document Depth', 'Context Length',
                        model_name]].copy()
        model_df.rename(columns={model_name: 'Score'}, inplace=True)
        pivot_table = pd.pivot_table(model_df,
                                    values='Score',
                                    index=['Document Depth'],
                                    columns=['Context Length'],
                                    aggfunc='mean')

        mean_scores = pivot_table.mean().values
        overall_score = mean_scores.mean()
        plt.figure(figsize=(10, 6))
        ax = plt.gca()
        cmap = LinearSegmentedColormap.from_list(
            'custom_cmap', ['#F0496E', '#EBB839', '#0CD79F'])

        sns.heatmap(pivot_table,
                    cmap=cmap,
                    ax=ax,
                    vmin=0,
                    vmax=100)
        cbar = ax.collections[0].colorbar
        x_data = [i + 0.5 for i in range(len(mean_scores))]
        y_data = mean_scores

        ax2 = ax.twinx()
        ax2.plot(x_data,
                y_data,
                color='white',
                marker='o',
                linestyle='-',
                linewidth=2,
                markersize=8,
                label='Average Depth Score'
                )
        for x_value, y_value in zip(x_data, y_data):
            ax2.text(x_value, y_value, f'{y_value:.2f}', ha='center', va='top')

        ax2.set_ylim(0, 100)

        ax2.set_yticklabels([])
        ax2.set_yticks([])