class_docstr: ## Sampler
  
  This class greedily samples from a model.
class_docstr_tok: ["#","#","Sampler","This","class","greedily","samples","from","a","model","."]
docstr: ### Sample text from the given prompt
docstr_tok: ["#","#","#","Sample","text","from","the","given","prompt"]
code: def sample(self, prompt: str, sample_len: int):
      """
      ### Sample text from the given prompt
      """
  
      # To store nearest neighbors as strings
      neighbors_str = []
  
      # Sampled text
      sampled = ""
  
      # Sample `sample_len` tokens
      for i in range(sample_len):
          # We need to retrieve neighbors,
          # if there are more sampled chunks than we have already retrieved for
          while len(neighbors_str) < len(prompt) // self.chunk_len:
              # Get the last chunk for which we haven't retrieved neighbors
              off = len(neighbors_str) * self.chunk_len
              chunk = prompt[off : off + self.chunk_len]
              # Retrieve nearest neighbors
              neighbors_str.append(
                  self.retrieve_nearest_neighbours(chunk)
              )
  
          # Tokenize the input
          src = self.tds.text_to_i(prompt)
          # Tokenize the retrieved neighbors
          neighbors = torch.stack(
              [
                  torch.stack([self.tds.text_to_i(n) for n in chunk])
                  for chunk in neighbors_str
              ]
          )
  
          # Move them to the same device as the model
          src = src.to(self.device)
          neighbors = neighbors.to(self.device)
  
          # Get model output
          res = self.model(src[None, :], neighbors[None, :, :, :])
  
          # Greedily sample the last token
          token = res[0, -1, :].argmax(dim=-1)
  
          # Add the sampled token text to the prompt and sample text
          prompt += self.tds.itos[token.item()]
          sampled += self.tds.itos[token.item()]
  
      #
      return sampled
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/transformers/retro/train.py
imports: import torch
  from torch import nn
  from torch.utils.data import DataLoader, RandomSampler
  from labml import monit, lab, tracker, experiment, logger
  from labml.logger import Text
  from labml_helpers.datasets.text import TextFileDataset
  from labml_nn.optimizers.noam import Noam
  from labml_nn.transformers.retro import model
  from labml_nn.transformers.retro.dataset import Dataset, RetroIndex
  from labml_nn.transformers.retro.model import RetroModel, NearestNeighborEncoder
label_desc: Use image node_id 3 for calling the Sampler obj's underlying member method code with example usage: obj.sample(prompt, sample_len) and returns: sampled