Spaces:
Sleeping
Sleeping
| """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 [] | |