FilesExpand file tree

 main

/

test_llama_speculative.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

16 lines (12 loc) · 696 Bytes

 main

/

test_llama_speculative.py

Copy path

Top

File metadata and controls

• Code
• Blame

16 lines (12 loc) · 696 Bytes

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

import numpy as np

from llama_cpp.llama_speculative import LlamaPromptLookupDecoding

def test_find_candidate_pred_tokens():

find_candidate_pred_tokens = LlamaPromptLookupDecoding.find_candidate_pred_tokens

# Test Case 1: Matching ngram is found

input_ids1 = np.array([1, 2, 3, 1, 2, 3, 1, 2, 3])

result1 = find_candidate_pred_tokens(input_ids1, max_ngram_size=3, num_pred_tokens=2)

assert np.array_equal(result1, np.array([1, 2]))

# Test Case 2: Matching ngram is not found

input_ids2 = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9])

result2 = find_candidate_pred_tokens(input_ids2, max_ngram_size=3, num_pred_tokens=2)

assert np.array_equal(result2, np.array([]))