docstr: Wizard custom batchify, which passes along the knowledge/title.
  
  Following the docstring of TorchAgent.batchify, it calls super, then
  uses an extended version of the torch_agent.Batch namedtuple.
  
  The purpose of extending the info is to keep track of some custom
  metrics.
docstr_tok: ["Wizard","custom","batchify",",","which","passes","along","the","knowledge\/title",".","Following","the","docstring","of","TorchAgent.batchify",",","it","calls","super",",","then","uses","an","extended","version","of","the","torch_agent.Batch","namedtuple",".","The","purpose","of","extending","the","info","is","to","keep","track","of","some","custom","metrics","."]
code: def batchify(self, obs_batch):
      """
      Wizard custom batchify, which passes along the knowledge/title.
  
      Following the docstring of TorchAgent.batchify, it calls super, then
      uses an extended version of the torch_agent.Batch namedtuple.
  
      The purpose of extending the info is to keep track of some custom
      metrics.
      """
      batch = super().batchify(obs_batch)
      reordered_observations = [
          obs_batch[i] for i in batch.valid_indices
      ]
      is_training = "labels" in reordered_observations[0]
  
      # first parse and compile all the knowledge together
      all_knowledges = (
          []
      )  # list-of-lists knowledge items for each observation
      knowledge_counts = []  # how much knowledge each observation gets
      for obs in reordered_observations:
          obs_know = self._parse_knowledge(obs)
          # downsample if desired
          if (
              is_training
              and self.max_knowledge
              and len(obs_know) > self.max_knowledge
          ):
              # offset by one so that we don't choose 0
              keepers = 1 + np.random.choice(
                  len(obs_know) - 1, self.max_knowledge, False
              )
              # correct answer is always the first one
              keepers[0] = 0
              obs_know = [obs_know[i] for i in keepers]
          all_knowledges.append(obs_know)
          knowledge_counts.append(len(obs_know))
  
      # now we want to actually pack this into a tensor, along with the mask
      N = len(reordered_observations)
      K = max(knowledge_counts)
      # round out the array so everything is equally sized
      for i in range(N):
          all_knowledges[i] += [""] * (K - knowledge_counts[i])
      flattened_knowledge = list(chain(*all_knowledges))
  
      knowledge_vec = [
          self._vectorize_text(
              # the beginning of the sentence is more useful
              k,
              truncate=self.knowledge_truncate,
              add_end=True,
              truncate_left=False,
          )
          for k in flattened_knowledge
      ]
      knowledge_vec, _ = padded_tensor(
          knowledge_vec, pad_idx=self.NULL_IDX, left_padded=True
      )
      knowledge_vec[:, -1] = self.END_IDX
      T = knowledge_vec.size(-1)
      knowledge_vec = knowledge_vec.view(N, K, T)
  
      # knowledge mask is a N x K tensor saying which items we're allowed to
      # attend over
      bsz = len(reordered_observations)
      ck_mask = th.zeros(bsz, K, dtype=th.uint8)
      for i, klen in enumerate(knowledge_counts):
          ck_mask[i, :klen] = 1
      ck_mask = (
          ck_mask != 0
      )  # for pytorch 1.0/1.2 uint8/bool compatibility
      # and the correct labels
      cs_ids = th.LongTensor(bsz).zero_()
  
      if self.use_cuda:
          knowledge_vec = knowledge_vec.cuda()
          ck_mask = ck_mask.cuda()
          cs_ids = cs_ids.cuda()
  
      batch["know_vec"] = knowledge_vec
      batch["ck_mask"] = ck_mask
      batch["cs_ids"] = cs_ids
      batch["use_cs_ids"] = is_training
      batch["knowledge"] = np.array(flattened_knowledge).reshape(N, K)
      return batch
  
imports: from typing import Optional
  from parlai.core.params import ParlaiParser
  from parlai.core.opt import Opt
  from itertools import chain
  from functools import lru_cache
  import torch
  import numpy
  from parlai.utils.torch import padded_tensor
  from parlai.utils.misc import round_sigfigs
  from parlai.agents.transformer.transformer import TransformerGeneratorAgent
  from .modules import EndToEndModel
  from parlai.tasks.wizard_of_wikipedia.agents import TOKEN_KNOWLEDGE
label_desc: Use image node_id 6 for calling the EndToEndAgent obj's underlying member method code with example usage: obj.batchify(obs_batch) and returns: batch