2️⃣ Black-box Analysis

def load_single_file(file_path: str):
    local_path = hf_hub_download(repo_id=DATASET_NAME, filename=file_path, repo_type="dataset")
    with open(local_path, "r") as f:
        return json.load(f)


def load_problem_data(problem_id: int, model_name: str = "deepseek-r1-distill-qwen-14b", verbose: bool = True):
    disable_progress_bars()

    problem_dir = "correct_base_solution"
    problem_dir_forced = "correct_base_solution_forced_answer"
    problem_path = f"{model_name}/temperature_0.6_top_p_0.95/{problem_dir}/problem_{problem_id}"
    problem_path_forced = f"{model_name}/temperature_0.6_top_p_0.95/{problem_dir_forced}/problem_{problem_id}"

    base_solution = load_single_file(f"{problem_path}/base_solution.json")
    problem = load_single_file(f"{problem_path}/problem.json")
    chunks_labeled = load_single_file(f"{problem_path}/chunks_labeled.json")

    n_chunks = len(chunks_labeled)
    chunk_solutions = [None] * n_chunks
    chunk_solutions_forced = [None] * n_chunks

    def load_chunk(chunk_idx):
        sol = load_single_file(f"{problem_path}/chunk_{chunk_idx}/solutions.json")
        forced = load_single_file(f"{problem_path_forced}/chunk_{chunk_idx}/solutions.json")
        return chunk_idx, sol, forced

    with ThreadPoolExecutor(max_workers=16) as executor:
        futures = [executor.submit(load_chunk, i) for i in range(n_chunks)]
        for future in tqdm(as_completed(futures), total=n_chunks, disable=not verbose, desc="Loading chunks"):
            idx, sol, forced = future.result()
            chunk_solutions[idx] = sol
            chunk_solutions_forced[idx] = forced

    enable_progress_bars()

    return {
        "problem": problem,
        "base_solution": base_solution,
        "chunks_labeled": chunks_labeled,
        "chunk_solutions_forced": chunk_solutions_forced,
        "chunk_solutions": chunk_solutions,
    }


# Load a problem
problem_data = load_problem_data(PROBLEM_ID)

# Inspect the problem structure
print("Problem:", problem_data["problem"]["problem"][:200], "...")
print(f"\nGround truth answer: {problem_data['problem']['gt_answer']}")
print(f"Number of chunks: {len(problem_data['chunks_labeled'])}")
print(f"Rollouts per chunk: {len(problem_data['chunk_solutions'][0])}")

# Show first few chunks
for i, c in enumerate(problem_data["chunks_labeled"][:5]):
    print(f"{i}. [{c['function_tags'][0]}] {c['chunk'][:80]}...")

def extract_answer_from_cot(cot: str) -> str:
    """Extract the numerical answer from a chain-of-thought solution by parsing the \\boxed{} expression."""
    ans = cot.split("\\boxed{")[-1].split("}")[0]
    return "".join(char for char in ans if char.isdigit() or char == ".")


def get_filtered_indices(
    chunk_removed: str, resampled: list[str], embedding_model: "SentenceTransformer", threshold: float = 0.7
) -> list[int]:
    """Return indices of resampled rollouts whose chunk is sufficiently dissimilar from the original."""
    emb_original = embedding_model.encode(chunk_removed)
    emb_resampled = embedding_model.encode(resampled)
    cos_sims = emb_original @ emb_resampled.T
    return np.where(cos_sims < threshold)[0]

def calculate_answer_importance(full_cot_list: list[list[str]], answer: str) -> list[float]:
    """
    Calculate importance for chunks based on accuracy differences.

    Args:
        full_cot_list: List of lists of rollouts. full_cot_list[i][j] is the j-th rollout
            generated by forcing an answer after the i-th chunk.
        answer: The ground truth answer.

    Returns:
        List of importance scores (one fewer than chunks, since we measure differences).
    """

    raise NotImplementedError("Implement answer importance calculation")


tests.test_calculate_answer_importance(calculate_answer_importance)

# Calculate forced answer importance
full_cot_list = [
    [rollout["full_cot"] for rollout in chunk_rollouts] for chunk_rollouts in problem_data["chunk_solutions_forced"]
]
answer = problem_data["problem"]["gt_answer"]

forced_importances = calculate_answer_importance(full_cot_list, answer)

# Get chunk texts for hover data
chunks_for_hover = [chunk["chunk"] for chunk in problem_data["chunks_labeled"][:-1]]

