Implicit Chain-of-Thought: рассуждения без слов

# Формат данных для поэтапного обучения
# Этап 1: полный CoT
stage_1_data = [
    {
        "input": "Задача: 23 × 17",
        "output": "Шаг 1: 23 × 10 = 230\nШаг 2: 23 × 7 = 161\nШаг 3: 230 + 161 = 391\nОтвет: 391"
    },
    # ...
]

# Этап 2: сокращённый CoT (оставляем ключевые шаги)
stage_2_data = [
    {
        "input": "Задача: 23 × 17",
        "output": "230 + 161 = 391\nОтвет: 391"
    },
]

# Этап 3: только ответ
stage_3_data = [
    {
        "input": "Задача: 23 × 17",
        "output": "Ответ: 391"
    },
]

from transformers import (
    AutoModelForCausalLM, AutoTokenizer,
    TrainingArguments, Trainer
)
from datasets import Dataset

model_name = "meta-llama/Llama-3.1-8B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

def progressive_finetune(
    model,
    stages: list[list[dict]],
    epochs_per_stage: int = 3
):
    """
    Поэтапное обучение с постепенным удалением CoT.
    Каждый этап начинается с чекпоинта предыдущего.
    """
    for i, stage_data in enumerate(stages):
        print(f"=== Stage {i+1}/{len(stages)} ===")
        print(f"  Samples: {len(stage_data)}")
        print(f"  Avg output tokens: {avg_tokens(stage_data)}")

        dataset = Dataset.from_list(stage_data)
        dataset = dataset.map(
            lambda x: tokenize(x, tokenizer),
            batched=True
        )

        args = TrainingArguments(
            output_dir=f"./implicit-cot-stage-{i+1}",
            num_train_epochs=epochs_per_stage,
            per_device_train_batch_size=4,
            learning_rate=2e-5 * (0.5 ** i),  # снижаем LR
            warmup_ratio=0.1,
            save_strategy="epoch",
            logging_steps=50,
            bf16=True,
        )

        trainer = Trainer(
            model=model,
            args=args,
            train_dataset=dataset,
        )
        trainer.train()

        print(f"  Stage {i+1} complete. Evaluating...")
        evaluate(model, tokenizer, stage=i+1)

    return model


# Запуск
stages = [stage_1_data, stage_2_data, stage_3_data]
model = progressive_finetune(model, stages, epochs_per_stage=3)

from openai import OpenAI

client = OpenAI()

def generate_cot_dataset(
    tasks: list[str],
    teacher_model: str = "gpt-4o"
) -> list[dict]:
    """
    Генерация CoT-рассуждений от модели-учителя
    для последующей дистилляции.
    """
    dataset = []

    for task in tasks:
        # Шаг 1: получить полное CoT от учителя
        cot_response = client.chat.completions.create(
            model=teacher_model,
            messages=[{
                "role": "user",
                "content": f"Реши задачу пошагово:\n{task}"
            }],
            temperature=0
        )
        cot_text = cot_response.choices[0].message.content

        # Шаг 2: извлечь только ответ
        answer_response = client.chat.completions.create(
            model=teacher_model,
            messages=[{
                "role": "user",
                "content": (
                    f"Вот решение задачи:\n{cot_text}\n\n"
                    "Извлеки только финальный ответ, "
                    "одно число или одно предложение."
                )
            }],
            temperature=0
        )
        answer = answer_response.choices[0].message.content

        dataset.append({
            "input": task,
            "cot": cot_text,       # для stage 1
            "answer": answer,       # для stage 3
        })

    return dataset

import torch
import numpy as np

def probe_hidden_reasoning(
    model,
    tokenizer,
    task: str,
    layer_idx: int = -1
) -> dict:
    """
    Анализ скрытых представлений: ищем следы
    рассуждений в hidden states модели.
    """
    inputs = tokenizer(task, return_tensors="pt")

    with torch.no_grad():
        outputs = model(
            **inputs,
            output_hidden_states=True
        )

    hidden = outputs.hidden_states[layer_idx][0]

    # Метрика 1: активация скрытых слоёв
    # У модели с Implicit CoT активация выше,
    # чем у модели без обучения на CoT
    activation_norm = torch.norm(hidden, dim=-1).mean().item()

    # Метрика 2: энтропия предсказаний
    # Низкая энтропия = модель уверена в ответе
    logits = outputs.logits[0, -1, :]
    probs = torch.softmax(logits, dim=-1)
    entropy = -(probs * torch.log(probs + 1e-10)).sum().item()

    return {
        "activation_norm": activation_norm,
        "entropy": entropy,
        "confident": entropy < 2.0,
    }


# Сравнение базовой модели и Implicit CoT
base_result = probe_hidden_reasoning(base_model, tokenizer, task)
icot_result = probe_hidden_reasoning(icot_model, tokenizer, task)

print(f"Base activation: {base_result['activation_norm']:.2f}")
print(f"ICoT activation: {icot_result['activation_norm']:.2f}")
# ICoT обычно показывает более высокую активацию —
# модель «работает» внутри, даже не генерируя токены

def calculate_savings(
    daily_requests: int = 100_000,
    avg_input_tokens: int = 50,
    explicit_cot_output: int = 90,
    implicit_cot_output: int = 8,
    input_price_per_m: float = 2.50,   # $/1M tokens
    output_price_per_m: float = 10.00,  # $/1M tokens
) -> dict:
    """Расчёт экономии при переходе на Implicit CoT."""

    def monthly_cost(output_tokens: int) -> float:
        input_cost = (
            daily_requests * avg_input_tokens / 1_000_000
            * input_price_per_m * 30
        )
        output_cost = (
            daily_requests * output_tokens / 1_000_000
            * output_price_per_m * 30
        )
        return input_cost + output_cost

    explicit = monthly_cost(explicit_cot_output)
    implicit = monthly_cost(implicit_cot_output)

    return {
        "explicit_cot_monthly": f"${explicit:,.0f}",
        "implicit_cot_monthly": f"${implicit:,.0f}",
        "savings_monthly": f"${explicit - implicit:,.0f}",
        "savings_percent": f"{(1 - implicit/explicit) * 100:.0f}%",
    }


print(calculate_savings())
# explicit_cot_monthly: $3,075
# implicit_cot_monthly: $615
# savings_monthly: $2,460
# savings_percent: 80%