Модуль 10. Поиск, проверка и вычислительный бюджет на этапе генерации

from __future__ import annotations
from pathlib import Path
import sys
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from IPython.display import display

PROFILE = "smoke"  # "smoke" or "full"
RUN_REAL_MODEL = False
SEEDS = range(40 if PROFILE == "smoke" else 100)

ROOT = Path.cwd()
if not (ROOT / "modules").exists():
    if (ROOT.parent / "modules").exists():
        ROOT = ROOT.parent
    else:
        raise RuntimeError("Run the notebook from the course root or notebooks directory")
MODULE = ROOT / "modules" / "module-10"
sys.path.insert(0, str(MODULE))

from module_10_reference import (
    pass_at_n, pass_curve, unbiased_pass_at_k,
    plurality_accuracy_exact, plurality_accuracy_mc,
    binary_verifier_success, process_score, prefix_success_estimate,
    wilson_interval, proxy_overoptimization_curve,
    allocate_budget_greedy, mean_coverage, pareto_frontier,
    break_even_queries, cost_regularized_utility,
)
from module_10_search import (
    ReasoningMaze, repeated_sampling, best_of_n_proxy, beam_search,
    best_first_search, mcts_search, evaluate_search_methods, maze_to_array,
)

LABELS = {'coverage': 'вероятность хотя бы одного успеха', 'budget': 'бюджет N', 'exact': 'точно', 'mc': 'Монте-Карло', 'vote': 'точность голосования', 'proxy': 'прокси-оценка', 'goal': 'независимая цель', 'score': 'оценка', 'success': 'доля успешных запусков', 'calls': 'вызовы модели', 'good': 'хороший оцениватель', 'biased': 'смещённый оцениватель', 'allocation': 'выделенные выборки', 'prompt': 'запрос', 'quality': 'качество', 'cost': 'стоимость', 'skipped': 'Реальная ветвь выключена. Включайте её только в подготовленной среде с доступом к модели.'}
np.set_printoptions(precision=4, suppress=True)
print(f"profile={PROFILE}, independent seeds={len(list(SEEDS))}")

budgets = np.array([1, 2, 4, 8, 16, 32, 64])
fig, ax = plt.subplots(figsize=(8.8, 5.0))
for p in (0.0, 0.05, 0.2, 0.5):
    ax.plot(budgets, pass_curve(p, budgets), marker="o", label=f"p={p:g}")
ax.set_xscale("log", base=2)
ax.set_xticks(budgets, labels=budgets)
ax.set_xlabel(LABELS["budget"])
ax.set_ylabel(LABELS["coverage"])
ax.set_ylim(-0.02, 1.02)
ax.grid(True, alpha=0.25)
ax.legend()
plt.show()

finite = pd.DataFrame({
    "n": [100, 100, 100, 20],
    "c": [1, 5, 20, 3],
    "k": [1, 10, 10, 5],
})
finite["pass@k"] = [unbiased_pass_at_k(int(r.n), int(r.c), int(r.k)) for r in finite.itertuples()]
display(finite.round(6))
assert np.isclose(pass_at_n(0.2, 8), 1 - 0.8**8)

cases = {
    "scattered_errors": np.array([0.40, 0.15, 0.15, 0.15, 0.15]),
    "systematic_error": np.array([0.30, 0.45, 0.25]),
}
rows = []
for name, probs in cases.items():
    for k in ((1, 4, 8, 16) if len(probs) == 5 else (1, 8, 32, 64)):
        exact = plurality_accuracy_exact(probs, k)
        mc = plurality_accuracy_mc(probs, k, trials=20_000 if PROFILE == "smoke" else 80_000, seed=10+k)
        rows.append({"case": name, "k": k, "exact": exact, "Monte Carlo": mc})
votes = pd.DataFrame(rows)
display(votes.round(6))

fig, ax = plt.subplots(figsize=(8.6, 4.8))
for name, group in votes.groupby("case"):
    ax.plot(group["k"], group["exact"], marker="o", label=name)
ax.set_xscale("log", base=2)
ax.set_xlabel("k")
ax.set_ylabel(LABELS["vote"])
ax.set_ylim(0, 1.02)
ax.grid(True, alpha=0.25)
ax.legend()
plt.show()
assert abs(votes["exact"].iloc[0] - 0.4) < 1e-12