# Plot with plotly
fig = go.Figure()
fig.add_trace(
    go.Bar(
        x=list(range(len(forced_importances))),
        y=forced_importances,
        opacity=0.7,
        hovertemplate="<b>Chunk %{x}</b><br>Importance: %{y:.4f}<br>Text: %{customdata}<extra></extra>",
        customdata=[chunk[:100] + "..." if len(chunk) > 100 else chunk for chunk in chunks_for_hover],
    )
)
fig.add_hline(y=0, line_color="black", line_width=0.5)
fig.update_layout(
    title="Forced Answer Importance by Chunk",
    xaxis_title="Chunk Index",
    yaxis_title="Forced Answer Importance",
    width=900,
    height=400,
)
fig.show()

# YOUR CODE HERE - compute `resampling_importances` using `calculate_answer_importance`
# on the resampled rollouts
resampling_importances = []

# Compare with precomputed values from dataset (they used a slightly different method to us, but
# we should get an answer within 1% of theirs)
precomputed = [chunk["resampling_importance_accuracy"] for chunk in problem_data["chunks_labeled"][:-1]]
avg_diff = np.abs(np.subtract(resampling_importances, precomputed)).mean()
assert avg_diff < 0.01, "Error above 1% threshold"

# Plot comparison between these two metrics
chunks_for_hover = [chunk["chunk"] for chunk in problem_data["chunks_labeled"][:-1]]
hover_texts = [chunk[:100] + "..." if len(chunk) > 100 else chunk for chunk in chunks_for_hover]

fig = make_subplots(
    rows=2,
    cols=1,
    shared_xaxes=True,
    vertical_spacing=0.1,
    subplot_titles=("Forced Answer Importance", "Resampling Importance"),
)

for row, color, y in [(1, "cornflowerblue", forced_importances), (2, "orange", resampling_importances)]:
    fig.add_trace(
        go.Bar(
            x=list(range(len(y))),
            y=y,
            opacity=0.7,
            name="Forced" if color == "cornflowerblue" else "Resampling",
            marker_color=color,
            hovertemplate="<b>Chunk %{x}</b><br>Importance: %{y:.4f}<br>Text: %{customdata}<extra></extra>",
            customdata=hover_texts,
        ),
        row=row,
        col=1,
    )

fig.add_hline(y=0, line_color="black", line_width=0.5, row=1, col=1)
fig.add_hline(y=0, line_color="black", line_width=0.5, row=2, col=1)
fig.update_layout(width=900, height=500, showlegend=False)
fig.update_xaxes(title_text="Chunk Index", row=2, col=1)
fig.update_yaxes(title_text="Importance", row=1, col=1)
fig.update_yaxes(title_text="Importance", row=2, col=1)
fig.show()

# Compute cosine similarity between original and resampled chunks
chunks_removed = [chunk["chunk"] for chunk in problem_data["chunks_labeled"]]
embeddings_original = embedding_model.encode(chunks_removed)

chunks_resampled = [
    [rollout["chunk_resampled"] for rollout in chunk_rollouts] for chunk_rollouts in problem_data["chunk_solutions"]
]
embeddings_resampled = np.stack([embedding_model.encode(r) for r in chunks_resampled])

# Compute similarities
cos_sims = einops.einsum(embeddings_original, embeddings_resampled, "chunk d, chunk resample d -> chunk resample")
cos_sims_mean = cos_sims.mean(axis=1)

# Plot by category with plotly
chunk_labels = [CATEGORIES[chunk["function_tags"][0]] for chunk in problem_data["chunks_labeled"]]
chunks_for_hover = [chunk["chunk"] for chunk in problem_data["chunks_labeled"]]
df = pd.DataFrame({"Label": chunk_labels, "Cosine Similarity": cos_sims_mean, "Chunk Text": chunks_for_hover})

fig = go.Figure()
for label in df["Label"].unique():
    subset = df[df["Label"] == label]
    hover_texts = [text[:100] + "..." if len(text) > 100 else text for text in subset["Chunk Text"]]
    fig.add_trace(
        go.Bar(
            x=subset.index.tolist(),
            y=subset["Cosine Similarity"].tolist(),
            name=label,
            marker_color=utils.CATEGORY_COLORS.get(label, "#9E9E9E"),
            hovertemplate="<b>Chunk %{x}</b><br>Category: "
            + label
            + "<br>Cosine Similarity: %{y:.4f}<br>Text: %{customdata}<extra></extra>",
            customdata=hover_texts,
        )
    )

