Add integrations tests for the vLLM integration

outlines/models/vllm.py

@@ -74,6 +74,8 @@ def generate(
         if max_tokens is not None:
             sampling_params.max_tokens = max_tokens
         if stop_at is not None:
+            if isinstance(stop_at, str):
+                stop_at = [stop_at]
             sampling_params.stop = stop_at
         if seed is not None:
             sampling_params.seed = seed

tests/generate/test_integration_llamacpp.py

@@ -181,7 +181,7 @@ def test_llamacpp_date(model):
     prompt = (
         "<|im_start|>user\nWhat day is it today?<|im_end|>\n<|im_start|>assistant\n"
     )
-    sequence = generate.format(model, datetime.date)(prompt, max_tokens=10)
+    sequence = generate.format(model, datetime.date)(prompt, max_tokens=20, seed=10)
     assert isinstance(sequence, datetime.date)
 
 

tests/generate/test_integration_vllm.py

@@ -1,8 +1,13 @@
+import datetime
+import re
+
 import pytest
 import torch
+from pydantic import BaseModel, constr
 from vllm.sampling_params import SamplingParams
 
 import outlines.generate as generate
+import outlines.grammars as grammars
 import outlines.models as models
 import outlines.samplers as samplers
 
@@ -114,7 +119,121 @@ def test_vllm_beam_search(model):
     assert res[0] != res[1]
 
 
-@pytest.mark.xfail(reason="CFG logits processor not available for vLLM")
-def test_cfg_simple(model):
-    generator = generate.cfg(model)
-    _ = generator("test")
+def test_vllm_text_stop(model):
+    prompt = "Write a short sentence containing 'You': "
+    sequence = generate.text(model)(prompt, max_tokens=100, seed=10)
+    assert sequence.find("news") != -1
+
+    sequence = generate.text(model)(prompt, stop_at="news", max_tokens=100, seed=10)
+    assert isinstance(sequence, str)
+    assert sequence.find("news") == -1
+
+
+def test_vllm_regex(model):
+    prompt = "Write an email address: "
+    regex_str = r"([a-z]{10})@([a-z]{5})\.([a-z]{3})"
+    generator = generate.regex(model, regex_str)
+
+    # One prompt
+    sequence = generator(prompts=prompt)
+    assert isinstance(sequence, str)
+    assert re.fullmatch(pattern=regex_str, string=sequence) is not None
+
+
+def test_vllm_integer(model):
+    prompt = "Give me an integer: "
+    sequence = generate.format(model, int)(prompt, max_tokens=10)
+    assert isinstance(sequence, int)
+    assert sequence != ""
+    int(sequence)
+
+
+def test_vllm_float(model):
+    prompt = "Give me a floating-point number: "
+    sequence = generate.format(model, float)(prompt, max_tokens=10)
+    assert isinstance(sequence, float)
+
+    assert sequence != ""
+    float(sequence)
+
+
+def test_vllm_bool(model):
+    prompt = "Is this True or False? "
+    sequence = generate.format(model, bool)(prompt, max_tokens=10)
+    assert isinstance(sequence, bool)
+
+    assert sequence != ""
+    bool(sequence)
+
+
+def test_vllm_date(model):
+    prompt = "What day is it today? "
+    sequence = generate.format(model, datetime.date)(prompt, max_tokens=10)
+    assert isinstance(sequence, datetime.date)
+
+
+def test_vllm_time(model):
+    prompt = "What time is it? "
+    sequence = generate.format(model, datetime.time)(prompt, max_tokens=10)
+    assert isinstance(sequence, datetime.time)
+
+
+def test_vllm_datetime(model):
+    prompt = "What time is it? "
+    sequence = generate.format(model, datetime.datetime)(prompt, max_tokens=20)
+    assert isinstance(sequence, datetime.datetime)
+
+
+def test_vllm_choice(model):
+    prompt = "Which one between 'test' and 'choice'? "
+    sequence = generate.choice(model, ["test", "choice"])(prompt)
+    assert sequence == "test" or sequence == "choice"
+
+
+def test_vllm_json_basic(model):
+    prompt = "Output some JSON. "
+
+    class Spam(BaseModel):
+        spam: constr(max_length=10)
+        fuzz: bool
+
+    sampling_params = SamplingParams(temperature=0)
+    result = generate.json(model, Spam, whitespace_pattern="")(
+        prompt, max_tokens=100, seed=1, sampling_params=sampling_params
+    )
+    assert isinstance(result, BaseModel)
+    assert isinstance(result.spam, str)
+    assert isinstance(result.fuzz, bool)
+    assert len(result.spam) <= 10
+
+
+def test_vllm_json_schema(model):
+    prompt = "Output some JSON. "
+
+    schema = """{
+      "title": "spam",
+      "type": "object",
+      "properties": {
+           "foo" : {"type": "boolean"},
+           "bar": {"type": "string", "maxLength": 4}
+        },
+      "required": ["foo", "bar"]
+      }
+    """
+
+    sampling_params = SamplingParams(temperature=0)
+    result = generate.json(model, schema, whitespace_pattern="")(
+        prompt, max_tokens=100, seed=10, sampling_params=sampling_params
+    )
+    assert isinstance(result, dict)
+    assert isinstance(result["foo"], bool)
+    assert isinstance(result["bar"], str)
+
+
+@pytest.mark.xfail(
+    reason="The CFG logits processor for vLLM has not been implemented yet."
+)
+def test_vllm_cfg(model):
+    prompt = "<|im_start|>user\nOutput a short and valid JSON object with two keys.<|im_end|>\n><|im_start|>assistant\n"
+    result = generate.cfg(model, grammars.arithmetic)(prompt, seed=11)
+    assert isinstance(result, str)