Автопоиск признаков для ML

дегустационная заметка
    |
    v
38 ответов TypeSafe
    |-- 29 вопросов score x 2 столбца = 58
    |     ожидаемый уровень рубрики + неопределенность ответа
    `--  9 вопросов noul  x 1 столбец =  9
          вероятность true
    |
    v
67 числовых столбцов --> CatBoost --> предсказанная оценка критика
                                      RMSE на отложенной выборке: 1.77 балла

from __future__ import annotations

import json
import os
import random
import textwrap
import urllib.request
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from time import perf_counter
from typing import NamedTuple

import matplotlib
import matplotlib.pyplot as plt
import numpy as np
from catboost import CatBoostRegressor
from cooksafe import JsonCache, make_playground_link
from IPython.display import Markdown, display
from typesafe_sdk import Noul, NoulCriteria, Score, TypeSafeClient

matplotlib.use("Agg")  # headless render

TYPESAFE_MODEL = "jev-1.12"
FOLDS, REPEATS = 5, 3  # repeats steady the error at this sample size
CATBOOST = dict(
    iterations=400,
    depth=4,
    learning_rate=0.05,
    loss_function="RMSE",
    verbose=0,
    random_seed=0,
    thread_count=1,
    allow_writing_files=False,
)

client = TypeSafeClient(
    # keyless kernels replay the cache
    api_key=os.environ.get("TYPESAFE_API_KEY", "cache-only"),
    base_url=os.environ.get("TYPESAFE_ENDPOINT"),
    timeout=120.0,
)
json_cache = JsonCache(Path("json_cache.json"))

# ----------------------------------------------------------------- the specification

INTENSITY_LEVELS = [
    "Not present in this note at all",
    "Barely present - mentioned once, in passing",
    "Present at a moderate level",
    "Present strongly - the note dwells on it",
    "Dominant - the note is largely about this",
]
PRESENCE_CRITERIA = NoulCriteria(
    true="The note states this or clearly implies it",
    false="The note gives no indication of this",
)

# Asking for the score outright: ten quality bands, rescaled onto the 80-100 critic scale.
SCORE_LEVELS = [
    "Faulty or unpleasant - the note is mostly criticism",
    "Barely acceptable - drinkable, with nothing to recommend it",
    "Simple and sound - correct, plain, forgettable",
    "Pleasant everyday wine - some appeal, little depth",
    "Good - clear varietal character, well made",
    "Very good - balanced, with something to say",
    "Excellent - complex and structured",
    "Outstanding - depth and length, built to age",
    "Superb - among the best of its type",
    "Profound - the note treats it as exceptional",
]

# Structured output requires every property in `required`, so unused fields come back empty.
PROPOSAL_SCHEMA = {
    "type": "object",
    "properties": {
        "actions": {
            "type": "array",
            "items": {
                "type": "object",
                "properties": {
                    "op": {"type": "string", "enum": ["add", "revise", "drop"]},
                    "target": {"type": "string"},
                    "name": {"type": "string"},
                    "kind": {"type": "string", "enum": ["intensity", "presence"]},
                    "question": {"type": "string"},
                },
                "required": ["op", "target", "name", "kind", "question"],
                "additionalProperties": False,
            },
        },
    },
    "required": ["actions"],
    "additionalProperties": False,
}

PROPOSER_TASK = """You are designing numeric features for a gradient-boosting model that
predicts the score a wine critic gave (an integer from 80 to 100) from the tasting note alone.
The model sees nothing but the features you design.

Return up to 18 actions. Each action is one of:
- add a new question
- revise an existing question to separate notes better (set `target` to the existing name)
- drop a question that is not helping (set `target` to the existing name; `question` is ignored)

Two question kinds:
- `intensity`: a Score question answered on a 5-point scale (0 = not present at all,
  1 = barely present, 2 = moderate, 3 = strong, 4 = dominant). Use this for any quality
  that comes in degrees (balance, oakiness, tannin, complexity, length of finish, acidity,
  fruit intensity).
- `presence`: a Noul question answered with a probability (true = stated or clearly implied,
  false = no indication). Use this for facts that either hold or do not (mentions a specific
  flaw, single-vineyard designation, reserve bottling, old vines, specific blend).

Write the question so an intelligent reader can answer it from the note alone. Do not ask for
the score directly. Focus on what separates high-scoring wines (92+) from average ones (85-88)
and flawed or dilute ones (<84)."""

# ------------------------------------------------------------- data loading and splits


class Split(NamedTuple):
    notes: list[str]
    scores: np.ndarray
    dev: np.ndarray
    test: np.ndarray


def load_split(n_dev: int, n_test: int, seed: int = 0) -> Split:
    """Download the wine-mag sample once, split it into dev and held-out rows."""
    path = Path("winemag_sample.json")
    if not path.exists():
        url = (
            "https://raw.githubusercontent.com/zackthoutt/wine-deep-learning/master/data/"
            "winemag-data_first150k.json"
        )
        with urllib.request.urlopen(url) as r:
            records = json.loads(r.read().decode("utf-8"))
        rows = [
            {"note": row["description"].strip(), "score": float(row["points"])}
            for row in records
            if row.get("description") and row.get("points") is not None
        ]
        rng = random.Random(seed)
        path.write_text(json.dumps(rng.sample(rows, n_dev + n_test)))

    rows = json.loads(path.read_text())[: n_dev + n_test]
    notes = [r["note"] for r in rows]
    scores = np.array([r["score"] for r in rows], dtype=float)
    rng = np.random.default_rng(seed)
    indices = rng.permutation(len(rows))
    return Split(
        notes=notes,
        scores=scores,
        dev=indices[:n_dev],
        test=indices[n_dev : n_dev + n_test],
    )


# ------------------------------------------------------------------ question to column


def as_question(feature: dict) -> Score | Noul:
    """Turn a feature record into a TypeSafe question."""
    if feature["kind"] == "intensity":
        return Score(instructions=feature["question"], criteria=INTENSITY_LEVELS)
    return Noul(instructions=feature["question"], criteria=PRESENCE_CRITERIA)


def feature_questions(features: list[dict]) -> dict[str, Score | Noul]:
    return {f["name"]: as_question(f) for f in features}


@json_cache
def fetch_answers(model: str, notes: tuple[str, ...], feature: dict) -> list[dict]:
    """Score one feature across all notes. Batched across threads, one request per note."""
    name, question = feature["name"], as_question(feature)

    def ask(note: str) -> dict:
        r = client.system_one(state=note, questions={name: question}, model=model)
        ans = r.answers[name]
        if feature["kind"] == "intensity":
            return {"expected": ans.score, "uncertainty": ans.uncertainty}
        return {"noul": ans.noul}

    with ThreadPoolExecutor(max_workers=8) as pool:
        return list(pool.map(ask, notes))


@json_cache
def ask_score(model: str, note: str) -> dict:
    """The direct baseline: ask for the score itself across ten bands."""
    q = Score(
        instructions="What score band does this tasting note describe?",
        criteria=SCORE_LEVELS,
    )
    ans = client.system_one(state=note, questions={"score": q}, model=model).answers[
        "score"
    ]
    # level 0 -> 80, level 9 -> 100
    rescaled = 80.0 + (ans.score / 9.0) * 20.0
    return {"expected": rescaled, "uncertainty": ans.uncertainty}


def encode_column(answers: list[dict], kind: str, encoding: str) -> np.ndarray:
    """Turn a feature's answers into one or two numeric columns."""
    if kind == "presence":
        return np.array([[a["noul"]] for a in answers], dtype=float)
    expected = np.array([a["expected"] for a in answers], dtype=float)
    if encoding == "mean_only":
        return expected[:, None]
    uncertainty = np.array([a["uncertainty"] for a in answers], dtype=float)
    return np.column_stack([expected, uncertainty])


def column_names(features: list[dict], encoding: str) -> list[str]:
    names = []
    for f in features:
        if f["kind"] == "presence" or encoding == "mean_only":
            names.append(f["name"])
        else:
            names.extend([f["name"], f"{f['name']}_uncertainty"])
    return names


def design(
    features: list[dict], answers_for: dict[str, list[dict]], encoding: str
) -> tuple[np.ndarray, list[str]]:
    """Assemble the design matrix from accepted features."""
    cols = [
        encode_column(answers_for[f["name"]], f["kind"], encoding) for f in features
    ]
    return np.hstack(cols), column_names(features, encoding)


# ------------------------------------------------------------- modelling and evaluation


def kfold_splits(
    n: int, k: int, repeats: int, seed: int = 0
) -> list[tuple[np.ndarray, np.ndarray]]:
    """Repeated stratified-style random k-fold."""
    rng = np.random.default_rng(seed)
    splits = []
    for _ in range(repeats):
        indices = rng.permutation(n)
        for fold in np.array_split(indices, k):
            val_idx = fold
            train_mask = np.ones(n, dtype=bool)
            train_mask[val_idx] = False
            splits.append((np.where(train_mask)[0], val_idx))
    return splits


def cross_validate(
    X_dev: np.ndarray, y_dev: np.ndarray, folds: int = FOLDS, repeats: int = REPEATS
) -> tuple[float, np.ndarray]:
    """Out-of-fold predictions and mean RMSE across repeated k-fold splits."""
    splits = kfold_splits(len(y_dev), folds, repeats)
    oof_accum = np.zeros(len(y_dev), dtype=float)
    oof_counts = np.zeros(len(y_dev), dtype=int)
    scores = []

    for train_idx, val_idx in splits:
        model = CatBoostRegressor(**CATBOOST)
        model.fit(X_dev[train_idx], y_dev[train_idx])
        preds = model.predict(X_dev[val_idx])
        oof_accum[val_idx] += preds
        oof_counts[val_idx] += 1
        scores.append(np.sqrt(np.mean((y_dev[val_idx] - preds) ** 2)))

    return float(np.mean(scores)), oof_accum / oof_counts


def fit_predict(X: np.ndarray, split: Split) -> np.ndarray:
    """Train on dev rows, predict held-out rows."""
    model = CatBoostRegressor(**CATBOOST)
    model.fit(X[split.dev], split.scores[split.dev])
    return model.predict(X[split.test])


def fit_predict_text(split: Split) -> np.ndarray:
    """The word-count baseline: pass the raw text through CatBoost's text handling."""
    model = CatBoostRegressor(
        **CATBOOST,
        text_features=[0],
        tokenizers=[dict(tokenizer_id="Space", separator_type="BySense", lowercasing="True")],
        dictionaries=[
            dict(
                dictionary_id="BiGram",
                max_dictionary_size="50000",
                occurrence_lower_bound="3",
                gram_order="2",
            )
        ],
        feature_calcers=["BoW:top_tokens_count=1000"],
    )
    X_text = np.array(split.notes, dtype=object)[:, None]
    model.fit(X_text[split.dev], split.scores[split.dev])
    return model.predict(X_text[split.test])


def rmse(y_true: np.ndarray, y_pred: np.ndarray) -> float:
    return float(np.sqrt(np.mean((y_true - y_pred) ** 2)))


def spearman(y_true: np.ndarray, y_pred: np.ndarray) -> float:
    # rank correlation without pulling in scipy
    def rank(a: np.ndarray) -> np.ndarray:
        s = np.argsort(a)
        r = np.empty_like(s, dtype=float)
        r[s] = np.arange(len(a))
        return r

    rt, rp = rank(y_true), rank(y_pred)
    return float(np.corrcoef(rt, rp)[0, 1])


def importances(X_dev: np.ndarray, y_dev: np.ndarray) -> np.ndarray:
    model = CatBoostRegressor(**CATBOOST)
    model.fit(X_dev, y_dev)
    imp = np.array(model.get_feature_importance())
    total = imp.sum()
    return imp / total if total > 0 else imp


def importance_per_feature(
    features: list[dict], labels: list[str], col_importances: np.ndarray
) -> dict[str, float]:
    """Sum importance across the one or two columns a feature expanded into."""
    per_feature = {f["name"]: 0.0 for f in features}
    for label, imp in zip(labels, col_importances):
        base = label.removesuffix("_uncertainty")
        if base in per_feature:
            per_feature[base] += imp * 100.0
    return per_feature


def polarity(
    feature: dict, answers_for: dict[str, list[dict]], split: Split
) -> float:
    """Does this feature move with or against the critic score?

    Returns Pearson r between the feature's mean answer and the critic score on dev rows.
    Positive means higher feature -> higher score; negative means it flags lower quality."""
    answers = answers_for[feature["name"]]
    if feature["kind"] == "presence":
        vals = np.array([a["noul"] for a in answers])[split.dev]
    else:
        vals = np.array([a["expected"] for a in answers])[split.dev]
    return float(np.corrcoef(vals, split.scores[split.dev])[0, 1])


# ------------------------------------------------------------- the LLM proposer


def prompt_llm(model: str, system: str, user: str, schema: dict) -> str:
    """One LLM call under structured output."""
    if model.startswith("claude"):
        import anthropic

        c = anthropic.Anthropic(
            api_key=os.environ.get("ANTHROPIC_API_KEY", "cache-only")
        )
        r = c.messages.create(
            model=model,
            max_tokens=4096,
            system=system,
            messages=[{"role": "user", "content": user}],
            extra_body={"output_format": schema},
        )
        return r.content[0].text
    from openai import OpenAI

    o = OpenAI(api_key=os.environ.get("OPENAI_API_KEY", "cache-only"))
    r = o.chat.completions.create(
        model=model,
        messages=[{"role": "system", "content": system}, {"role": "user", "content": user}],
        response_format={
            "type": "json_schema",
            "json_schema": {"name": "actions", "strict": True, "schema": schema},
        },
    )
    return r.choices[0].message.content


@json_cache
def propose(
    model: str, round_idx: int, task: str, schema: dict, context: str
) -> dict:
    started = perf_counter()
    reply = prompt_llm(model, task, context, schema)
    return {"reply": reply, "seconds": round(perf_counter() - started, 2)}


# ------------------------------------------------------------- the autoresearch loop


def build_context(
    round_idx: int,
    accepted: list[dict],
    labels: list[str],
    col_importances: np.ndarray,
    oof_preds: np.ndarray,
    split: Split,
    n_examples: int,
) -> str:
    """The briefing for the proposer: what exists, what matters, and where the model fails."""
    lines = [f"=== ROUND {round_idx} ==="]
    if not accepted:
        lines.append(
            "No features designed yet. Propose an initial set of 15-18 diverse questions."
        )
    else:
        per_f = importance_per_feature(accepted, labels, col_importances)
        ranked = sorted(accepted, key=lambda f: -per_f[f["name"]])
        lines.append(f"Current features ({len(accepted)} active):")
        for f in ranked:
            lines.append(
                f"- {f['name']} ({f['kind']}, importance {per_f[f['name']]:.1f}%): \"{f['question']}\""
            )

    dev_err = np.abs(split.scores[split.dev] - oof_preds)
    worst = split.dev[np.argsort(-dev_err)[: n_examples // 2]]
    best = split.dev[np.argsort(dev_err)[: n_examples // 2]]

    lines.append(
        f"\nSample of {n_examples} training notes with target score and current prediction:"
    )
    for idx in list(worst) + list(best):
        note = " ".join(split.notes[idx].split())
        lines.append(
            f"[{split.scores[idx]:.0f} pts | pred {oof_preds[idx]:.1f}] {note[:200]}"
        )
    return "\n".join(lines)


class LoopResult(NamedTuple):
    accepted: list[dict]
    answers_for: dict[str, list[dict]]
    snapshots: list[list[dict]]
    history: list[dict]


def owner_of(col_name: str, features: list[dict]) -> dict:
    base = col_name.removesuffix("_uncertainty")
    return next(f for f in features if f["name"] == base)


def run_loop(
    split: Split,
    proposer_model: str,
    rounds: int,
    n_examples: int,
    encoding: str,
    min_spread: float,
    change_tolerance: float,
) -> LoopResult:
    accepted: list[dict] = []
    answers_for: dict[str, list[dict]] = {}
    snapshots: list[list[dict]] = []
    history: list[dict] = []

    for round_idx in range(1, rounds + 1):
        # 1. Summarise the state so the proposer can read its own predictions.
        if round_idx == 1:
            context = build_context(
                round_idx,
                [],
                [],
                np.array([]),
                np.full(len(split.notes), split.scores[split.dev].mean()),
                split,
                n_examples,
            )
        else:
            X_curr, labels = design(accepted, answers_for, encoding)
            col_imp = importances(X_curr[split.dev], split.scores[split.dev])
            _, oof_preds = cross_validate(X_curr[split.dev], split.scores[split.dev])
            context = build_context(
                round_idx,
                accepted,
                labels,
                col_imp,
                oof_preds,
                split,
                n_examples,
            )

        # 2. Ask the LLM for actions.
        raw = propose(proposer_model, round_idx, PROPOSER_TASK, PROPOSAL_SCHEMA, context)
        actions = json.loads(raw["reply"])["actions"]

        tally = {op: sum(1 for a in actions if a["op"] == op) for op in ("add", "revise", "drop")}
        print(
            f"round {round_idx}: {tally['add']} add, {tally['revise']} revise, {tally['drop']} drop"
        )

        # 3. Additions: fetch answers for every new question; keep it if it has any variance.
        added_this_round = []
        for a in actions:
            if a["op"] != "add":
                continue
            name = a["name"].strip().replace(" ", "_")
            feature = {"name": name, "kind": a["kind"], "question": a["question"].strip()}
            answers = fetch_answers(TYPESAFE_MODEL, tuple(split.notes), feature)
            answers_for[name] = answers
            # drop columns that returned a flat line
            col = encode_column(answers, feature["kind"], encoding)
            if float(np.std(col[:, 0])) < min_spread:
                continue
            accepted.append(feature)
            added_this_round.append(name)
        if added_this_round:
            wrapped = textwrap.fill(
                ", ".join(added_this_round),
                width=80,
                initial_indent="  added  ",
                subsequent_indent="          ",
            )
            print(wrapped)

        # 4. Score baseline after additions.
        X_curr, _ = design(accepted, answers_for, encoding)
        current_rmse, _ = cross_validate(X_curr[split.dev], split.scores[split.dev])

        # 5. Revisions and drops: keep only if dev error improves.
        for a in actions:
            if a["op"] == "revise":
                target = a["target"]
                curr = next((f for f in accepted if f["name"] == target), None)
                if not curr:
                    continue
                name = a["name"].strip().replace(" ", "_")
                revised = {"name": name, "kind": a["kind"], "question": a["question"].strip()}
                answers = fetch_answers(TYPESAFE_MODEL, tuple(split.notes), revised)
                answers_for[name] = answers
                candidate = [revised if f["name"] == target else f for f in accepted]
                X_cand, _ = design(candidate, answers_for, encoding)
                cand_rmse, _ = cross_validate(X_cand[split.dev], split.scores[split.dev])
                if cand_rmse < current_rmse - change_tolerance:
                    print(
                        f"  revise {name:<32} was {target}, CV {current_rmse:.3f} -> {cand_rmse:.3f}"
                    )
                    accepted = candidate
                    current_rmse = cand_rmse
                else:
                    print(
                        f"  reject {name:<32} was {target}, would cost +{cand_rmse - current_rmse:.3f}"
                    )

            elif a["op"] == "drop":
                target = a["target"]
                if not any(f["name"] == target for f in accepted):
                    continue
                candidate = [f for f in accepted if f["name"] != target]
                if not candidate:
                    continue
                X_cand, _ = design(candidate, answers_for, encoding)
                cand_rmse, _ = cross_validate(X_cand[split.dev], split.scores[split.dev])
                if cand_rmse < current_rmse - change_tolerance:
                    print(f"  drop   {target:<32} CV {current_rmse:.3f} -> {cand_rmse:.3f}")
                    accepted = candidate
                    current_rmse = cand_rmse
                else:
                    print(f"  keep   {target:<32} would cost +{cand_rmse - current_rmse:.3f}")

        snapshots.append(list(accepted))
        history.append({"round": round_idx, "rmse": current_rmse, "n_features": len(accepted)})
        print(f"  -> {len(accepted)} features, dev CV RMSE {current_rmse:.3f}\n")

    return LoopResult(
        accepted=accepted,
        answers_for=answers_for,
        snapshots=snapshots,
        history=history,
    )


# ----------------------------------------------------------------- chart styling


SURFACE, INK, INK2, MUTED = "#fcfcfb", "#0b0b0b", "#52514e", "#898781"
GRID, AXIS, BLUE, ORANGE = "#e1e0d9", "#c3c2b7", "#2a78d6", "#eb6834"
HEATMAP_HIGH, HEATMAP_LOW = "#4338ca", "#fbbf24"


def style_ax(ax):
    ax.set_facecolor(SURFACE)
    for side in ("top", "right"):
        ax.spines[side].set_visible(False)
    for side in ("left", "bottom"):
        ax.spines[side].set_color(AXIS)
    ax.tick_params(colors=MUTED, labelcolor=INK2, labelsize=9)
    ax.set_axisbelow(True)


def reviews_heatmap(plt, questions, answers_for, split, review_rows):
    """Plot features across five held-out reviews of increasing critic score."""
    q_names = [q["name"] for q in questions]
    grid = np.zeros((len(q_names), len(review_rows)), dtype=float)

    for i, q in enumerate(questions):
        answers = answers_for[q["name"]]
        if q["kind"] == "intensity":
            # intensity is 0-4; rescale to 0-1 so it matches nouls
            grid[i] = [answers[r]["expected"] / 4.0 for r in review_rows]
        else:
            grid[i] = [answers[r]["noul"] for r in review_rows]

    fig, ax = plt.subplots(figsize=(10.5, 9.5), facecolor=SURFACE)
    style_ax(ax)
    cmap = matplotlib.colors.LinearSegmentedColormap.from_list(
        "ts", [SURFACE, "#e0e7ff", HEATMAP_HIGH]
    )
    im = ax.imshow(grid, aspect="auto", cmap=cmap, vmin=0, vmax=1)

    ax.set_xticks(range(len(review_rows)))
    scores_at = [f"{split.scores[r]:.0f} pts" for r in review_rows]
    ax.set_xticklabels(
        [f"review {i + 1}\n{s}" for i, s in enumerate(scores_at)],
        fontsize=9.5,
        color=INK,
    )
    ax.xaxis.tick_top()
    ax.xaxis.set_label_position("top")

    # wrap question text so the y-labels stay readable
    ylabels = []
    for q in questions:
        wrapped = textwrap.shorten(q["question"], width=72, placeholder="...")
        tag = "[score]" if q["kind"] == "intensity" else "[noul] "
        ylabels.append(f"{tag} {wrapped}")

    ax.set_yticks(range(len(q_names)))
    ax.set_yticklabels(ylabels, fontsize=8.5, color=INK)

    # overlay numbers on each cell
    for i in range(len(q_names)):
        for j in range(len(review_rows)):
            val = grid[i, j]
            # show intensity on its original 0-4 scale, nouls as 0-1 probabilities
            displayed = (
                f"{val * 4:.1f}"
                if questions[i]["kind"] == "intensity"
                else f"{val:.2f}"
            )
            text_color = SURFACE if val > 0.65 else INK
            ax.text(
                j,
                i,
                displayed,
                ha="center",
                va="center",
                color=text_color,
                fontsize=8.5,
            )

    # find the flip point: the first question with negative correlation
    pols = [polarity(q, answers_for, split) for q in questions]
    div = next((i for i, p in enumerate(pols) if p < 0), None)
    if div is not None:
        ax.axhline(div - 0.5, color=ORANGE, linewidth=1.2, linestyle="--")
        ax.text(
            len(review_rows) - 0.45,
            div - 0.65,
            "correlates positively with score  ▲",
            ha="right",
            va="bottom",
            color=INK2,
            fontsize=8,
        )
        ax.text(
            len(review_rows) - 0.45,
            div - 0.35,
            "correlates negatively with score  ▼",
            ha="right",
            va="top",
            color=ORANGE,
            fontsize=8,
        )

    ax.set_title(
        "Fifteen discovered questions across five held-out reviews of increasing critic score",
        color=INK,
        fontsize=11,
        loc="left",
        pad=30,
    )
    fig.tight_layout()
    return fig


def paired_gain(
    y_test: np.ndarray, round1_preds: np.ndarray, round5_preds: np.ndarray
) -> tuple[float, float, float]:
    """Bootstrap CI for the error difference between round 1 and round 5 on held-out rows."""
    rng = np.random.default_rng(0)
    diffs = []
    n = len(y_test)
    for _ in range(2000):
        idx = rng.integers(0, n, size=n)
        r1 = np.sqrt(np.mean((y_test[idx] - round1_preds[idx]) ** 2))
        r5 = np.sqrt(np.mean((y_test[idx] - round5_preds[idx]) ** 2))
        diffs.append(r5 - r1)  # negative means round 5 improved
    return (
        float(np.mean(diffs)),
        float(np.percentile(diffs, 2.5)),
        float(np.percentile(diffs, 97.5)),
    )


def rounds_chart(plt, curve, history, n_heldout, gain):
    fig, ax = plt.subplots(figsize=(7.5, 4.0), facecolor=SURFACE)
    style_ax(ax)
    ax.grid(color=GRID, linewidth=0.8)

    rounds = [h["round"] for h in history]
    dev_rmse = [h["rmse"] for h in history]
    test_rmse = [c[1] for c in curve]

    ax.plot(
        rounds,
        dev_rmse,
        marker="o",
        color=BLUE,
        linewidth=1.8,
        linestyle="--",
        label="dev CV RMSE (what the loop steers by)",
    )
    ax.plot(
        rounds,
        test_rmse,
        marker="s",
        color=ORANGE,
        linewidth=1.8,
        label=f"held-out RMSE ({n_heldout} unseen reviews)",
    )

    for r, n, t in zip(rounds, [h["n_features"] for h in history], test_rmse):
        ax.annotate(
            f"{t:.3f}\n({n} features)",
            (r, t),
            textcoords="offset points",
            xytext=(0, -22),
            ha="center",
            fontsize=8,
            color=INK2,
        )

    ax.set_xticks(rounds)
    ax.set_xlabel("round", color=INK2, fontsize=9)
    ax.set_ylabel("RMSE (critic score points, lower is better)", color=INK2, fontsize=9)
    ax.set_title(
        "Prediction error by round: the loop vs the held-out rows",
        color=INK,
        fontsize=11,
        loc="left",
    )
    mean_gain, low_ci, high_ci = gain
    ax.text(
        0.02,
        0.30,
        f"round 1 -> 5 on held-out: {mean_gain:+.3f} pts\n"
        f"95% CI [{low_ci:+.3f}, {high_ci:+.3f}]",
        transform=ax.transAxes,
        color=MUTED,
        fontsize=9,
    )
    ax.legend(frameon=False, labelcolor=INK2, fontsize=9, loc="lower left")
    return fig

pip install anthropic openai catboost numpy matplotlib ipython "typesafe-sdk>=0.5.7" cooksafe --extra-index-url https://pypi.typesafe.ai/

N_DEV, N_TEST = 1200, 800  # the loop reads dev labels only; test is scored once
ROUNDS = 5  # a round answers questions for all 2,000 rows: 2,000 requests
PROPOSER = "claude-sonnet-5"  # or "gpt-5.6-luna"; the cache holds the Anthropic run
EXAMPLES = 60  # dev notes the proposer reads per round, half of them its worst misses
MIN_SPREAD = 0.05  # a column this flat cannot separate anything, so it is not kept
CHANGE_TOLERANCE = 0.0  # a revision or drop has to improve dev error, not just not hurt
ENCODING = "mean_spread"  # a score answer becomes two columns: its mean and spread

split = load_split(N_DEV, N_TEST, seed=0)
NOTES, SCORES, DEV, TEST = split.notes, split.scores, split.dev, split.test

print(
    f"{len(DEV)} dev rows, {len(TEST)} held out; scores run "
    f"{SCORES.min():.0f}-{SCORES.max():.0f}, mean {SCORES.mean():.2f}, sd {SCORES.std():.2f}"
)
print(f"\none of the notes:\n{NOTES[0]}")

1200 dev rows, 800 held out; scores run 80-98, mean 88.73, sd 3.17

one of the notes:
A Champagne that is very much wine. The structure and the richness are just right for a food wine, showing ripe acidity, flavors of plums and apricots, and balancing these primary fruits with a dense, complex structure that takes in yeast, maturity and a tight apple skin finish.

questions <- {}
repeat for each round:
    notes  <- round 1 ? 60 dev notes across the score range
                      : the 30 worst-predicted dev notes + the 30 best,
                        each with its score, this prediction and the last
    actions <- LLM(brief, questions, notes, importance and error so far)
    answers[q] <- TypeSafe(note, all new questions of this round) for every row
    for each added q:      keep it unless its column is flat
    for each revised q:    refit; keep the change only if dev error drops
    for each dropped q:    refit; drop it only if dev error drops
    out_of_fold <- k-fold CatBoost on the columns   # judges, and picks next round's notes

print("every intensity question is graded on these five levels:\n")
for i, level in enumerate(INTENSITY_LEVELS):
    print(f"  {i}. {level}")
print("\nevery presence question is judged true or false against these:\n")
print(f"  true:  {PRESENCE_CRITERIA['true']}")
print(f"  false: {PRESENCE_CRITERIA['false']}")
print("\nthe brief the proposer works from:\n")
print("\n".join(PROPOSER_TASK.splitlines()[:6]) + "\n  ...")

every intensity question is graded on these five levels:

  0. Not present in this note at all
  1. Barely present - mentioned once, in passing
  2. Present at a moderate level
  3. Present strongly - the note dwells on it
  4. Dominant - the note is largely about this

every presence question is judged true or false against these:

  true:  The note states this or clearly implies it
  false: The note gives no indication of this

the brief the proposer works from:

You are designing numeric features for a gradient-boosting model that
predicts the score a wine critic gave (an integer from 80 to 100) from the tasting note alone.
The model sees nothing but the features you design.

Return up to 18 actions. Each action is one of:

  ...

run = run_loop(
    split, PROPOSER, ROUNDS, EXAMPLES, ENCODING, MIN_SPREAD, CHANGE_TOLERANCE
)
accepted, answers_for = run.accepted, run.answers_for
snapshots, history = run.snapshots, run.history

round 1: 18 add, 0 revise, 0 drop
  added  complexity, fruit_intensity, tannin_structure, acidity_intensity,
          oak_intensity, finish_length, balance_harmony, aging_potential,
          positive_superlative_language, negative_critical_language,
          drinkability_easiness, body_richness, sweetness_level, texture_descriptors,
          earthy_savory_notes, flaw_or_defect_mentioned,
          single_vineyard_or_prestige_signal, varietal_blend_detail
  -> 18 features, dev CV RMSE 1.903

round 2: 5 add, 3 revise, 3 drop
  added  power_concentration_language, flavor_distinctiveness, generic_fruit_language,
          candied_artificial_flavor, rustic_authentic_character
  reject oak_dominance                     was oak_intensity, would cost +0.005
  revise negative_critical_language        was negative_critical_language, CV 1.897 -> 1.894
  revise single_vineyard_or_prestige_signalwas single_vineyard_or_prestige_signal, CV 1.894 -> 1.881
  keep   finish_length                     would cost +0.009
  keep   texture_descriptors               would cost +0.001
  keep   varietal_blend_detail             would cost +0.023
  -> 23 features, dev CV RMSE 1.881

round 3: 7 add, 2 revise, 1 drop
  added  elegance_finesse_language, minerality_precision_language,
          hedged_qualified_praise, underripe_green_character,
          reviewer_overall_verdict_strength, unusual_or_funky_descriptor_valence,
          botrytis_or_special_winemaking_signal
  revise negative_critical_language        was negative_critical_language, CV 1.868 -> 1.864
  revise finish_quality                    was finish_length, CV 1.864 -> 1.861
  keep   candied_artificial_flavor         would cost +0.014
  -> 30 features, dev CV RMSE 1.861

round 4: 5 add, 2 revise, 3 drop
  added  excess_or_imbalance_signal, descriptive_detail_density,
          critic_enthusiasm_confidence, savory_food_wine_seriousness,
          note_overall_tone_positivity
  revise rustic_authentic_character        was rustic_authentic_character, CV 1.843 -> 1.838
  reject hedged_qualified_praise           was hedged_qualified_praise, would cost +0.014
  keep   botrytis_or_special_winemaking_signalwould cost +0.011
  keep   candied_artificial_flavor         would cost +0.009
  keep   unusual_or_funky_descriptor_valencewould cost +0.010
  -> 35 features, dev CV RMSE 1.838

round 5: 4 add, 2 revise, 8 drop
  added  structural_seriousness, youthful_tension_signal, surface_prettiness_vs_depth,
          price_value_signal
  reject unconventional_character_as_virtuewas rustic_authentic_character, would cost +0.010
  revise flavor_distinctiveness            was flavor_distinctiveness, CV 1.849 -> 1.843
  keep   candied_artificial_flavor         would cost +0.002
  keep   botrytis_or_special_winemaking_signalwould cost +0.003
  keep   hedged_qualified_praise           would cost +0.006
  keep   excess_or_imbalance_signal        would cost +0.005
  drop   underripe_green_character         CV 1.843 -> 1.840
  keep   unusual_or_funky_descriptor_valencewould cost +0.002
  keep   texture_descriptors               would cost +0.001
  keep   generic_fruit_language            would cost +0.000
  -> 38 features, dev CV RMSE 1.840

X, labels = design(accepted, answers_for, ENCODING)
column_importances = importances(X[DEV], SCORES[DEV])
# an encoding gives a feature more than one column, so add a feature's columns back up
feature_importances = importance_per_feature(accepted, labels, column_importances)
ranked = sorted(accepted, key=lambda f: -feature_importances[f["name"]])
score_questions = [f for f in ranked if f["kind"] == "intensity"][:8]
noul_questions = [f for f in ranked if f["kind"] == "presence"][:7]
# ordered by which way the answer moves with the score, so the map flips halfway down
heatmap_questions = sorted(
    score_questions + noul_questions,
    key=lambda f: -polarity(f, answers_for, split),
)
ordered_test = TEST[np.argsort(SCORES[TEST], kind="stable")]
positions = np.linspace(0, len(ordered_test) - 1, 5).round().astype(int)
review_rows = tuple(ordered_test[positions])

print("the five held-out heatmap columns:\n")
for i, row in enumerate(review_rows, 1):
    excerpt = " ".join(NOTES[row].split())
    print(f"  {i}. {SCORES[row]:.0f} points: {excerpt[:100]}...")

fig = reviews_heatmap(plt, heatmap_questions, answers_for, split, review_rows)
display(fig)
plt.close(fig)

the five held-out heatmap columns:

  1. 80 points: Raw cherry and plum aromas are resiny and suggest wet cement. This is shearing and so jacked up with...
  2. 86 points: A slight spritz brightens the mouthfeel of this lemony wine. Aromas are a bit musky, but flavors of ...
  3. 89 points: This is a European-style Syrah, cofermented with 2% Viognier. It's soft and round, medium in body, a...
  4. 91 points: From the producer's dry-farmed estate vineyard, and supported by small amounts of Merlot and Caberne...
  5. 97 points: A thoroughly elegant, serious and yet immensely enjoyable wine that stays lively many days after ope...

predicted = fit_predict(X, split)
text_predicted = fit_predict_text(split)

# ask TypeSafe for the score itself, one request per row
with ThreadPoolExecutor(max_workers=8) as pool:
    direct = list(pool.map(lambda note: ask_score(TYPESAFE_MODEL, note), NOTES))
asked = np.array([d["expected"] for d in direct])
shift = float(SCORES[DEV].mean() - asked[DEV].mean())  # one number, from the dev labels

# what one proposal call gets you, before any feedback: the set round 1 ended with
first_round, _ = design(snapshots[0], answers_for, ENCODING)

print(f"{'arm':<46}{'RMSE':>7}{'spearman':>10}")
for label, p in (
    ("predict the mean of the dev rows", np.full(len(TEST), SCORES[DEV].mean())),
    ("the note as word counts, same CatBoost", text_predicted),
    (f"ask for the score itself, shifted {shift:+.2f}", asked[TEST] + shift),
    (
        f"{len(snapshots[0])} questions from round 1, no loop",
        fit_predict(first_round, split),
    ),
    (f"{len(accepted)} questions after all {ROUNDS} rounds", predicted),
):
    print(f"{label:<46}{rmse(SCORES[TEST], p):>7.3f}{spearman(SCORES[TEST], p):>10.3f}")

arm                                              RMSE  spearman
predict the mean of the dev rows                3.088    -0.014
the note as word counts, same CatBoost          2.466     0.605
ask for the score itself, shifted -1.71         2.145     0.761
18 questions from round 1, no loop              1.869     0.778
38 questions after all 5 rounds                 1.772     0.799

curve, per_round = [], []
for features in snapshots:
    X_round, _ = design(features, answers_for, ENCODING)
    per_round.append(fit_predict(X_round, split))
    curve.append((len(features), rmse(SCORES[TEST], per_round[-1])))

# the same held-out rows resampled 2,000 times, both arms scored on each resample
gain = paired_gain(SCORES[TEST], per_round[0], per_round[-1])
print(
    f"round 1 -> round {ROUNDS} on the held-out rows: {gain[0]:+.3f} points, "
    f"95% CI [{gain[1]:+.3f}, {gain[2]:+.3f}]"
)

fig = rounds_chart(plt, curve, history, len(TEST), gain)
display(fig)
plt.close(fig)

round 1 -> round 5 on the held-out rows: -0.097 points, 95% CI [-0.147, -0.050]

kinds = {f["name"]: f["kind"] for f in accepted}
print("feature importance share: % of total CatBoost importance across all questions")
print(f"{'feature':<38}{'asked as':<10}{'importance share':>16}")
for name, importance_share in sorted(feature_importances.items(), key=lambda p: -p[1])[
    :12
]:
    kind = "score" if kinds[name] == "intensity" else "noul"
    print(
        f"{name[:36]:<38}{kind:<10}{importance_share:>8.1f}%  "
        f"{'#' * round(importance_share)}"
    )
counts = f"{sum(1 for k in kinds.values() if k == 'intensity')} score"
counts += f", {sum(1 for k in kinds.values() if k == 'presence')} noul"
print(f"\nthe {len(accepted)} questions the loop kept: {counts}")
top = max(feature_importances, key=feature_importances.get)
print(
    f'the question behind the top row:\n  {top}: "{owner_of(top, accepted)["question"]}"'
)

feature importance share: % of total CatBoost importance across all questions
feature                               asked as  importance share
note_overall_tone_positivity          score         17.4%  #################
savory_food_wine_seriousness          score          8.7%  #########
positive_superlative_language         score          8.4%  ########
single_vineyard_or_prestige_signal    noul           7.2%  #######
descriptive_detail_density            score          5.7%  ######
elegance_finesse_language             score          5.0%  #####
complexity                            score          5.0%  #####
aging_potential                       score          5.0%  #####
balance_harmony                       score          2.9%  ###
drinkability_easiness                 score          2.9%  ###
critic_enthusiasm_confidence          score          2.7%  ###
flavor_distinctiveness                score          2.6%  ###

the 38 questions the loop kept: 29 score, 9 noul
the question behind the top row:
  note_overall_tone_positivity: "Setting aside specific descriptors, how positive is the overall emotional tone and word choice of the note taken as a whole (warm, admiring language throughout vs. flat, neutral, or lukewarm phrasing)?"

playground_link = make_playground_link(
    NOTES[0], feature_questions(accepted), models=[TYPESAFE_MODEL]
)
display(
    Markdown(
        f"🔗 [Open the note + questions in the TypeSafe playground]({playground_link})"
    )
)