"tests/vscode:/vscode.git/clone" did not exist on "2991dd3d2241cb3b188a047282e58111684e0201"
test_tokenizer_group.py 2.7 KB
Newer Older
1
2
# SPDX-License-Identifier: Apache-2.0

zhuwenwen's avatar
zhuwenwen committed
3
import os
4
5
import pytest
from transformers import AutoTokenizer, PreTrainedTokenizerBase
6

7
8
from vllm.lora.request import LoRARequest
from vllm.transformers_utils.tokenizer import get_lora_tokenizer
zhuwenwen's avatar
zhuwenwen committed
9

10
from vllm.transformers_utils.tokenizer_group import TokenizerGroup
zhuwenwen's avatar
zhuwenwen committed
11
from ..utils import models_path_prefix
12
13
14
15
16
17


@pytest.mark.asyncio
@pytest.mark.parametrize("tokenizer_group_type", [None, "ray"])
async def test_tokenizer_group_lora(sql_lora_files, tokenizer_group_type):
    reference_tokenizer = AutoTokenizer.from_pretrained(sql_lora_files)
18
    tokenizer_group = TokenizerGroup(
19
        tokenizer_id=os.path.join(models_path_prefix,"gpt2"),
20
21
        enable_lora=True,
        max_num_seqs=1,
22
        max_loras=1,
23
24
25
26
        max_input_length=None,
    )
    lora_request = LoRARequest("1", 1, sql_lora_files)
    assert reference_tokenizer.encode("prompt") == tokenizer_group.encode(
27
        prompt="prompt", lora_request=lora_request)
28
29
    assert reference_tokenizer.encode(
        "prompt") == await tokenizer_group.encode_async(
30
            prompt="prompt", lora_request=lora_request)
31
32
33
34
35
36
37
38
39
40
41
42
43
44
    assert isinstance(tokenizer_group.get_lora_tokenizer(None),
                      PreTrainedTokenizerBase)
    assert tokenizer_group.get_lora_tokenizer(
        None) == await tokenizer_group.get_lora_tokenizer_async(None)

    assert isinstance(tokenizer_group.get_lora_tokenizer(lora_request),
                      PreTrainedTokenizerBase)
    assert tokenizer_group.get_lora_tokenizer(
        lora_request) != tokenizer_group.get_lora_tokenizer(None)
    assert tokenizer_group.get_lora_tokenizer(
        lora_request) == await tokenizer_group.get_lora_tokenizer_async(
            lora_request)


45
def test_get_lora_tokenizer(sql_lora_files, tmp_path):
46
47
48
49
50
51
52
53
    lora_request = None
    tokenizer = get_lora_tokenizer(lora_request)
    assert not tokenizer

    lora_request = LoRARequest("1", 1, sql_lora_files)
    tokenizer = get_lora_tokenizer(lora_request)
    assert tokenizer.get_added_vocab()

54
    lora_request = LoRARequest("1", 1, str(tmp_path))
55
56
    tokenizer = get_lora_tokenizer(lora_request)
    assert not tokenizer
57
58
59
60
61
62


@pytest.mark.parametrize("enable_lora", [True, False])
@pytest.mark.parametrize("max_num_seqs", [1, 2])
@pytest.mark.parametrize("max_loras", [1, 2])
def test_lora_tokenizers(enable_lora, max_num_seqs, max_loras):
63
    tokenizer_group = TokenizerGroup(
64
65
66
67
68
69
70
71
72
73
74
        tokenizer_id="gpt2",
        enable_lora=enable_lora,
        max_num_seqs=max_num_seqs,
        max_loras=max_loras,
        max_input_length=None,
    )
    if enable_lora:
        assert tokenizer_group.lora_tokenizers.capacity == max(
            max_num_seqs, max_loras)
    else:
        assert tokenizer_group.lora_tokenizers.capacity == 0