code: def drm_get_batch_context(self, batch: Batch, orig_fun: Callable):
      """
      Override the get context method for the DRM.
  
      This override removes the knowledge response from the set of tokens being blocked in
      the dialogue response.
      """
      ctxts = orig_fun(batch)
      knowledge_start_id = self.dict.txt2vec(TOKEN_KNOWLEDGE)
      knowledge_end_id = self.dict.txt2vec(TOKEN_END_KNOWLEDGE)
  
      def mask_ctxttensor_between_sublists(
          ctxts: torch.Tensor, sub1: List[int], sub2: List[int]
      ) -> torch.Tensor:
          """
          Generate a mask that masks out the context between sub1 and sub2.
          """
          mask = []
          for ctxt in ctxts:
              mask_idxs = []
              should_copy = False
              idx_pointer = 0
              id_to_match = sub1
              for j, token in enumerate(ctxt.cpu().numpy()):
                  if token == id_to_match[idx_pointer]:
                      idx_pointer += 1
                      if idx_pointer == 1 and id_to_match == sub1:
                          mask_idxs.append([j])
                      elif idx_pointer >= len(id_to_match):
                          should_copy = id_to_match == sub1
                          idx_pointer = 0
                          id_to_match = (
                              sub2 if (id_to_match == sub1) else sub1
                          )
                          mask_idxs[-1].append(j)
                      else:
                          mask_idxs[-1].append(j)
                  elif should_copy:
                      assert isinstance(mask_idxs[-1], list)
                      mask_idxs[-1].append(j)
                  elif idx_pointer > 0:
                      idx_pointer = 0
                      del mask_idxs[-1]
              mask.append(
                  [
                      0
                      if idx in [i for sl in mask_idxs for i in sl]
                      else 1
                      for idx in range(len(ctxt))
                  ]
              )
          return torch.LongTensor(mask).to(ctxts.device)
  
      ctxts *= mask_ctxttensor_between_sublists(
          ctxts, knowledge_start_id, knowledge_end_id
      )
      return ctxts
  
file_path: ParlAI/projects/seeker/utils.py
imports: from collections import Counter
  from enum import Enum
  import torch
  from typing import Callable, List, Tuple
  from parlai.core.torch_agent import Batch
  from parlai.tasks.wizard_of_wikipedia.agents import TOKEN_KNOWLEDGE, TOKEN_END_KNOWLEDGE
  import parlai.utils.logging
label_desc: Use image node_id 1 for calling a global function with example usage: drm_get_batch_context(batch, orig_fun) and returns: ctxts, torch