Expose sample from generation utils

docs/source/api_ref_generation.rst

@@ -11,3 +11,4 @@ torchtune.generation
     :nosignatures:
 
     generate
+    sample

torchtune/generation/__init__.py

@@ -4,6 +4,6 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-from ._generation import generate, generate_next_token
+from ._generation import generate, generate_next_token, sample
 
-__all__ = ["generate", "generate_next_token"]
+__all__ = ["generate", "generate_next_token", "sample"]

torchtune/generation/_generation.py

@@ -17,9 +17,18 @@ def multinomial_sample_one(probs: torch.Tensor) -> torch.Tensor:
 
 
 def sample(
-    logits: torch.Tensor, temperature: float = 1.0, top_k: int = None
+    logits: torch.Tensor, temperature: float = 1.0, top_k: Optional[int] = None
 ) -> torch.Tensor:
-    """Generic sample from a probability distribution."""
+    """Generic sample from a probability distribution.
+
+    Args:
+        logits (torch.Tensor): logits from which to sample
+        temperature (float): value to scale the predicted logits by, default 1.0.
+        top_k (Optional[int]): If specified, we prune the sampling to only token ids within the top_k probabilities
+
+    Returns:
+        torch.Tensor: sampled token id
+    """
     # scale the logits based on temperature
     logits = logits / max(temperature, 1e-5)
     if top_k is not None: