Activate CFG test for VLLM integration

dottxt-ai · Apr 11, 2024 · 778b581 · 778b581
1 parent 274e2f3
commit 778b581
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 4 deletions.
diff --git a/outlines/generate/cfg.py b/outlines/generate/cfg.py
@@ -2,11 +2,12 @@
 
 from outlines.fsm.guide import CFGGuide
 from outlines.generate.api import SequenceGenerator, SequenceGeneratorAdapter
+from outlines.integrations.vllm import CFGLogitsProcessor
 from outlines.models import OpenAI
 from outlines.models.llamacpp import LlamaCpp
 from outlines.models.vllm import VLLM
 from outlines.samplers import Sampler, multinomial
-from outlines.integrations.vllm import CFGLogitsProcessor
+
 
 @singledispatch
 def cfg(model, cfg_str: str, sampler: Sampler = multinomial()) -> SequenceGenerator:

diff --git a/tests/generate/test_integration_vllm.py b/tests/generate/test_integration_vllm.py
@@ -230,9 +230,6 @@ def test_vllm_json_schema(model):
     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)