Skip to content

Samplers

BeamSearchSampler

Beam Search sampling algorithm.

Attributes

samples The number of samples taken for each input sequence. Equivalent to the number of beams.

Source code in outlines/samplers.py 
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
class BeamSearchSampler:
    """Beam Search sampling algorithm.

    Attributes
    ----------
    samples
        The number of samples taken for each input sequence. Equivalent to the
        number of beams.
    """

    def __init__(self, beams: int = 1):
        self.samples = beams

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

        Parameters
        ----------
        next_token_logits
            A tensor of shape ``(n_seqs, vocab_size,)`` that represents the
            probability distribution of the next token over the vocabulary.
        sequence_weights
            A tensor of shape ``(n_seqs,)`` that represents the cumulative
            weight of each sequence.
        rng
            A random number generator.

        Returns
        -------
        A tuple with an array that contains the ids of the sampled tokens of
        shape ``(n_seqs, 1)``, an array that contains the ancestors of each
        sampled id of shape ``(n_seqs,)`` and an array that contains the updated
        cumulative weights of each sequence of shape ``(n_seqs,)``.

        """
        import torch

        logprobs = torch.nn.functional.log_softmax(next_token_logits, dim=-1)
        weights = logprobs + sequence_weights.unsqueeze(1).expand_as(next_token_logits)

        # Flatten scores to (n_batch, n_samples * vocab_size)
        # and find the top-k weights for each batch.
        batch_size = next_token_logits.shape[0] // self.samples
        vocab_size = next_token_logits.shape[-1]
        weights = weights.view(batch_size, self.samples * vocab_size)

        # If the weights are all equal to 0 we are at the beginning of the search
        # and thus only need to sample from one set of token logits for each
        # batch.
        if torch.all(sequence_weights == 0):
            weights = weights[:, :vocab_size]

        weights, indices = torch.topk(
            weights, self.samples, dim=1, largest=True, sorted=True
        )

        ancestors = torch.div(indices, vocab_size, rounding_mode="floor")
        next_token_ids = indices % vocab_size

        # Re-shape the weights, next_token_ids and ancestors to (n_batch * n_samples, 1)
        first_batch_idx = torch.arange(
            0, batch_size * self.samples, self.samples, device=next_token_logits.device
        ).unsqueeze(1)
        ancestors = ancestors + first_batch_idx

        ancestors = ancestors.view(self.samples * batch_size)
        weights = weights.view(self.samples * batch_size)
        next_token_ids = next_token_ids.view(self.samples * batch_size, 1)

        return next_token_ids, ancestors, weights

__call__(next_token_logits, sequence_weights, _)

Call the beam search sampler.

Parameters

next_token_logits A tensor of shape (n_seqs, vocab_size,) that represents the probability distribution of the next token over the vocabulary. sequence_weights A tensor of shape (n_seqs,) that represents the cumulative weight of each sequence. rng A random number generator.

Returns

A tuple with an array that contains the ids of the sampled tokens of shape (n_seqs, 1), an array that contains the ancestors of each sampled id of shape (n_seqs,) and an array that contains the updated cumulative weights of each sequence of shape (n_seqs,).

Source code in outlines/samplers.py 
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
def __call__(
    self,
    next_token_logits: "torch.DoubleTensor",
    sequence_weights: "torch.DoubleTensor",
    _,
) -> Tuple["torch.DoubleTensor", "torch.DoubleTensor", "torch.DoubleTensor"]:
    """Call the beam search sampler.

    Parameters
    ----------
    next_token_logits
        A tensor of shape ``(n_seqs, vocab_size,)`` that represents the
        probability distribution of the next token over the vocabulary.
    sequence_weights
        A tensor of shape ``(n_seqs,)`` that represents the cumulative
        weight of each sequence.
    rng
        A random number generator.

    Returns
    -------
    A tuple with an array that contains the ids of the sampled tokens of
    shape ``(n_seqs, 1)``, an array that contains the ancestors of each
    sampled id of shape ``(n_seqs,)`` and an array that contains the updated
    cumulative weights of each sequence of shape ``(n_seqs,)``.

    """
    import torch

    logprobs = torch.nn.functional.log_softmax(next_token_logits, dim=-1)
    weights = logprobs + sequence_weights.unsqueeze(1).expand_as(next_token_logits)

    # Flatten scores to (n_batch, n_samples * vocab_size)
    # and find the top-k weights for each batch.
    batch_size = next_token_logits.shape[0] // self.samples
    vocab_size = next_token_logits.shape[-1]
    weights = weights.view(batch_size, self.samples * vocab_size)

    # If the weights are all equal to 0 we are at the beginning of the search
    # and thus only need to sample from one set of token logits for each
    # batch.
    if torch.all(sequence_weights == 0):
        weights = weights[:, :vocab_size]

    weights, indices = torch.topk(
        weights, self.samples, dim=1, largest=True, sorted=True
    )

    ancestors = torch.div(indices, vocab_size, rounding_mode="floor")
    next_token_ids = indices % vocab_size

    # Re-shape the weights, next_token_ids and ancestors to (n_batch * n_samples, 1)
    first_batch_idx = torch.arange(
        0, batch_size * self.samples, self.samples, device=next_token_logits.device
    ).unsqueeze(1)
    ancestors = ancestors + first_batch_idx

    ancestors = ancestors.view(self.samples * batch_size)
    weights = weights.view(self.samples * batch_size)
    next_token_ids = next_token_ids.view(self.samples * batch_size, 1)

    return next_token_ids, ancestors, weights

GreedySampler

Greedy Sampling algorithm.

Greedy sampling consists in choosing the token with the largest likelihood at every step.

We don't allow more than one sample. We could attribute this a meaning, for instance the k-th sample represents the k-th most likely token. In which case it would be equivalent to beam search without the sequence weights.

Attributes

samples The number of samples taken for each input sequence.

Source code in outlines/samplers.py 
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
class GreedySampler:
    """Greedy Sampling algorithm.

    Greedy sampling consists in choosing the token with the largest
    likelihood at every step.

    We don't allow more than one sample. We could attribute this a meaning, for
    instance the k-th sample represents the k-th most likely token. In which
    case it would be equivalent to beam search without the sequence weights.

    Attributes
    ----------
    samples
        The number of samples taken for each input sequence.

    """

    def __init__(self):
        self.samples = 1

    def __call__(
        self,
        next_token_logits: "torch.DoubleTensor",
        sequence_weights: "torch.DoubleTensor",
        _,
    ) -> "torch.DoubleTensor":
        """Call the greedy sampler.

        Parameters
        ----------
        next_token_logits
            A tensor of shape ``(n_seqs, vocab_size,)`` that represents the
            probability distribution of the next token over the vocabulary.
        sequence_weights
            A tensor of shape ``(n_seqs,)`` that represents the cumulative
            weight of each sequence.
        rng
            A random number generator.

        Returns
        -------
        A tuple with an array that contains the ids of the sampled tokens of
        shape ``(n_seqs, 1)``, an array that contains the ancestors of each
        sampled id of shape ``(n_seqs,)`` and an array that contains the updated
        cumulative weights of each sequence of shape ``(n_seqs,)``.

        """
        import torch

        logprobs = torch.nn.functional.log_softmax(next_token_logits, dim=-1)
        next_token_ids = torch.argmax(logprobs, dim=-1, keepdim=True)

        ancestors = torch.arange(
            next_token_logits.shape[0], device=next_token_logits.device
        )
        weights = sequence_weights + torch.gather(logprobs, 1, next_token_ids).squeeze()

        return next_token_ids, ancestors, weights

__call__(next_token_logits, sequence_weights, _)

Call the greedy sampler.

Parameters

next_token_logits A tensor of shape (n_seqs, vocab_size,) that represents the probability distribution of the next token over the vocabulary. sequence_weights A tensor of shape (n_seqs,) that represents the cumulative weight of each sequence. rng A random number generator.

Returns

A tuple with an array that contains the ids of the sampled tokens of shape (n_seqs, 1), an array that contains the ancestors of each sampled id of shape (n_seqs,) and an array that contains the updated cumulative weights of each sequence of shape (n_seqs,).

Source code in outlines/samplers.py 
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
def __call__(
    self,
    next_token_logits: "torch.DoubleTensor",
    sequence_weights: "torch.DoubleTensor",
    _,
) -> "torch.DoubleTensor":
    """Call the greedy sampler.

    Parameters
    ----------
    next_token_logits
        A tensor of shape ``(n_seqs, vocab_size,)`` that represents the
        probability distribution of the next token over the vocabulary.
    sequence_weights
        A tensor of shape ``(n_seqs,)`` that represents the cumulative
        weight of each sequence.
    rng
        A random number generator.

    Returns
    -------
    A tuple with an array that contains the ids of the sampled tokens of
    shape ``(n_seqs, 1)``, an array that contains the ancestors of each
    sampled id of shape ``(n_seqs,)`` and an array that contains the updated
    cumulative weights of each sequence of shape ``(n_seqs,)``.

    """
    import torch

    logprobs = torch.nn.functional.log_softmax(next_token_logits, dim=-1)
    next_token_ids = torch.argmax(logprobs, dim=-1, keepdim=True)

    ancestors = torch.arange(
        next_token_logits.shape[0], device=next_token_logits.device
    )
    weights = sequence_weights + torch.gather(logprobs, 1, next_token_ids).squeeze()

    return next_token_ids, ancestors, weights

MultinomialSampler

Multinomial sampling algorithm.

Multinomial sampling consists in randomly sampling the next token assuming its distribution is a Categorical distribution parametrized by the next-token logits.

Attributes

samples The number of samples taken for each input sequence.

Source code in outlines/samplers.py 
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
class MultinomialSampler:
    """Multinomial sampling algorithm.

    Multinomial sampling consists in randomly sampling the next token assuming
    its distribution is a Categorical distribution parametrized by the
    next-token logits.


    Attributes
    ----------
    samples
        The number of samples taken for each input sequence.

    """

    def __init__(
        self,
        samples: int = 1,
        *,
        top_k: Optional[int] = None,
        top_p: Optional[float] = None,
        temperature: Optional[float] = None,
    ):
        self.samples = samples
        self.top_k = top_k
        self.top_p = top_p
        self.temperature = temperature

        self.logits_processors = []
        if top_k is not None:
            self.logits_processors.append(keep_top_k_logits(top_k))
        elif top_p is not None:
            self.logits_processors.append(keep_top_p_logits(top_p))

        if temperature is not None:
            self.logits_processors.append(rescale_logits(temperature))

    def __call__(
        self,
        next_token_logits: "torch.DoubleTensor",
        sequence_weights: "torch.DoubleTensor",
        rng: "torch.Generator",
    ) -> Tuple["torch.DoubleTensor", "torch.DoubleTensor", "torch.DoubleTensor"]:
        """Call the multinomial sampler.

        Parameters
        ----------
        next_token_logits
            A tensor of shape ``(n_seqs, vocab_size,)`` that represents the
            probability distribution of the next token over the vocabulary.
        sequence_weights
            A tensor of shape ``(n_seqs,)`` that represents the cumulative
            weight of each sequence.
        rng
            A random number generator.

        Returns
        -------
        A tuple with an array that contains the ids of the sampled tokens of
        shape ``(n_seqs, 1)``, an array that contains the ancestors of each
        sampled id of shape ``(n_seqs,)`` and an array that contains the updated
        cumulative weights of each sequence of shape ``(n_seqs,)``.

        """
        import torch

        altered_next_token_logits = next_token_logits
        for logit_processor in self.logits_processors:
            altered_next_token_logits = logit_processor(next_token_logits)

        probs = torch.nn.functional.softmax(altered_next_token_logits, dim=-1)
        next_token_ids = torch.multinomial(probs, num_samples=1, generator=rng)

        logprobs = torch.nn.functional.log_softmax(altered_next_token_logits, dim=-1)
        ancestors = torch.arange(
            altered_next_token_logits.shape[0], device=next_token_logits.device
        )
        weights = sequence_weights + torch.gather(logprobs, 1, next_token_ids).squeeze()

        return next_token_ids, ancestors, weights

__call__(next_token_logits, sequence_weights, rng)

Call the multinomial sampler.

Parameters

next_token_logits A tensor of shape (n_seqs, vocab_size,) that represents the probability distribution of the next token over the vocabulary. sequence_weights A tensor of shape (n_seqs,) that represents the cumulative weight of each sequence. rng A random number generator.

Returns

A tuple with an array that contains the ids of the sampled tokens of shape (n_seqs, 1), an array that contains the ancestors of each sampled id of shape (n_seqs,) and an array that contains the updated cumulative weights of each sequence of shape (n_seqs,).

Source code in outlines/samplers.py 
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
def __call__(
    self,
    next_token_logits: "torch.DoubleTensor",
    sequence_weights: "torch.DoubleTensor",
    rng: "torch.Generator",
) -> Tuple["torch.DoubleTensor", "torch.DoubleTensor", "torch.DoubleTensor"]:
    """Call the multinomial sampler.

    Parameters
    ----------
    next_token_logits
        A tensor of shape ``(n_seqs, vocab_size,)`` that represents the
        probability distribution of the next token over the vocabulary.
    sequence_weights
        A tensor of shape ``(n_seqs,)`` that represents the cumulative
        weight of each sequence.
    rng
        A random number generator.

    Returns
    -------
    A tuple with an array that contains the ids of the sampled tokens of
    shape ``(n_seqs, 1)``, an array that contains the ancestors of each
    sampled id of shape ``(n_seqs,)`` and an array that contains the updated
    cumulative weights of each sequence of shape ``(n_seqs,)``.

    """
    import torch

    altered_next_token_logits = next_token_logits
    for logit_processor in self.logits_processors:
        altered_next_token_logits = logit_processor(next_token_logits)

    probs = torch.nn.functional.softmax(altered_next_token_logits, dim=-1)
    next_token_ids = torch.multinomial(probs, num_samples=1, generator=rng)

    logprobs = torch.nn.functional.log_softmax(altered_next_token_logits, dim=-1)
    ancestors = torch.arange(
        altered_next_token_logits.shape[0], device=next_token_logits.device
    )
    weights = sequence_weights + torch.gather(logprobs, 1, next_token_ids).squeeze()

    return next_token_ids, ancestors, weights

keep_top_k_logits(k)

Build a function that masks logits values smaller than the top k ones.

Parameters

k The ranking below which logit values are replaced by -math.inf.

Source code in outlines/samplers.py 
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
def keep_top_k_logits(k: int) -> Callable[["torch.Tensor"], "torch.Tensor"]:
    """Build a function that masks logits values smaller than the top `k` ones.

    Parameters
    ----------
    k
        The ranking below which logit values are replaced by `-math.inf`.

    """
    import torch

    if not isinstance(k, int) or k < 1:
        raise ValueError(f"`k` must be a strictly positive integers, got {k} instead.")

    def logits_processor(logits: torch.Tensor) -> torch.Tensor:
        num_to_keep = min(k, logits.size(-1))
        mask_idx = logits < torch.topk(logits, num_to_keep)[0][..., -1, None]
        return logits.masked_fill(mask_idx, -math.inf)

    return logits_processor

keep_top_p_logits(p)

Build a function that masks the lowest probability tokens whose cumulative probability is below a certain threshold.

Parameters

p The value of the threshold. We keep the highest probability tokens whose cumulative distribution is greater than or equal to p and mask the others. Its value must be between 0 (excluded) and 1 (included).

Source code in outlines/samplers.py 
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
def keep_top_p_logits(p: float) -> Callable[["torch.Tensor"], "torch.Tensor"]:
    """Build a function that masks the lowest probability tokens whose
    cumulative probability is below a certain threshold.

    Parameters
    ----------
    p
        The value of the threshold. We keep the highest probability tokens whose
        cumulative distribution is greater than or equal to `p` and mask the
        others. Its value must be between 0 (excluded) and 1 (included).

    """
    import torch

    if p <= 0.0 or p > 1.0:
        raise ValueError(
            f"`p` must be a floating point number between 0 (excluded) and 1 (included), got {p} instead."
        )

    def logits_processor(logits: torch.Tensor) -> torch.Tensor:
        sorted_logits, sorted_idx = torch.sort(logits, descending=False)
        cumulative_probabilties = torch.nn.functional.softmax(
            sorted_logits, dim=-1
        ).cumsum(dim=-1)

        sorted_masked_idx = cumulative_probabilties <= (1 - p)
        mask_idx = torch.scatter(sorted_masked_idx, 1, sorted_idx, sorted_masked_idx)
        return logits.masked_fill(mask_idx, -math.inf)

    return logits_processor

rescale_logits(temperature)

Build a function that rescales the token probabilities exponentially.

Parameters

temperature The value by which we rescale the logits.

Source code in outlines/samplers.py 
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
def rescale_logits(temperature: float) -> Callable[["torch.Tensor"], "torch.Tensor"]:
    """Build a function that rescales the token probabilities exponentially.

    Parameters
    ----------
    temperature
        The value by which we rescale the logits.

    """

    if not isinstance(temperature, float) or temperature < 0.0:
        raise ValueError(
            f"`temperature` must be a strictly positive floating point number, got {temperature} instead."
        )
    elif temperature == 0.0:
        raise ValueError(
            "Please use the greedy sampler instead of setting the temperature to 0."
        )

    def logits_processor(logits: "torch.Tensor") -> "torch.Tensor":
        return logits / temperature

    return logits_processor