Can a Language Model Pick Your Features?

import json
import os
import re
import warnings
from pathlib import Path

# keep the output to results: no download progress bars or library notices
os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"
os.environ["HF_DATASETS_DISABLE_PROGRESS_BARS"] = "1"
warnings.filterwarnings("ignore")

import numpy as np
import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
import plotly.io as pio
from datasets import load_dataset
from huggingface_hub.utils import logging as hf_logging
from sklearn.ensemble import HistGradientBoostingClassifier
from sklearn.feature_selection import mutual_info_classif
from sklearn.inspection import permutation_importance
from sklearn.metrics import roc_auc_score
from sklearn.model_selection import RepeatedStratifiedKFold, cross_val_score, train_test_split

hf_logging.set_verbosity_error()
RANDOM_STATE = 42

# language model settings (any OpenAI-compatible server)
LLM_BASE_URL = os.environ.get("LLM_BASE_URL", "http://127.0.0.1:8080/v1")
LLM_API_KEY = os.environ.get("LLM_API_KEY", "not-needed")
LLM_MODEL = os.environ.get("LLM_MODEL", "minicpm-v")
LLM_CACHE = Path("llm_feature_scores.json")
# reuse the saved scores so results are reproducible; set True to score again with your own model
REFRESH_LLM = os.environ.get("REFRESH_LLM") == "1"

# one look for every chart: dark, transparent, and colorblind-friendly
COLORS = {"Mutual information": "#38bdf8", "Permutation importance": "#a78bfa", "Language model": "#34d399", "Random": "#94a3b8"}
pio.templates["post"] = go.layout.Template(
    layout=dict(
        paper_bgcolor="rgba(0,0,0,0)",
        plot_bgcolor="rgba(0,0,0,0)",
        font=dict(family="Inter, system-ui, sans-serif", color="#cbd5e1", size=13),
        title=dict(font=dict(size=17, color="#f1f5f9"), x=0, xanchor="left"),
        xaxis=dict(gridcolor="rgba(148,163,184,0.15)", zerolinecolor="rgba(148,163,184,0.3)"),
        yaxis=dict(gridcolor="rgba(148,163,184,0.15)", zerolinecolor="rgba(148,163,184,0.3)"),
        colorway=list(COLORS.values()),
        margin=dict(l=10, r=10, t=60, b=10),
        legend=dict(bgcolor="rgba(0,0,0,0)"),
    )
)
pio.templates.default = "plotly_dark+post"
# show charts in JupyterLab and VS Code, and keep their data in the saved notebook
pio.renderers.default = "plotly_mimetype+notebook_connected"

raw = load_dataset("scikit-learn/adult-census-income", split="train").to_pandas()

df = raw.replace("?", np.nan)
CATEGORICAL = ["workclass", "education", "marital.status", "occupation", "relationship", "race", "sex", "native.country"]
NUMERIC = ["age", "fnlwgt", "education.num", "capital.gain", "capital.loss", "hours.per.week"]
FEATURES = [f for f in raw.columns if f != "income"]

for col in CATEGORICAL:
    df[col] = df[col].fillna("Unknown").astype("category")
y = (df["income"] == ">50K").astype(int)
X = df[FEATURES]

print(f"{len(df):,} people, {len(FEATURES)} features")
print(f"{y.mean():.1%} earn more than $50K")
X.head()

# count ages into 2-year bins per group, so the chart carries 70 numbers instead of 32,000 rows
bins = np.arange(16, 92, 2)
age_mix = pd.concat([
    pd.DataFrame({"age": bins[:-1], "upper": bins[:-1] + 1, "share": np.histogram(df.loc[y == flag, "age"], bins=bins)[0] / (y == flag).sum(), "Income": label})
    for flag, label in [(0, "$50K or less"), (1, "More than $50K")]
])
age_fig = px.bar(
    age_mix, x="age", y="share", color="Income", barmode="overlay", opacity=0.75, custom_data=["upper"],
    color_discrete_sequence=["#94a3b8", "#34d399"],
    title="Higher earners skew older, peaking in their 40s and 50s",
    labels={"age": "Age", "share": "Share of group"},
)
age_fig.update_traces(hovertemplate="Ages %{x} to %{customdata[0]}: %{y:.1%} of the group<extra></extra>")
age_fig.update_layout(yaxis_tickformat=".0%", legend_title_text="", bargap=0.05)
age_fig.show()