fig.update_layout(
    title="How Similar are Resampled Chunks to Originals?",
    xaxis_title="Chunk Index",
    yaxis_title="Mean Cosine Similarity to Resamples",
    width=900,
    height=400,
    legend=dict(x=1.02, y=1, xanchor="left"),
    bargap=0,
)
fig.show()

def calculate_counterfactual_importance(
    chunks_removed: list[str],
    chunks_resampled: list[list[str]],
    rollout_data: list[list],
    score_fn: Callable,
    embedding_model: SentenceTransformer,
    threshold: float = 0.8,
    min_samples: int = 5,
) -> list[float]:
    """
    Calculate counterfactual importance by filtering for low-similarity resamples.

    This is a generic function: `rollout_data[i]` is a list of per-rollout values at chunk
    position i, and `score_fn(rollout_data_item)` returns True/False for whether that rollout
    counts as a "success". For math accuracy, rollout_data contains CoT strings and score_fn
    checks if the extracted answer matches the ground truth. For blackmail rate, rollout_data
    contains boolean labels and score_fn is just the identity.

    Args:
        chunks_removed: Original chunks that were removed
        chunks_resampled: List of resampled chunks for each position
        rollout_data: Per-rollout data for each chunk position (passed to score_fn)
        score_fn: Function mapping a single rollout data item to True/False
        threshold: Maximum cosine similarity to count as "different"
        min_samples: Minimum samples needed to compute probability
        embedding_model: Sentence embedding model

    Returns:
        List of counterfactual importance scores
    """

    raise NotImplementedError("Implement counterfactual importance calculation")

# Calculate counterfactual importance
chunks_removed = [chunk["chunk"] for chunk in problem_data["chunks_labeled"]]
chunks_resampled = [
    [rollout["chunk_resampled"] for rollout in chunk_rollouts] for chunk_rollouts in problem_data["chunk_solutions"]
]
full_cot_list = [
    [rollout["full_cot"] for rollout in chunk_rollouts] for chunk_rollouts in problem_data["chunk_solutions"]
]

counterfactual_importances = calculate_counterfactual_importance(
    chunks_removed,
    chunks_resampled,
    rollout_data=full_cot_list,
    score_fn=lambda cot: extract_answer_from_cot(cot) == answer,
    embedding_model=embedding_model,
)

# Compare with precomputed
# (We flip the sign because the authors store the negative of the counterfactual metric in the dataset)
precomputed_cf = [-chunk["counterfactual_importance_accuracy"] for chunk in problem_data["chunks_labeled"][:-1]]
avg_diff = np.abs(np.subtract(counterfactual_importances, precomputed_cf)).mean()
assert avg_diff < 0.025, "Error above 2.5% threshold"
print("Precomputed comparison passed!")

tests.test_calculate_counterfactual_importance(calculate_counterfactual_importance)

# Plot comparison of all three metrics with subplots (like previous bar chart)
chunks_for_hover = [chunk["chunk"] for chunk in problem_data["chunks_labeled"][:-1]]
hover_texts = [chunk[:100] + "..." if len(chunk) > 100 else chunk for chunk in chunks_for_hover]

fig = make_subplots(
    rows=3,
    cols=1,
    shared_xaxes=True,
    vertical_spacing=0.08,
    subplot_titles=("Forced Answer Importance", "Resampling Importance", "Counterfactual Importance"),
)

for row, (name, importances, color) in enumerate(
    [
        ("Forced", forced_importances, "cornflowerblue"),
        ("Resampling", resampling_importances, "orange"),
        ("Counterfactual", counterfactual_importances, "seagreen"),
    ],
    start=1,
):
    fig.add_trace(
        go.Bar(
            x=list(range(len(importances))),
            y=importances,
            name=name,
            opacity=0.8,
            marker_color=color,
            hovertemplate="<b>Chunk %{x}</b><br>"
            + name
            + " Importance: %{y:.4f}<br>Text: %{customdata}<extra></extra>",
            customdata=hover_texts,
        ),
        row=row,
        col=1,
    )
    fig.add_hline(y=0, line_color="black", line_width=0.5, row=row, col=1)

fig.update_layout(
    title="Comparison of Importance Metrics",
    width=1000,
    height=700,
    showlegend=False,
)
fig.show()

