Fix output shape of greedy sampler with batch > 1

rlouf · Dec 22, 2023 · 1e9b510 · 1e9b510
1 parent e611c9b
commit 1e9b510
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/outlines/generate/samplers.py b/outlines/generate/samplers.py
@@ -33,7 +33,7 @@ def greedy(logits: torch.DoubleTensor, samples: int, *_) -> torch.DoubleTensor:
 
  """
  if samples == 1:
- next_token_ids = torch.argmax(logits, dim=-1, keepdim=True).T
+ next_token_ids = torch.argmax(logits, dim=-1, keepdim=True)
  else:
  next_token_ids = torch.topk(
  logits, samples, dim=-1, largest=True, sorted=True

diff --git a/tests/generate/test_samplers.py b/tests/generate/test_samplers.py
@@ -15,7 +15,7 @@ def test_greedy():
 
  logits = torch.tensor([[10.0, 0.0, 3.0], [-math.inf, 2.0, 5.0]])
  next_token_ids = greedy(logits, samples=1)
- assert next_token_ids.equal(torch.tensor([[0, 2]]))
+ assert next_token_ids.equal(torch.tensor([[0], [2]]))
 
  next_token_ids = greedy(logits, samples=2)
  assert next_token_ids.equal(torch.tensor([[0, 2], [2, 1]]))