p, alpha, beta = 0.2, 0.9, 0.1
ns = np.arange(1, 65)
rejected = np.array([
    binary_verifier_success(int(n), p, alpha, beta, fallback_mode="rejected_pool").success_probability
    for n in ns
])
fresh = np.array([
    binary_verifier_success(int(n), p, alpha, beta, fallback_mode="fresh_unverified").success_probability
    for n in ns
])
example_rej = binary_verifier_success(8, p, alpha, beta, fallback_mode="rejected_pool")
example_fresh = binary_verifier_success(8, p, alpha, beta, fallback_mode="fresh_unverified")
display(pd.DataFrame([example_rej.__dict__, example_fresh.__dict__]).round(6))

fig, ax = plt.subplots(figsize=(8.7, 4.8))
ax.plot(ns, rejected, label="fallback: rejected pool")
ax.plot(ns, fresh, label="fallback: fresh draw")
ax.axhline(example_rej.accepted_precision, linestyle="--", label="accepted-pool precision")
ax.set_xscale("log", base=2)
ax.set_xlabel(LABELS["budget"])
ax.set_ylabel(LABELS["success"])
ax.grid(True, alpha=0.25)
ax.legend()
plt.show()
assert np.isclose(example_rej.success_probability, 0.632486004809728)
assert np.isclose(example_fresh.success_probability, 0.6480396435591987)

budgets_proxy = np.array([1, 2, 4, 8, 16, 32, 64, 128, 256, 512])
proxy, independent = proxy_overoptimization_curve(
    budgets_proxy, trials=4_000 if PROFILE == "smoke" else 20_000, seed=202_610
)
proxy_table = pd.DataFrame({"N": budgets_proxy, "proxy": proxy, "independent_objective": independent})
display(proxy_table.round(4))

fig, ax = plt.subplots(figsize=(8.8, 4.9))
ax.plot(budgets_proxy, proxy, marker="o", label=LABELS["proxy"])
ax.plot(budgets_proxy, independent, marker="o", label=LABELS["goal"])
ax.set_xscale("log", base=2)
ax.set_xlabel(LABELS["budget"])
ax.set_ylabel(LABELS["score"])
ax.grid(True, alpha=0.25)
ax.legend()
plt.show()
peak_n = int(budgets_proxy[np.argmax(independent)])
print("independent-objective peak N =", peak_n)

successes = np.arange(9)
rows = []
for s in successes:
    lo, hi = wilson_interval(int(s), 8)
    rows.append({"successes": s, "estimate": prefix_success_estimate(int(s), 8), "Wilson_low": lo, "Wilson_high": hi})
prefix_table = pd.DataFrame(rows)
display(prefix_table.round(4))

fig, ax = plt.subplots(figsize=(8.6, 4.8))
mid = prefix_table["estimate"].to_numpy()
lo = prefix_table["Wilson_low"].to_numpy()
hi = prefix_table["Wilson_high"].to_numpy()
ax.errorbar(successes, mid, yerr=np.vstack([mid-lo, hi-mid]), fmt="o", capsize=4)
ax.set_xlabel("successes out of 8")
ax.set_ylabel("prefix-success estimate and 95% Wilson interval")
ax.set_ylim(-0.03, 1.03)
ax.grid(True, alpha=0.25)
plt.show()

step_scores = [0.92, 0.78, 0.61, 0.85]
aggregation = {m: process_score(step_scores, m) for m in ("min", "mean", "product", "last", "log_product")}
display(pd.Series(aggregation, name="aggregate").to_frame())

env = ReasoningMaze.course_maze()
methods = {
    "repeated": repeated_sampling(env, 48, seed=7),
    "best_of_n_proxy": best_of_n_proxy(env, 48, seed=7),
    "beam": beam_search(env, 48, beam_width=3, score_mode="good", seed=7),
    "best_first": best_first_search(env, 48, score_mode="good", seed=7),
    "mcts": mcts_search(env, 48, score_mode="good", seed=7),
}
summary = pd.DataFrame([
    {"method": name, "success": r.success, "model_calls": r.model_calls,
     "generated_trajectories": r.generated_trajectories, "path_length": r.selected_state.depth}
    for name, r in methods.items()
])
display(summary)