edu = (
    df.assign(high=y)
    .groupby("education", observed=True)
    .agg(rate=("high", "mean"), people=("high", "size"), level=("education.num", "first"))
    .sort_values("level")
    .reset_index()
)
edu_fig = px.bar(
    edu, x="education", y="rate", custom_data=["people"],
    title="The chance of earning over $50K climbs with every step of education",
    labels={"education": "", "rate": "Share earning over $50K"},
    color_discrete_sequence=["#34d399"],
)
edu_fig.update_traces(hovertemplate="%{x}<br>%{y:.0%} earn over $50K<br>%{customdata[0]:,} people<extra></extra>")
edu_fig.update_layout(yaxis_tickformat=".0%", xaxis_tickangle=-40)
edu_fig.show()

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, stratify=y, random_state=RANDOM_STATE)


def make_model():
    return HistGradientBoostingClassifier(categorical_features="from_dtype", random_state=RANDOM_STATE)


baseline = make_model().fit(X_train, y_train)
baseline_auc = roc_auc_score(y_test, baseline.predict_proba(X_test)[:, 1])
print(f"Test ROC AUC with all {len(FEATURES)} features: {baseline_auc:.3f}")

X_codes = X_train.copy()
for col in CATEGORICAL:
    X_codes[col] = X_codes[col].cat.codes

mi = pd.Series(
    mutual_info_classif(X_codes, y_train, discrete_features=[c in CATEGORICAL for c in FEATURES], random_state=RANDOM_STATE),
    index=FEATURES,
).sort_values()

mi_fig = px.bar(
    x=mi.values, y=mi.index, orientation="h",
    title="Mutual information: relationship and marital status lead",
    labels={"x": "Mutual information with income", "y": ""},
    color_discrete_sequence=[COLORS["Mutual information"]],
)
mi_fig.update_traces(hovertemplate="%{y}: %{x:.3f}<extra></extra>")
mi_fig.show()

perm = permutation_importance(
    baseline, X_test, y_test, scoring="roc_auc", n_repeats=10, random_state=RANDOM_STATE, n_jobs=-1
)
pi = pd.DataFrame({"mean": perm.importances_mean, "std": perm.importances_std}, index=FEATURES).sort_values("mean")

pi_fig = px.bar(
    pi, x="mean", y=pi.index, orientation="h", error_x="std",
    title="Permutation importance: what the trained model actually relies on",
    labels={"mean": "Drop in ROC AUC when shuffled", "index": ""},
    color_discrete_sequence=[COLORS["Permutation importance"]],
)
pi_fig.update_traces(hovertemplate="%{y}: %{x:.4f}<extra></extra>")
pi_fig.update_layout(yaxis_title="")
pi_fig.show()

DESCRIPTIONS = {
    "age": "Age in years",
    "workclass": "Employer type: private, self-employed, government, and so on",
    "fnlwgt": "Census sampling weight: how many people in the population this row represents",
    "education": "Highest level of education, as a label",
    "education.num": "The same education level, as a number from 1 to 16",
    "marital.status": "Married, never married, divorced, and so on",
    "occupation": "Broad job category",
    "relationship": "Role in the household: husband, wife, own child, and so on",
    "race": "Race as recorded by the census",
    "sex": "Sex as recorded by the census",
    "capital.gain": "Income from investment gains",
    "capital.loss": "Losses from investments",
    "hours.per.week": "Usual hours worked per week",
    "native.country": "Country of origin",
}

PROMPT = (
    "You are helping select features for a model that predicts whether a person's annual income is above $50,000, "
    "using records from the 1994 U.S. Census.\n\n"
    "Features:\n" + "\n".join(f"- {name}: {desc}" for name, desc in DESCRIPTIONS.items()) + "\n\n"
    "For each feature, give an importance score from 0 (useless for predicting income) to 1 (extremely useful), "
    "based on your knowledge of what drives income. Score features independently; several can share a score.\n"
    'Reply with only a JSON object mapping each feature name to {"score": number, "reason": "one short sentence"}.'
)
print(PROMPT)

