code: def sample(
      model,
      ds,
      sampler: Sampler,
      n_samples: int,
      n_tokens: int,
      seq_len: int,
      prompt: str,
  ):
      with torch.no_grad():
          data = torch.tile(
              ds.text_to_i(prompt)[:, None], (1, n_samples)
          )
  
          # Collect output for printing
          logs = [[(prompt, Text.meta)] for _ in range(n_samples)]
          # Sample 25 tokens
          for i in monit.iterate("Sample", n_tokens):
              # Tokenize the prompt
              data = data[-seq_len:]
              # Get the model output
              logits, *_ = model(data)
              logits = logits[-1]
              # Get the model prediction (greedy)
              res = sampler(logits)
              data = torch.cat([data, res[None, :]], dim=0)
              # Add the prediction for logging
              for j in range(n_samples):
                  logs[j] += [("" + ds.itos[res[j]], Text.value)]
  
      # Print the sampled output
      for j in range(n_samples):
          logger.log(logs[j])
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/sampling/experiment_tiny.py
imports: from typing import Tuple
  import torch
  from labml import experiment, monit
  from labml import logger
  from labml.logger import Text
  from labml_helpers.datasets.text import TextDataset
  from labml_nn.sampling import Sampler
  from labml_nn.sampling.greedy import GreedySampler
  from labml_nn.sampling.nucleus import NucleusSampler
  from labml_nn.sampling.temperature import TemperatureSampler
  from labml_nn.sampling.top_k import TopKSampler
  from labml_nn.transformers.basic.autoregressive_experiment import Configs, AutoregressiveTransformer
label_desc: Use image node_id 2 for calling a global function with example usage: sample(model, ds, sampler, n_samples, n_tokens, seq_len, prompt) without return types