code: def forward(
      self, src_tokens, know_tokens, ck_mask, cs_ids, use_cs_ids
  ):
      # encode the context, pretty basic
      context_encoded, context_mask = self.transformer(src_tokens)
  
      # make all the knowledge into a 2D matrix to encode
      N, K, Tk = know_tokens.size()
      know_flat = know_tokens.reshape(-1, Tk)
      know_encoded, know_mask = self.transformer(know_flat)
  
      # compute our sentence embeddings for context and knowledge
      context_use = universal_sentence_embedding(
          context_encoded, context_mask
      )
      know_use = universal_sentence_embedding(know_encoded, know_mask)
  
      # remash it back into the shape we need
      know_use = know_use.reshape(
          N, know_tokens.size(1), self.embed_dim
      )
      context_use /= np.sqrt(self.embed_dim)
      know_use /= np.sqrt(self.embed_dim)
  
      ck_attn = th.bmm(know_use, context_use.unsqueeze(-1)).squeeze(-1)
      # fill with near -inf
      ck_attn.masked_fill_(~ck_mask, neginf(context_encoded.dtype))
  
      if not use_cs_ids:
          # if we're not given the true chosen_sentence (test time), pick our
          # best guess
          _, cs_ids = ck_attn.max(1)
  
      # pick the true chosen sentence. remember that TransformerEncoder outputs
      #   (batch, time, embed)
      # but because know_encoded is a flattened, it's really
      #   (N * K, T, D)
      # We need to compute the offsets of the chosen_sentences
      cs_offsets = th.arange(N, device=cs_ids.device) * K + cs_ids
      cs_encoded = know_encoded[cs_offsets]
      # but padding is (N * K, T)
      cs_mask = know_mask[cs_offsets]
  
      # finally, concatenate it all
      full_enc = th.cat([cs_encoded, context_encoded], dim=1)
      full_mask = th.cat([cs_mask, context_mask], dim=1)
  
      # also return the knowledge selection mask for the loss
      return full_enc, full_mask, ck_attn
  
file_path: ParlAI/projects/wizard_of_wikipedia/generator/modules.py
imports: from typing import Dict
  import numpy
  import torch
  import torch.nn
  from parlai.utils.torch import neginf
  from parlai.agents.transformer.modules import TransformerGeneratorModel
label_desc: Use image node_id 2 for calling the ContextKnowledgeEncoder obj's underlying member method code with example usage: obj.forward(src_tokens, know_tokens, ck_mask, cs_ids, use_cs_ids) and returns: full_enc,  full_mask,  ck_attn