N_RUNS = 5


def ask_llm(n_runs=N_RUNS):
    from openai import OpenAI

    client = OpenAI(base_url=LLM_BASE_URL, api_key=LLM_API_KEY, timeout=180)
    runs = []
    for i in range(n_runs):
        reply = client.chat.completions.create(
            model=LLM_MODEL,
            messages=[{"role": "user", "content": PROMPT}],
            temperature=0.7,
            max_tokens=1500,
        ).choices[0].message.content
        parsed = json.loads(re.search(r"\{.*\}", reply, re.S).group(0))  # tolerate text around the JSON
        runs.append({f: parsed[f] for f in FEATURES})
        print(f"run {i + 1}/{n_runs} done")
    return runs


CACHE_URL = "https://raw.githubusercontent.com/towardsuffering/rotunno.io/main/src/data/blog/en/llm-feature-selection/llm_feature_scores.json"


def load_cache():
    # opened from GitHub in Colab, the saved scores aren't next to the notebook, so fetch them
    if not LLM_CACHE.exists():
        import urllib.request

        LLM_CACHE.write_bytes(urllib.request.urlopen(CACHE_URL).read())
    return json.loads(LLM_CACHE.read_text())


if REFRESH_LLM or not (LLM_CACHE.exists() or os.environ.get("COLAB_RELEASE_TAG")):
    llm_runs = ask_llm()
    LLM_CACHE.write_text(json.dumps({"model": LLM_MODEL, "prompt": PROMPT, "runs": llm_runs}, indent=1))
    llm_model_used = LLM_MODEL
    print(f"Scored by {LLM_MODEL}; saved to {LLM_CACHE}")
else:
    cached = load_cache()
    llm_runs, llm_model_used = cached["runs"], cached["model"]
    print(f"Using saved scores from {llm_model_used} ({len(llm_runs)} runs)")

scores = pd.DataFrame([{f: float(run[f]["score"]) for f in FEATURES} for run in llm_runs])
llm = pd.DataFrame({"mean": scores.mean(), "std": scores.std()}).sort_values("mean")
reasons = {f: llm_runs[0][f]["reason"] for f in FEATURES}

llm_fig = px.bar(
    llm, x="mean", y=llm.index, orientation="h", error_x="std",
    custom_data=[[reasons[f] for f in llm.index]],
    title=f"Language model scores ({llm_model_used}, average of {len(llm_runs)} runs)",
    labels={"mean": "Importance score (0 to 1)", "index": ""},
    color_discrete_sequence=[COLORS["Language model"]],
)
llm_fig.update_traces(hovertemplate="<b>%{y}</b>: %{x:.2f}<br>%{customdata[0]}<extra></extra>")
llm_fig.update_layout(yaxis_title="", xaxis_range=[0, 1.05])
llm_fig.show()

pd.DataFrame({"score": llm["mean"].round(2), "reason": [reasons[f] for f in llm.index]}, index=llm.index).iloc[::-1]

rankings = pd.DataFrame({
    "Mutual information": mi.rank(ascending=False),
    "Permutation importance": pi["mean"].rank(ascending=False),
    "Language model": llm["mean"].rank(ascending=False, method="first"),
}).loc[FEATURES]

methods = list(rankings.columns)
bump = go.Figure()
for feature, row in rankings.iterrows():
    highlight = feature in ("fnlwgt", "education", "education.num", "capital.gain")
    bump.add_trace(go.Scatter(
        x=methods, y=row.values, mode="lines+markers+text", name=feature,
        text=["", "", feature], textposition="middle right",
        line=dict(width=3 if highlight else 1.5),
        opacity=1 if highlight else 0.55,
        hovertemplate=f"<b>{feature}</b><br>%{{x}}: rank %{{y:.0f}}<extra></extra>",
    ))
