1️⃣ Setup & visualizing truth representations

def extract_activations(
    statements: list[str],
    model: AutoModelForCausalLM,
    tokenizer: AutoTokenizer,
    layers: list[int],
    batch_size: int = 25,
) -> dict[int, Float[Tensor, "n_statements d_model"]]:
    """
    Extract last-token hidden state activations from specified layers for a list of statements.

    Args:
        statements: List of text statements to process.
        model: A HuggingFace causal language model.
        tokenizer: The corresponding tokenizer.
        layers: List of layer indices (0-indexed) to extract activations from.
        batch_size: Number of statements to process at once.

    Returns:
        Dictionary mapping layer index to tensor of activations, shape [n_statements, d_model].
    """
    raise NotImplementedError()


tests.test_extract_activations(extract_activations, model, tokenizer, PROBE_LAYER, D_MODEL)

# Extract activations at the probe layer for all datasets
activations = {}
labels_dict = {}

for name in DATASET_NAMES:
    df = datasets[name]
    statements = df["statement"].tolist()
    labs = t.tensor(df["label"].values, dtype=t.float32)

    acts = extract_activations(statements, model, tokenizer, [PROBE_LAYER])
    activations[name] = acts[PROBE_LAYER]
    labels_dict[name] = labs

# Show summary table
summary = pd.DataFrame(
    {
        "Dataset": DATASET_NAMES,
        "N statements": [len(datasets[n]) for n in DATASET_NAMES],
        "N true": [int(datasets[n]["label"].sum()) for n in DATASET_NAMES],
        "N false": [int((1 - datasets[n]["label"]).sum()) for n in DATASET_NAMES],
        "Act shape": [str(tuple(activations[n].shape)) for n in DATASET_NAMES],
        "Mean norm": [f"{activations[n].norm(dim=-1).mean():.1f}" for n in DATASET_NAMES],
    }
)
display(summary)

def get_pca_components(
    activations: Float[Tensor, "n d_model"],
    k: int = 2,
) -> Float[Tensor, "d_model k"]:
    """
    Compute the top-k principal components of the activation matrix.

    Args:
        activations: Activation matrix, shape [n_samples, d_model].
        k: Number of principal components to return.

    Returns:
        Matrix of top-k eigenvectors as columns, shape [d_model, k].
    """
    raise NotImplementedError()


tests.test_get_pca_components(get_pca_components, activations["cities"], D_MODEL)

fig = make_subplots(rows=1, cols=3, subplot_titles=DATASET_NAMES)

for i, name in enumerate(DATASET_NAMES):
    acts = activations[name]
    labs = labels_dict[name]
    pcs = get_pca_components(acts, k=2)
    X_centered = acts - acts.mean(dim=0)
    projected = (X_centered @ pcs).numpy()

    # Compute variance explained
    total_var = X_centered.var(dim=0).sum().item()
    pc_var = t.tensor(projected).var(dim=0)
    pct_explained = (pc_var / total_var * 100).tolist()

    colors = ["blue" if l == 1 else "red" for l in labs.tolist()]
    fig.add_trace(
        go.Scatter(
            x=projected[:, 0],
            y=projected[:, 1],
            mode="markers",
            marker=dict(color=colors, size=3, opacity=0.5),
            name=name,
            showlegend=False,
        ),
        row=1,
        col=i + 1,
    )
    fig.update_xaxes(title_text=f"PC1 ({pct_explained[0]:.1f}%)", row=1, col=i + 1)
    fig.update_yaxes(title_text=f"PC2 ({pct_explained[1]:.1f}%)", row=1, col=i + 1)

# Add a legend manually
fig.add_trace(go.Scatter(x=[None], y=[None], mode="markers", marker=dict(color="blue", size=8), name="True"))
fig.add_trace(go.Scatter(x=[None], y=[None], mode="markers", marker=dict(color="red", size=8), name="False"))

fig.update_layout(
    title="PCA of Truth Representations (Layer 14, Last Token)",
    height=400,
    width=1200,
)
fig.show()

def layer_sweep_accuracy(
    statements: list[str],
    labels: Float[Tensor, " n"],
    model: AutoModelForCausalLM,
    tokenizer: AutoTokenizer,
    layers: list[int],
    train_frac: float = 0.8,
    batch_size: int = 25,
) -> dict[str, list[float]]:
    """
    For each layer, train a difference-of-means classifier and compute train/test accuracy.

    Args:
        statements: List of statements.
        labels: Binary labels (1=true, 0=false).
        model: The language model.
        tokenizer: The tokenizer.
        layers: List of layer indices to sweep over.
        train_frac: Fraction of data for training.
        batch_size: Batch size for activation extraction.

    Returns:
        Dict with keys "train_acc" and "test_acc", each a list of accuracies per layer.
    """
    raise NotImplementedError()


t.manual_seed(42)
all_layers = list(range(NUM_LAYERS))
cities_statements = datasets["cities"]["statement"].tolist()
cities_labels = t.tensor(datasets["cities"]["label"].values, dtype=t.float32)

sweep_results = layer_sweep_accuracy(cities_statements, cities_labels, model, tokenizer, all_layers)

# Print results as a table
sweep_df = pd.DataFrame(
    {
        "Layer": all_layers,
        "Train Acc": [f"{a:.3f}" for a in sweep_results["train_acc"]],
        "Test Acc": [f"{a:.3f}" for a in sweep_results["test_acc"]],
    }
)
display(sweep_df)

# Plot
fig = go.Figure()
fig.add_trace(go.Scatter(x=all_layers, y=sweep_results["train_acc"], mode="lines+markers", name="Train"))
fig.add_trace(go.Scatter(x=all_layers, y=sweep_results["test_acc"], mode="lines+markers", name="Test"))
fig.add_vline(x=PROBE_LAYER, line_dash="dash", line_color="gray", annotation_text=f"Probe layer ({PROBE_LAYER})")
fig.update_layout(
    title="Layer Sweep: Difference-of-Means Accuracy on Cities Dataset",
    xaxis_title="Layer",
    yaxis_title="Accuracy",
    yaxis_range=[0.4, 1.05],
    height=400,
    width=800,
)
fig.show()

best_layer = all_layers[int(np.argmax(sweep_results["test_acc"]))]
print(f"\nBest layer by test accuracy: {best_layer} ({max(sweep_results['test_acc']):.3f})")
print(f"Configured probe layer: {PROBE_LAYER} ({sweep_results['test_acc'][PROBE_LAYER]:.3f})")