# YOUR CODE HERE - replicate figure 3b!

# Demo: visualize a causal graph with sample data
n_chunks = 5
sample_labels = ["Problem Setup", "Plan Generation", "Active Computation", "Self Checking", "Final Answer Emission"]
sample_texts = [
    "Let me start by understanding the problem...",
    "I notice that this is equivalent to...",
    "Computing the value: 3 * 7 = 21",
    "Let me verify: 21 / 7 = 3, correct",
    "Therefore the answer is 21",
]

# Create a sample importance matrix (i, j) = how much does chunk i causally influence chunk j
# We expect later chunks to be influenced by earlier ones. In our demo, each step influences
# the next one, plus plan generation influences everything!
sample_importance = np.zeros((n_chunks, n_chunks))
for i in range(4):
    sample_importance[i, i + 1] += 0.5
for i in range(2, 5):
    sample_importance[1, i] += 0.5

html_str = utils.chunk_graph_html(
    edge_weights=sample_importance,
    chunk_labels=sample_labels,
    chunk_texts=sample_texts,
)
display(HTML(html_str))

def precompute_rollout_embeddings(
    chunk_solutions: list[list[dict]],
    embedding_model: SentenceTransformer,
    batch_size: int = 128,
) -> tuple[list[list[np.ndarray]], list[list[list[str]]]]:
    """
    Precompute embeddings for all sentences in all rollouts using batched encoding. Does this by
    first collecting every single rollout from each chunk, then batch encode them all.

    Args:
        chunk_solutions: List of chunk solutions, where chunk_solutions[k] contains
                        rollouts from resampling after chunk k
        embedding_model: Sentence embedding model
        batch_size: Batch size for embedding model encoding

    Returns:
        rollout_embeddings: rollout_embeddings[i][j] = array of shape (n_sentences, embed_dim)
        rollout_sentences: rollout_sentences[i][j] = list of sentence strings
    """
    # First pass: collect all sentences and track their positions
    all_sentences: list[str] = []
    # position_map[i][j] = (start_idx, end_idx) into all_sentences, or None if empty
    position_map: list[list[tuple[int, int] | None]] = []
    rollout_sentences: list[list[list[str]]] = []

    # YOUR CODE HERE!

    return rollout_embeddings, rollout_sentences


def precompute_target_embeddings(
    chunks: list[dict],
    embedding_model: SentenceTransformer,
    n_chunks: int,
) -> Float[np.ndarray, "n_chunks embed_dim"]:
    """
    Precompute embeddings for all target chunk sentences.

    Args:
        chunks: List of chunk dictionaries with "chunk" key
        embedding_model: Sentence embedding model
        n_chunks: Number of chunks to process

    Returns:
        Array of shape (n_chunks, embed_dim) with target embeddings
    """
    # YOUR CODE HERE!
    return embedding_model.encode(chunk_texts)


def compute_match_rate_from_embeddings(
    target_embedding: Float[np.ndarray, " embed_dim"],
    rollout_embeddings_list: list[Float[np.ndarray, "n_sentences embed_dim"]],
    similarity_threshold: float = 0.7,
) -> float:
    """
    Compute fraction of rollouts containing a sentence similar to target.

    Args:
        target_embedding: Embedding of target sentence, shape (embed_dim,)
        rollout_embeddings_list: List of arrays, each shape (n_sentences, embed_dim)
        similarity_threshold: Minimum cosine similarity for a match

    Returns:
        Fraction of rollouts with a matching sentence (0 to 1)
    """
    # Filter out empty embeddings and track rollout boundaries
    valid_embeddings = []
    rollout_boundaries = []  # (start, end) indices into concatenated array

    current_idx = 0
    for embeddings in rollout_embeddings_list:
        valid_embeddings.append(embeddings)
        rollout_boundaries.append((current_idx, current_idx + len(embeddings)))
        current_idx += len(embeddings)

    # YOUR CODE HERE!
    matches = ...  # compute number of rollouts with a match using cosine similarity and threshold

    return matches / len(rollout_boundaries)