chosen = methods["best_first"]
grid = maze_to_array(env, chosen.selected_state)
fig, ax = plt.subplots(figsize=(6.3, 5.5))
ax.imshow(grid)
for i in range(env.height):
    for j in range(env.width):
        if (i, j) not in env.walls:
            ax.text(j, i, "·", ha="center", va="center", fontsize=10)
ax.set_title(f"best-first: success={chosen.success}, calls={chosen.model_calls}")
ax.set_xticks(range(env.width)); ax.set_yticks(range(env.height))
plt.show()

search_budgets = np.array([12, 16, 24, 32, 48])
good = evaluate_search_methods(search_budgets, seeds=SEEDS, score_mode="good")
biased = evaluate_search_methods(search_budgets, seeds=SEEDS, score_mode="biased")

fig, axes = plt.subplots(1, 2, figsize=(13.8, 4.9), sharey=True)
for ax, curves, title in zip(axes, (good, biased), (LABELS["good"], LABELS["biased"])):
    for method, values in curves.items():
        ax.plot(search_budgets, values, marker="o", label=method)
    ax.set_title(title)
    ax.set_xlabel(LABELS["calls"])
    ax.grid(True, alpha=0.25)
axes[0].set_ylabel(LABELS["success"])
axes[0].legend(fontsize=8)
plt.show()

comparison = pd.DataFrame({
    ("good", method): values for method, values in good.items()
} | {
    ("biased", method): values for method, values in biased.items()
}, index=search_budgets)
comparison.index.name = "model_calls"
display(comparison.round(3))

ps = np.array([0.02, 0.08, 0.20, 0.50])
total = 20
uniform = np.full(len(ps), total // len(ps), dtype=int)
greedy = allocate_budget_greedy(ps, total, minimum_per_prompt=1)
allocation = pd.DataFrame({"p": ps, "uniform": uniform, "greedy": greedy})
display(allocation)
print("uniform mean coverage:", round(mean_coverage(ps, uniform), 6))
print("greedy mean coverage:", round(mean_coverage(ps, greedy), 6))

costs = np.array([8, 12, 16, 24, 32, 48, 64], dtype=float)
quality = np.array([0.42, 0.55, 0.69, 0.76, 0.79, 0.80, 0.795])
front = pareto_frontier(quality, costs)
pareto_table = pd.DataFrame({"cost": costs, "quality": quality, "Pareto": front})
pareto_table["utility(lambda=0.004)"] = [cost_regularized_utility(q, c, 0.004) for q, c in zip(quality, costs)]
display(pareto_table)
print("break-even queries:", break_even_queries(12_000, 0.12, 0.07))

fig, ax = plt.subplots(figsize=(7.6, 4.8))
ax.scatter(costs[~front], quality[~front], label="dominated")
ax.scatter(costs[front], quality[front], label="Pareto")
ax.plot(costs[front], quality[front], linestyle="--")
ax.set_xlabel(LABELS["cost"]); ax.set_ylabel(LABELS["quality"])
ax.grid(True, alpha=0.25); ax.legend(); plt.show()

if not RUN_REAL_MODEL:
    print('Реальная ветвь выключена. Включайте её только в подготовленной среде с доступом к модели.')
else:
    # Установите Transformers в отдельной среде для реальной модели и зафиксируйте
    # фактически разрешённую версию перед воспроизводимым запуском.
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed

    model_id = "Qwen/Qwen3-0.6B"
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        torch_dtype="auto",
        device_map="auto",
    )
    prompt = "Solve carefully and end with a boxed integer: 17 * 23"
    messages = [{"role": "user", "content": prompt}]
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer([text], return_tensors="pt").to(model.device)
    set_seed(10)
    generated = model.generate(
        **inputs,
        do_sample=True,
        temperature=0.8,
        top_p=0.95,
        max_new_tokens=192,
        num_return_sequences=8,
        return_dict_in_generate=True,
        output_scores=True,
    )
    transition_scores = model.compute_transition_scores(
        generated.sequences,
        generated.scores,
        normalize_logits=True,
    )
    prompt_len = inputs.input_ids.shape[1]
    responses = tokenizer.batch_decode(generated.sequences[:, prompt_len:], skip_special_tokens=True)
    rows = []
    for i, response in enumerate(responses):
        rows.append({
            "candidate": i,
            "response": response,
            "sum_logprob": float(transition_scores[i].sum().cpu()),
            "mean_logprob": float(transition_scores[i].mean().cpu()),
        })
    display(pd.DataFrame(rows))