def prompt_graph_predict(emb, support_idx, support_lbl, query_idx, n_way):
    """Explicit PRODIGY-style task graph with a parameter-free task GNN (mean aggregation, cosine readout)."""
    # task-graph edges: data node i -> label node support_lbl[i]
    src = torch.arange(len(support_idx))
    dst = torch.as_tensor(support_lbl)
    label_nodes = torch.zeros(n_way, emb.size(1)).index_add_(
        0, dst, emb[support_idx][src]
    )
    label_nodes = F.normalize(
        label_nodes / torch.bincount(dst, minlength=n_way).unsqueeze(1),
        dim=1,
    )
    return (emb[query_idx] @ label_nodes.T).argmax(
        1
    )  # query–label "link" scores

def run_episodes(
    emb, n_way=3, k_shot=3, n_query=20, n_episodes=300, seed=SEED
):
    r = np.random.default_rng(seed)
    y = cora.y.numpy()
    accs = []
    for _ in range(n_episodes):
        classes = r.choice(7, n_way, replace=False)
        sup, sup_l, qry, qry_l = [], [], [], []
        for j, c in enumerate(classes):
            idx = r.choice(
                np.flatnonzero(y == c), k_shot + n_query, replace=False
            )
            sup += list(idx[:k_shot])
            sup_l += [j] * k_shot
            qry += list(idx[k_shot:])
            qry_l += [j] * n_query
        pred = prompt_graph_predict(
            emb, torch.tensor(sup), sup_l, torch.tensor(qry), n_way
        )
        accs.append((pred.numpy() == np.array(qry_l)).mean())
    return float(np.mean(accs))

icl_rows = []
for _k in [1, 3, 5, 10]:
    icl_rows.append(
        {
            "k-shot (3-way)": _k,
            "raw bag-of-words features": run_episodes(raw_feat, k_shot=_k),
            "pretrained GNN embeddings (ICL)": run_episodes(
                node_emb, k_shot=_k
            ),
        }
    )
icl_results = pd.DataFrame(icl_rows)
icl_results.round(3)

_rev = reveal(
    r"""
Measured (300 random 3-way episodes, 20 queries per class): with **one** example per class the label-free GNN embeddings
already reach **0.70** (chance 0.33), and **0.85** at 10 shots, against 0.48 → 0.70 for the raw features. No weights were updated
for any task, and no labels were used in pretraining. PRODIGY goes further: it *learns* the task-graph GNN during pretraining on
self-supervised prompt tasks (neighbour matching, multi-task), and transfers to graphs and tasks it never saw (up to +48% over
contrastive pretraining). Ours is the parameter-free special case: the label node is the mean of its examples, i.e. a prototype.

""",
    predict="With one labelled example per class and no training at all, how far above chance (0.33) can prototypes on label-free GNN embeddings get?",
)
_post = mo.md(r"""
## 2 · Relational Transformer: cells as tokens

RT tokenises **every cell** of every table as (value, column name, table name), so any schema lands in one input space.
A task is a table whose label cell is **masked**; earlier rows of the task table act as in-context examples. Attention is
restricted by **relational masks**: *column* (same column), *feature* (same row + its FK-parent rows), *neighbour* (rows linked
by a foreign key) and *full*. A tiny database:
""")
mo.vstack([_rev, _post])

tiny_db = {
    "users": pd.DataFrame(
        {"user_id": [0, 1], "age": [34, 51], "country": ["IR", "US"]}
    ),
    "products": pd.DataFrame(
        {
            "product_id": [0, 1],
            "category": ["books", "toys"],
            "price": [12.0, 30.0],
        }
    ),
    "orders": pd.DataFrame(
        {
            "order_id": [0, 1, 2],
            "user_id": [0, 0, 1],
            "product_id": [1, 0, 1],
            "amount": [30.0, 24.0, 60.0],
        }
    ),
    "task_churn": pd.DataFrame(
        {"user_id": [0, 1], "churn": [0, None]}
    ),  # label of user 1 is the [MASK] to predict
}
FOREIGN_KEYS = {
    ("orders", "user_id"): "users",
    ("orders", "product_id"): "products",
    ("task_churn", "user_id"): "users",
}

def cell_tokens(db):
    toks = []
    for table, df in db.items():
        for i, row in df.iterrows():
            for col, val in row.items():
                if (
                    col.endswith("_id")
                    and not (table, col) in FOREIGN_KEYS
                    and col == f"{table[:-1]}_id"
                ):
                    continue  # primary keys are not content tokens
                if (table, col) in FOREIGN_KEYS:
                    continue  # FKs become links, not tokens
                toks.append(
                    {
                        "table": table,
                        "row": i,
                        "column": col,
                        "value": "[MASK]" if pd.isna(val) else val,
                    }
                )
    return pd.DataFrame(toks)

rt_tokens = cell_tokens(tiny_db)
rt_tokens

class Classifier(nn.Module):
    def __init__(self, use_graph, d=64):
        super().__init__()
        self.use_graph = use_graph
        self.l1 = (
            GCNConv(cora.num_features, d)
            if use_graph
            else nn.Linear(cora.num_features, d)
        )
        self.l2 = GCNConv(d, 7) if use_graph else nn.Linear(d, 7)

    def forward(self, x, ei):
        h = F.relu(self.l1(x, ei) if self.use_graph else self.l1(x))
        h = F.dropout(h, 0.5, self.training)
        return self.l2(h, ei) if self.use_graph else self.l2(h)

