Class PromptLookupDecoding

examples/server/server.py:1208–1242

class PromptLookupDecoding(DraftProvider):
    def __init__(self, max_ngram_size: int = 2, num_pred_tokens: int = 10) -> None:
        self._max_ngram_size = max_ngram_size
        self._num_pred_tokens = num_pred_tokens

    def draft(
        self,
        input_ids: np.ndarray,
        /,
        *,
        seq_id: int,
        max_tokens: Optional[int],
    ) -> np.ndarray:
        input_length = input_ids.shape[0]
        if input_length < 2:
            return np.array([], dtype=np.intc)
        num_pred_tokens = self._num_pred_tokens
        if max_tokens is not None:
            num_pred_tokens = min(num_pred_tokens, max_tokens)
        if num_pred_tokens <= 0:
            return np.array([], dtype=np.intc)
        max_ngram_size = min(self._max_ngram_size, input_length - 1)
        for ngram_size in range(max_ngram_size, 0, -1):
            windows = np.lib.stride_tricks.sliding_window_view(input_ids, (ngram_size,))
            ngram = input_ids[-ngram_size:]
            matches = np.all(windows == ngram, axis=1)
            match_indices = np.nonzero(matches)[0]
            for index in match_indices:
                start = index + ngram_size
                if start >= input_length:
                    continue
                end = min(start + num_pred_tokens, input_length)
                if start < end:
                    return input_ids[start:end].astype(np.intc, copy=False)
        return np.array([], dtype=np.intc)


class MTPDraftProvider(DraftProvider):