Spaces:
Sleeping
Sleeping
File size: 1,597 Bytes
74679d4 | 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 | """Prompt lookup decoding (PLD): drafts copied from the context (prompt + output).
At each cycle, the longest suffix of the context of length n_max..n_min that occurred
earlier is matched to its MOST RECENT previous occurrence, and the m tokens that followed
it are proposed as drafts. The drafts are verified like any other (exact, lossless).
"""
class NgramIndex:
"""Incremental index: for n in [n_min, n_max], the last END position of every n-gram,
excluding the gram that ends at the current last token (so a match is always earlier)."""
def __init__(self, n_max: int = 4, n_min: int = 1):
self.n_max, self.n_min = n_max, n_min
self.seq: list[int] = []
self.last = {n: {} for n in range(n_min, n_max + 1)}
def extend(self, toks) -> None:
for t in toks:
self.seq.append(t)
e = len(self.seq) - 2 # index grams ending one before the new last token
if e < 0:
continue
for n in range(self.n_min, self.n_max + 1):
if e - n + 1 >= 0:
self.last[n][tuple(self.seq[e - n + 1:e + 1])] = e
def propose(self, m: int) -> list[int]:
"""Up to m copied tokens; self.last_n = length of the matched suffix (0 if none)."""
s = self.seq
self.last_n = 0
for n in range(self.n_max, self.n_min - 1, -1):
if len(s) < n:
continue
e = self.last[n].get(tuple(s[-n:]))
if e is not None:
self.last_n = n
return s[e + 1:e + 1 + m]
return []
|