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 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]],
    full_cot_list: list[list[str]],
    answer: str,
    threshold: float = 0.8,
    min_samples: int = 5,
    embedding_model: SentenceTransformer | None = None,
) -> list[float]:
    """
    Calculate counterfactual importance by filtering for low-similarity resamples.

    Args:
        chunks_removed: Original chunks that were removed
        chunks_resampled: List of resampled chunks for each position
        full_cot_list: Full CoT rollouts for each position
        answer: Ground truth answer
        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, full_cot_list, answer, embedding_model=embedding_model
)

# Compare with precomputed
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("All tests for `calculate_counterfactual_importance` passed!")

# 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]]] = []

    for i in tqdm(range(len(chunk_solutions)), desc="Precomputing embeddings"):
        chunk_positions = []
        chunk_sentences = []

        for rollout_dict in chunk_solutions[i]:
            rollout_text = rollout_dict["rollout"]
            sentences = split_solution_into_chunks(rollout_text)
            assert sentences, "Expected at least one sentence per rollout"

            all_sentences.extend(sentences)
            chunk_positions.append((len(all_sentences) - len(sentences), len(all_sentences)))
            chunk_sentences.append(sentences)

        position_map.append(chunk_positions)
        rollout_sentences.append(chunk_sentences)

    # Batch embed all sentences at once
    all_embeddings = embedding_model.encode(all_sentences, batch_size=batch_size, show_progress_bar=True)

    # Reconstruct per-rollout embeddings using position map
    rollout_embeddings: list[list[np.ndarray]] = []

    for i in range(len(chunk_solutions)):
        chunk_embeddings = []
        for position in position_map[i]:
            start_idx, end_idx = position
            chunk_embeddings.append(all_embeddings[start_idx:end_idx])
        rollout_embeddings.append(chunk_embeddings)

    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
    """
    chunk_texts = [chunks[i]["chunk"] for i in range(n_chunks)]
    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)

    # Concatenate all embeddings: (total_sentences, embed_dim)
    all_embeddings = np.vstack(valid_embeddings)

    # Normalize target and all embeddings for cosine similarity
    target_norm = target_embedding / (np.linalg.norm(target_embedding) + 1e-8)
    all_norms = np.linalg.norm(all_embeddings, axis=1, keepdims=True) + 1e-8
    all_embeddings_norm = all_embeddings / all_norms

    # Compute all cosine similarities at once: (total_sentences,)
    all_similarities = all_embeddings_norm @ target_norm

    # Count matches per rollout using boundaries
    matches = sum(1 for start, end in rollout_boundaries if np.max(all_similarities[start:end]) >= similarity_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)
    """
    if source_idx >= target_idx:
        return 0.0

    target_embedding = target_embeddings[target_idx]

    # Rollouts where source was kept
    include_embeddings = rollout_embeddings[source_idx + 1]
    # Rollouts where source was removed
    exclude_embeddings = rollout_embeddings[source_idx]

    include_rate = compute_match_rate_from_embeddings(target_embedding, include_embeddings, similarity_threshold)
    exclude_rate = compute_match_rate_from_embeddings(target_embedding, exclude_embeddings, similarity_threshold)

    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).
    """
    raise NotImplementedError()

    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, torch_dtype=torch.bfloat16, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME_8B)

resampled_rollouts = get_resampled_rollouts(
    prompt=problem_data["base_solution"]["prompt"],
    model=model,
    tokenizer=tokenizer,
    num_resamples_per_chunk=2,
    batch_size=2,
    max_new_tokens=2048,
    up_to_n_chunks=4,
)

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