hymenjj/llama-cpp-python-prebuilt

1import abc2 3from typing import Any4 5import numpy as np6import numpy.typing as npt7 8 9class LlamaDraftModel(abc.ABC):10    @abc.abstractmethod11    def __call__(12        self, input_ids: npt.NDArray[np.intc], /, **kwargs: Any13    ) -> npt.NDArray[np.intc]:14        raise NotImplementedError()15 16 17class LlamaPromptLookupDecoding(LlamaDraftModel):18    """Based on https://github.com/apoorvumang/prompt-lookup-decoding"""19 20    def __init__(self, max_ngram_size: int = 2, num_pred_tokens: int = 10):21        self.max_ngram_size = max_ngram_size22        self.num_pred_tokens = num_pred_tokens23 24    @staticmethod25    def find_candidate_pred_tokens(26        input_ids: npt.NDArray[np.intc],27        max_ngram_size: int,28        num_pred_tokens: int,29    ):30        input_length = input_ids.shape[0]31 32        for ngram_size in range(min(max_ngram_size, input_length - 1), 0, -1):33            # Create sliding windows of size ngram_size34            windows = np.lib.stride_tricks.sliding_window_view(input_ids, (ngram_size,))35 36            # Convert ngram to an array for comparison37            ngram_array = input_ids[-ngram_size:]38 39            # Find where the windows match the ngram40            matches = np.all(windows == ngram_array, axis=1)41 42            # Get the indices of matches43            match_indices = np.nonzero(matches)[0]44 45            # Iterate through match indices to find a valid continuation46            for idx in match_indices:47                start_idx = idx + ngram_size48                end_idx = start_idx + num_pred_tokens49                end_idx = min(end_idx, input_length)50 51                if start_idx < end_idx:52                    return input_ids[start_idx:end_idx]53 54        # If no match is found, return an empty array55        return np.array([], dtype=np.intc)56 57    def __call__(58        self, input_ids: npt.NDArray[np.intc], /, **kwargs: Any59    ) -> npt.NDArray[np.intc]:60        return self.find_candidate_pred_tokens(61            input_ids=input_ids,62            max_ngram_size=self.max_ngram_size,63            num_pred_tokens=self.num_pred_tokens,64        )65