Spaces:
Sleeping
Sleeping
| """Sparse, smoothed trigram corrections and exact second-order chain decoding.""" | |
| from pathlib import Path | |
| import math | |
| import torch | |
| MAX_TRIGRAM_LOOKUP_ELEMENTS = 2_000_000 | |
| class TrigramScorer(torch.nn.Module): | |
| """Correction log P(c|a,b) - log P(c|b), with bigram backoff. | |
| All context totals precede singleton filtering during training. Unseen or | |
| low-count triples give zero correction; clipping limits corpus preferences. | |
| """ | |
| def __init__( | |
| self, | |
| vocab_size, | |
| triple_keys, | |
| triple_counts, | |
| prefix_keys, | |
| prefix_counts, | |
| *, | |
| alpha=50.0, | |
| metadata=None, | |
| ): | |
| super().__init__() | |
| if ( | |
| vocab_size < 1 | |
| or vocab_size**3 >= 2**63 | |
| or not math.isfinite(alpha) | |
| or alpha <= 0 | |
| ): | |
| raise ValueError("Invalid trigram vocabulary or smoothing") | |
| self.vocab_size = int(vocab_size) | |
| self.alpha = float(alpha) | |
| self.metadata = metadata or {} | |
| for name, keys, counts, limit in [ | |
| ("triple", triple_keys, triple_counts, vocab_size**3), | |
| ("prefix", prefix_keys, prefix_counts, vocab_size**2), | |
| ]: | |
| keys = torch.as_tensor(keys, dtype=torch.long) | |
| counts = torch.as_tensor(counts, dtype=torch.float32) | |
| if ( | |
| keys.ndim != 1 | |
| or counts.shape != keys.shape | |
| or (keys < 0).any() | |
| or (keys >= limit).any() | |
| or (keys[1:] <= keys[:-1]).any() | |
| ): | |
| raise ValueError("Invalid sorted sparse trigram keys") | |
| if not torch.isfinite(counts).all() or (counts <= 0).any(): | |
| raise ValueError("Invalid trigram counts") | |
| self.register_buffer(name + "_keys", keys) | |
| self.register_buffer(name + "_counts", counts) | |
| def _lookup(keys, counts, query): | |
| if not keys.numel(): | |
| return torch.zeros_like(query, dtype=torch.float32) | |
| index = torch.searchsorted(keys, query.contiguous()) | |
| safe = index.clamp_max(len(keys) - 1) | |
| return torch.where( | |
| (index < len(keys)) & (keys[safe] == query), counts[safe], 0.0 | |
| ) | |
| def correction(self, a, b, c, bigram, *, min_count=2, clip=2.0): | |
| if bigram.vocab_size != self.vocab_size: | |
| raise ValueError("Bigram/trigram vocabulary mismatch") | |
| if ( | |
| not isinstance(min_count, int) | |
| or min_count < 1 | |
| or not math.isfinite(clip) | |
| or clip < 0 | |
| ): | |
| raise ValueError("Invalid trigram minimum count or clipping") | |
| if min_count < self.metadata.get("storage_min_count", 1): | |
| raise ValueError( | |
| "Requested count threshold is below checkpoint storage threshold" | |
| ) | |
| a, b, c = torch.broadcast_tensors(a, b, c) | |
| for ids in (a, b, c): | |
| if ids.device != self.triple_keys.device: | |
| raise ValueError("Move trigram scorer to candidate device") | |
| if ((ids < 0) | (ids >= self.vocab_size)).any(): | |
| raise ValueError("Token outside trigram vocabulary") | |
| context = a * self.vocab_size + b | |
| counts = self._lookup( | |
| self.triple_keys, self.triple_counts, context * self.vocab_size + c | |
| ) | |
| totals = self._lookup(self.prefix_keys, self.prefix_counts, context) | |
| base = bigram.log_prob(b, c) | |
| smoothed = ( | |
| torch.logaddexp(counts.log(), math.log(self.alpha) + base) | |
| - (totals + self.alpha).log() | |
| ) | |
| return torch.where( | |
| counts >= min_count, (smoothed - base).clamp(-clip, clip), 0.0 | |
| ) | |
| def lattice_corrections(self, ids, bigram, *, min_count=2, clip=2.0): | |
| k = ids.shape[1] | |
| if len(ids) < 3: | |
| return torch.zeros((0, k, k, k), device=ids.device) | |
| if bigram.vocab_size != self.vocab_size: | |
| raise ValueError("Bigram/trigram vocabulary mismatch") | |
| if ( | |
| not isinstance(min_count, int) | |
| or min_count < 1 | |
| or not math.isfinite(clip) | |
| or clip < 0 | |
| ): | |
| raise ValueError("Invalid trigram minimum count or clipping") | |
| if min_count < self.metadata.get("storage_min_count", 1): | |
| raise ValueError( | |
| "Requested count threshold is below checkpoint storage threshold" | |
| ) | |
| if ids.device != self.triple_keys.device: | |
| raise ValueError("Move trigram scorer to candidate device") | |
| if ((ids < 0) | (ids >= self.vocab_size)).any(): | |
| raise ValueError("Token outside trigram vocabulary") | |
| # At ordinary K=16 this is one vectorized lookup across all positions. | |
| # Large K uses window chunks so sparse-lookup temporaries do not grow | |
| # into multiple gigabytes beyond the required K^3 potential table. | |
| chunk = max(1, MAX_TRIGRAM_LOOKUP_ELEMENTS // (k**3)) | |
| output = torch.empty( | |
| (len(ids) - 2, k, k, k), device=ids.device, dtype=torch.float32 | |
| ) | |
| for start in range(0, len(ids) - 2, chunk): | |
| end = min(len(ids) - 2, start + chunk) | |
| # Prefix totals depend on (a,b), and backoff on (b,c). Looking them | |
| # up on the full K^3 broadcast repeats each sparse search K times. | |
| a, b, c = ids[start:end], ids[start + 1 : end + 1], ids[start + 2 : end + 2] | |
| context = a[:, :, None] * self.vocab_size + b[:, None, :] | |
| totals = self._lookup(self.prefix_keys, self.prefix_counts, context)[ | |
| :, :, :, None | |
| ] | |
| base = bigram.log_prob(b[:, :, None], c[:, None, :])[:, None, :, :] | |
| counts = self._lookup( | |
| self.triple_keys, | |
| self.triple_counts, | |
| context[:, :, :, None] * self.vocab_size + c[:, None, None, :], | |
| ) | |
| smoothed = ( | |
| torch.logaddexp(counts.log(), math.log(self.alpha) + base) | |
| - (totals + self.alpha).log() | |
| ) | |
| output[start:end] = torch.where( | |
| counts >= min_count, (smoothed - base).clamp(-clip, clip), 0.0 | |
| ) | |
| return output | |
| def load(cls, path, *, bigram, device="cpu"): | |
| state = torch.load(Path(path), map_location="cpu", weights_only=True) | |
| if state.pop("version", None) != 1: | |
| raise ValueError("Unsupported trigram checkpoint") | |
| scorer = cls(**state) | |
| if scorer.vocab_size != bigram.vocab_size or scorer.metadata.get( | |
| "tokenizer_sha256" | |
| ) != bigram.metadata.get("tokenizer_sha256"): | |
| raise ValueError("Bigram and trigram tokenizer mismatch") | |
| if scorer.metadata.get("tokens") != bigram.metadata.get("tokens"): | |
| raise ValueError("Bigram and trigram training corpora mismatch") | |
| return scorer.to(device) | |
| def _potentials(emissions, edges, triples, temperature=1.0): | |
| length, k = emissions.shape | |
| if edges.shape != (length - 1, k, k) or triples.shape != ( | |
| max(0, length - 2), | |
| k, | |
| k, | |
| k, | |
| ): | |
| raise ValueError("Invalid second-order lattice shapes") | |
| if length < 1 or k < 1: | |
| raise ValueError("Empty second-order lattice") | |
| temperature = float(temperature) | |
| if not math.isfinite(temperature) or temperature <= 0: | |
| raise ValueError("Temperature must be positive") | |
| scale = max(temperature, 1e-6) | |
| unary = emissions.float() | |
| unary = (unary - unary.amax(-1, keepdim=True)) / scale | |
| edges = edges.float() | |
| triples = triples.float() | |
| if edges.numel(): | |
| edges = (edges - edges.amax((1, 2), keepdim=True)) / scale | |
| if triples.numel(): | |
| triples = (triples - triples.amax((1, 2, 3), keepdim=True)) / scale | |
| return unary, edges, triples | |
| def _forward(unary, edges, triples, *, maximum=False): | |
| length, k = unary.shape | |
| rows = [] | |
| back = [] | |
| if length < 2: | |
| return rows, back | |
| row = unary[0, :, None] + unary[1, None, :] + edges[0] | |
| rows.append(row - row.max()) | |
| for t in range(2, length): | |
| score = rows[-1][:, :, None] + triples[t - 2] | |
| if maximum: | |
| row, index = score.max(0) | |
| back.append(index) | |
| else: | |
| row = torch.logsumexp(score, 0) | |
| row = row + edges[t - 1] + unary[t, None, :] | |
| rows.append(row - row.max()) | |
| return rows, back | |
| def decode_second_order(ids, emissions, edges, triples): | |
| # A small chain launches hundreds of tiny GPU reductions. A single host | |
| # transfer and vectorized CPU reductions can cost less than those launches. | |
| # Keep the GPU route for large lattices, and sampling on its existing route. | |
| if ids.device.type == "cuda" and len(ids) <= 256 and ids.shape[1] <= 16: | |
| return _decode_second_order_host(ids, emissions, edges, triples) | |
| unary, edges, triples = _potentials(emissions, edges, triples) | |
| if len(ids) == 1: | |
| return ids.gather(1, unary.argmax(-1, keepdim=True)).squeeze(1) | |
| forward, back = _forward(unary, edges, triples, maximum=True) | |
| k = ids.shape[1] | |
| if ids.device.type != "cuda": | |
| # Scalar indexing is inexpensive on CPU; retain its lean traceback. | |
| path = torch.empty(len(ids), device=ids.device, dtype=torch.long) | |
| last = forward[-1].reshape(-1).argmax() | |
| path[-2] = last // k | |
| path[-1] = last % k | |
| for t in range(len(ids) - 3, -1, -1): | |
| path[t] = back[t][path[t + 1], path[t + 2]] | |
| else: | |
| path = _traceback_device(back, forward[-1].reshape(-1).argmax(), k) | |
| return ids.gather(1, path[:, None]).squeeze(1) | |
| def _decode_second_order_host(ids, emissions, edges, triples): | |
| """Exact MAP recurrence on CPU, preserving float32 operations/tie order.""" | |
| import numpy as np | |
| # Normalize on the original device to retain the original score arithmetic. | |
| unary, edges, triples = _potentials(emissions, edges, triples) | |
| length, k = unary.shape | |
| # One synchronous copy instead of one per predecessor or lattice row. | |
| packed = torch.cat((unary.reshape(-1), edges.reshape(-1), triples.reshape(-1))) | |
| host = packed.detach().cpu().numpy() | |
| cut = length * k | |
| end = cut + (length - 1) * k * k | |
| u = host[:cut].reshape(length, k) | |
| e = host[cut:end].reshape(length - 1, k, k) | |
| tri = host[end:].reshape(max(0, length - 2), k, k, k) | |
| if length == 1: | |
| path = np.array([u[0].argmax()], dtype=np.int64) | |
| else: | |
| row = u[0, :, None] + u[1, None, :] + e[0] | |
| row = row - row.max() | |
| back = [] | |
| for t in range(2, length): | |
| score = row[:, :, None] + tri[t - 2] | |
| index = score.argmax(axis=0) | |
| row = score.max(axis=0) + e[t - 1] + u[t, None, :] | |
| row = row - row.max() | |
| back.append(index) | |
| path = np.empty(length, dtype=np.int64) | |
| last = row.reshape(-1).argmax() | |
| path[-2] = last // k | |
| path[-1] = last % k | |
| for t in range(length - 3, -1, -1): | |
| path[t] = back[t][path[t + 1], path[t + 2]] | |
| indices = torch.from_numpy(path).to(device=ids.device) | |
| return ids.gather(1, indices[:, None]).squeeze(1) | |
| def _traceback_device(back, last, k): | |
| """Trace pair states without converting GPU indices to host scalars.""" | |
| pair = last.reshape(1) | |
| reverse = [(pair % k).squeeze(0), (pair // k).squeeze(0)] | |
| for predecessors in reversed(back): | |
| # A one-element tensor gather keeps the index on the GPU. Scalar | |
| # tensor indexing can otherwise synchronize with the host each hop. | |
| previous = predecessors.reshape(-1).gather(0, pair) | |
| reverse.append(previous.squeeze(0)) | |
| pair = previous * k + pair // k | |
| return torch.stack(reverse[::-1]) | |
| def sample_second_order(ids, emissions, edges, triples, temperature=1.0): | |
| if temperature == 0: | |
| return decode_second_order(ids, emissions, edges, triples) | |
| unary, edges, triples = _potentials(emissions, edges, triples, temperature) | |
| if len(ids) == 1: | |
| index = torch.multinomial(unary[0].softmax(-1), 1) | |
| return ids[0, index] | |
| forward, _ = _forward(unary, edges, triples) | |
| k = ids.shape[1] | |
| path = torch.empty(len(ids), device=ids.device, dtype=torch.long) | |
| last = torch.multinomial(forward[-1].reshape(-1).softmax(-1), 1).squeeze(0) | |
| path[-2] = last // k | |
| path[-1] = last % k | |
| for t in range(len(ids) - 3, -1, -1): | |
| scores = forward[t][:, path[t + 1]] + triples[t, :, path[t + 1], path[t + 2]] | |
| path[t] = torch.multinomial(scores.softmax(-1), 1).squeeze(0) | |
| return ids.gather(1, path[:, None]).squeeze(1) | |
| def second_order_marginals( | |
| emissions, edges, triples, temperature=1.0, *, maximum=False | |
| ): | |
| """Return pair log-potentials (max) or normalized joint pair probabilities.""" | |
| unary, edges, triples = _potentials(emissions, edges, triples, temperature or 1.0) | |
| if len(unary) < 2: | |
| return edges.clone() | |
| forward, _ = _forward(unary, edges, triples, maximum=maximum) | |
| backward = [None] * (len(unary) - 1) | |
| backward[-1] = torch.zeros_like(edges[-1]) | |
| for t in range(len(unary) - 3, -1, -1): | |
| scores = ( | |
| triples[t] | |
| + edges[t + 1][None, :, :] | |
| + unary[t + 2, None, None, :] | |
| + backward[t + 1][None, :, :] | |
| ) | |
| row = scores.max(-1).values if maximum else torch.logsumexp(scores, -1) | |
| backward[t] = row - row.max() | |
| pair_scores = torch.stack([a + b for a, b in zip(forward, backward)]) | |
| return ( | |
| pair_scores | |
| if maximum | |
| else pair_scores.flatten(1).softmax(-1).reshape_as(pair_scores) | |
| ) | |
| def second_order_path_margins(ids, emissions, edges, triples, selected): | |
| if len(ids) == 1: | |
| from .bigram import BigramViterbiSelector | |
| return BigramViterbiSelector.path_margins(ids, emissions, edges, selected, 1.0) | |
| pairs = second_order_marginals(emissions, edges, triples, maximum=True) | |
| scores = torch.cat((pairs[0].max(1).values[None, :], pairs.max(1).values), 0) | |
| same = ids == selected[:, None] | |
| return ( | |
| scores.masked_fill(~same, -torch.inf).max(-1).values | |
| - scores.masked_fill(same, -torch.inf).max(-1).values | |
| ).clamp_min(0.0) | |