test_prompt_validation.py 763 Bytes
Newer Older
1
import pytest
2
import os
3
4

from vllm import LLM
5
from ...utils import models_path_prefix
6
7


Joe Runde's avatar
Joe Runde committed
8
9
10
11
12
13
14
15
@pytest.fixture(autouse=True)
def v1(run_with_both_engines):
    # Simple autouse wrapper to run both engines for each test
    # This can be promoted up to conftest.py to run for every
    # test in a package
    pass


16
def test_empty_prompt():
zhuwenwen's avatar
zhuwenwen committed
17
    llm = LLM(model=os.path.join(models_path_prefix, "gpt2"), enforce_eager=True)
18
19
    with pytest.raises(ValueError, match='Prompt cannot be empty'):
        llm.generate([""])
20
21


Joe Runde's avatar
Joe Runde committed
22
@pytest.mark.skip_v1
23
def test_out_of_vocab_token():
zhuwenwen's avatar
zhuwenwen committed
24
    llm = LLM(model=os.path.join(models_path_prefix, "gpt2"), enforce_eager=True)
25
26
    with pytest.raises(ValueError, match='out of vocabulary'):
        llm.generate({"prompt_token_ids": [999999]})