AI Papers: Test-Time Scaling 核心论文梳理 — 用推理预算提升性能的方法

from dataclasses import dataclass, field
from typing import Callable
import asyncio
import time

@dataclass
class ScoredCandidate:
    text: str
    score: float
    latency_ms: float
    token_count: int

@dataclass
class BestOfNConfig:
    n_samples: int = 4
    temperature: float = 0.7
    max_tokens: int = 1024
    timeout_seconds: float = 15.0
    cost_budget_usd: float = 0.05

class BestOfNPipeline:
    """生产环境用 Best-of-N 流水线。

    将候选生成、评分、预算控制整合在一个流程中管理。
    """

    def __init__(
        self,
        generator: Callable,
        scorers: list[Callable],
        config: BestOfNConfig = BestOfNConfig(),
    ):
        self.generator = generator
        self.scorers = scorers
        self.config = config

    async def generate_candidates(self, prompt: str) -> list[ScoredCandidate]:
        """异步生成 N 个候选并打分。"""
        candidates = []
        tasks = [
            self._generate_one(prompt)
            for _ in range(self.config.n_samples)
        ]

        results = await asyncio.gather(*tasks, return_exceptions=True)

        for result in results:
            if isinstance(result, Exception):
                continue  # 跳过 timeout 或 API 错误
            text, latency_ms, token_count = result
            score = self._compute_score(text)
            candidates.append(ScoredCandidate(
                text=text,
                score=score,
                latency_ms=latency_ms,
                token_count=token_count,
            ))

        return sorted(candidates, key=lambda c: c.score, reverse=True)

    def _compute_score(self, text: str) -> float:
        """计算多个 scorer 的加权平均值。"""
        if not self.scorers:
            return 0.0
        scores = [scorer(text) for scorer in self.scorers]
        return sum(scores) / len(scores)

    async def _generate_one(self, prompt: str):
        """生成单个候选。超过 timeout 时抛出 asyncio.TimeoutError。"""
        start = time.monotonic()
        result = await asyncio.wait_for(
            self.generator(prompt, temperature=self.config.temperature),
            timeout=self.config.timeout_seconds,
        )
        elapsed_ms = (time.monotonic() - start) * 1000
        return result.text, elapsed_ms, result.token_count

    async def run(self, prompt: str) -> ScoredCandidate | None:
        """返回最优候选。没有候选时返回 None。"""
        candidates = await self.generate_candidates(prompt)
        return candidates[0] if candidates else None

from collections import Counter
import re
from typing import Optional


def extract_final_answer(text: str) -> Optional[str]:
    """从 CoT 输出中提取最终答案。

    依次尝试多种模式，如果全部都不匹配，
    则使用最后一行作为答案。
    """
    # 模式 1: "정답: X"、"Answer: X" 形式
    patterns = [
        r"(?:정답|답|answer|result)\s*[:：=]\s*(.+)",
        r"\\boxed\{(.+?)\}",              # LaTeX boxed 形式
        r"####\s*(.+)",                     # GSM8K 形式
        r"(?:따라서|그러므로|therefore)\s*(.+)",
    ]
    for pattern in patterns:
        matches = re.findall(pattern, text, re.IGNORECASE)
        if matches:
            return matches[-1].strip()

    # 所有模式都失败时，返回最后一个非空行
    lines = [line.strip() for line in text.strip().splitlines() if line.strip()]
    return lines[-1] if lines else None


def normalize_answer(answer: str) -> str:
    """用于答案比较的归一化处理。

    去除空格、标点、单位表记等细微差异，
    以便把相同的答案识别为同一个答案。
    """
    s = answer.strip().lower()
    # 去除货币/单位符号
    s = re.sub(r"[$€₩\\,]", "", s)
    # 去除结尾的句点/逗号
    s = s.rstrip(".,;")
    # 把多个空格合并为一个空格
    s = re.sub(r"\s+", " ", s)
    return s


