class_docstr: The overall arch described in the BERT4Rec paper: (https://arxiv.org/abs/1904.06690)
  the encoder_layer was described in the section of 3.3, the output_layer was described in the
  section of 3.5
  
  Args:
      vocab_size (int): the item count including mask and padding
      max_len (int): the max length
      emb_dim (int): embedding dimension
      nhead (int): number of the transformation headers
      num_layers (int): number of the transformation layers
      dropout (float): the dropout probability
      device: (Optional[torch.device]).
  
  Example::
  
      input_kjt = KeyedJaggedTensor.from_lengths_sync(
          keys=["item"],
          values=torch.tensor([2, 4, 3, 4, 5]),
          lengths=torch.tensor([2, 3]),
      )
      device = (
          torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
      )
      input_kjt = input_kjt.to(device)
      bert4rec = BERT4Rec(
          vocab_size=6, max_len=3, emb_dim=4, nhead=4, num_layers=4, device=device
      )
      logits = bert4rec(input_kjt)
class_docstr_tok: ["The","overall","arch","described","in","the","BERT4Rec","paper",":","(","https",":","\/\/arxiv.org\/abs\/1904.06690",")","the","encoder_layer","was","described","in","the","section","of","3.3",",","the","output_layer","was","described","in","the","section","of","3.5","Args",":","vocab_size","(","int",")",":","the","item","count","including","mask","and","padding","max_len","(","int",")",":","the","max","length","emb_dim","(","int",")",":","embedding","dimension","nhead","(","int",")",":","number","of","the","transformation","headers","num_layers","(","int",")",":","number","of","the","transformation","layers","dropout","(","float",")",":","the","dropout","probability","device",":","(","Optional","[","torch.device","]",")",".","Example",":",":","input_kjt","=","KeyedJaggedTensor.from_lengths_sync","(","keys=","[","``","item","''","]",",","values=torch.tensor","(","[","2",",","4",",","3",",","4",",","5","]",")",",","lengths=torch.tensor","(","[","2",",","3","]",")",",",")","device","=","(","torch.device","(","``","cuda","''",")","if","torch.cuda.is_available","(",")","else","torch.device","(","``","cpu","''",")",")","input_kjt","=","input_kjt.to","(","device",")","bert4rec","=","BERT4Rec","(","vocab_size=6",",","max_len=3",",","emb_dim=4",",","nhead=4",",","num_layers=4",",","device=device",")","logits","=","bert4rec","(","input_kjt",")"]
code: def __init__(
      self,
      vocab_size: int,
      max_len: int,
      emb_dim: int,
      nhead: int,
      num_layers: int,
      dropout: float = 0.1,
      device: Optional[torch.device] = None,
  ) -> None:
      super().__init__()
      self.vocab_size = vocab_size
      self.emb_dim = emb_dim
      self.max_len = max_len
      self.history = HistoryArch(
          vocab_size, max_len, emb_dim, dropout=dropout, device=device
      )
  
      self.transformer_blocks = nn.ModuleList(
          [
              TransformerBlock(
                  emb_dim, nhead, emb_dim * 4, dropout, device=device
              )
              for _ in range(num_layers)
          ]
      )
      # this linear layer is different from the original paper which used the
      # embedding layer weight to get the final output by matmul, from our experimentation,
      # this implementation has higher performance
      self.out = nn.Linear(self.emb_dim, self.vocab_size, device=device)
  
imports: import copy
  import math
  from typing import Callable, Optional, Tuple
  import torch
  import torch.nn
  from torchrec.modules.embedding_configs import EmbeddingConfig
  from torchrec.modules.embedding_modules import EmbeddingCollection
  from torchrec.sparse.jagged_tensor import KeyedJaggedTensor
label_desc: Use image node_id 1 to create a new BERT4Rec object from inherited base classes: nn with example: obj = BERT4Rec(vocab_size, max_len, emb_dim, nhead, num_layers, dropout, device)