Added generate.probabilities for BeamSearch

docs/reference/probabilities.md

@@ -0,0 +1,25 @@
+# Multiple choices with Probabilities
+
+Outlines allows you to generate probabilities for different options, giving you insights into the model's confidence for each choice.
+
+```python
+from outlines import models, generate
+
+model = models.transformers("mistralai/Mistral-7B-v0.1")
+probabilities = generate.probabilities(model, ["skirt", "dress", "pen", "jacket"])
+answer = probabilities("Pick the odd word out: skirt, dress, pen, jacket")
+print(answer)
+```
+
+!!! Warning "Compatibility"
+
+    `generate.probabilities` uses a beam search sampler. It is not compatible with other samplers. Ensure that no other samplers are used in conjunction with this method.
+
+## How It Works
+
+
+Beam search is a heuristic search algorithm used to explore the most promising sequences in a limited set. In text generation, it maintains the top `k` sequences (beams) at each step based on their cumulative probabilities. Each sequence has a weight, which is the product of the probabilities of its tokens, representing the likelihood of the sequence according to the model.
+
+!!! Warning "Probabilities Summation"
+
+    The probabilities returned by `generate.probabilities` might not sum to one because the `topk` limitation only keeps the best sequences. This means other sequences with potentially non-negligible probabilities are not taken into account, leading to an incomplete probability distribution.

outlines/generate/__init__.py

@@ -4,5 +4,6 @@
 from .format import format
 from .fsm import fsm
 from .json import json
+from .probabilities import probabilities
 from .regex import regex
 from .text import text

outlines/generate/api.py

@@ -14,19 +14,19 @@
 
 
 class SequenceGenerator:
-    def __init__(
-        self,
-        fsm,
-        model,
-        sampler,
-        device,
-    ):
+    def __init__(self, fsm, model, sampler, device, probabilities=None):
         self.fsm = fsm
         self.model = model
         self.sampler = sampler
         self.tokenizer = model.tokenizer
         self.device = device
         self.num_samples = sampler.samples
+        # The choices for which we want to compute the probabilities
+        self.probabilities = probabilities
+        if self.probabilities:
+            assert isinstance(
+                self.sampler, BeamSearchSampler
+            ), "Probabilities are only supported with a beam search sampler"
 
     def get_generated_token_ids(
         self,
@@ -132,7 +132,9 @@ def __call__(
         max_tokens: Optional[int] = None,
         stop_at: Optional[Union[str, List[str]]] = None,
         rng: Optional["torch.Generator"] = None,
-    ) -> Union[FormattedOutput, List[FormattedOutput], List[List[FormattedOutput]]]:
+    ) -> Union[
+        FormattedOutput, List[FormattedOutput], List[List[FormattedOutput]], tuple
+    ]:
         """Generate the full text sequence.
 
         Since `SequenceGenerator.stream` calls the tokenizer at every step this
@@ -231,9 +233,33 @@ def __call__(
 
         # We reshape the output to (batch_size, sample_size)
         output: List[List[FormattedOutput]] = list()
+
         for i in range(batch_size):
             output.append(formatted[i : i + num_samples])
 
+        if self.probabilities:
+            logprobs = last_state.weights
+            probs = torch.exp(logprobs)
+            output_probs = [
+                {choice: 0.0 for choice in self.probabilities}
+                for _ in range(batch_size)
+            ]
+            for i in range(batch_size):
+                for choice in self.probabilities:
+                    for output_index, output_item in enumerate(output[i]):
+                        if choice == output_item:
+                            output_probs[i][choice] += float(
+                                probs[i * num_samples + output_index]
+                            )
+            if batch_size == 1 and num_samples == 1:
+                return output[0][0], output_probs[0]
+            elif batch_size == 1:
+                return output[0], output_probs[0]
+            elif num_samples == 1:
+                return [samples[0] for samples in output], output_probs
+            else:
+                return output, output_probs
+
         # We remove leading dimensions for the output
         if batch_size == 1 and num_samples == 1:
             return output[0][0]

outlines/generate/probabilities.py

@@ -0,0 +1,20 @@
+from typing import List
+
+from outlines.generate.api import SequenceGenerator
+from outlines.samplers import BeamSearchSampler, Sampler, beam_search
+
+from .regex import regex
+
+
+def probabilities(
+    model, choices: List[str], sampler: Sampler = beam_search()
+) -> SequenceGenerator:
+    regex_str = r"(" + r"|".join(choices) + r")"
+    assert isinstance(
+        sampler, BeamSearchSampler
+    ), "Only BeamSearchSampler is supported for probabilities"
+    generator = regex(model, regex_str, sampler)
+    generator.format_sequence = lambda x: x
+    generator.probabilities = choices
+
+    return generator

outlines/samplers.py

@@ -261,7 +261,7 @@ def __call__(
         self,
         next_token_logits: "torch.DoubleTensor",
         sequence_weights: "torch.DoubleTensor",
-        _,
+        rng: "torch.Generator",
     ) -> Tuple["torch.DoubleTensor", "torch.DoubleTensor", "torch.DoubleTensor"]:
         """Call the beam search sampler.
 

tests/generate/test_integration_transformers.py

@@ -320,6 +320,21 @@ def test_transformers_integration_choice():
     assert sequence == "test" or sequence == "choice"
 
 
+def test_transformers_integration_probabilities():
+    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, probs = generate.probabilities(
+        model, ["test", "choice"], sampler=beam_search(beams=5)
+    )(prompt, rng=rng)
+    assert probs.keys() == {"test", "choice"}
+    assert probs["test"] > 0
+    assert probs["choice"] > 0
+
+
 def test_transformers_integration_with_pad_token():
     model_name = "hf-internal-testing/tiny-random-XLMRobertaXLForCausalLM"
     model = models.transformers(model_name, device="meta")