def self_consistency_vote(
    samples: list[str],
    min_agreement: int = 2,
) -> tuple[Optional[str], int, int]:
    """执行 Self-Consistency 多数表决。

    Returns:
        (best_answer, vote_count, total_valid):
        - best_answer: 出现次数最多的答案（提取失败时为 None）
        - vote_count: 该答案获得的票数
        - total_valid: 成功提取出答案的样本总数
    """
    extracted = []
    for sample in samples:
        ans = extract_final_answer(sample)
        if ans is not None:
            extracted.append(normalize_answer(ans))

    if not extracted:
        return None, 0, 0

    counter = Counter(extracted)
    best_answer, vote_count = counter.most_common(1)[0]

    if vote_count < min_agreement:
        return None, vote_count, len(extracted)

    return best_answer, vote_count, len(extracted)

from dataclasses import dataclass, field
from typing import Optional
import heapq

@dataclass
class ThoughtNode:
    """推理树中的一个节点。"""
    content: str
    depth: int
    score: float = 0.0
    parent: Optional['ThoughtNode'] = None
    children: list['ThoughtNode'] = field(default_factory=list)

    def path(self) -> list[str]:
        """返回从根节点到当前节点的推理路径。"""
        nodes = []
        current = self
        while current is not None:
            nodes.append(current.content)
            current = current.parent
        return list(reversed(nodes))


class BeamSearchReasoner:
    """基于 Beam Search 的推理搜索器。

    在每个 depth 维持 beam_width 个候选，
    搜索到 max_depth 后返回得分最高的路径。
    """

    def __init__(
        self,
        expand_fn,       # (context: str) -> list[str]: 下一推理步骤的候选
        evaluate_fn,     # (partial_solution: str) -> float: 中间评估分数
        beam_width: int = 3,
        max_depth: int = 5,
    ):
        self.expand_fn = expand_fn
        self.evaluate_fn = evaluate_fn
        self.beam_width = beam_width
        self.max_depth = max_depth

    def search(self, problem: str) -> list[str]:
        """用 beam search 寻找最优推理路径。"""
        root = ThoughtNode(content=problem, depth=0)
        beam = [root]

        for depth in range(1, self.max_depth + 1):
            all_candidates = []

            for node in beam:
                context = "\n".join(node.path())
                expansions = self.expand_fn(context)

                for thought in expansions:
                    child = ThoughtNode(
                        content=thought,
                        depth=depth,
                        parent=node,
                    )
                    partial = "\n".join(child.path())
                    child.score = self.evaluate_fn(partial)
                    node.children.append(child)
                    all_candidates.append(child)

            # 只保留得分最高的 beam_width 个
            beam = heapq.nlargest(
                self.beam_width,
                all_candidates,
                key=lambda n: n.score,
            )

            if not beam:
                break

        # 从整个搜索结果中返回得分最高的路径
        best = max(beam, key=lambda n: n.score)
        return best.path()

import subprocess
import tempfile
import textwrap
from pathlib import Path
from dataclasses import dataclass


@dataclass
class CodeEvalResult:
    code: str
    passed: bool
    tests_passed: int
    tests_total: int
    error_output: str
    execution_time_ms: float