bump.update_layout(
    title="How each feature ranks under each method (1 = most important)",
    yaxis=dict(autorange="reversed", title="Rank", dtick=1),
    xaxis=dict(title=""), showlegend=False, height=560, margin=dict(r=120),
)
bump.show()

corr = rankings.corr(method="spearman")
corr_fig = px.imshow(
    corr, text_auto=".2f", zmin=-1, zmax=1, color_continuous_scale="Tealgrn",
    title="Rank agreement between methods (Spearman correlation)",
)
corr_fig.update_layout(coloraxis_showscale=False, height=420)
corr_fig.show()

# grade every model with 5-fold cross-validation, repeated twice, so one lucky split can't decide the result
cv = RepeatedStratifiedKFold(n_splits=5, n_repeats=2, random_state=RANDOM_STATE)


def auc_with(features):
    return cross_val_score(make_model(), X_train[features], y_train, cv=cv, scoring="roc_auc", n_jobs=-1).mean()


orders = {method: list(rankings[method].sort_values().index) for method in methods}
ks = range(1, len(FEATURES) + 1)
curve_df = pd.DataFrame({method: [auc_with(order[:k]) for k in ks] for method, order in orders.items()}, index=list(ks))

# 30 random orders show how much luck alone can do
N_RANDOM = 30
rng = np.random.default_rng(RANDOM_STATE)
random_orders = [list(rng.permutation(FEATURES)) for _ in range(N_RANDOM)]
random_curves = np.array([[auc_with(order[:k]) for k in ks] for order in random_orders])
low, mid, high = np.percentile(random_curves, [10, 50, 90], axis=0)
full_auc = curve_df.iloc[-1, 0]

curve_fig = go.Figure()
curve_fig.add_trace(go.Scatter(
    x=list(ks) + list(ks)[::-1], y=list(high) + list(low)[::-1], fill="toself",
    fillcolor="rgba(148,163,184,0.18)", line=dict(width=0), hoverinfo="skip",
    name=f"Random: middle 80% of {N_RANDOM} orders",
))
curve_fig.add_trace(go.Scatter(
    x=list(ks), y=mid, mode="lines", name="Random: median",
    line=dict(color=COLORS["Random"], width=2, dash="dot"),
    hovertemplate="Random median<br>top %{x} features: AUC %{y:.3f}<extra></extra>",
))
for method in methods:
    curve_fig.add_trace(go.Scatter(
        x=curve_df.index, y=curve_df[method], mode="lines+markers", name=method,
        line=dict(color=COLORS[method], width=3 if method == "Language model" else 2),
        hovertemplate=f"{method}<br>top %{{x}} features: AUC %{{y:.3f}}<extra></extra>",
    ))
curve_fig.add_hline(y=full_auc, line_dash="dash", line_color="rgba(148,163,184,0.5)", annotation_text="all features")
curve_fig.update_layout(
    title="Model accuracy using only each method's top k features",
    xaxis=dict(title="Number of features used (k)", dtick=1),
    yaxis=dict(title="Cross-validated ROC AUC"),
    legend=dict(orientation="h", y=-0.2), height=520,
)
curve_fig.show()

# two ways to score a whole curve: how few features reach near-full accuracy, and the average AUC across every k
target = full_auc - 0.01


def needed(curve):
    return int(np.argmax(np.asarray(curve) >= target)) + 1


random_needed = np.array([needed(c) for c in random_curves])
random_avg = random_curves.mean(axis=1)

scoreboard = pd.DataFrame({
    "Features needed": {m: needed(curve_df[m]) for m in methods},
    "Average AUC (k = 1 to 14)": curve_df[methods].mean().round(3),
    "Random orders it beats": {m: f"{(random_avg < curve_df[m].mean()).mean():.0%}" for m in methods},
})
scoreboard.loc[f"Random (median of {N_RANDOM})"] = [int(np.median(random_needed)), round(float(np.median(random_avg)), 3), ""]
print(f"Full model AUC: {full_auc:.3f}. 'Features needed' = fewest top features within 0.01 of it.")
print(f"Random orders needed anywhere from {random_needed.min()} to {random_needed.max()} features.")
scoreboard