test - validate llamacpp cfg outputs

outlines/integrations/llamacpp.py

@@ -116,7 +116,10 @@ def __call__(
             last_token = input_ids[-1]
             self._fsm_state = self.fsm.get_next_state(self._fsm_state, last_token)
 
-        allowed_tokens = self.fsm.get_next_instruction(self._fsm_state).tokens
+        instruction = self.fsm.get_next_instruction(self._fsm_state)
+        print(type(instruction))
+        print(instruction)
+        allowed_tokens = instruction.tokens
 
         mask = torch.full((scores.shape[-1],), -math.inf, device="cpu").numpy()
         mask[allowed_tokens] = 0

tests/generate/test_integration_llamacpp.py

@@ -1,4 +1,6 @@
+import ast
 import datetime
+import json
 import re
 
 import pytest
@@ -243,10 +245,24 @@ def test_llamacpp_json_schema(model):
     assert isinstance(result["bar"], str)
 
 
-def test_llamacpp_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)
+@pytest.mark.parametrize(
+    "grammar,validator,query",
+    [
+        (grammars.json, json.loads, "short and valid JSON object with two keys."),
+        (
+            grammars.arithmetic,
+            ast.parse,
+            "a mathematical expression with plus minus times and divide.",
+        ),
+    ],
+)
+def test_llamacpp_cfg(model, grammar, validator, query):
+    prompt = f"<|im_start|>user\nOutput a {query}<|im_end|>\n><|im_start|>assistant\n"
+    result = generate.cfg(model, grammar)(prompt, seed=11)
+    try:
+        validator(result)
+    except Exception:
+        raise Exception(f"Failed to validate. output = `{result}`")
 
 
 @pytest.mark.parametrize(