Skip to content

Commit

Permalink
Fix test failing with vllm version upgrade
Browse files Browse the repository at this point in the history
  • Loading branch information
plaguss committed Oct 24, 2024
1 parent 40e408d commit 6c2e1fd
Showing 1 changed file with 16 additions and 4 deletions.
20 changes: 16 additions & 4 deletions tests/unit/llms/test_vllm.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
from openai.types.completion_choice import CompletionChoice
from openai.types.completion_usage import CompletionUsage
from pydantic import BaseModel
from transformers import AutoTokenizer

from distilabel.llms import vLLM
from distilabel.llms.vllm import ClientvLLM
Expand Down Expand Up @@ -101,10 +102,10 @@ class Animal(BaseModel):
]


# Just a mock to avoid loading the model
class DummyTokenizer:
# chat_template = None
chat_template = "template"
vocabulary = {"I'm": 1, "fine": 2, "thank": 3, "you": 4, "sir": 5}

def __init__(self) -> None:
pass
Expand All @@ -115,6 +116,12 @@ def apply_chat_template(self, input, **kwargs):
def encode(self, text: str):
return [1, 2, 3, 4, 5]

def convert_token_to_string(self, token: str) -> str:
return "token"

def get_vocab(self):
return self.vocabulary


class TestvLLM:
@pytest.mark.parametrize("multi_structured_output", (False, True))
Expand Down Expand Up @@ -148,8 +155,11 @@ def test_generate(
expected_result: List[Dict[str, Any]],
) -> None:
llm = vLLM(model="dummy")
llm._tokenizer = DummyTokenizer()
tokenizer = AutoTokenizer.from_pretrained(
"distilabel-internal-testing/tiny-random-mistral"
)
vllm_mock = mock.MagicMock()
vllm_mock.get_tokenizer = mock.MagicMock(return_value=tokenizer)
# mock the import by hacking sys.modules
# https://stackoverflow.com/questions/60919705/how-to-mock-in-a-python-unittest-a-library-not-installed-locally
import sys
Expand Down Expand Up @@ -192,8 +202,10 @@ def test_generate(
},
],
{
"format": "json",
"schema": Character.model_json_schema(),
# "format": "json",
"format": "regex",
"schema": r".*",
# "schema": Character.model_json_schema(),
},
)
]
Expand Down

0 comments on commit 6c2e1fd

Please sign in to comment.