Chain-of-Thought Decoding

from openai import OpenAI
import math

client = OpenAI()


def cot_decoding(
    question: str,
    k: int = 10,
    model: str = "gpt-4o",
    max_tokens: int = 512,
) -> dict:
    """
    CoT-Decoding: генерация k альтернативных путей
    и выбор наиболее уверенного.

    1. Получаем top-k первых токенов через logprobs
    2. Для каждого токена генерируем полный ответ
    3. Оцениваем уверенность каждого пути
    4. Возвращаем путь с максимальной уверенностью
    """
    # Шаг 1: получить top-k первых токенов
    first_token_response = client.chat.completions.create(
        model=model,
        messages=[{"role": "user", "content": question}],
        max_tokens=1,
        logprobs=True,
        top_logprobs=k,
    )

    top_tokens = (
        first_token_response
        .choices[0]
        .logprobs
        .content[0]
        .top_logprobs
    )

    # Шаг 2: для каждого первого токена сгенерировать
    # полный ответ
    paths = []
    for token_info in top_tokens:
        prefix = token_info.token
        prefix_logprob = token_info.logprob

        # Генерируем продолжение, начиная с этого
        # токена, через assistant prefill
        response = client.chat.completions.create(
            model=model,
            messages=[
                {"role": "user", "content": question},
                {
                    "role": "assistant",
                    "content": prefix,
                },
            ],
            max_tokens=max_tokens,
            logprobs=True,
            top_logprobs=5,
            temperature=0,
        )

        choice = response.choices[0]
        full_text = prefix + (choice.message.content or "")

        # Шаг 3: вычислить уверенность пути
        confidence = compute_path_confidence(choice)

        paths.append({
            "first_token": prefix,
            "first_token_logprob": prefix_logprob,
            "full_text": full_text,
            "confidence": confidence,
        })

    # Шаг 4: выбрать путь с максимальной уверенностью
    best = max(paths, key=lambda p: p["confidence"])

    return {
        "answer": best["full_text"],
        "confidence": best["confidence"],
        "paths_explored": len(paths),
        "all_paths": paths,
    }


def compute_path_confidence(choice) -> float:
    """
    Метрика уверенности пути.
    Считаем среднюю дельту между top-1 и top-2
    логарифмическими вероятностями по всем токенам.
    Высокая дельта = модель уверена в каждом токене.
    """
    if not choice.logprobs or not choice.logprobs.content:
        return 0.0

    deltas = []
    for token_lp in choice.logprobs.content:
        top = token_lp.top_logprobs
        if len(top) >= 2:
            delta = top[0].logprob - top[1].logprob
            deltas.append(delta)

    return sum(deltas) / len(deltas) if deltas else 0.0


# Пример использования
result = cot_decoding(
    "Джон купил 3 пакета яблок по 4 штуки в каждом. "
    "Он отдал 5 яблок другу. Сколько яблок осталось?"
)
print(f"Ответ: {result['answer']}")
print(f"Уверенность: {result['confidence']:.3f}")
print(f"Путей исследовано: {result['paths_explored']}")

import anthropic

client = anthropic.Anthropic()

# Типичные CoT-стартеры, которые запускают
# рассуждения в моделях
COT_PREFIXES = [
    "",                     # без префикса (baseline)
    "Давайте разберёмся",   # русский CoT-стартер
    "Для решения",          # русский CoT-стартер
    "Let me think",         # английский CoT-стартер
    "Step 1:",              # пошаговый стартер
    "First,",               # последовательный стартер
]


def cot_decoding_anthropic(
    question: str,
    prefixes: list[str] = COT_PREFIXES,
    model: str = "claude-sonnet-4-20250514",
    max_tokens: int = 1024,
) -> dict:
    """
    CoT-Decoding для Claude через assistant prefill.
    Вместо top-k токенов используем заранее
    определённые CoT-стартеры.
    """
    paths = []

    for prefix in prefixes:
        messages = [{"role": "user", "content": question}]

        # Используем assistant prefill, если есть
        # префикс
        if prefix:
            messages.append({
                "role": "assistant",
                "content": prefix,
            })

        response = client.messages.create(
            model=model,
            max_tokens=max_tokens,
            messages=messages,
            temperature=0,
        )

        full_text = prefix + response.content[0].text

        # Эвристика уверенности: длина рассуждений
        # коррелирует с качеством ответа
        has_reasoning = any(
            marker in full_text.lower()
            for marker in [
                "шаг", "step", "значит", "итого",
                "получаем", "therefore", "thus",
                "сначала", "во-первых",
            ]
        )

        paths.append({
            "prefix": prefix or "(greedy)",
            "full_text": full_text,
            "has_reasoning": has_reasoning,
            "length": len(full_text),
        })

    # Предпочитаем пути с рассуждениями
    reasoning_paths = [p for p in paths if p["has_reasoning"]]
    best = (
        max(reasoning_paths, key=lambda p: p["length"])
        if reasoning_paths
        else paths[0]
    )

    return {
        "answer": best["full_text"],
        "prefix_used": best["prefix"],
        "has_reasoning": best["has_reasoning"],
        "paths_explored": len(paths),
    }

import re


def extract_final_answer(text: str) -> str | None:
    """Извлечь финальный числовой ответ из текста."""
    patterns = [
        r"(?:ответ|answer)[:\s]*(\d+[\.,]?\d*)",
        r"(?:итого|total|=)\s*(\d+[\.,]?\d*)",
        r"(\d+[\.,]?\d*)\s*$",
    ]
    for pattern in patterns:
        match = re.search(pattern, text, re.IGNORECASE)
        if match:
            return match.group(1).replace(",", ".")
    return None


def cot_decoding_with_voting(
    question: str,
    k: int = 10,
    model: str = "gpt-4o",
) -> dict:
    """
    CoT-Decoding с голосованием по ответам.
    Комбинирует CoT-Decoding с Self-Consistency:
    - Генерируем k путей через top-k
    - Извлекаем финальный ответ из каждого
    - Взвешиваем ответы по confidence
    - Выбираем ответ с максимальным взвешенным счётом
    """
    result = cot_decoding(question, k=k, model=model)
    paths = result["all_paths"]

    # Группируем по ответам с взвешиванием
    answer_scores: dict[str, float] = {}
    answer_paths: dict[str, list] = {}

    for path in paths:
        answer = extract_final_answer(path["full_text"])
        if answer is None:
            continue

        weight = max(path["confidence"], 0.01)

        if answer not in answer_scores:
            answer_scores[answer] = 0
            answer_paths[answer] = []
        answer_scores[answer] += weight
        answer_paths[answer].append(path)

    if not answer_scores:
        return result

    best_answer = max(
        answer_scores, key=answer_scores.get
    )

    return {
        "answer": best_answer,
        "total_confidence": answer_scores[best_answer],
        "votes": len(answer_paths[best_answer]),
        "all_answers": answer_scores,
        "best_path": answer_paths[best_answer][0],
    }