class_docstr: ## Nucleus Sampler
class_docstr_tok: ["#","#","Nucleus","Sampler"]
docstr: Sample from logits with Nucleus Sampling
docstr_tok: ["Sample","from","logits","with","Nucleus","Sampling"]
code: def __call__(self, logits: torch.Tensor):
      """
      Sample from logits with Nucleus Sampling
      """
  
      # Get probabilities $P(x_i | x_{1:i-1})$
      probs = self.softmax(logits)
  
      # Sort probabilities in descending order
      sorted_probs, indices = torch.sort(probs, dim=-1, descending=True)
      # Get the cumulative sum of probabilities in the sorted order
      cum_sum_probs = torch.cumsum(sorted_probs, dim=-1)
      # Find the cumulative sums less than $p$.
      nucleus = cum_sum_probs < self.p
      # Prepend ones so that we add one token after the minimum number
      # of tokens with cumulative probability less that $p$.
      nucleus = torch.cat(
          [
              nucleus.new_ones(nucleus.shape[:-1] + (1,)),
              nucleus[..., :-1],
          ],
          dim=-1,
      )
  
      # Get log probabilities and mask out the non-nucleus
      sorted_log_probs = torch.log(sorted_probs)
      sorted_log_probs[~nucleus] = float("-inf")
  
      # Sample from the sampler
      sampled_sorted_indexes = self.sampler(sorted_log_probs)
  
      # Get the actual indexes
      res = indices.gather(-1, sampled_sorted_indexes.unsqueeze(-1))
  
      #
      return res.squeeze(-1)
  
imports: import torch
  from torch import nn
  from labml_nn.sampling import Sampler
label_desc: Use image node_id 2 for calling the NucleusSampler obj's underlying member method code with example usage: obj.__call__(logits) and returns: res