Class Sampling

java.lang.Object
smile.llm.engine.Sampling
public final class Sampling extends Object
Shared sampling helpers for greedy / temperature + nucleus decoding.

  • Method Details

    • sampleGreedyTokenId

      public static int sampleGreedyTokenId(Tensor logits)
      Greedy argmax on device; synchronizes a single scalar to the host via Tensor.getLong(long...) (works on views; avoids a full-vocab D2H).
      Parameters:
      logits - last-step logits [batch, vocab] (or [1, vocab]).
      Returns:
      sampled token id for batch row 0.

    • sampleGreedyTokenIds

      public static int[] sampleGreedyTokenIds(Tensor logits)
      Batched greedy argmax; one kernel and one host copy for the whole batch.
      Parameters:
      logits - last-step logits [batch, vocab].
      Returns:
      token id per batch row.

    • sampleTokenIds

      public static int[] sampleTokenIds(Tensor logits, double temperature, double topp)
      Batched sample: greedy when temperature <= 0, otherwise one temperature + nucleus draw for every row of [batch, vocab].
      Parameters:
      logits - logits [batch, vocab].
      temperature - sampling temperature (<= 0 → greedy).
      topp - nucleus sampling threshold.
      Returns:
      sampled token ids [batch].

    • sampleNext

      public static Tensor sampleNext(Tensor logits, double temperature, double topp)
      Samples the next token from last-position logits [batch, vocab] (or [batch, 1, vocab] squeezed by the caller).
      Parameters:
      logits - last-step logits; not closed by this method.
      temperature - <= 0 selects greedy argmax.
      topp - nucleus threshold when temperature > 0.
      Returns:
      owned token ids shaped [batch] (caller must close).

    • mergeWithPromptMask

      public static Tensor mergeWithPromptMask(Tensor textMask, Tensor currentTokens, Tensor nextToken)
      Merges sampled tokens with prompt-overlay positions using the input mask.
      Parameters:
      textMask - true where the sequence still holds prompt tokens.
      currentTokens - tokens currently at this position (prompt or pad).
      nextToken - freshly sampled tokens.
      Returns:
      owned merged tokens (caller must close).