Spaces:
Running on Zero
Running on Zero
File size: 4,393 Bytes
1473488 24666c5 1473488 24666c5 1473488 24666c5 1473488 24666c5 1473488 24666c5 1473488 942fb38 1473488 942fb38 1473488 942fb38 24666c5 1473488 942fb38 1473488 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 | """Optional confidence-based token commitment, independent of token scoring."""
import math
import torch
def confidence_threshold(step: int, steps: int, aggressiveness: float = 0.0) -> float:
"""CAI-dLLM cosine shape, 15% warmup, endpoints 0.90 and 0.50.
Normalize the discrete block's first/last steps to 0/1 so its final step
reaches the requested endpoint. A one-step block uses the final threshold.
Aggressiveness scales both endpoints by (1-a): 0 preserves the old schedule,
1 admits every finite-confidence candidate. Commitment guards still apply.
"""
aggressiveness = float(aggressiveness)
if not math.isfinite(aggressiveness) or not 0.0 <= aggressiveness <= 1.0:
raise ValueError("APD aggressiveness must be finite and between 0 and 1")
if steps < 1 or not 0 <= step < steps:
raise ValueError("Expected a valid zero-based step within the block")
if steps == 1:
return 0.50 * (1.0 - aggressiveness)
progress = step / (steps - 1)
phase = max(0.0, (progress - 0.15) / 0.85)
original = 0.50 + 0.20 * (1.0 + math.cos(math.pi * phase))
return original if aggressiveness == 0 else original * (1.0 - aggressiveness)
def selected_token_probabilities(logits, token_ids):
"""Full-vocabulary model probabilities, without top-k/temperature scaling."""
scores = logits.float()
return (scores.gather(1, token_ids[:, None]).squeeze(1) - scores.logsumexp(-1)).exp()
def resolve_mask_predictions(logits, selected, confidence, pending, mask_token_id):
"""Output-validity fallback if a decoder selects MASK itself.
Use the same forward's highest-scoring non-mask token only for those rows.
Ordinary sampled/Viterbi-selected tokens and their confidence stay unchanged.
"""
invalid = pending & (selected == mask_token_id)
if not bool(invalid.any()):
return selected, confidence
if logits.shape[-1] < 2:
raise ValueError("APD cannot resolve a vocabulary containing only MASK")
rows = invalid.nonzero(as_tuple=True)[0]
top_ids = logits[rows].topk(2, dim=-1).indices
alternative = torch.where(top_ids[:, 0] == mask_token_id, top_ids[:, 1], top_ids[:, 0])
selected, confidence = selected.clone(), confidence.clone()
selected[rows] = alternative
confidence[rows] = selected_token_probabilities(logits[rows], alternative)
return selected, confidence
def commitment_mask(confidence, pending, step, steps, *, endings=None, delay_endings=False, aggressiveness=0.0, minimum_tokens=1):
"""Commit all qualifying pending positions; force progress and final cleanup.
With ending retention delay, postpone endings while ordinary predictions
would remain uncommitted. The fallback prefers an ordinary token in that
case. Fill a minimum per-step quota using highest-confidence pending rows.
Ordinary rows precede endings when ending retention is delayed.
"""
if not isinstance(minimum_tokens, int) or minimum_tokens < 1:
raise ValueError("APD minimum tokens must be a positive integer")
threshold = confidence_threshold(step, steps, aggressiveness)
if step == steps - 1:
return pending.clone()
finite = torch.isfinite(confidence)
chosen = pending & finite & (confidence >= threshold)
ordinary = pending if endings is None else pending & ~endings
if delay_endings and endings is not None and bool((ordinary & ~chosen).any()):
chosen &= ~endings
required = min(minimum_tokens, int(pending.sum()))
needed = max(0, required - int(chosen.sum()))
if needed:
# Preserve the old ending delay while satisfying the mandatory quota:
# exhaust ordinary rows before filling from endings. Nonfinite scores
# rank last; stable ordering retains the old first-index fallback.
pools = [ordinary, pending & ~ordinary] if delay_endings and endings is not None else [pending]
for pool in pools:
indices = (pool & ~chosen).nonzero(as_tuple=True)[0]
count = min(needed, len(indices))
if count:
scores = confidence[indices].masked_fill(~finite[indices], -torch.inf)
ranked = scores.argsort(descending=True, stable=True)
chosen[indices[ranked[:count]]] = True
needed -= count
if needed == 0:
break
return chosen
|