def test_probabilities(use_graph, seed=SEED):
    torch.manual_seed(seed)
    m = Classifier(use_graph)
    opt = torch.optim.Adam(m.parameters(), lr=0.01, weight_decay=5e-4)
    for _ in range(200):
        m.train()
        opt.zero_grad()
        F.cross_entropy(
            m(cora.x, cora.edge_index)[cora.train_mask],
            cora.y[cora.train_mask],
        ).backward()
        opt.step()
    m.eval()
    with torch.no_grad():
        return F.softmax(m(cora.x, cora.edge_index), 1)[
            cora.test_mask
        ].numpy(), cora.y[cora.test_mask].numpy()

conformal_probs = {
    "GCN": test_probabilities(True),
    "MLP": test_probabilities(False),
}

def split_conformal(probs, y, alpha, n_splits=200, seed=SEED):
    r = np.random.default_rng(seed)
    n = len(y)
    cover, size = [], []
    for _ in range(n_splits):
        perm = r.permutation(n)
        cal, ev = perm[: n // 2], perm[n // 2 :]
        scores = 1 - probs[cal, y[cal]]
        level = min(1.0, np.ceil((len(cal) + 1) * (1 - alpha)) / len(cal))
        q = np.quantile(scores, level, method="higher")
        sets = (1 - probs[ev]) <= q
        cover.append(sets[np.arange(len(ev)), y[ev]].mean())
        size.append(sets.sum(1).mean())
    return float(np.mean(cover)), float(np.mean(size))

conf_alpha = mo.ui.slider(
    0.01,
    0.3,
    step=0.01,
    value=0.1,
    label="α (target miscoverage)",
    show_value=True,
)
mo.vstack(
    [
        mo.md(
            f"Test accuracy: GCN {(conformal_probs['GCN'][0].argmax(1) == conformal_probs['GCN'][1]).mean():.3f}, "
            f"MLP {(conformal_probs['MLP'][0].argmax(1) == conformal_probs['MLP'][1]).mean():.3f}. "
            "**Try it:** lower α demands more coverage, so sets grow."
        ),
        conf_alpha,
    ]
)

_rev = reveal(
    r"""
Measured (200 random calibration/evaluation splits): both models hit the target **exactly** (0.90 → 0.902 empirical), as the theory
promises, *whatever their accuracy*. The difference is **efficiency**: at 90% coverage the GCN's sets contain ≈ **1.5** labels on
average, the MLP's ≈ **3.6** of 7. At 95% coverage: 2.3 vs 4.5; at 80%: 1.0 vs 2.4. The guarantee is free; a good model makes it
*useful*. CF-GNN adds a second, topology-aware GNN that shrinks the sets further using neighbours' uncertainty.

""",
    predict="Both models are calibrated for 90% coverage. Will their empirical coverage differ? Will their prediction-set sizes?",
)
_post = mo.md(r"""
---

## 4 · Exercises

**Exercise 1 (conformal quantile).** 9 calibration scores: 0.05, 0.1, 0.12, 0.2, 0.3, 0.35, 0.5, 0.7, 0.9; $\alpha=0.2$.
$\lceil(n+1)(1-\alpha)\rceil = 8$, so $\hat q$ is the 8th smallest score. Which classes enter the set for a node with
$\hat p=(0.6, 0.25, 0.1, 0.05)$? → `ex1_set` (a Python set of class indices).

**Exercise 2 (ICL by hand).** Label-node embeddings (unit vectors) $\ell_A=(1,0)$, $\ell_B=(0,1)$, $\ell_C=(-0.6,0.8)$;
query $(0.8,0.6)$. Which label does the prompt graph predict? → `ex2_label`.

**Exercise 3 (code).** Write `aps_scores(probs, y)`: the **adaptive prediction set** score, i.e. the total probability mass
of all classes ranked *at or above* the true class (sort probabilities descending, cumulative sum up to and including $y$).
""")
mo.vstack([_rev, _post])

ex1_set = None
ex2_label = None

def aps_scores(probs, y):
    return None

def _verdict(ok, answered):
    if not answered:
        return mo.callout(mo.md("Not answered yet."), kind="neutral")
    return mo.callout(
        mo.md("Correct!" if ok else "Not yet — try again or open a hint."),
        kind="success" if ok else "warn",
    )

def _check_aps():
    try:
        p = np.array([[0.6, 0.25, 0.1, 0.05], [0.1, 0.2, 0.3, 0.4]])
        y = np.array([1, 3])
        got = aps_scores(p, y)
        if got is None:
            return _verdict(False, False)
        return _verdict(np.allclose(got, [0.85, 0.4]), True)
    except Exception as e:
        return mo.callout(mo.md(f"Error: `{e}`"), kind="danger")

mo.vstack(
    [
        mo.md("**Exercise 1**"),
        _verdict(ex1_set == {0}, ex1_set is not None),
        mo.md("**Exercise 2**"),
        _verdict(ex2_label == "A", ex2_label is not None),
        mo.md("**Exercise 3**"),
        _check_aps(),
        mo.accordion(
            {
                "Solution 1": mo.md(
                    r"$\hat q$ = 8th smallest score = **0.7**. Scores $1-\hat p$: 0.4, 0.75, 0.9, 0.95 → only class 0 has score ≤ 0.7, so the set is **{0}**. Class 1 just misses (0.75)."
                ),
                "Solution 2": mo.md(
                    "Cosines: A 0.8, B 0.6, C $-0.48+0.48=0$ → **A**."
                ),
                "Solution 3": mo.md("""```python
def aps_scores(probs, y):
    order = np.argsort(-probs, axis=1)
    cum = np.cumsum(np.take_along_axis(probs, order, 1), axis=1)
    rank = np.argmax(order == y[:, None], axis=1)
    return cum[np.arange(len(y)), rank]
```
APS sets adapt their size to how spread out the model's probabilities are, unlike the simple $1-\\hat p_y$ score."""),
            }
        ),
    ]
)