def evaluate_code_candidate(
    code: str,
    test_code: str,
    timeout_seconds: int = 30,
) -> CodeEvalResult:
    """用测试套件评估代码候选。

    实际运行 pytest 来判定通过/失败，
    并处理 timeout、import error 等异常情况。
    """
    import time

    with tempfile.TemporaryDirectory() as tmpdir:
        solution_path = Path(tmpdir) / "solution.py"
        test_path = Path(tmpdir) / "test_solution.py"

        solution_path.write_text(code, encoding="utf-8")

        # 构造测试文件，以便能从中 import solution
        full_test = textwrap.dedent(f"""\
            import sys
            sys.path.insert(0, "{tmpdir}")
            from solution import *

            {test_code}
        """)
        test_path.write_text(full_test, encoding="utf-8")

        start = time.monotonic()
        try:
            result = subprocess.run(
                ["python", "-m", "pytest", str(test_path), "-v", "--tb=short"],
                capture_output=True,
                text=True,
                timeout=timeout_seconds,
                cwd=tmpdir,
            )
            elapsed_ms = (time.monotonic() - start) * 1000

            # 从 pytest 输出中解析 passed/failed 数量
            passed = result.returncode == 0
            lines = result.stdout.splitlines()
            summary_line = [l for l in lines if "passed" in l or "failed" in l]

            return CodeEvalResult(
                code=code,
                passed=passed,
                tests_passed=result.stdout.count("PASSED"),
                tests_total=result.stdout.count("PASSED") + result.stdout.count("FAILED"),
                error_output=result.stderr if not passed else "",
                execution_time_ms=elapsed_ms,
            )
        except subprocess.TimeoutExpired:
            elapsed_ms = (time.monotonic() - start) * 1000
            return CodeEvalResult(
                code=code,
                passed=False,
                tests_passed=0,
                tests_total=0,
                error_output=f"Execution timed out after {timeout_seconds}s",
                execution_time_ms=elapsed_ms,
            )


def pick_best_code(
    candidates: list[str],
    test_code: str,
) -> CodeEvalResult | None:
    """在多个代码候选中选出通过测试数最多的一个。

    如果有多个候选全部通过，则优先选择执行时间更短的。
    """
    results = [evaluate_code_candidate(c, test_code) for c in candidates]

    # 第 1 优先级：通过测试数，第 2 优先级：执行时间（越低越好）
    results.sort(key=lambda r: (r.tests_passed, -r.execution_time_ms), reverse=True)

    return results[0] if results else None

# inference_policy.yaml
# 按问题难度分配 TTS 预算的策略

policies:
  easy:
    description: '简单 QA、FAQ、分类'
    strategy: single_pass
    samples: 1
    temperature: 0.0
    max_tokens: 256
    estimated_cost_per_query_usd: 0.001

  medium:
    description: '摘要、翻译、简单推理'
    strategy: best_of_n
    samples: 3
    temperature: 0.5
    max_tokens: 512
    scorer: format_and_length
    estimated_cost_per_query_usd: 0.005

  hard:
    description: '数学、编程、多步推理'
    strategy: self_consistency
    samples: 8
    temperature: 0.7
    max_tokens: 1024
    min_agreement: 3
    estimated_cost_per_query_usd: 0.015

  critical:
    description: '代码审查、安全判断、医疗/法律'
    strategy: verified_best_of_n
    samples: 5
    temperature: 0.6
    max_tokens: 2048
    verifier: domain_specific_verifier
    fallback_on_low_confidence: human_review
    estimated_cost_per_query_usd: 0.025

routing:
  classifier: difficulty_classifier_v2
  fallback_policy: medium

guardrails:
  max_total_tokens_per_query: 16384
  timeout_seconds: 30
  daily_budget_usd: 500
  alert_threshold_usd: 400

-- 用于 TTS 运营看板的查询
-- 一眼看清按天、按策略划分的性能/成本/延迟

SELECT
    date_trunc('day', created_at)     AS day,
    tts_strategy,
    COUNT(*)                          AS total_queries,
    AVG(correct::int)                 AS accuracy,
    PERCENTILE_CONT(0.50) WITHIN GROUP (ORDER BY latency_ms) AS p50_latency_ms,
    PERCENTILE_CONT(0.95) WITHIN GROUP (ORDER BY latency_ms) AS p95_latency_ms,
    AVG(samples_generated)            AS avg_samples,
    SUM(cost_usd)                     AS total_cost_usd,
    AVG(cost_usd)                     AS avg_cost_per_query,
    -- 效率指标：每个正确答案的成本
    SUM(cost_usd) / NULLIF(SUM(correct::int), 0) AS cost_per_correct_answer
FROM reasoning_eval_runs
WHERE created_at >= CURRENT_DATE - INTERVAL '14 days'
GROUP BY 1, 2
ORDER BY 1 DESC, 2;