File size: 6,329 Bytes
4b0402b | 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 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 | """Response-only penalties for Clover block LowConfidence generation.
Presence, frequency and repetition use committed response history.
Length and thinking controls provide gradual stopping biases.
All state belongs to one generate request; prompt tokens are never accumulated.
"""
import math
import torch
PENALTY_DEFAULTS = {
"presence_penalty": 0.0,
"frequency_penalty": 0.0,
"repetition_penalty": 1.0,
"dllm_length_penalty": 0.0,
"dllm_length_penalty_start": 8192,
"dllm_length_penalty_window": 8192,
"dllm_think_end_bias": 0.0,
"dllm_think_end_start": 4096,
"dllm_think_end_window": 4096,
"dllm_think_end_token_id": None,
"stop_token_ids": [],
}
def validate_penalties(gen, vocab_size):
for name, default in PENALTY_DEFAULTS.items():
if getattr(gen, name, None) is None:
setattr(gen, name, default)
for name in ("presence_penalty", "frequency_penalty"):
value = getattr(gen, name)
if not math.isfinite(value) or not -2 <= value <= 2:
raise ValueError(f"{name} must be finite and in [-2, 2]")
if not math.isfinite(gen.repetition_penalty) or not 0 < gen.repetition_penalty <= 2:
raise ValueError("repetition_penalty must be finite and in (0, 2]")
for prefix, strength in (
("dllm_length_penalty", "dllm_length_penalty"),
("dllm_think_end", "dllm_think_end_bias"),
):
value = getattr(gen, strength)
if not math.isfinite(value) or value < 0:
raise ValueError(f"{strength} must be finite and nonnegative")
for suffix, minimum in (("start", 0), ("window", 1)):
name = f"{prefix}_{suffix}"
value = getattr(gen, name)
if not isinstance(value, int) or isinstance(value, bool) or value < minimum:
raise ValueError(f"{name} must be an integer >= {minimum}")
close = gen.dllm_think_end_token_id
if close is not None and (
not isinstance(close, int) or not 0 <= close < vocab_size
):
raise ValueError("dllm_think_end_token_id must be a valid token ID")
if gen.dllm_think_end_bias and close is None:
raise ValueError("dllm_think_end_bias requires dllm_think_end_token_id")
if not isinstance(gen.stop_token_ids, (list, tuple)) or any(
not isinstance(i, int) or not 0 <= i < vocab_size for i in gen.stop_token_ids
):
raise ValueError("stop_token_ids must be a list of valid token IDs")
class CloverPenaltyState:
def __init__(self, gen, vocab_size, device, stop_ids):
self.gen = gen
self.length = 0
self.active = bool(
gen.presence_penalty
or gen.frequency_penalty
or gen.repetition_penalty != 1
or gen.dllm_length_penalty
or gen.dllm_think_end_bias
)
self.counts = (
torch.zeros(vocab_size, dtype=torch.int64, device=device)
if self.active
else None
)
self.stop_ids = sorted({i for i in stop_ids if 0 <= i < vocab_size})
def commit(self, tokens):
"""Count each returned response token once, after full block resolution."""
self.length += tokens.numel()
if self.active:
self.counts.index_add_(
0, tokens, torch.ones_like(tokens, dtype=self.counts.dtype)
)
def apply(self, logits, block, mask_id):
"""Mutate fresh [block, vocab] logits before sampling and confidence."""
if not self.active:
return logits
gen = self.gen
# Uncommitted block tokens must not penalize themselves or other slots.
# This snapshot changes only when the returned block is committed.
seen = self.counts > 0
# Keep SGLang's additive operation order and FP32 rounding boundaries.
if gen.frequency_penalty:
logits.sub_(self.counts.float() * gen.frequency_penalty)
if gen.presence_penalty:
logits.sub_(seen * gen.presence_penalty)
if gen.repetition_penalty != 1:
values = logits[:, seen]
logits[:, seen] = torch.where(
values < 0,
values * gen.repetition_penalty,
values / gen.repetition_penalty,
)
positions = self.length + torch.arange(
1, logits.shape[0] + 1, device=logits.device
)
def ramp(start, window, strength):
return ((positions.float() - start) / window).clamp(0, 1) * strength
if gen.dllm_length_penalty and self.stop_ids:
logits[:, self.stop_ids] += ramp(
gen.dllm_length_penalty_start,
gen.dllm_length_penalty_window,
gen.dllm_length_penalty,
)[:, None]
if gen.dllm_think_end_bias:
close = gen.dllm_think_end_token_id
# In LowConfidence the whole block receives the bias until a close
# token is visible or committed.
visible_close = ((block != mask_id) & (block == close)).any()
enabled = ~(seen[close] | visible_close)
logits[:, close] += enabled * ramp(
gen.dllm_think_end_start,
gen.dllm_think_end_window,
gen.dllm_think_end_bias,
)
return logits
def predict_with_confidence(logits, gen):
"""Sample with configured filters; confidence always uses T=1 logits."""
if gen.do_sample:
scaled = logits / gen.temperature
if gen.top_k and 0 < gen.top_k < scaled.shape[-1]:
cutoff = scaled.topk(gen.top_k, -1).values[:, -1:]
scaled = scaled.masked_fill(scaled < cutoff, -float("inf"))
if gen.top_p is not None and gen.top_p < 1:
sorted_logits, order = scaled.sort(-1, descending=True)
probs = sorted_logits.softmax(-1)
remove = (probs.cumsum(-1) - probs) > gen.top_p
sorted_logits = sorted_logits.masked_fill(remove, -float("inf"))
scaled = torch.empty_like(scaled).scatter(-1, order, sorted_logits)
prediction = torch.multinomial(scaled.softmax(-1), 1).squeeze(-1)
else:
prediction = logits.argmax(-1)
confidence = logits.softmax(-1).gather(-1, prediction[:, None]).squeeze(-1)
return prediction, confidence
|