code: def generate():
      """
      ## Generate text
      """
  
      # Setup [cache](../utils/cache.html) to cache intermediate key/value pairs for faster generation
      cache = get_cache()
      cache.set("use_cache", True)
  
      # Device
      device = torch.device("cuda:0")
  
      # Load layers
      layers = list(
          LayerGenerator(
              is_clone_layers=True,
              filter_layers=LAYERS,
              dtype=torch.float16,
              device=device,
          ).load()
      )
  
      model = nn.Sequential(*layers)
  
      # Get token ids
      ids = get_tokens(PROMPT)
  
      # Run the model
      cache.set("state_ids", (None, 1))
      with monit.section("Infer"):
          next_token = infer(model, ids, device)[-1]
  
      # Append the predicted token
      ids += [next_token]
  
      # Predict 100 tokens
      for i in range(1, 100):
          # Set the state to use cached activations
          cache.set("state_ids", (i, i + 1))
          # Get next token. Note that we only feed the last token to the model because
          # we cache the key/value pairs of previous tokens.
          with monit.section("Infer"):
              next_token = infer(model, [next_token], device)[-1]
          # Append the predicted token
          ids += [next_token]
          # Print
          print_tokens(ids, [ids])
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/neox/samples/generate.py
imports: from typing import List
  import torch
  from torch import nn
  from labml import monit
  from labml_nn.neox.model import LayerGenerator
  from labml_nn.neox.utils import get_tokens, print_tokens
  from labml_nn.neox.utils.cache import get_cache
label_desc: Use image node_id 2 for calling a global function with example usage: generate() without return types