class_docstr: The QueryGenerator is a wrapper around a generator model.
  
  This model can be trained for both dataset classification and search query
  generation.
class_docstr_tok: ["The","QueryGenerator","is","a","wrapper","around","a","generator","model",".","This","model","can","be","trained","for","both","dataset","classification","and","search","query","generation","."]
docstr: Classify input and get retrieval type.
  
  Here, we classify which "type" of retrieval to do for each input batch item.
  
  In the case of "search", we additionally return search queries.
  
  :param input:
      input to classify
  :param num_memories:
      how many memories each example has.
      we override classification if there are no mems for the example.
  :param generated_memories:
      the generated memories from a memory decoder.
  
  :return (retrieval_type, searches):
      retrieval_type: a bsz-length tensor indicating which "type" of retrieval
                      we're doing (see RetrievalType above)
      searches: For batch items classified as search, we return the search queries
                as well.
docstr_tok: ["Classify","input","and","get","retrieval","type",".","Here",",","we","classify","which","``","type","''","of","retrieval","to","do","for","each","input","batch","item",".","In","the","case","of","``","search","''",",","we","additionally","return","search","queries",".",":","param","input",":","input","to","classify",":","param","num_memories",":","how","many","memories","each","example","has",".","we","override","classification","if","there","are","no","mems","for","the","example",".",":","param","generated_memories",":","the","generated","memories","from","a","memory","decoder",".",":","return","(","retrieval_type",",","searches",")",":","retrieval_type",":","a","bsz-length","tensor","indicating","which","``","type","''","of","retrieval","we","'re","doing","(","see","RetrievalType","above",")","searches",":","For","batch","items","classified","as","search",",","we","return","the","search","queries","as","well","."]
code: def classify_retrieval(
      self,
      input: torch.LongTensor,
      num_memories: torch.LongTensor,
      generated_memories: Optional[List[List[str]]],
      skip_search: Optional[torch.BoolTensor],
  ) -> Tuple[torch.LongTensor, List[str]]:
      """
      Classify input and get retrieval type.
  
      Here, we classify which "type" of retrieval to do for each input batch item.
  
      In the case of "search", we additionally return search queries.
  
      :param input:
          input to classify
      :param num_memories:
          how many memories each example has.
          we override classification if there are no mems for the example.
      :param generated_memories:
          the generated memories from a memory decoder.
  
      :return (retrieval_type, searches):
          retrieval_type: a bsz-length tensor indicating which "type" of retrieval
                          we're doing (see RetrievalType above)
          searches: For batch items classified as search, we return the search queries
                    as well.
      """
      self.retrieval_type = torch.LongTensor(input.size(0))
      self.retrieval_type.fill_(0)
      assert self.agent_dict is not None
      texts = [
          self.agent_dict.vec2txt(self.clean_input(i)) for i in input
      ]
      if (
          self.knowledge_access_method
          is KnowledgeAccessMethod.MEMORY_ONLY
      ):
          search_queries = [MEMORY_STRINGS[-1]] * len(texts)
      else:
          search_queries = self._batch_generate(texts)
      logging.debug(f"search queries: {search_queries}")
      logging.verbose(f"Search: {search_queries[0]}")
      searches = []
      if not generated_memories:
          generated_memories = [[] for _ in range(input.size(0))]
      for i, s in enumerate(search_queries):
          if (
              (strip_punc(s) in MEMORY_STRINGS)
              or any(ms in s for ms in MEMORY_STRINGS)
          ) and (
              (num_memories is not None and num_memories[i] > 0)
              or generated_memories[i]
          ):
              self.retrieval_type[i] = RetrievalType.MEMORY.value
          elif strip_punc(s) in NONE_STRINGS + MEMORY_STRINGS:
              self.retrieval_type[i] = RetrievalType.NONE.value
          elif skip_search is not None and skip_search[i]:
              self.retrieval_type[i] = RetrievalType.NONE.value
          else:
              self.retrieval_type[i] = RetrievalType.SEARCH.value
              searches.append(s)
  
      return self.retrieval_type, searches
  
imports: from enum import Enum, auto
  import os
  import string
  import time
  import torch
  import torch.nn
  from typing import List, Tuple, Dict, Optional, Any
  from parlai.agents.rag.retrievers import clean_vec
  from parlai.core.agents import create_agent_from_model_file, create_agent_from_shared
  from parlai.core.build_data import modelzoo_path
  from parlai.core.dict import DictionaryAgent
  from parlai.core.message import Message
  from parlai.core.opt import Opt
  from parlai.core.torch_agent import TorchAgent
  from parlai.tasks.msc.agents import NOPERSONA
  import parlai.utils.logging
label_desc: Use image node_id 2 for calling the QueryGenerator obj's underlying member method code with example usage: obj.classify_retrieval(input, num_memories, generated_memories, skip_search) and returns: self,  searches