modilify/Modilify-Mk1

1# Copyright 2026 Modilify2# SPDX-License-Identifier: LicenseRef-Modilify-Open-Model-1.03"""Confidence-and-entropy commit policy for inference."""4 5from __future__ import annotations6 7from collections.abc import Sequence8from dataclasses import dataclass9import math10 11import torch12 13from .latent_deliberation import (14    advance_trajectory_clocks,15    should_force_trajectory_jump,16)17 18FUSED_EPS = 1e-619 20 21def fused_commit_confidence(22    proposal_confidence: torch.Tensor,23    token_entropy: torch.Tensor,24    *,25    vocab_size: int = 256000,26    eps: float = FUSED_EPS,27) -> torch.Tensor:28    """Fuse proposal confidence with token entropy.29 30    Effective confidence uses an excess-entropy sigmoid:31 32        p = clamp(proposal_confidence, eps, 1 - eps)33        h2 = -p * log(p) - (1 - p) * log(1 - p)34        excess = max(token_entropy - h2, 0)35        fused = sigmoid(logit(p) - excess) ** 236 37    When token entropy equals the binary entropy of ``p``, fused confidence38    equals ``p ** 2``. Entropy above that binary entropy reduces confidence.39 40    Args:41        proposal_confidence: Sampled-token probabilities, shape ``[batch, canvas]``.42        token_entropy: Token-level entropy, shape ``[batch, canvas]``.43        vocab_size: Unused; retained so callers can pass the model vocabulary.44        eps: Clamp that keeps logits finite.45 46    Returns:47        Fused commit confidence in ``(eps, 1 - eps)``.48    """49 50    del vocab_size51    p = proposal_confidence.float().clamp(min=eps, max=1.0 - eps)52    entropy = token_entropy.float().clamp(min=0.0)53    binary_entropy = -p * torch.log(p) - (1.0 - p) * torch.log1p(-p)54    excess = (entropy - binary_entropy).clamp(min=0.0)55    logit_p = torch.log(p) - torch.log1p(-p)56    fused = torch.sigmoid(logit_p - excess).square()57    return fused.clamp(min=eps, max=1.0 - eps)58 59 60def fused_commit_failure_rate(61    proposal_confidence: torch.Tensor,62    token_entropy: torch.Tensor,63    **kwargs: object,64) -> torch.Tensor:65    """Return ``1 - fused_commit_confidence``."""66 67    return 1.0 - fused_commit_confidence(68        proposal_confidence, token_entropy, **kwargs69    )70 71 72@dataclass(frozen=True)73class CommitPolicyDecision:74    """One inference transition from proposal to committed prefix."""75 76    normal_lengths: torch.LongTensor77    commit_lengths: torch.LongTensor78    commit_token_ids: torch.LongTensor79    jump_rows: torch.BoolTensor80    ponder_steps: torch.IntTensor81    stagnation_steps: torch.IntTensor82 83 84def prefix_failure_commit_lengths(85    failure_rate: torch.Tensor,86    *,87    failure_budget: float,88    valid_mask: torch.BoolTensor | None = None,89) -> torch.LongTensor:90    """Return the longest prefix with ``cumsum(failure_rate) < budget``.91 92    Args:93        failure_rate: Per-token failure rates, shape ``[batch, canvas]``.94        failure_budget: Strict cumulative risk limit.95        valid_mask: Optional canvas mask with the same shape.96 97    Returns:98        Commit lengths of shape ``[batch]``.99    """100 101    if failure_rate.ndim != 2:102        raise ValueError("Failure rate must have shape [batch, canvas].")103    if not math.isfinite(failure_budget) or failure_budget <= 0:104        raise ValueError("Commit failure budget must be finite and positive.")105    if valid_mask is None:106        valid_mask = torch.ones_like(failure_rate, dtype=torch.bool)107    if valid_mask.shape != failure_rate.shape:108        raise ValueError("Commit validity mask must match failure rate.")109 110    risk = failure_rate.float().clamp(0.0, 1.0) * valid_mask.to(torch.float32)111    cumulative_risk = risk.cumsum(dim=-1)112    contiguous_valid = valid_mask.long().cumprod(dim=-1).bool()113    allowed = cumulative_risk.lt(float(failure_budget)) & contiguous_valid114    return allowed.long().cumprod(dim=-1).sum(dim=-1)115 116 117def first_committed_token_lengths(118    proposal: torch.LongTensor,119    commit_lengths: torch.LongTensor,120    token_id: int | Sequence[int],121) -> torch.LongTensor:122    """Clip each prefix immediately after its first stop token.123 124    Args:125        proposal: Token IDs, shape ``[batch, canvas]``.126        commit_lengths: Unclipped prefix lengths, shape ``[batch]``.127        token_id: One stop ID or a sequence of stop IDs.128 129    Returns:130        Clipped commit lengths of shape ``[batch]``.131    """132 133    if proposal.ndim != 2 or commit_lengths.shape != proposal.shape[:1]:134        raise ValueError("Proposal and commit lengths must share a batch dimension.")135    positions = torch.arange(proposal.shape[1], device=proposal.device).unsqueeze(0)136    committed = positions.lt(commit_lengths[:, None])137    stop_token_ids = (138        (int(token_id),)139        if isinstance(token_id, int)140        else tuple(dict.fromkeys(int(value) for value in token_id))141    )142    if not stop_token_ids:143        raise ValueError("At least one stop token ID is required.")144    matches = proposal.eq(stop_token_ids[0])145    for value in stop_token_ids[1:]:146        matches |= proposal.eq(value)147    matches &= committed148    sentinel = torch.full_like(positions, proposal.shape[1])149    first = torch.where(matches, positions, sentinel).min(dim=-1).values150    clipped = torch.where(first.lt(proposal.shape[1]), first + 1, commit_lengths)151    return torch.minimum(clipped, commit_lengths)152 153 154def bounded_prefix_failure_commit_lengths(155    committed_token_ids: torch.LongTensor,156    failure_rate: torch.Tensor,157    *,158    failure_budget: float,159    remaining_lengths: torch.LongTensor,160    stop_token_id: int | Sequence[int],161    valid_mask: torch.BoolTensor | None = None,162) -> torch.LongTensor:163    """Apply remaining-length and stop-token bounds to the prefix policy."""164 165    if committed_token_ids.shape != failure_rate.shape:166        raise ValueError("Committed token IDs and failure rate must share [batch, canvas].")167    if remaining_lengths.shape != committed_token_ids.shape[:1]:168        raise ValueError("Remaining lengths must have shape [batch].")169    commit_lengths = prefix_failure_commit_lengths(170        failure_rate,171        failure_budget=failure_budget,172        valid_mask=valid_mask,173    )174    commit_lengths = torch.minimum(commit_lengths, remaining_lengths.clamp_min(0))175    return first_committed_token_lengths(176        committed_token_ids,177        commit_lengths,178        stop_token_id,179    )180 181 182def select_commit_lengths(183    sampled_token_ids: torch.LongTensor,184    normal_failure_rate: torch.Tensor,185    previous_failure_rate: torch.Tensor,186    greedy_token_ids: torch.LongTensor,187    jump_failure_rate: torch.Tensor,188    *,189    ponder_steps: torch.Tensor,190    stagnation_steps: torch.Tensor,191    active_rows: torch.BoolTensor,192    remaining_lengths: torch.LongTensor,193    failure_budget: float,194    jump_failure_budget: float,195    stop_token_id: int | Sequence[int],196    max_ponder_steps: int,197    stagnation_threshold: int,198    min_progress: float,199    valid_mask: torch.BoolTensor | None = None,200) -> CommitPolicyDecision:201    """Select sampled commits or a greedy jump after stagnation.202 203    Progress is the signed change in fused failure rate over the union of the204    previous and current prefixes plus one blocking position.205 206    Args:207        sampled_token_ids: Temperature-sampled canvas tokens.208        normal_failure_rate: Fused failure rates for the sampled tokens.209        previous_failure_rate: Fused failure rates from the previous step.210        greedy_token_ids: Greedy canvas tokens used for jumps.211        jump_failure_rate: Fused failure rates for the greedy tokens.212        ponder_steps: Per-row useful-ponder clocks.213        stagnation_steps: Per-row stagnation clocks.214        active_rows: Rows that are still generating.215        remaining_lengths: Tokens still allowed on each row.216        failure_budget: Normal commit budget.217        jump_failure_budget: Forced-jump budget.218        stop_token_id: Turn or EOS stop IDs.219        max_ponder_steps: Watchdog on useful pondering.220        stagnation_threshold: Watchdog on true stagnation.221        min_progress: Minimum signed improvement counted as progress.222        valid_mask: Optional canvas mask.223 224    Returns:225        Commit lengths, token IDs, jump flags, and updated clocks.226    """227 228    if not (229        sampled_token_ids.shape230        == normal_failure_rate.shape231        == previous_failure_rate.shape232        == greedy_token_ids.shape233        == jump_failure_rate.shape234    ):235        raise ValueError("Sampled and greedy statistics must share [batch, canvas].")236 237    normal = bounded_prefix_failure_commit_lengths(238        sampled_token_ids,239        normal_failure_rate,240        failure_budget=failure_budget,241        remaining_lengths=remaining_lengths,242        stop_token_id=stop_token_id,243        valid_mask=valid_mask,244    )245    canvas_length = normal_failure_rate.shape[1]246    previous_prefix_length = prefix_failure_commit_lengths(247        previous_failure_rate,248        failure_budget=failure_budget,249        valid_mask=valid_mask,250    )251    frontier_length = torch.maximum(previous_prefix_length, normal) + 1252    valid_lengths = (253        valid_mask.long().sum(dim=-1)254        if valid_mask is not None255        else torch.full_like(frontier_length, canvas_length)256    )257    frontier_length = torch.minimum(frontier_length, valid_lengths)258    positions = torch.arange(canvas_length, device=normal_failure_rate.device)[None, :]259    progress_mask = positions < frontier_length[:, None]260    if valid_mask is not None:261        progress_mask &= valid_mask262    progress_mask &= active_rows[:, None]263    signed_improvement = previous_failure_rate.float() - normal_failure_rate.float()264    weights = progress_mask.float()265    progress = (signed_improvement * weights).sum(dim=-1) / weights.sum(dim=-1).clamp_min(266        1.0267    )268    next_ponder, next_stagnation = advance_trajectory_clocks(269        ponder_steps,270        stagnation_steps,271        commit_lengths=normal,272        active_rows=active_rows,273        progress_scores=progress,274        min_progress=min_progress,275    )276    jump_rows = normal.eq(0) & active_rows & should_force_trajectory_jump(277        next_ponder,278        next_stagnation,279        max_ponder_steps=max_ponder_steps,280        stagnation_threshold=stagnation_threshold,281    )282    jump_commit = bounded_prefix_failure_commit_lengths(283        greedy_token_ids,284        jump_failure_rate,285        failure_budget=jump_failure_budget,286        remaining_lengths=remaining_lengths,287        stop_token_id=stop_token_id,288        valid_mask=valid_mask,289    )290    committed = torch.where(jump_rows, jump_commit, normal)291    commit_token_ids = torch.where(292        jump_rows[:, None],293        greedy_token_ids,294        sampled_token_ids,295    )296    committed = first_committed_token_lengths(297        commit_token_ids,298        committed,299        stop_token_id,300    )301    committed = torch.where(active_rows, committed, 0)302    jump_rows &= committed.gt(0)303    next_ponder = torch.where(committed.gt(0), 0, next_ponder).to(torch.int32)304    next_stagnation = torch.where(committed.gt(0), 0, next_stagnation).to(torch.int32)305    return CommitPolicyDecision(306        normal_lengths=normal,307        commit_lengths=committed,308        commit_token_ids=commit_token_ids,309        jump_rows=jump_rows,310        ponder_steps=next_ponder,311        stagnation_steps=next_stagnation,312    )313 314 315__all__ = [316    "CommitPolicyDecision",317    "bounded_prefix_failure_commit_lengths",318    "first_committed_token_lengths",319    "fused_commit_confidence",320    "fused_commit_failure_rate",321    "prefix_failure_commit_lengths",322    "select_commit_lengths",323]324