Generate a choice between different strings

Author: rlouf

outlines/text/generate/__init__.py

@@ -1,2 +1,2 @@
 from .continuation import continuation
-from .regex import float, integer, regex
+from .regex import choice, float, integer, regex

outlines/text/generate/regex.py

@@ -200,3 +200,9 @@ def float(model, max_tokens: Optional[int] = None):
 
     """
     return Regex(model, r"([+-]?((0|[1-9]+)([.][0-9]*)?)|([.][0-9]+))", max_tokens)
+
+
+def choice(model, choices: List[str], max_tokens: Optional[int] = None):
+    """Choose between different sequences."""
+    regex_str = r"(" + r"|".join(choices) + r")"
+    return Regex(model, regex_str, max_tokens)

tests/text/generate/test_integration_transfomers.py

@@ -98,6 +98,19 @@ def test_transformers_integration_float():
     float(generated)
 
 
+def test_transformers_integration_choice():
+    rng = torch.Generator()
+    rng.manual_seed(0)
+
+    model_name = "hf-internal-testing/tiny-random-GPTJForCausalLM"
+    model = models.transformers(model_name, device="cpu")
+    prompt = "Write a short sentence "
+    sequence = generate.choice(model, ["test", "choice"])(prompt, rng=rng)
+
+    generated = sequence[len(prompt) :]
+    assert generated == "test" or generated == "choice"
+
+
 def test_transformers_integration_with_pad_token():
     model_name = "hf-internal-testing/tiny-random-XLMRobertaXLForCausalLM"
     model = models.transformers(model_name, device="cpu")

tests/text/generate/test_regex.py

@@ -106,6 +106,35 @@ def test_integer_proposal(input_ids, proposal):
     )
 
 
+def test_choice_proposal():
+    model = Model()
+    generator = generate.choice(model, ["1", "431a", "431A-"])
+    logits = torch.ones(len(model.tokenizer.vocabulary))
+    result = generator.create_proposal(torch.tensor([[]]), logits)
+    assert torch.equal(
+        result,
+        torch.tensor(
+            [[-math.inf, -math.inf, 1.0, -math.inf, 1.0, -math.inf, -math.inf]]
+        ),
+    )
+
+    result = generator.create_proposal(torch.tensor([[4]]), logits)
+    assert torch.equal(
+        result,
+        torch.tensor(
+            [[-math.inf, -math.inf, -math.inf, -math.inf, -math.inf, 1.0, 1.0]]
+        ),
+    )
+
+    result = generator.create_proposal(torch.tensor([[4, 6]]), logits)
+    assert torch.equal(
+        result,
+        torch.tensor(
+            [[-math.inf, 1.0, -math.inf, -math.inf, -math.inf, -math.inf, -math.inf]]
+        ),
+    )
+
+
 @pytest.mark.parametrize(
     "input_ids, proposal",
     [