Модуль 7. Модели награды и данные о предпочтениях

from __future__ import annotations

from dataclasses import dataclass
import math
import random

import numpy as np
import torch

SEED = 7
RUN_REAL_MODEL = False
SMOKE_MODE = True

random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)

print({"torch": torch.__version__, "device": "cuda" if torch.cuda.is_available() else "cpu"})

@dataclass(frozen=True)
class PairBatch:
    chosen: torch.Tensor
    rejected: torch.Tensor
    prompt_ids: list[str]
    first: torch.Tensor
    second: torch.Tensor
    first_preferred: torch.Tensor


def make_synthetic_pairs(n_prompts: int = 800, pairs_per_prompt: int = 2, seed: int = SEED):
    rng = np.random.default_rng(seed)
    true_w = np.array([1.5, 0.8, -0.9, 0.4], dtype=np.float32)
    rows = []
    for prompt_index in range(n_prompts):
        prompt_id = f"p{prompt_index:05d}"
        prompt_shift = rng.normal(0.0, 0.5)
        for _ in range(pairs_per_prompt):
            a = rng.normal(size=4).astype(np.float32)
            b = rng.normal(size=4).astype(np.float32)
            utility_a = float(a @ true_w + prompt_shift + rng.normal(0.0, 0.35))
            utility_b = float(b @ true_w + prompt_shift + rng.normal(0.0, 0.35))
            chosen, rejected = (a, b) if utility_a >= utility_b else (b, a)
            rows.append((prompt_id, chosen, rejected, a, b, float(utility_a >= utility_b)))
    return rows


def split_by_prompt(rows, eval_fraction: float = 0.2, seed: int = SEED):
    prompt_ids = sorted({row[0] for row in rows})
    rng = np.random.default_rng(seed)
    rng.shuffle(prompt_ids)
    cut = int(round(len(prompt_ids) * (1.0 - eval_fraction)))
    train_ids = set(prompt_ids[:cut])
    eval_ids = set(prompt_ids[cut:])
    assert train_ids.isdisjoint(eval_ids)

    def pack(selected_ids):
        chosen = np.stack([row[1] for row in rows if row[0] in selected_ids])
        rejected = np.stack([row[2] for row in rows if row[0] in selected_ids])
        ids = [row[0] for row in rows if row[0] in selected_ids]
        first = np.stack([row[3] for row in rows if row[0] in selected_ids])
        second = np.stack([row[4] for row in rows if row[0] in selected_ids])
        labels = [row[5] for row in rows if row[0] in selected_ids]
        return PairBatch(torch.tensor(chosen), torch.tensor(rejected), ids,
                         torch.tensor(first), torch.tensor(second), torch.tensor(labels))

    return pack(train_ids), pack(eval_ids)


rows = make_synthetic_pairs()
train_batch, eval_batch = split_by_prompt(rows)
print({"train_pairs": len(train_batch.prompt_ids), "eval_pairs": len(eval_batch.prompt_ids)})

class TinyRewardModel(torch.nn.Module):
    def __init__(self, n_features: int = 4):
        super().__init__()
        self.score = torch.nn.Linear(n_features, 1)

    def forward(self, features: torch.Tensor) -> torch.Tensor:
        return self.score(features).squeeze(-1)


def pair_metrics(model: TinyRewardModel, batch: PairBatch) -> dict[str, float]:
    model.eval()
    with torch.no_grad():
        delta = model(batch.chosen) - model(batch.rejected)
        probability = torch.sigmoid(delta)
        target = torch.ones_like(probability)
        loss = torch.nn.functional.binary_cross_entropy(probability, target)
        brier = torch.mean((probability - target) ** 2)
        return {
            "accuracy": float((delta > 0).float().mean()),
            "log_loss": float(loss),
            "brier": float(brier),
            "mean_margin": float(delta.mean()),
        }


model = TinyRewardModel()
optimizer = torch.optim.AdamW(model.parameters(), lr=0.03, weight_decay=1e-3)
steps = 120 if SMOKE_MODE else 600
batch_size = 256
rng = np.random.default_rng(SEED)

for step in range(steps):
    index = torch.tensor(rng.integers(0, len(train_batch.prompt_ids), size=batch_size))
    chosen_score = model(train_batch.chosen[index])
    rejected_score = model(train_batch.rejected[index])
    loss = torch.nn.functional.softplus(-(chosen_score - rejected_score)).mean()
    # Optional gauge convention: discourage a drifting common score origin.
    center = 0.5 * (chosen_score.mean() + rejected_score.mean())
    total_loss = loss + 1e-2 * center.square()
    optimizer.zero_grad(set_to_none=True)
    total_loss.backward()
    optimizer.step()

