LLM-as-a-Judge - Iteratively Improve Models, AI Apps, and Agents Offline

import matplotlib.pyplot as plt
import seaborn as sns

# read the parquet file
df = pl.read_parquet('arxiv-metadata-id-categories-abstract-100000.snappy.parquet')
# cast categories to list by separating on space
df = df.with_columns(pl.col('categories').str.split(' ').alias('categories_list'))
# create high-level categories list by splitting on '.' for each item in categories_list
def get_high_level_categories(categories_list):
    return list(set([item.split('.')[0] for item in categories_list]))
df = df.with_columns(pl.col('categories_list').map_elements(get_high_level_categories).alias('high_level_categories_list'))
# explode the high_level_categories_list
df = df.explode('high_level_categories_list')
# aggregate by high_level_categories_list and count the number of papers
df = df.group_by('high_level_categories_list').agg(pl.count()).sort('count', descending=True)

# create a better looking bar chart of the high_level_categories_list
plt.figure(figsize=(12, 8))
sns.set_theme(style="whitegrid")

# Convert to pandas for seaborn
df_pandas = df.to_pandas()

# Create a more attractive bar plot
ax = sns.barplot(
    data=df_pandas, 
    x='high_level_categories_list', 
    y='count',
    palette='viridis',
    order=df_pandas.sort_values('count', ascending=False)['high_level_categories_list']
)

# Improve the styling
plt.xticks(rotation=45, ha='right')
plt.title('ArXiv High-Level Categories Distribution', fontsize=16, fontweight='bold', pad=20)
plt.xlabel('High-Level Categories', fontsize=12, fontweight='bold')
plt.ylabel('Number of Papers', fontsize=12, fontweight='bold')
# Improve layout
plt.tight_layout()

# Save with higher DPI for better quality
plt.savefig('high_level_categories_distribution.png', dpi=300, bbox_inches='tight')

... previous code ...

jobs = so.list_jobs()
for model, eval_models in eval_pairs.items():
    for eval_model in eval_models:
        job_name = eval_model + '_evals_' + model
        job = next((job for job in jobs if job['name'] == job_name), None)
        if job:
            job_id = job['job_id']
            results = so.get_job_results(job_id)
            if eval_model == 'gpt-oss-120b': # automatically unpacks final response to content field b/c it's a reasoning model, so we handle this one differently
                results = results.with_columns(pl.col('content').alias('score'))
                results = results.with_columns(pl.col('score').struct.field('score').alias(job_name))
                results = results.drop(['content', 'reasoning_content', 'score'])
            else:
                results = results.with_columns(pl.col('score').alias(job_name))
                # drop all columns except job_name
                results = results.drop([col for col in results.columns if col != job_name])
            df = pl.concat([df, results], how='horizontal')

import pandas as pd
df = pl.read_parquet('arxiv-metadata-id-categories-abstract-100000-explanations-evals.parquet')

i_vals = sorted(list(set([eval_model for eval_models in eval_pairs.values() for eval_model in eval_models])))
j_vals = sorted(list(eval_pairs.keys()))

matrix = np.zeros((len(i_vals), len(j_vals)))
for model, eval_models in eval_pairs.items():
    for eval_model in eval_models:
        job_name = eval_model + '_evals_' + model
        scores = df[job_name]
        matrix[i_vals.index(eval_model), j_vals.index(model)] = scores.mean()

# wrap in DataFrame with labels
df_matrix = pd.DataFrame(matrix, index=i_vals, columns=j_vals)

plt.figure(figsize=(10, 8))
sns.heatmap(df_matrix, annot=True, fmt=".2f", cmap="RdYlGn", cbar=True, square=True)
plt.title("Score Matrix")
plt.xlabel("ELI5 Model")
plt.ylabel("Evaluator Model")
plt.savefig("llm-as-a-judge-score-matrix.png", dpi=300, bbox_inches='tight')
plt.show()

df = pl.read_parquet('arxiv-metadata-id-categories-abstract-100000-explanations-evals.parquet')

