| """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 |
| |
| |
| seen = self.counts > 0 |
| |
| 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 |
| |
| |
| 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 |
|
|