print("train", pair_metrics(model, train_batch))
print("eval ", pair_metrics(model, eval_batch))

def reliability_bins(probability: np.ndarray, label: np.ndarray, n_bins: int = 10):
    probability = np.asarray(probability, dtype=np.float64).reshape(-1)
    label = np.asarray(label, dtype=np.float64).reshape(-1)
    if probability.shape != label.shape or probability.size == 0:
        raise ValueError("aligned non-empty arrays required")
    indices = np.minimum((probability * n_bins).astype(int), n_bins - 1)
    bins = []
    for bin_id in range(n_bins):
        mask = indices == bin_id
        if np.any(mask):
            bins.append({
                "bin": bin_id,
                "count": int(mask.sum()),
                "confidence": float(probability[mask].mean()),
                "frequency": float(label[mask].mean()),
            })
    return bins


with torch.no_grad():
    eval_delta = model(eval_batch.first) - model(eval_batch.second)
    eval_probability = torch.sigmoid(eval_delta).numpy()
labels = eval_batch.first_preferred.numpy()
bins = reliability_bins(eval_probability, labels, n_bins=5)
ece = sum(item["count"] / len(labels) * abs(item["confidence"] - item["frequency"]) for item in bins)
print({"ece_5_bins": ece, "bins": bins})

def toy_overoptimization(pool_sizes=(1, 2, 4, 8, 16, 32, 64, 128, 256, 512), trials=4000, seed=7):
    rng = np.random.default_rng(seed)
    proxy_means = []
    target_means = []
    for pool_size in pool_sizes:
        useful = rng.normal(size=(trials, pool_size))
        nuisance = rng.normal(size=(trials, pool_size))
        proxy = useful + nuisance
        target = useful - np.maximum(nuisance - 1.0, 0.0) ** 2
        selected = np.argmax(proxy, axis=1)
        rows = np.arange(trials)
        proxy_means.append(proxy[rows, selected].mean())
        target_means.append(target[rows, selected].mean())
    return np.asarray(proxy_means), np.asarray(target_means)


pool_sizes = np.array([1, 2, 4, 8, 16, 32, 64, 128, 256, 512])
proxy, target = toy_overoptimization(pool_sizes)
# Finite Monte Carlo estimates need not be monotone for every seed/budget.
assert np.isfinite(proxy).all() and np.isfinite(target).all()
print({"best_target_pool": int(pool_sizes[np.argmax(target)]), "final_proxy": proxy[-1], "final_target": target[-1]})

if RUN_REAL_MODEL:
    import torch
    from datasets import load_dataset
    from peft import LoraConfig, TaskType
    from trl import RewardConfig, RewardTrainer

    MODEL_ID = "Qwen/Qwen3-0.6B"
    DATASET_ID = "trl-lib/ultrafeedback_binarized"
    TRAIN_PAIRS = 512 if SMOKE_MODE else 10_000
    EVAL_PAIRS = 256 if SMOKE_MODE else 2_000
    MAX_LENGTH = 512 if SMOKE_MODE else 1024

    dataset = load_dataset(DATASET_ID)
    train_dataset = dataset["train"].select(range(min(TRAIN_PAIRS, len(dataset["train"]))))
    eval_split_name = "test" if "test" in dataset else next(name for name in dataset if name != "train")
    eval_dataset = dataset[eval_split_name].select(range(min(EVAL_PAIRS, len(dataset[eval_split_name]))))

    peft_config = LoraConfig(
        task_type=TaskType.SEQ_CLS,
        r=16,
        lora_alpha=32,
        lora_dropout=0.05,
        target_modules=["q_proj", "v_proj"],
        modules_to_save=["score"],
    )

    args = RewardConfig(
        output_dir="module-07-rm-smoke" if SMOKE_MODE else "module-07-rm-full",
        max_length=MAX_LENGTH,
        learning_rate=1e-4,
        per_device_train_batch_size=2,
        per_device_eval_batch_size=2,
        gradient_accumulation_steps=8,
        max_steps=40 if SMOKE_MODE else -1,
        num_train_epochs=1.0,
        eval_strategy="steps",
        eval_steps=20 if SMOKE_MODE else 200,
        save_strategy="no" if SMOKE_MODE else "steps",
        logging_steps=5,
        fp16=True,
        gradient_checkpointing=True,
        center_rewards_coefficient=1e-2,
        report_to="none",
        seed=SEED,
    )

    trainer = RewardTrainer(
        model=MODEL_ID,
        args=args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        peft_config=peft_config,
    )
    trainer.train()
    metrics = trainer.evaluate()
    print(metrics)
else:
    print("Real-model branch skipped. Set RUN_REAL_MODEL=True in a reviewed Colab runtime.")