Модуль 6. LLM как политика: токены, логарифмы вероятностей и KL

import numpy as np


def stable_log_softmax(logits: np.ndarray, axis: int = -1) -> np.ndarray:
    z = np.asarray(logits, dtype=np.float64)
    if z.ndim == 0:
        raise ValueError("logits must have at least one dimension")
    if not np.all(np.isfinite(z)):
        raise ValueError("logits must be finite")
    if isinstance(axis, bool) or not isinstance(axis, (int, np.integer)):
        raise TypeError("axis must be an integer")
    axis = int(axis)
    if axis < 0:
        axis += z.ndim
    if not 0 <= axis < z.ndim:
        raise ValueError("axis outside array dimensions")
    maximum = np.max(z, axis=axis, keepdims=True)
    shifted = z - maximum
    log_normalizer = np.log(np.exp(shifted).sum(axis=axis, keepdims=True))
    return shifted - log_normalizer

import numpy as np


def _stable_log_softmax(logits: np.ndarray, axis: int = -1) -> np.ndarray:
    z = np.asarray(logits, dtype=np.float64)
    if not np.all(np.isfinite(z)):
        raise ValueError("logits must be finite")
    maximum = np.max(z, axis=axis, keepdims=True)
    shifted = z - maximum
    return shifted - np.log(np.exp(shifted).sum(axis=axis, keepdims=True))


def causal_response_logprobs(
    model_logits: np.ndarray,
    input_ids: np.ndarray,
    response_token_mask: np.ndarray,
) -> tuple[np.ndarray, np.ndarray]:
    logits = np.asarray(model_logits, dtype=np.float64)
    ids = np.asarray(input_ids, dtype=np.int64)
    mask = np.asarray(response_token_mask, dtype=bool)
    if logits.ndim != 3:
        raise ValueError("model_logits must have shape [B, L, V]")
    if ids.ndim != 2 or ids.shape != logits.shape[:2] or mask.shape != ids.shape:
        raise ValueError("input_ids and response_token_mask must match logits[:2]")
    if logits.shape[1] < 2:
        raise ValueError("the sequence must contain at least two positions")
    vocabulary_size = logits.shape[-1]
    if np.any(ids < 0) or np.any(ids >= vocabulary_size):
        raise ValueError("token id outside vocabulary")

    shifted_logprobs = _stable_log_softmax(logits[:, :-1, :], axis=-1)
    targets = ids[:, 1:]
    selected = np.take_along_axis(shifted_logprobs, targets[..., None], axis=-1)[..., 0]
    return selected, mask[:, 1:]

from dataclasses import dataclass
import numpy as np


@dataclass(frozen=True)
class SequenceMetrics:
    token_count: np.ndarray
    sum_logprob: np.ndarray
    mean_nll: np.ndarray
    perplexity: np.ndarray


def masked_sequence_metrics(
    selected_logprobs: np.ndarray,
    mask: np.ndarray,
) -> SequenceMetrics:
    logprobs = np.asarray(selected_logprobs, dtype=np.float64)
    active = np.asarray(mask, dtype=bool)
    if logprobs.shape != active.shape:
        raise ValueError("selected_logprobs and mask must have the same shape")
    if logprobs.ndim == 1:
        logprobs = logprobs[None, :]
        active = active[None, :]
    elif logprobs.ndim != 2:
        raise ValueError("expected shape [B, T] or [T]")
    if np.any(~np.isfinite(logprobs[active])):
        raise ValueError("active log-probabilities must be finite")

    token_count = active.sum(axis=1).astype(np.int64)
    if np.any(token_count == 0):
        raise ValueError("empty sequence mask")
    sum_logprob = np.where(active, logprobs, 0.0).sum(axis=1)
    mean_nll = -sum_logprob / token_count
    return SequenceMetrics(
        token_count=token_count,
        sum_logprob=sum_logprob,
        mean_nll=mean_nll,
        perplexity=np.exp(mean_nll),
    )

import numpy as np


def _stable_log_softmax(logits: np.ndarray, axis: int = -1) -> np.ndarray:
    z = np.asarray(logits, dtype=np.float64)
    if not np.all(np.isfinite(z)):
        raise ValueError("logits must be finite")
    maximum = np.max(z, axis=axis, keepdims=True)
    shifted = z - maximum
    return shifted - np.log(np.exp(shifted).sum(axis=axis, keepdims=True))


def exact_kl_from_logits(p_logits: np.ndarray, q_logits: np.ndarray) -> np.ndarray:
    p_logits = np.asarray(p_logits, dtype=np.float64)
    q_logits = np.asarray(q_logits, dtype=np.float64)
    if p_logits.shape != q_logits.shape:
        raise ValueError("p_logits and q_logits must have the same shape")
    log_p = _stable_log_softmax(p_logits, axis=-1)
    log_q = _stable_log_softmax(q_logits, axis=-1)
    p = np.exp(log_p)
    return np.sum(p * (log_p - log_q), axis=-1)


def kl_sample_values(log_q: np.ndarray, log_p: np.ndarray) -> dict[str, np.ndarray]:
    log_q = np.asarray(log_q, dtype=np.float64)
    log_p = np.asarray(log_p, dtype=np.float64)
    if log_q.shape != log_p.shape or not (np.all(np.isfinite(log_q)) and np.all(np.isfinite(log_p))):
        raise ValueError("log_q and log_p must be finite arrays of the same shape")
    log_ratio = log_p - log_q
    ratio = np.exp(log_ratio)
    return {
        "k1": -log_ratio,
        "k2": 0.5 * log_ratio**2,
        "k3": ratio - 1.0 - log_ratio,
    }

