FilesExpand file tree

 main

/

llama_speculative.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

64 lines (48 loc) · 2.04 KB

 main

/

llama_speculative.py

Copy path

Top

File metadata and controls

• Code
• Blame

64 lines (48 loc) · 2.04 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

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

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

import abc

from typing import Any

import numpy as np

import numpy.typing as npt

class LlamaDraftModel(abc.ABC):

@abc.abstractmethod

def __call__(

self, input_ids: npt.NDArray[np.intc], /, **kwargs: Any

) -> npt.NDArray[np.intc]:

raise NotImplementedError()

class LlamaPromptLookupDecoding(LlamaDraftModel):

"""Based on https://github.com/apoorvumang/prompt-lookup-decoding"""

def __init__(self, max_ngram_size: int = 2, num_pred_tokens: int = 10):

self.max_ngram_size = max_ngram_size

self.num_pred_tokens = num_pred_tokens

@staticmethod

def find_candidate_pred_tokens(

input_ids: npt.NDArray[np.intc],

max_ngram_size: int,

num_pred_tokens: int,

):

input_length = input_ids.shape[0]

for ngram_size in range(min(max_ngram_size, input_length - 1), 0, -1):

# Create sliding windows of size ngram_size

windows = np.lib.stride_tricks.sliding_window_view(input_ids, (ngram_size,))

# Convert ngram to an array for comparison

ngram_array = input_ids[-ngram_size:]

# Find where the windows match the ngram

matches = np.all(windows == ngram_array, axis=1)

# Get the indices of matches

match_indices = np.nonzero(matches)[0]

# Iterate through match indices to find a valid continuation

for idx in match_indices:

start_idx = idx + ngram_size

end_idx = start_idx + num_pred_tokens

end_idx = min(end_idx, input_length)

if start_idx < end_idx:

return input_ids[start_idx:end_idx]

# If no match is found, return an empty array

return np.array([], dtype=np.intc)

def __call__(

self, input_ids: npt.NDArray[np.intc], /, **kwargs: Any

) -> npt.NDArray[np.intc]:

return self.find_candidate_pred_tokens(

input_ids=input_ids,

max_ngram_size=self.max_ngram_size,

num_pred_tokens=self.num_pred_tokens,

)