relative_eval_jobs = {
    "qwen-3-32b": 'job-8f460128-6426-4078-b355-7790a00c2dfa',
    "gemma-3-27b-it": 'job-5a1eed4a-44b2-4d74-a6cd-d2424e0df113',
    "gpt-oss-120b": 'job-64b2a381-5f25-4917-b2c9-78140da3c475',
    "llama-3.3-70b": 'job-2860b900-31e5-4160-b8f1-4162767215c0',
}

jobs = so.list_jobs()
for eval_model, job_id in relative_eval_jobs.items():
    eval_job_name = eval_model + '_relative_evals'
    results = so.get_job_results(job_id)
    if eval_job_name == 'gpt-oss-120b_relative_evals': # automatically unpacks final response to content field b/c it's a reasoning model, so we handle this one differently
        results = results.with_columns(pl.col('content').struct.field('ranking').alias(eval_job_name))
        results = results.drop(['content', 'reasoning_content'])
    elif eval_job_name == 'llama-3.3-70b_relative_evals': # 2 of the results had problematic json formatting, so we handle this one differently
        results = results.with_columns(
            pl.when(pl.col("inference_result").str.contains('"C", "B", "A'))
            .then(pl.lit('{"ranking": ["C", "B", "A"]}'))
            .otherwise(pl.col("inference_result"))
            .alias("inference_result")
        )
        results = results.with_columns(pl.col('inference_result').str.json_decode().alias('inference_result'))
        results = results.with_columns(pl.col('inference_result').struct.field('ranking').alias(eval_job_name))
        results = results.drop(['inference_result'])
    else:
        results = results.drop([col for col in results.columns if col != 'ranking'])
        results = results.with_columns(pl.col('ranking').alias(eval_job_name))
        results = results.drop(['ranking'])
    eli5_map = {"A": inverse_eval_pairs[eval_model][0], "B": inverse_eval_pairs[eval_model][1], "C": inverse_eval_pairs[eval_model][2]}
    results = results.with_columns(pl.col(eval_job_name).map_elements(lambda x: [eli5_map[y] for y in x]).alias(eval_job_name))
    df = pl.concat([df, results], how='horizontal')

df.write_parquet('arxiv-metadata-id-categories-abstract-100000-explanations-relative-evals.parquet')

df = pl.read_parquet('arxiv-metadata-id-categories-abstract-100000-explanations-relative-evals.parquet')

eval_models = inverse_eval_pairs.keys()
eli5_models = list(set([model for eval_models in inverse_eval_pairs.values() for model in eval_models]))
eli5_model_matrix = {}
for i in eli5_models:
    for j in eli5_models:
        if i != j:
            eli5_model_matrix[(i, j)] = 0

for eval_model in eval_models:
    eval_job_name = eval_model + '_relative_evals'
    eval_results = df[eval_job_name].to_list()
    for eval_result in eval_results:
        if eval_result is not None:
            eli5_model_matrix[(eval_result[0], eval_result[1])] += 1
            eli5_model_matrix[(eval_result[1], eval_result[2])] += 1
            eli5_model_matrix[(eval_result[0], eval_result[2])] += 1

win_rate_pct_matrix = {}
for i in eli5_models:
    for j in eli5_models:
        if i != j:
            num = eli5_model_matrix[(i, j)]
            den = num + eli5_model_matrix[(j, i)]
            win_rate_pct_matrix[(i, j)] = num / den if den else float('nan')

df = pd.DataFrame(np.nan, index=eli5_models, columns=eli5_models, dtype=float)
for (i, j), v in win_rate_pct_matrix.items():
    df.loc[i, j] = v
np.fill_diagonal(df.values, np.nan)  # optional: blank diagonal

plt.figure(figsize=(10, 8))
sns.heatmap(df, annot=True, fmt=".2f", cmap="RdYlGn", cbar=True, square=True)
plt.title("Win Rate Percentage Matrix (cell = P(row beats column))")
plt.xlabel("ELI5 Model (column)")
plt.ylabel("ELI5 Model (row)")
plt.savefig("llm-as-a-judge-win-rate-percentage-matrix.png", dpi=300, bbox_inches='tight')
plt.tight_layout()
plt.show()