import numpy as np


def _softmax_1d(logits: np.ndarray, temperature: float = 1.0) -> np.ndarray:
    z = np.asarray(logits, dtype=np.float64)
    if z.ndim != 1 or z.size == 0 or not np.all(np.isfinite(z)):
        raise ValueError("logits must be a non-empty finite one-dimensional array")
    if not np.isfinite(temperature) or temperature <= 0.0:
        raise ValueError("temperature must be a positive finite number")
    scaled = z / temperature
    shifted = scaled - np.max(scaled)
    probabilities = np.exp(shifted)
    return probabilities / probabilities.sum()


def _renormalize_support(probabilities: np.ndarray, keep: np.ndarray) -> np.ndarray:
    out = np.zeros_like(probabilities, dtype=np.float64)
    out[keep] = probabilities[keep]
    mass = float(out.sum())
    if mass <= 0.0:
        raise ValueError("the retained support has zero probability mass")
    return out / mass


def temperature_distribution(logits: np.ndarray, temperature: float) -> np.ndarray:
    return _softmax_1d(logits, temperature)


def top_k_distribution(logits: np.ndarray, k: int) -> np.ndarray:
    probabilities = _softmax_1d(logits)
    if isinstance(k, bool) or not isinstance(k, (int, np.integer)) or not (1 <= int(k) <= probabilities.size):
        raise ValueError("k must be an integer between 1 and the vocabulary size")
    order = np.argsort(-probabilities, kind="stable")
    return _renormalize_support(probabilities, order[: int(k)])


def top_p_distribution(
    logits: np.ndarray,
    top_p: float,
    *,
    min_tokens_to_keep: int = 1,
) -> np.ndarray:
    probabilities = _softmax_1d(logits)
    if not np.isfinite(top_p) or not (0.0 < top_p <= 1.0):
        raise ValueError("top_p must lie in (0, 1]")
    if (
        isinstance(min_tokens_to_keep, bool)
        or not isinstance(min_tokens_to_keep, (int, np.integer))
        or not (1 <= int(min_tokens_to_keep) <= probabilities.size)
    ):
        raise ValueError("min_tokens_to_keep must be a valid positive integer")

    order = np.argsort(-probabilities, kind="stable")
    cumulative = np.cumsum(probabilities[order])
    threshold_count = int(np.searchsorted(cumulative, top_p, side="left")) + 1
    keep_count = max(int(min_tokens_to_keep), threshold_count)
    return _renormalize_support(probabilities, order[:keep_count])

from dataclasses import dataclass
from typing import Any, Mapping
import numpy as np


@dataclass(frozen=True)
class RolloutBatchSummary:
    batch_size: int
    max_response_tokens: int
    active_tokens: int
    terminated_sequences: int
    truncated_sequences: int
    mean_response_length: float


def validate_rollout_batch(batch: Mapping[str, Any]) -> RolloutBatchSummary:
    required = {
        "response_ids", "response_mask", "old_logprobs", "ref_logprobs",
        "sequence_rewards", "terminated", "truncated", "stop_reason",
    }
    missing = sorted(required - set(batch))
    if missing:
        raise ValueError("missing fields: " + ", ".join(missing))

    ids = np.asarray(batch["response_ids"])
    mask = np.asarray(batch["response_mask"], dtype=bool)
    old = np.asarray(batch["old_logprobs"], dtype=np.float64)
    ref = np.asarray(batch["ref_logprobs"], dtype=np.float64)
    rewards = np.asarray(batch["sequence_rewards"], dtype=np.float64)
    terminated = np.asarray(batch["terminated"], dtype=bool)
    truncated = np.asarray(batch["truncated"], dtype=bool)
    reasons = np.asarray(batch["stop_reason"], dtype=object)

    if ids.ndim != 2 or not (ids.shape == mask.shape == old.shape == ref.shape):
        raise ValueError("token-level arrays must align")
    batch_size, max_tokens = ids.shape
    if not (rewards.shape == terminated.shape == truncated.shape == reasons.shape == (batch_size,)):
        raise ValueError("sequence-level arrays must have shape (batch,)")

    lengths = mask.sum(axis=1).astype(int)
    if np.any(lengths <= 0):
        raise ValueError("every response must contain an active token")
    expected = np.arange(max_tokens)[None, :] < lengths[:, None]
    if not np.array_equal(mask, expected):
        raise ValueError("response mask must be a contiguous right-padded prefix")
    if np.any(~np.isfinite(old[mask])) or np.any(~np.isfinite(ref[mask])):
        raise ValueError("active log-probabilities must be finite")
    if np.any(np.isfinite(old[~mask])) or np.any(np.isfinite(ref[~mask])):
        raise ValueError("padding log-probabilities must be NaN")
    if np.any(~np.isfinite(rewards)):
        raise ValueError("rewards must be finite")
    if any(not isinstance(x, str) or not x.strip() for x in reasons.tolist()):
        raise ValueError("every stop_reason must be a non-empty string")

    return RolloutBatchSummary(
        batch_size=int(batch_size),
        max_response_tokens=int(max_tokens),
        active_tokens=int(mask.sum()),
        terminated_sequences=int(terminated.sum()),
        truncated_sequences=int(truncated.sum()),
        mean_response_length=float(lengths.mean()),
    )