Move samplers module to root module

docs/api/samplers.md

@@ -1 +1 @@
-::: outlines.generate.samplers
+::: outlines.samplers

outlines/generate/api.py

@@ -14,7 +14,7 @@
     sequence_generator,
     token_generator,
 )
-from outlines.generate.samplers import MultinomialSampler, Sampler
+from outlines.samplers import MultinomialSampler, Sampler
 
 
 class SequenceGenerator:

outlines/generate/generator.py

@@ -8,8 +8,8 @@
 
 if TYPE_CHECKING:
     from outlines.fsm.fsm import FSM
-    from outlines.generate.samplers import Sampler
     from outlines.models.tokenizer import Tokenizer
+    from outlines.samplers import Sampler
 
 
 @dataclasses.dataclass(frozen=True)

outlines/generate/samplers.py → outlines/samplers.py

@@ -18,15 +18,12 @@ class GreedySampler:
     Greedy sampling consists in choosing the token with the largest
     likelihood at every step.
 
-    Attributes
-    ----------
-    particles
-        The number of samples taken for each input sequence.
+    We don't allow more than one sample as this does not really make sense.
 
     """
 
-    def __init__(self, samples: int = 1):
-        self.particles = samples
+    def __init__(self):
+        self.particles = 1
 
     def __call__(self, logits: torch.DoubleTensor, *_) -> torch.DoubleTensor:
         """Call the greedy sampler.
@@ -41,15 +38,10 @@ def __call__(self, logits: torch.DoubleTensor, *_) -> torch.DoubleTensor:
 
         Returns
         -------
-        The ids of the sampled tokens having shape ``(samples, n_seqs)``.
+        The ids of the sampled tokens, of shape ``(n_seqs, 1)``
 
         """
-        if self.particles == 1:
-            next_token_ids = torch.argmax(logits, dim=-1, keepdim=True)
-        else:
-            next_token_ids = torch.topk(
-                logits, self.particles, dim=-1, largest=True, sorted=True
-            ).indices.T
+        next_token_ids = torch.argmax(logits, dim=-1, keepdim=True)
 
         return next_token_ids
 

outlines/text/generate/api.py

@@ -2,7 +2,7 @@
 from typing import Callable, List, Optional, Union
 
 import outlines
-from outlines.generate.samplers import MultinomialSampler, Sampler
+from outlines.samplers import MultinomialSampler, Sampler
 
 
 def json(

tests/generate/test_integration_transfomers.py

@@ -10,8 +10,8 @@
 import outlines.generate as generate
 import outlines.models as models
 from outlines.fsm.regex import reduced_vocabulary
-from outlines.generate.samplers import multinomial
 from outlines.models.transformers import TransformerTokenizer
+from outlines.samplers import multinomial
 
 
 def test_deprecation():

tests/generate/test_samplers.py → tests/test_samplers.py

@@ -2,12 +2,7 @@
 
 import torch
 
-from outlines.generate.samplers import (
-    GreedySampler,
-    MultinomialSampler,
-    greedy,
-    multinomial,
-)
+from outlines.samplers import GreedySampler, MultinomialSampler, greedy, multinomial
 
 
 def test_aliases():
@@ -21,19 +16,11 @@ def test_greedy():
     next_token_ids = sampler(logits)
     assert next_token_ids.equal(torch.tensor([[2]]))
 
-    sampler = GreedySampler(2)
-    next_token_ids = sampler(logits)
-    assert next_token_ids.equal(torch.tensor([[2], [1]]))
-
     sampler = GreedySampler()
     logits = torch.tensor([[10.0, 0.0, 3.0], [-math.inf, 2.0, 5.0]])
     next_token_ids = sampler(logits)
     assert next_token_ids.equal(torch.tensor([[0], [2]]))
 
-    sampler = GreedySampler(2)
-    next_token_ids = sampler(logits)
-    assert next_token_ids.equal(torch.tensor([[0, 2], [2, 1]]))
-
 
 def test_multinomial():
     rng = torch.Generator()