code: def _test():
      """
      ### Code to test BERT embeddings
      """
      from labml.logger import inspect
  
      # Initialize
      device = torch.device("cuda:0")
      bert = BERTChunkEmbeddings(device)
  
      # Sample
      text = ["Replace me by any text you'd like.", "Second sentence"]
  
      # Check BERT tokenizer
      encoded_input = bert.tokenizer(
          text,
          return_tensors="pt",
          add_special_tokens=False,
          padding=True,
      )
  
      inspect(encoded_input, _expand=True)
  
      # Check BERT model outputs
      output = bert.model(
          input_ids=encoded_input["input_ids"].to(device),
          attention_mask=encoded_input["attention_mask"].to(device),
          token_type_ids=encoded_input["token_type_ids"].to(device),
      )
  
      inspect(
          {
              "last_hidden_state": output["last_hidden_state"],
              "pooler_output": output["pooler_output"],
          },
          _expand=True,
      )
  
      # Check recreating text from token ids
      inspect(
          bert.tokenizer.convert_ids_to_tokens(
              encoded_input["input_ids"][0]
          ),
          _n=-1,
      )
      inspect(
          bert.tokenizer.convert_ids_to_tokens(
              encoded_input["input_ids"][1]
          ),
          _n=-1,
      )
  
      # Get chunk embeddings
      inspect(bert(text))
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/transformers/retro/bert_embeddings.py
imports: from typing import List
  import torch
  from transformers import BertTokenizer, BertModel
  from labml import lab, monit
label_desc: Use image node_id 1 for calling a global function with example usage: _test() without return types