def compute_pairwise_importance(
    source_idx: int,
    target_idx: int,
    target_embeddings: Float[np.ndarray, "n_chunks embed_dim"],
    rollout_embeddings: list[list[Float[np.ndarray, "n_sentences embed_dim"]]],
    similarity_threshold: float = 0.7,
) -> float:
    """
    Compute causal importance of sentence source_idx on sentence target_idx.
    Uses precomputed embeddings for efficiency.

    Compares:
    - Rollouts from chunk_{source_idx + 1} (source was KEPT)
    - Rollouts from chunk_{source_idx} (source was REMOVED)
    """
    # YOUR CODE HERE!
    include_rate = ...  # compute match rate for rollouts where source was kept
    exclude_rate = ...  # compute match rate for rollouts where source was removed

    return include_rate - exclude_rate


# Precompute all embeddings upfront
n_chunks = 74
target_embeddings = precompute_target_embeddings(problem_data["chunks_labeled"], embedding_model, n_chunks=n_chunks)
rollout_embeddings, rollout_sentences = precompute_rollout_embeddings(
    problem_data["chunk_solutions"], embedding_model
)

# Compute importance matrix
importance_matrix = np.zeros((n_chunks, n_chunks))

for i in tqdm(range(n_chunks - 1), desc="Computing pairwise importance"):
    for j in range(i + 1, n_chunks):
        importance_matrix[i, j] = compute_pairwise_importance(i, j, target_embeddings, rollout_embeddings)

# Plot the importance matrix with plotly
# Mask upper triangle of transposed matrix (we want to show lower triangle which has the values)
mask = np.triu(np.ones_like(importance_matrix, dtype=bool))
importance_matrix = np.where(mask, importance_matrix, 0.0)

chunk_labels = [CATEGORIES[chunk["function_tags"][0]] for chunk in problem_data["chunks_labeled"][:n_chunks]]
chunk_texts = [chunk["chunk"] for chunk in problem_data["chunks_labeled"][:n_chunks]]

html_str = utils.chunk_graph_html(
    edge_weights=importance_matrix,
    chunk_labels=chunk_labels,
    chunk_texts=chunk_texts,
    n_top_edges_per_direction=3,
)
display(HTML(html_str))

class StopOnThink(StoppingCriteria):
    """Helper class for stopping generation when the <think>...</think> tags are closed."""

    def __init__(self, tokenizer):
        self.tokenizer = tokenizer
        self.think_token_id = tokenizer.encode("</think>", add_special_tokens=False)[0]

    def __call__(self, input_ids, scores, **kwargs):
        return input_ids[0, -1] == self.think_token_id


def get_resampled_rollouts(
    prompt: str,
    model: transformers.models.llama.modeling_llama.LlamaForCausalLM,
    tokenizer: transformers.models.llama.tokenization_llama.LlamaTokenizer,
    num_resamples_per_chunk: int = 100,
    batch_size: int = 4,
    max_new_tokens: int = 2048,
    up_to_n_chunks: int | None = None,
) -> tuple[str, list[str], list[list[dict]]]:
    """
    After each chunk in `chunks`, computes a number of resampled rollouts.

    Args:
        prompt: The initial problem prompt (which ends with a <think> tag).
        num_resamples_per_chunk: Number of resamples to compute for each chunk.

    Returns:
        Tuple of (full_answer, chunks, resampled_rollouts) where the latter is a list of lists of
        dicts (one for each chunk & resample on that chunk).
    """

    @t.inference_mode()
    def generate(inputs):
        return model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            temperature=0.6,
            top_p=0.95,
            do_sample=True,
            stopping_criteria=[StopOnThink(tokenizer=tokenizer)],
            pad_token_id=tokenizer.eos_token_id,
        )

    # YOUR CODE HERE - implement resampling procedure as described in the text!

    return full_answer, chunks, chunk_rollouts


# Load the model for generation (only if you want to try this)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME_8B, dtype=dtype, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME_8B)

t0 = time.time()
full_answer_resampled, chunks_resampled, chunk_rollouts_resampled = get_resampled_rollouts(
    prompt=problem_data["base_solution"]["prompt"],
    model=model,
    tokenizer=tokenizer,
    num_resamples_per_chunk=4,
    batch_size=2,
    max_new_tokens=2048,
    up_to_n_chunks=5,
)
print(f"Got resampled rollouts in {time.time() - t0:.2f} seconds")

chunk_rollouts_resampled[0]

for i, resamples in enumerate(chunk_rollouts_resampled):
    print("Replaced chunk: " + repr(resamples[0]["chunk_replaced"]))
    for j, r in enumerate(resamples):
        print(f"    Resample {j}: " + repr(r["chunk_resampled"]))
    print()