import math
import numpy as np
import pandas as pd

def bt_elo_from_pair_counts(
    pair_counts: dict[tuple[str, str], int],
    ties: dict[tuple[str, str], int] | None = None,
    laplace: float = 0.5,
    max_iter: int = 1000,
    tol: float = 1e-8,
    elo_mean: float = 1500.0,
):
    """
    pair_counts: { (winner, loser): wins } for all observed directed pairs.
    ties: optional { (a, b): tie_count } counted once per unordered pair (a,b).
          If provided, each tie contributes 0.5 win to both directions.
    laplace: additive smoothing to each *directed* count (prevents zeros).
    """

    # ---- Build model list ----
    models = sorted(set([k[0] for k in pair_counts] + [k[1] for k in pair_counts]))
    m = len(models)
    idx = {name: i for i, name in enumerate(models)}

    # ---- Build directed wins matrix W[i,j] = times i beat j ----
    W = np.zeros((m, m), dtype=float)
    for (w, l), c in pair_counts.items():
        if w == l: 
            continue
        W[idx[w], idx[l]] += float(c)

    # ---- Optional ties: add 0.5 to both directions for each tie ----
    if ties:
        for (a, b), t in ties.items():
            if a == b: 
                continue
            i, j = idx[a], idx[b]
            W[i, j] += 0.5 * t
            W[j, i] += 0.5 * t

    # ---- Laplace smoothing on directed edges (excluding diagonal) ----
    if laplace and laplace > 0:
        W += laplace
        np.fill_diagonal(W, 0.0)

    # Unordered totals N_ij = W_ij + W_ji
    N = W + W.T
    np.fill_diagonal(N, 0.0)

    # Guard: drop models with zero matches
    active = (N.sum(axis=1) > 0)
    if not np.all(active):
        keep = np.where(active)[0]
        models = [models[i] for i in keep]
        idx = {name: i for i, name in enumerate(models)}
        W = W[np.ix_(keep, keep)]
        N = N[np.ix_(keep, keep)]
        m = len(models)

    # ---- Bradley–Terry via MM updates (Hunter 2004) ----
    s = np.ones(m, dtype=float)  # abilities (positive)
    for _ in range(max_iter):
        s_old = s.copy()
        w_i = W.sum(axis=1)  # total (smoothed) wins per model
        # denom_i = sum_j N_ij / (s_i + s_j)
        denom = (N / (s.reshape(-1,1) + s.reshape(1,-1) + 1e-12)).sum(axis=1)
        upd = denom > 0
        s[upd] = w_i[upd] / denom[upd]
        # normalize to keep scale stable (geometric mean = 1)
        s /= np.prod(s) ** (1.0 / m)
        if np.max(np.abs(np.log(s + 1e-12) - np.log(s_old + 1e-12))) < tol:
            break

    # ---- Convert to beta and Elo-like ratings ----
    beta = np.log(s + 1e-12)
    elo = (400.0 / math.log(10.0)) * beta
    elo = elo - np.mean(elo) + elo_mean  # center

    # ---- Summaries and expected probabilities ----
    wins = W.sum(axis=1)
    losses = W.sum(axis=0)
    matches = N.sum(axis=1)  # unordered total vs all opponents

    ratings = pd.DataFrame({
        "ability": s,
        "beta": beta,
        "elo": elo,
        "wins": wins,
        "losses": losses,
        "matches": matches,
    }, index=models).sort_values("elo", ascending=False)

    P = s.reshape(-1,1) / (s.reshape(-1,1) + s.reshape(1,-1))
    np.fill_diagonal(P, np.nan)
    p_matrix = pd.DataFrame(P, index=models, columns=models)

    return ratings, p_matrix

ratings, p_exp = bt_elo_from_pair_counts(eli5_model_matrix, ties=None)
print(ratings[["elo","wins","losses","matches"]].to_markdown())