Spaces:
Running on Zero
Running on Zero
| """Sparse count bigrams and exact O(B K²) Viterbi token selection. | |
| Candidate-only scores are relative emissions (no full-vocabulary normalization | |
| can be recovered). Use ``select_block`` when full diffusion logits are available. | |
| No transitions are inserted across document boundaries during count building. | |
| """ | |
| from __future__ import annotations | |
| import hashlib | |
| import json | |
| import math | |
| from collections import Counter | |
| from pathlib import Path | |
| from typing import Iterable | |
| import torch | |
| def tokenizer_fingerprint(tokenizer) -> str: | |
| """Fingerprint the token-to-ID mapping, including special tokens.""" | |
| payload = json.dumps(sorted(tokenizer.get_vocab().items()), ensure_ascii=False).encode() | |
| return hashlib.sha256(payload).hexdigest() | |
| class BigramScorer(torch.nn.Module): | |
| """P(next|prev) = (count(prev,next) + alpha * unigram(next)) / (count(prev) + alpha). | |
| Unigrams have additive beta smoothing, so unseen/special tokens remain | |
| finite. Storage is O(observed pairs + vocabulary), never vocabulary squared. | |
| Pair keys are sorted int64 values for vectorized searchsorted lookups. | |
| """ | |
| def __init__(self, vocab_size, pair_keys, pair_counts, unigram_counts, *, alpha=10.0, | |
| beta=0.1, metadata=None, successor_counts=None): | |
| super().__init__() | |
| if int(vocab_size) < 1 or not math.isfinite(alpha) or alpha <= 0 or not math.isfinite(beta) or beta <= 0: | |
| raise ValueError("vocab_size, alpha, and beta must be positive and finite") | |
| self.vocab_size = int(vocab_size) | |
| self.alpha, self.beta = float(alpha), float(beta) | |
| self.metadata = metadata or {} | |
| keys = torch.as_tensor(pair_keys, dtype=torch.long) | |
| counts = torch.as_tensor(pair_counts, dtype=torch.float32) | |
| unigrams = torch.as_tensor(unigram_counts, dtype=torch.float32) | |
| if keys.ndim != 1 or counts.shape != keys.shape or unigrams.shape != (self.vocab_size,): | |
| raise ValueError("Invalid sparse bigram checkpoint shapes") | |
| if (keys < 0).any() or (keys >= self.vocab_size**2).any() or (keys[1:] <= keys[:-1]).any(): | |
| raise ValueError("Pair keys must be sorted, unique, and within the vocabulary") | |
| if not torch.isfinite(counts).all() or not torch.isfinite(unigrams).all() or (counts < 0).any() or (unigrams < 0).any(): | |
| raise ValueError("Counts must be finite and nonnegative") | |
| totals = torch.zeros(self.vocab_size) | |
| totals.scatter_add_(0, keys // self.vocab_size, counts) | |
| successors = torch.zeros(self.vocab_size) | |
| successors.scatter_add_(0, keys % self.vocab_size, counts) | |
| if successor_counts is not None: | |
| saved_successors = torch.as_tensor(successor_counts, dtype=torch.float32) | |
| if saved_successors.shape != successors.shape or not torch.equal(saved_successors, successors): | |
| raise ValueError("Successor counts do not match sparse bigram counts") | |
| self.register_buffer("successor_counts", successors) | |
| self.register_buffer("successor_unigram", (successors + beta) / (successors.sum() + beta * self.vocab_size)) | |
| self.register_buffer("pair_keys", keys) | |
| self.register_buffer("pair_counts", counts) | |
| self.register_buffer("row_totals", totals) | |
| self.register_buffer("unigram_counts", unigrams) | |
| self.register_buffer("unigram", (unigrams + beta) / (unigrams.sum() + beta * self.vocab_size)) | |
| def from_sequences(cls, sequences: Iterable[Iterable[int]], vocab_size: int, **kwargs): | |
| pairs, unigrams = Counter(), Counter() | |
| for sequence in sequences: | |
| ids = [int(i) for i in sequence] | |
| if any(i < 0 or i >= vocab_size for i in ids): | |
| raise ValueError("Training token outside tokenizer vocabulary") | |
| unigrams.update(ids) | |
| pairs.update(a * vocab_size + b for a, b in zip(ids, ids[1:])) | |
| keys = sorted(pairs) | |
| return cls(vocab_size, keys, [pairs[k] for k in keys], | |
| [unigrams[i] for i in range(vocab_size)], **kwargs) | |
| def _lookup_counts(self, previous: torch.Tensor, following: torch.Tensor) -> torch.Tensor: | |
| if previous.device != self.pair_keys.device or following.device != self.pair_keys.device: | |
| raise ValueError("Move the bigram scorer to the candidate device before scoring") | |
| # Checkpoints cover the backbone vocabulary. Out-of-range IDs are errors, | |
| # rather than aliases for unrelated pairs through packed-key overflow. | |
| if ((previous < 0) | (previous >= self.vocab_size)).any() or ((following < 0) | (following >= self.vocab_size)).any(): | |
| raise ValueError("Candidate token outside bigram vocabulary") | |
| query = previous * self.vocab_size + following | |
| if self.pair_keys.numel(): | |
| index = torch.searchsorted(self.pair_keys, query.contiguous()) | |
| safe = index.clamp_max(self.pair_keys.numel() - 1) | |
| counts = torch.where((index < self.pair_keys.numel()) & (self.pair_keys[safe] == query), | |
| self.pair_counts[safe], 0.0) | |
| else: | |
| counts = torch.zeros_like(query, dtype=torch.float32) | |
| return counts | |
| def _conditional(self, counts, previous, following): | |
| # Keep the original all-token unigram backoff for conditional compatibility. | |
| return (counts + self.alpha * self.unigram[following]).log() - (self.row_totals[previous] + self.alpha).log() | |
| def log_prob(self, previous: torch.Tensor, following: torch.Tensor) -> torch.Tensor: | |
| """Original smoothed conditional probability, independent of scoring mode.""" | |
| return self._conditional(self._lookup_counts(previous, following), previous, following) | |
| def validate_scoring(scoring_method, min_count, pmi_clip): | |
| if scoring_method not in {"conditional", "pmi"}: | |
| raise ValueError("Scoring method must be 'conditional' or 'pmi'") | |
| if not isinstance(min_count, int) or isinstance(min_count, bool) or min_count < 0: | |
| raise ValueError("PMI minimum count must be a nonnegative integer") | |
| if not math.isfinite(pmi_clip) or pmi_clip < 0: | |
| raise ValueError("PMI clipping bound must be finite and nonnegative") | |
| def transition_score(self, previous, following, *, scoring_method="conditional", min_count=5, pmi_clip=5.0): | |
| """Conditional log-probability or guarded PMI using successor marginals. | |
| Unseen pairs always score zero in PMI mode, including min_count=0. | |
| Only the PMI correction is clipped; diffusion emissions are untouched. | |
| """ | |
| self.validate_scoring(scoring_method, min_count, pmi_clip) | |
| counts = self._lookup_counts(previous, following) | |
| conditional = self._conditional(counts, previous, following) | |
| if scoring_method == "conditional": | |
| return conditional | |
| pmi = conditional - self.successor_unigram[following].log() | |
| return torch.where((counts > 0) & (counts >= min_count), pmi.clamp(-pmi_clip, pmi_clip), 0.0) | |
| def forward(self, token_ids: torch.Tensor, logits: torch.Tensor, *, scoring_method="conditional", min_count=5, pmi_clip=5.0) -> torch.Tensor: | |
| """Return all adjacent transition scores [B-1,K,K] in one tensor operation.""" | |
| if token_ids.ndim != 2 or logits.shape != token_ids.shape: | |
| raise ValueError("Expected token_ids and logits [B,K]") | |
| return self.transition_score(token_ids[:-1, :, None], token_ids[1:, None, :], | |
| scoring_method=scoring_method, min_count=min_count, pmi_clip=pmi_clip) | |
| def save(self, path: str | Path): | |
| torch.save(dict(version=2, vocab_size=self.vocab_size, pair_keys=self.pair_keys.cpu(), | |
| pair_counts=self.pair_counts.cpu(), unigram_counts=self.unigram_counts.cpu(), | |
| successor_counts=self.successor_counts.cpu(), | |
| alpha=self.alpha, beta=self.beta, metadata=self.metadata), path) | |
| def load(cls, path: str | Path, *, tokenizer=None, device="cpu"): | |
| if not Path(path).is_file(): | |
| raise FileNotFoundError(f"Bigram checkpoint unavailable: {path}. Build it with scripts/build_bigram.py.") | |
| state = torch.load(path, map_location="cpu", weights_only=True) | |
| if state.pop("version", None) not in {1, 2}: | |
| raise ValueError("Unsupported bigram checkpoint version") | |
| scorer = cls(**state) | |
| if tokenizer is not None and scorer.metadata.get("tokenizer_sha256") != tokenizer_fingerprint(tokenizer): | |
| raise ValueError("Bigram checkpoint tokenizer does not match the inference tokenizer") | |
| return scorer.to(device) | |
| def expected_bigram_repetition_penalties(token_ids, probabilities, vocab_size, strength=1.0, | |
| excluded_token_ids=(), count_mode="expected"): | |
| """Return transition penalties gamma*C(a,b), shape [B-1,K,K]. | |
| Inputs are the legal candidate lattice [B,K] and *full-vocabulary* | |
| probabilities (fixed rows use a single point mass). A sparse tensor-indexed | |
| dictionary keeps C(a,b) on the candidate device. Each edge reads C before | |
| updating it, so its first occurrence contributes zero. State is local to | |
| this call/block; boundary-context transitions are never counted. | |
| Pairs containing excluded tokens (e.g. EOS/EOT) are neither scored nor counted. | |
| count_mode="candidate_occurrences" counts candidate-set membership once per | |
| edge instead of probability products; dummy fixed-row columns count once. | |
| """ | |
| if count_mode not in {"expected", "candidate_occurrences"}: | |
| raise ValueError("Unknown repetition count mode") | |
| strength = float(strength) | |
| if not math.isfinite(strength) or strength < 0: | |
| raise ValueError("Repetition penalty strength must be finite and nonnegative") | |
| if token_ids.ndim != 2 or token_ids.shape != probabilities.shape: | |
| raise ValueError("Expected candidate IDs and probabilities [B,K]") | |
| probabilities = probabilities.float() | |
| penalties = probabilities.new_zeros((max(0, token_ids.shape[0] - 1), token_ids.shape[1], token_ids.shape[1])) | |
| if strength == 0 or token_ids.shape[0] < 2: | |
| return penalties | |
| # Bound the scale before float32 arithmetic, including the DP's sum over B | |
| # rows. Reject impractical strengths instead of silently producing infinities. | |
| if strength > torch.finfo(torch.float32).max / (4 * token_ids.shape[0] ** 2): | |
| raise ValueError("Repetition penalty strength is too large for this block") | |
| keys = token_ids[:-1, :, None] * vocab_size + token_ids[1:, None, :] | |
| unique, inverse = torch.unique(keys.reshape(-1), return_inverse=True) | |
| indices = inverse.reshape_as(keys) | |
| expected = probabilities[:-1, :, None] * probabilities[1:, None, :] | |
| if count_mode == "candidate_occurrences": | |
| # Set membership, independent of probability. Duplicate dummy columns | |
| # in fixed rows represent one candidate, not additional occurrences. | |
| legal = ~((token_ids[:, :, None] == token_ids[:, None, :]) & | |
| torch.ones((token_ids.shape[1], token_ids.shape[1]), | |
| dtype=torch.bool, device=token_ids.device).tril(-1)).any(-1) | |
| expected = (legal[:-1, :, None] & legal[1:, None, :]).float() | |
| if excluded_token_ids: | |
| excluded = torch.as_tensor(tuple(excluded_token_ids), device=token_ids.device, dtype=torch.long) | |
| eligible = ~torch.isin(token_ids, excluded) | |
| # Endings neither incur a correction nor seed counts affecting later rows. | |
| expected = expected * (eligible[:-1, :, None] & eligible[1:, None, :]) | |
| counts = probabilities.new_zeros(unique.numel()) | |
| for t in range(len(expected)): | |
| previous = counts[indices[t]] | |
| # Penalize the particular transition selected by Viterbi. The current | |
| # q only updates future counts; it does not dilute this correction. | |
| penalties[t] = strength * previous | |
| counts.scatter_add_(0, indices[t].reshape(-1), expected[t].reshape(-1)) | |
| return penalties | |
| class BigramViterbiSelector: | |
| def __init__(self, scorer, weight=0.5, *, scoring_method="conditional", min_count=5, pmi_clip=5.0, remasking_confidence="diffusion", sequence_sampling=False, independent_temperature=False, | |
| trigram_scorer=None, trigram_weight=0.5, trigram_min_count=2, trigram_clip=2.0, | |
| randomized_viterbi=False): | |
| if not math.isfinite(weight) or weight < 0: | |
| raise ValueError("Bigram weight must be finite and nonnegative") | |
| BigramScorer.validate_scoring(scoring_method, min_count, pmi_clip) | |
| if remasking_confidence not in {"diffusion", "viterbi_margin"}: | |
| raise ValueError("Remasking confidence must be 'diffusion' or 'viterbi_margin'") | |
| self.remasking_confidence = remasking_confidence | |
| self.sequence_sampling = bool(sequence_sampling) | |
| self.randomized_viterbi = bool(randomized_viterbi) | |
| if self.sequence_sampling and self.randomized_viterbi: | |
| raise ValueError("Choose either exact sequence sampling or randomized Viterbi") | |
| self.independent_temperature = bool(independent_temperature) | |
| if not math.isfinite(trigram_weight) or trigram_weight < 0 or not isinstance(trigram_min_count,int) or trigram_min_count < 1 or not math.isfinite(trigram_clip) or trigram_clip < 0: | |
| raise ValueError("Invalid trigram correction settings") | |
| if trigram_scorer is not None and trigram_scorer.vocab_size != scorer.vocab_size: | |
| raise ValueError("Trigram vocabulary mismatch") | |
| self.trigram_scorer = trigram_scorer | |
| self.trigram_weight = float(trigram_weight) | |
| self.trigram_min_count = trigram_min_count | |
| self.trigram_clip = float(trigram_clip) | |
| self.last_path_margins = None | |
| self.scorer, self.weight = scorer, float(weight) | |
| self.scoring_method, self.min_count, self.pmi_clip = scoring_method, min_count, float(pmi_clip) | |
| def _score(self, previous, following): | |
| return self.scorer.transition_score(previous, following, scoring_method=self.scoring_method, | |
| min_count=self.min_count, pmi_clip=self.pmi_clip) | |
| def decode(token_ids, logits, transitions, weight=0.5): | |
| """Exact DP; each emission is added once. Only the time axis is sequential.""" | |
| if token_ids.ndim != 2 or token_ids.shape != logits.shape or min(token_ids.shape) < 1: | |
| raise ValueError("Expected nonempty token_ids and logits [B,K]") | |
| length, k = token_ids.shape | |
| if transitions.shape != (length - 1, k, k): | |
| raise ValueError("Expected transitions [B-1,K,K]") | |
| if weight == 0 or length == 1: | |
| return token_ids.gather(1, logits.argmax(-1, keepdim=True)).squeeze(1) | |
| scores = logits[0].float() | |
| back = torch.empty((length - 1, k), dtype=torch.long, device=logits.device) | |
| for t in range(1, length): | |
| scores, back[t - 1] = (scores[:, None] + weight * transitions[t - 1]).max(dim=0) | |
| scores = scores + logits[t].float() | |
| path = torch.empty(length, dtype=torch.long, device=logits.device) | |
| path[-1] = scores.argmax() | |
| for t in range(length - 2, -1, -1): | |
| path[t] = back[t, path[t + 1]] | |
| return token_ids.gather(1, path[:, None]).squeeze(1) | |
| def sample_sequence(token_ids, emissions, transitions, weight=1.0, temperature=1.0): | |
| """Exact forward-filter/backward sampling from exp(total path score). | |
| No score perturbations or model calls. Centering each forward row | |
| preserves probabilities while avoiding accumulation of large constants. | |
| Uses torch's seeded RNG, like the existing diffusion sampler. | |
| """ | |
| if token_ids.ndim != 2 or token_ids.shape != emissions.shape or min(token_ids.shape) < 1: | |
| raise ValueError("Expected nonempty candidate lattice [B,K]") | |
| length, k = emissions.shape | |
| if transitions.shape != (length - 1, k, k): | |
| raise ValueError("Expected transitions [B-1,K,K]") | |
| temperature = float(temperature) | |
| if not math.isfinite(temperature) or temperature < 0: | |
| raise ValueError("Sequence temperature must be finite and nonnegative") | |
| if temperature == 0: | |
| return BigramViterbiSelector.decode(token_ids, emissions, transitions, weight) | |
| # Global path temperature changes entropy without changing the relative | |
| # diffusion/PMI/repetition weights. Tiny positive T approaches argmax. | |
| scale = max(temperature, 1e-6) | |
| edges = transitions.float() * weight if weight else torch.zeros_like(transitions, dtype=torch.float32) | |
| if edges.numel(): | |
| edges = (edges - edges.amax(dim=(1, 2), keepdim=True)) / scale | |
| emissions = (emissions.float() - emissions.float().amax(-1, keepdim=True)) / scale | |
| forward = torch.empty_like(emissions, dtype=torch.float32) | |
| row = emissions[0].float() | |
| forward[0] = row - row.max() | |
| for t in range(1, length): | |
| row = emissions[t].float() + torch.logsumexp(forward[t-1, :, None] + edges[t-1], dim=0) | |
| forward[t] = row - row.max() | |
| path = torch.empty(length, device=token_ids.device, dtype=torch.long) | |
| path[-1] = torch.multinomial(forward[-1].softmax(-1), 1).squeeze(0) | |
| for t in range(length - 2, -1, -1): | |
| scores = forward[t] + edges[t, :, path[t+1]] | |
| path[t] = torch.multinomial(scores.softmax(-1), 1).squeeze(0) | |
| return token_ids.gather(1, path[:, None]).squeeze(1) | |
| def __call__(self, token_ids, logits): | |
| if self.weight == 0 or token_ids.shape[0] == 1: | |
| return token_ids.gather(1, logits.argmax(-1, keepdim=True)).squeeze(1) | |
| transitions = self.scorer(token_ids, logits, scoring_method=self.scoring_method, | |
| min_count=self.min_count, pmi_clip=self.pmi_clip) | |
| return self.decode(token_ids, logits, transitions, self.weight) | |
| def path_margins(token_ids, emissions, transitions, selected_ids, weight=0.5): | |
| """Best complete path through selected ID minus best through another ID. | |
| Forward/backward max-sum passes cost O(B K²), without enumerating paths. | |
| Centering each DP row removes shared constants for numerical stability. | |
| Fixed rows and K=1 have infinite margins (no alternative in the lattice). | |
| These are score gaps, not probabilities or full-vocabulary uncertainty. | |
| """ | |
| length, k = emissions.shape | |
| if token_ids.shape != emissions.shape or selected_ids.shape != (length,): | |
| raise ValueError("Expected candidates/emissions [B,K] and selected IDs [B]") | |
| if length < 1 or k < 1 or transitions.shape != (length - 1, k, k): | |
| raise ValueError("Expected nonempty lattice and transitions [B-1,K,K]") | |
| unary = emissions.float() - emissions.float().max(-1, keepdim=True).values | |
| edges = transitions.float() * weight if weight else torch.zeros_like(transitions, dtype=torch.float32) | |
| forward = torch.empty_like(unary) | |
| backward = torch.empty_like(unary) | |
| forward[0] = unary[0] | |
| for t in range(1, length): | |
| row = unary[t] + (forward[t - 1, :, None] + edges[t - 1]).max(0).values | |
| forward[t] = row - row.max() | |
| backward[-1] = 0. | |
| for t in range(length - 2, -1, -1): | |
| row = (edges[t] + unary[t + 1, None, :] + backward[t + 1, None, :]).max(1).values | |
| backward[t] = row - row.max() | |
| scores = forward + backward | |
| same_id = token_ids == selected_ids[:, None] | |
| best_selected = scores.masked_fill(~same_id, -torch.inf).max(-1).values | |
| best_alternative = scores.masked_fill(same_id, -torch.inf).max(-1).values | |
| return (best_selected - best_alternative).clamp_min(0.) | |
| def select_block(self, logits, current_ids, editable, *, top_k=16, temperature=1.0, left_id=None, right_id=None, | |
| repetition_penalty_strength=0.0, repetition_excluded_token_ids=(), | |
| repetition_count_mode="expected"): | |
| """Decode a contiguous block; fixed rows and external neighbors constrain it. | |
| Full-vocabulary log_softmax precedes top-k extraction. Fixed rows have | |
| one possible ID and zero emission (a path-independent constant). | |
| Boundary context adds transition terms without extra emission terms. | |
| Temperature scales emissions relative to bigrams. For T > 0, maximizing | |
| log_softmax(logits / T) + weight * transitions is equivalent to using | |
| raw log probabilities + (weight * T) * transitions. The latter avoids | |
| division by tiny T and has the exact greedy limit at T=0. Margins use | |
| these common rescaled units, preserving their retention ranking. | |
| With independent_temperature=True, weight directly multiplies PMI and | |
| temperature only sets the sequence sampling distribution exp(S/T). | |
| Maximum-path decoding is temperature-independent; T=0 samples its argmax. | |
| Returned diffusion confidence always uses raw full-vocabulary p_t. | |
| Optional repetition penalties subtract gamma*C(a,b) directly from each | |
| edge, independently of the PMI weight and temperature. Expected counts | |
| use original full-vocabulary probabilities, never adjusted scores. | |
| """ | |
| if logits.ndim != 2 or current_ids.shape != logits.shape[:1] or editable.shape != current_ids.shape: | |
| raise ValueError("Expected logits [B,V], current_ids/editable [B]") | |
| if top_k < 1: | |
| raise ValueError("top_k must be positive") | |
| repetition_penalty_strength = float(repetition_penalty_strength) | |
| if not math.isfinite(repetition_penalty_strength) or repetition_penalty_strength < 0: | |
| raise ValueError("Repetition penalty strength must be finite and nonnegative") | |
| temperature = float(temperature) | |
| if not math.isfinite(temperature) or temperature < 0: | |
| raise ValueError("temperature must be finite and nonnegative") | |
| weight = self.weight if self.independent_temperature else self.weight * temperature | |
| if not math.isfinite(weight): | |
| raise ValueError("temperature times bigram weight must be finite") | |
| self.last_path_margins = None | |
| log_probs = logits.float().log_softmax(-1) | |
| randomize = self.randomized_viterbi and temperature > 0 | |
| if weight == 0 and self.remasking_confidence == "diffusion" and repetition_penalty_strength == 0 and not self.sequence_sampling and not randomize: | |
| # topk need not break tied maxima like argmax (common for BF16). | |
| # The zero-weight baseline must exactly match full-vocabulary argmax. | |
| selected = torch.where(editable, logits.argmax(-1), current_ids) | |
| return selected, log_probs.gather(1, selected[:, None]).squeeze(1).exp() | |
| emissions, candidates = log_probs.topk(min(int(top_k), logits.shape[-1]), dim=-1) | |
| if weight == 0 or top_k == 1: | |
| # Keep full-vocabulary argmax tie behavior even if topk omits that | |
| # tied ID. Replacing one tied maximum leaves a valid top-k lattice. | |
| argmax_ids = logits.argmax(-1) | |
| absent = ~(candidates == argmax_ids[:, None]).any(-1) | |
| candidates[absent, 0] = argmax_ids[absent] | |
| emissions[absent, 0] = log_probs.gather(1, argmax_ids[:, None]).squeeze(1)[absent] | |
| candidates = torch.where(editable[:, None], candidates, current_ids[:, None]) | |
| fixed_scores = torch.full_like(emissions, -torch.inf) | |
| fixed_scores[:, 0] = 0.0 | |
| emissions = torch.where(editable[:, None], emissions, fixed_scores) | |
| if repetition_count_mode not in {"expected", "candidate_occurrences", "risk"}: | |
| raise ValueError("Unknown repetition count mode") | |
| if repetition_penalty_strength and repetition_count_mode != "risk": | |
| probabilities = log_probs.gather(1, candidates).exp() | |
| # Truly fixed positions contribute exactly one observed token, not K | |
| # duplicate candidates. Visible LAD positions remain fully editable. | |
| fixed_probabilities = torch.zeros_like(probabilities) | |
| fixed_probabilities[:, 0] = 1.0 | |
| probabilities = torch.where(editable[:, None], probabilities, fixed_probabilities) | |
| penalties = expected_bigram_repetition_penalties( | |
| candidates, probabilities, self.scorer.vocab_size, repetition_penalty_strength, | |
| excluded_token_ids=repetition_excluded_token_ids, count_mode=repetition_count_mode, | |
| ) | |
| if weight: | |
| if left_id is not None: | |
| left = torch.as_tensor(left_id, device=logits.device, dtype=torch.long) | |
| emissions[0] += weight * self._score(left, candidates[0]) | |
| if right_id is not None: | |
| right = torch.as_tensor(right_id, device=logits.device, dtype=torch.long) | |
| emissions[-1] += weight * self._score(candidates[-1], right) | |
| use_trigrams = self.trigram_scorer is not None and self.trigram_weight > 0 and weight > 0 and len(emissions) >= 3 | |
| if repetition_penalty_strength or self.sequence_sampling or use_trigrams or randomize: | |
| transitions = ( | |
| self.scorer(candidates, emissions, scoring_method=self.scoring_method, | |
| min_count=self.min_count, pmi_clip=self.pmi_clip) | |
| if weight else emissions.new_zeros((len(emissions)-1, emissions.shape[1], emissions.shape[1])) | |
| ) | |
| # Combine in raw log-score units so gamma is never multiplied by | |
| # the PMI weight or temperature (and still applies when either is 0). | |
| edges = weight * transitions | |
| triple_scores = None | |
| if use_trigrams: | |
| triple_scores = weight * self.trigram_weight * self.trigram_scorer.lattice_corrections( | |
| candidates, self.scorer, min_count=self.trigram_min_count, clip=self.trigram_clip, | |
| ) | |
| if repetition_penalty_strength: | |
| if repetition_count_mode == "risk": | |
| from .repetition_risk import repetition_risk_penalties | |
| risk_marginals = None | |
| if use_trigrams: | |
| from .trigram import second_order_marginals | |
| risk_marginals = second_order_marginals( | |
| emissions, edges, triple_scores, | |
| temperature=(temperature if self.sequence_sampling and self.independent_temperature else 1.0), | |
| ) | |
| penalties = repetition_risk_penalties( | |
| candidates, emissions, edges, self.scorer.vocab_size, | |
| repetition_penalty_strength, repetition_excluded_token_ids, | |
| temperature=(temperature if self.sequence_sampling and self.independent_temperature else 1.0), | |
| marginals=risk_marginals, | |
| ) | |
| edges = edges - penalties | |
| if randomize: | |
| emissions = self.perturb_unaries(emissions, editable, temperature) | |
| if use_trigrams: | |
| from .trigram import decode_second_order, sample_second_order, second_order_path_margins | |
| if self.sequence_sampling: | |
| selected = sample_second_order(candidates, emissions, edges, triple_scores, | |
| temperature=temperature if self.independent_temperature else 1.0) | |
| else: | |
| selected = decode_second_order(candidates, emissions, edges, triple_scores) | |
| if self.remasking_confidence == "viterbi_margin": | |
| self.last_path_margins = second_order_path_margins(candidates, emissions, edges, triple_scores, selected) | |
| else: | |
| decoder = self.sample_sequence if self.sequence_sampling else self.decode | |
| if self.sequence_sampling and self.independent_temperature: | |
| selected = decoder(candidates, emissions, edges, 1.0, temperature=temperature) | |
| else: | |
| selected = decoder(candidates, emissions, edges, 1.0) | |
| if self.remasking_confidence == "viterbi_margin": | |
| self.last_path_margins = self.path_margins(candidates, emissions, edges, selected, 1.0) | |
| elif self.remasking_confidence == "viterbi_margin": | |
| transitions = ( | |
| self.scorer(candidates, emissions, scoring_method=self.scoring_method, | |
| min_count=self.min_count, pmi_clip=self.pmi_clip) | |
| if weight else emissions.new_zeros((len(emissions) - 1, emissions.shape[1], emissions.shape[1])) | |
| ) | |
| selected = ( | |
| self.decode(candidates, emissions, transitions, weight) | |
| if weight or repetition_penalty_strength else torch.where(editable, logits.argmax(-1), current_ids) | |
| ) | |
| self.last_path_margins = self.path_margins(candidates, emissions, transitions, selected, weight) | |
| else: | |
| transitions = self.scorer(candidates, emissions, scoring_method=self.scoring_method, | |
| min_count=self.min_count, pmi_clip=self.pmi_clip) | |
| selected = self.decode(candidates, emissions, transitions, weight) | |
| confidence = log_probs.gather(1, selected[:, None]).squeeze(1).exp() | |
| return selected, confidence | |
| def perturb_unaries(emissions, editable, temperature): | |
| """Low-dimensional perturb-and-MAP; not exact joint sequence sampling. | |
| Temperature scales independent unit Gumbels in raw log-score units. | |
| Candidate support, transition scores and fixed positions stay unchanged. | |
| """ | |
| if temperature == 0: | |
| return emissions | |
| eps = torch.finfo(emissions.dtype).eps | |
| uniform = torch.rand_like(emissions).clamp(min=torch.finfo(emissions.dtype).tiny, max=1-eps) | |
| noise = -torch.log(-torch.log(uniform)) | |
| return emissions + torch.where(editable[:, None], temperature * noise, 0.) | |
| class RepetitionPenalizedBigramSelector: | |
| """Add block-local transition penalties to an existing selector interface.""" | |
| def __init__(self, selector, strength=1.0, *, excluded_token_ids=(), count_mode="expected", decay_steps=0): | |
| strength = float(strength) | |
| if not math.isfinite(strength) or strength < 0: | |
| raise ValueError("Repetition penalty strength must be finite and nonnegative") | |
| if count_mode not in {"expected", "candidate_occurrences", "risk"}: | |
| raise ValueError("Unknown repetition count mode") | |
| from .repetition_schedule import repetition_strength_factor | |
| repetition_strength_factor(0, decay_steps) # Validate configuration. | |
| self.decay_steps = decay_steps | |
| self.block_step = 0 | |
| self.count_mode = count_mode | |
| self.selector = selector | |
| self.strength = strength | |
| self.excluded_token_ids = tuple(sorted(set(excluded_token_ids))) | |
| def set_step_context(self, block_step, block_steps): | |
| from .repetition_schedule import repetition_strength_factor | |
| repetition_strength_factor(block_step, self.decay_steps) | |
| self.block_step = block_step | |
| def effective_strength(self): | |
| from .repetition_schedule import repetition_strength_factor | |
| return self.strength * repetition_strength_factor(self.block_step, self.decay_steps) | |
| def select_block(self, logits, current_ids, editable, **kwargs): | |
| return self.selector.select_block( | |
| logits, current_ids, editable, | |
| repetition_penalty_strength=self.effective_strength, | |
| repetition_excluded_token_ids=self.excluded_token_ids, | |
| repetition_count_mode=self.count_mode, | |
| **kwargs, | |
| ) | |
| def last_path_margins(self): | |
| return self.selector.last_path_margins | |