text: #!/usr/bin/env python3
  
  # Copyright (c) Facebook, Inc. and its affiliates.
  # This source code is licensed under the MIT license found in the
  # LICENSE file in the root directory of this source tree.
  
  """
  Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks.
  
  Original Paper: https://arxiv.org/abs/2005.11401
  
  As used in ParlAI: https://arxiv.org/abs/2104.07567
  """
  from abc import ABC, abstractstaticmethod
  import os
  import torch
  import torch.nn
  import torch.cuda
  from typing import Any, Dict, List, Optional, Tuple, Union, Type
  
  from parlai.agents.bart.bart import BartAgent
  import parlai.agents.hugging_face.hugging_face  # noqa: F401
  from parlai.agents.hugging_face.t5 import T5Agent
  from parlai.agents.transformer.polyencoder import PolyencoderAgent
  from parlai.agents.transformer.transformer import TransformerGeneratorAgent
  from parlai.core.build_data import modelzoo_path
  from parlai.core.dict import DictionaryAgent
  from parlai.core.message import Message
  from parlai.core.metrics import AverageMetric, normalize_answer, F1Metric
  from parlai.core.opt import Opt
  from parlai.core.params import ParlaiParser
  from parlai.core.torch_agent import History, Batch, Output
  from parlai.core.torch_generator_agent import PPLMetric, TorchGeneratorAgent, TreeSearch
  from parlai.utils.distributed import sync_parameters
  from parlai.utils.io import PathManager
  import parlai.utils.logging as logging
  from parlai.utils.misc import recursive_getattr
  import parlai.utils.pickle
  from parlai.utils.torch import total_parameters, trainable_parameters, PipelineHelper
  from parlai.utils.typing import TShared
  
  from parlai.agents.rag.args import setup_rag_args
  from parlai.agents.rag.model_types import (
      RagModelInterface,
      RagSequence,
      RagToken,
      RagTurn,
  )
  from parlai.agents.rag.modules import RagModel, T5RagModel
  from parlai.agents.rag.retrievers import Document
  from parlai.utils.fsdp import is_fsdp
  
  
  class BaseGenerationAgentMixin(ABC):
      """
      A Base Generation Agent Mixin.
      """
  
      @abstractstaticmethod
      def build_rag_model(opt: Opt, dictionary: DictionaryAgent) -> RagModel:
          """
          Build and return a RAG Model.
          """
  
  
  class TransformerGeneratorRagAgent(TransformerGeneratorAgent, BaseGenerationAgentMixin):
      @staticmethod
      def build_rag_model(opt: Opt, dictionary: DictionaryAgent) -> RagModel:
          return RagModel(opt, dictionary)
  
  
  class BartRagAgent(BartAgent, BaseGenerationAgentMixin):
      @staticmethod
      def build_rag_model(opt: Opt, dictionary: DictionaryAgent) -> RagModel:
          return RagModel(opt, dictionary)
  
  
  class T5RagAgent(T5Agent, BaseGenerationAgentMixin):
      @staticmethod
      def build_rag_model(opt: Opt, dictionary: DictionaryAgent) -> T5RagModel:
          return T5RagModel(opt, dictionary)
  
      def _generate(
          self,
          batch: Batch,
          beam_size: int,
          max_ts: int,
          prefix_tokens: Optional[torch.LongTensor] = None,
      ) -> Tuple[
          List[Tuple[torch.LongTensor, torch.Tensor, Optional[Dict]]], List[TreeSearch]
      ]:
          """
          Override since T5 needs to call TGA generate.
          """
          return TorchGeneratorAgent._generate(
              self, batch, beam_size, max_ts, prefix_tokens
          )
  
  
  GENERATION_AGENTS = {
      'transformer/generator': TransformerGeneratorRagAgent,
      'bart': BartRagAgent,
      't5': T5RagAgent,
  }
  
  RAG_MODELS = {'sequence': RagSequence, 'token': RagToken, 'turn': RagTurn}
  
  
  class RagAgent(TransformerGeneratorRagAgent, BartRagAgent, T5RagAgent):
      """
      RagAgent.
  
      The RAG Agent interacts with the RAG model mostly via it's RAG Model interface.
      """
  
      _generation_agent: Union[
          Type[T5RagAgent], Type[BartRagAgent], Type[TransformerGeneratorRagAgent]
      ]
      _rag_model_interface: RagModelInterface
  
      @classmethod
      def add_cmdline_args(
          cls, parser: ParlaiParser, partial_opt: Optional[Opt] = None
      ) -> ParlaiParser:
          """
          Add RAG Args.
          """
          PolyencoderAgent.add_cmdline_args(parser, partial_opt=None)
          TransformerGeneratorRagAgent.add_cmdline_args(parser, partial_opt)
          parser = setup_rag_args(parser)
          RagTurn.add_cmdline_args(parser, partial_opt)
          if partial_opt and partial_opt.get('generation_model') == 'bart':
              BartRagAgent.add_cmdline_args(parser, partial_opt=partial_opt)
          T5RagAgent.add_cmdline_args(parser, partial_opt=partial_opt)
          # BART Agent sets these to True; doesn't let you set anything else
          parser.set_defaults(fp16=False, force_fp16_tokens=False)
          return parser
  
      @property
      def generation_model(self) -> str:
          return self._generation_model
  
      @property
      def model_api(self) -> RagModel:
          if hasattr(self.model, 'module') and not is_fsdp(self.model):
              return self.model.module
          else:
              return self.model
  
      @generation_model.setter
      def generation_model(self, model: str):
          self._generation_model = model
          self._generation_agent = GENERATION_AGENTS[model]
  
      @property
      def rag_model_type(self) -> str:
          return self._rag_model_type
  
      @rag_model_type.setter
      def rag_model_type(self, model: str):
          self._rag_model_type = model
          self._rag_model_interface = RAG_MODELS[model](self.opt, self.NULL_IDX)
  
      @property
      def retriever_query(self) -> str:
          return self._retriever_query
  
      @retriever_query.setter
      def retriever_query(self, query: str):
          self._retriever_query = query
          if query == 'one_turn':
              self._query_key = 'text'
          elif query == 'full_history':
              self._query_key = 'full_text'
  
      def __init__(self, opt: Opt, shared: TShared = None):
          self.opt = opt
          if torch.cuda.is_available():
              torch.cuda.set_device('cuda:0')
          self.generation_model = opt['generation_model']
          self.regret = opt['regret']
          self.regret_intermediate_maxlen = opt['regret_intermediate_maxlen']
          self.retriever_query = opt['rag_retriever_query']
  
          # Super call
          self._generation_agent.__init__(self, opt, shared)  # type: ignore
          self.rag_model_type = opt['rag_model_type']
  
          if not shared and self.regret:
              self.regret_model = self.build_regret_model()
          elif shared:
              self.regret_model = shared.get('regret_model')
  
      def build_regret_model(self) -> RagModel:
          """
          Build and return regret RagModel.
          """
          model_file = modelzoo_path(self.opt['datapath'], self.opt['regret_model_file'])
          if model_file:
              assert os.path.exists(
                  model_file
              ), f'specify correct path for --regret-model-file (currently {model_file})'
              regret_opt = Opt.load(f'{model_file}.opt')
              regret_opt['n_docs'] = self.opt['n_docs']  # Urgent that this is the same
              # add keys that were not in this model when originally trained
              regret_opt.update(
                  {k: v for k, v in self.opt.items() if k not in regret_opt}
              )
              retriever_shared = None
              if all(
                  [
                      regret_opt[k] == self.opt[k]
                      for k in [
                          'rag_retriever_type',
                          'path_to_index',
                          'path_to_dpr_passages',
                      ]
                  ]
              ):
                  logging.warning('Sharing retrievers between model and regret model!')
                  retriever_shared = self.model.retriever.share()
              elif self.opt['regret_override_index']:
                  # Sharing Index Path & Passages only; not the full retriever
                  logging.warning('Overriding initial ReGReT model index')
                  regret_opt['path_to_index'] = self.opt['path_to_index']
                  regret_opt['path_to_dpr_passages'] = self.opt['path_to_dpr_passages']
  
              if self.opt['regret_dict_file']:
                  regret_opt['dict_file'] = self.opt['regret_dict_file']
  
              regret_dict = self.dictionary_class()(regret_opt)
              model = RagModel(regret_opt, regret_dict, retriever_shared=retriever_shared)
              with PathManager.open(model_file, 'rb') as f:
                  states = torch.load(
                      f,
                      map_location=lambda cpu, _: cpu,
                      pickle_module=parlai.utils.pickle,
                  )
              assert 'model' in states
              model.load_state_dict(states['model'])
              if self.model_parallel:
                  ph = PipelineHelper()
                  ph.check_compatibility(self.opt)
                  model = ph.make_parallel(model)
              elif self.use_cuda:
                  model.cuda()
              if self.fp16:
                  model = model.half()
  
              sync_parameters(model)
              train_params = trainable_parameters(model)
              total_params = total_parameters(model)
              logging.info(
                  f"Total regret parameters: {total_params:,d} ({train_params:,d} trainable)"
              )
          else:
              model = self.model
  
          return model
  
      def share(self):
          shared = super().share()
          if self.regret:
              shared['regret_model'] = self.regret_model
          return shared
  
      ########################################
      # BART-Specific Overrides              #
      ########################################
  
      def _convert_model(self, opt: Opt) -> Opt:
          """
          Override BartAgent._convert_model to use RagConversionScript.
          """
          return self._generation_agent._convert_model(self, opt)  # type: ignore
  
      def _initialize_bart(self, opt: Opt) -> Opt:
          return self._generation_agent._initialize_bart(self, opt)  # type: ignore
  
      ########################################
      # TorchGeneratorAgent Overrides        #
      ########################################
  
      ###### 0. Act/Observe ######
      def observe(self, observation: Union[Dict, Message]) -> Message:
          """
          Overrides TA.observe to tokenize the retriever query.
          """
          observation = Message(observation)
          observation = self._generation_agent.observe(self, observation)
          if observation.is_padding():
              return observation
          if 'query_vec' not in observation and self._query_key in observation:
              self._set_query_vec(observation)
          if 'input_turn_cnt_vec' not in observation and self._query_key in observation:
              self._set_input_turn_cnt_vec(observation)
          return observation
  
      def eval_step(self, batch: Batch) -> Optional[Output]:
          output = super().eval_step(batch)
          if output is None or not hasattr(self.model, 'retriever'):
              return output
          assert isinstance(self.model, RagModel)
          if hasattr(self.model.retriever, 'top_docs'):
              output.top_docs = self.model.retriever.top_docs  # type: ignore
          return output
  
      ###### 1. Model Inputs ######
  
      def _model_input(
          self, batch: Batch
      ) -> Tuple[torch.LongTensor, torch.LongTensor, torch.LongTensor, torch.LongTensor]:
          """
          Get input for RagModel.
  
          :param batch:
              batch to process
  
          :return (text_vec, text_lengths, query_vec, input_turn_cnt_vec):
              text_vec - tokenized batch
              text_lengths - length of each item in the batch
              query_vec - tokenized vectors for retrieval mechanism
              input_turn_cnt_vec - count of input turns for each batch item
          """
          return (
              batch.text_vec,
              batch.text_vec.ne(self.NULL_IDX).sum(1),
              batch.query_vec,
              batch.input_turn_cnt_vec,
          )
  
      def _encoder_input(
          self, batch: Batch
      ) -> Tuple[
          torch.LongTensor,
          Optional[torch.LongTensor],
          Optional[torch.LongTensor],
          Optional[torch.LongTensor],
      ]:
          """
          Called directly when generating.
  
          Some RAG Model Types require different encoder inputs.
  
          :param batch:
              batch to process
  
          :return (text_vec, text_lengths, query_vec, input_turn_cnt_vec):
              text_vec - tokenized batch
              text_lengths - length of each item in the batch
              query_vec - tokenized vectors for retrieval mechanism
              input_turn_cnt_vec - count of input turns for each batch item
          """
          if hasattr(self._rag_model_interface, 'get_generation_input'):
              return self._rag_model_interface.get_generation_input(batch)  # type: ignore
          return self._model_input(batch)
  
      ##### 2. Standard TGA Function Overrides #####
      def build_model(self) -> RagModel:
          """
          Build and return RagModel.
          """
          model = self._generation_agent.build_rag_model(self.opt, self.dict)
          if self.opt['embedding_type'] != 'random':
              self._copy_embeddings(
                  model.encoder.embeddings.weight, self.opt['embedding_type']
              )
          return model
  
      def _resize_token_embeddings(self, state_dict, msg=None):
          """
          Resize the token embeddings when are adding extra special tokens.
  
          Modify TGA._resize_token_embeddings to access correct modules within RAG.
          """
          # map extra special tokens carefully
          new_size = self.model.embeddings.weight.size()[0]
          orig_size = state_dict['embeddings.weight'].size()[0]
          logging.info(f'Resizing token embeddings from {orig_size} to {new_size}')
          if new_size <= orig_size:
              # new size should be greater than original size,
              # as we are adding special tokens
              raise RuntimeError(msg)
  
          for emb_weights in [
              'embeddings.weight',
              'seq2seq_encoder.embeddings.weight',
              'seq2seq_decoder.embeddings.weight',
          ]:
              # get new_embs
              old_embs = state_dict[emb_weights]
              new_embs = recursive_getattr(self.model, emb_weights).to(old_embs.device)
              # copy over old weights
              new_embs.data[:orig_size, :] = old_embs.data[:orig_size, :]
              # reset in state dict
              state_dict[emb_weights] = new_embs
  
          return state_dict
  
      def build_dictionary(self) -> DictionaryAgent:
          """
          Build and return dictionary.
          """
          return self._generation_agent.build_dictionary(self)
  
      @staticmethod
      def update_state_dict(
          opt: Opt, state_dict: Dict[str, torch.Tensor], model: torch.nn.Module
      ):
          """
          Update the given state dict to be RAG-ified.
  
          :param opt:
              options
          :param state_dict:
              weights to load
          :param model:
              underlying model that will load the state_dict
  
          :return updated_state_dict:
              return state_dict with appropriate keys/values
          """
          # 1. Substitute all "encoder" and "decoder" keys with "seq2seq_encoder" and "seq2seq_decoder"
          if not [k for k in state_dict if k.startswith('seq2seq')]:
              for k in list(state_dict.keys()):
                  if k.startswith('encoder') or k.startswith('decoder'):
                      weights = state_dict.pop(k)
                      state_dict[f'seq2seq_{k}'] = weights
          # 2. Retriever state
          if not [k for k in state_dict if 'retriever' in k]:
              retriever_state = {
                  f"retriever.{k}": v
                  for k, v in model.retriever.state_dict().items()  # type: ignore
              }
              state_dict.update(retriever_state)
          # 3. Handle n_positional difference
          if opt.get('n_extra_positions', 0) > 0:
              key = 'seq2seq_encoder.position_embeddings.weight'
              init_weight = (
                  model.seq2seq_encoder.position_embeddings.weight  # type: ignore
              )
              if state_dict[key].size(0) < opt['n_positions'] + opt['n_extra_positions']:
                  # Make sure we're not adding more positions to a model trained
                  # with extra positions
                  state_dict[key] = torch.cat(
                      [
                          state_dict[key].to(init_weight),  # type: ignore
                          init_weight[-opt['n_extra_positions'] :, :],  # type: ignore
                      ],
                      dim=0,
                  )
          return state_dict
  
      def _should_override_dpr_model_weights(self, opt: Opt):
          """
          Determine if we need to override the DPR Model weights.
  
          Under certain circumstances, one may wish to specify a different `--dpr-model-
          file` for a pre-trained, RAG model. Thus, we additionally check to make sure
          that the loaded DPR model weights are not overwritten by the state loading.
          """
          override_dpr = False
          overrides = opt.get('override', {})
          if overrides.get('dpr_model_file') and os.path.exists(
              overrides['dpr_model_file']
          ):
              override_dpr = True
              logging.warning(
                  f"Overriding DPR Model with {modelzoo_path(opt['datapath'], opt['dpr_model_file'])}"
              )
          return override_dpr
  
      def load_state_dict(self, state_dict: Dict[str, torch.Tensor]):
          """
          Potentially update state dict with relevant RAG components.
  
          Useful when initializing from a normal seq2seq model.
          """
          try:
              if self._should_override_dpr_model_weights(self.opt):
                  state_dict.update(
                      {
                          f"retriever.{k}": v
                          for k, v in self.model.retriever.state_dict().items()  # type: ignore
                      }
                  )
              self.model.load_state_dict(state_dict)
          except RuntimeError as msg:
              state_dict = self.update_state_dict(self.opt, state_dict, self.model)
              msg_ = str(msg)
              if 'size mismatch' in msg_ and 'embedding' in msg_:
                  if hasattr(self, 'special_toks') and len(self.special_toks) > 0:
                      state_dict = self._resize_token_embeddings(state_dict, msg_)
                      self.resized_embeddings = True  # make note that we resized here
                  else:
                      raise RuntimeError(
                          f'{msg_}\n'
                          '-----------------\n'
                          'Could not load the model due to a size mismatch in the '
                          'embeddings. A common reason for this is trying to load '
                          'a model trained with fp16 but loaded without fp16. Try '
                          'adding --fp16 true or --force-fp16-tokens true.'
                      )
              self.model.load_state_dict(state_dict)
  
      def batchify(self, obs_batch: List[Message], sort: bool = False) -> Batch:
          """
          Override TA.batchify to incorporate query and input turn vecs.
          """
          assert not sort
          if len(obs_batch) == 0:
              return Batch(batchsize=0)
  
          valid_exs = [ex for ex in obs_batch if self.is_valid(ex)]
  
          if len(valid_exs) == 0:
              return Batch(batchsize=0)
  
          batch = self._generation_agent.batchify(self, obs_batch, sort)
  
          if any(ex.get('query_vec') is not None for ex in valid_exs):
              _qs = []
              for ex in valid_exs:
                  q = ex.get('query_vec', self.EMPTY)
                  if type(q) is list and type(q[0]) is list:
                      # handling input turns
                      _qs += q
                  else:
                      _qs.append(q)
              self.set_batch_query(batch, _qs)
  
          batch.input_turn_cnt_vec = None
          if any(ex.get('input_turn_cnt_vec') is not None for ex in valid_exs):
              batch.input_turn_cnt_vec = torch.cat(
                  [
                      ex.get('input_turn_cnt_vec', torch.LongTensor([1]))
                      for ex in valid_exs
                  ],
                  dim=0,
              )
          return batch
  
      def vectorize(self, *args, **kwargs):
          """
          Override TA.vectorize to ensure proper super class handles vectorization.
          """
          return self._generation_agent.vectorize(self, *args, **kwargs)
  
      def _set_text_vec(
          self, obs: Message, history: History, truncate: Optional[int]
      ) -> Message:
          """
          Override TA._set_text_vec to ensure proper gen agent handles _set_text_vec.
          """
          return self._generation_agent._set_text_vec(self, obs, history, truncate)
  
      ##### 3. Custom Batch handling #####
  
      def set_batch_query(self, batch: Batch, queries: List[torch.LongTensor]) -> Batch:
          """
          Put the queries in the batch.
  
          :param batch:
              batch to put queries in
          :param queries:
              list of query tokens, presumably
  
          :return batch:
              return the batch, with queries.
          """
          qs, q_lens = self._pad_tensor(queries)
          batch.query_vec = qs
          batch.query_lengths = torch.LongTensor(q_lens)
          return batch
  
      def _set_query_vec(self, observation: Message) -> Message:
          """
          Tokenize the query for retrieval.
  
          :param observation:
              observation with input text.
  
          :return observation:
              return observation with query vec.
          """
          query_str = observation[self._query_key]
          if hasattr(self.model, 'module'):
              observation['query_vec'] = self.model.module.tokenize_query(query_str)
          else:
              observation['query_vec'] = self.model.tokenize_query(query_str)
          return observation
  
      def _set_input_turn_cnt_vec(self, observation: Message) -> Message:
          """
          Compute the number of turns of context for the observation query.
  
          :param observation:
              observation with input text.
  
          :return observation:
              return observation with input turn vec.
          """
          observation = self._rag_model_interface.set_input_turn_cnt_vec(
              observation, self.model, observation[self._query_key]
          )
          return observation
  
      ##### 4. Decoder Input Functions #####
  
      def _get_initial_decoder_input(
          self, bsz: int, beam_size: int, dev: torch.device
      ) -> torch.LongTensor:
          """
          Override TGA._get_initial_decoder_input to seed decoder with proper amount of
          inputs.
  
          This is called during generation, so additional inputs are necessary only for
          RAG Token (which retrieves during the generation step).
          """
          dec_input = self._generation_agent._get_initial_decoder_input(
              self, bsz, beam_size, dev
          )
          rag_dec_input = self._rag_model_interface.get_initial_decoder_input(dec_input)
          return rag_dec_input
  
      def _get_next_decoder_input(
          self,
          prev_input: torch.LongTensor,
          selection: torch.LongTensor,
          incr_state_inds: torch.LongTensor,
      ) -> torch.LongTensor:
          """
          Override TGA._get_next_decoder_input to repeat decoder input appropriately.
          """
          if hasattr(self._rag_model_interface, 'get_next_decoder_input'):
              dec_input = self._rag_model_interface.get_next_decoder_input(  # type: ignore
                  prev_input, selection, incr_state_inds
              )
          else:
              dec_input = self._generation_agent._get_next_decoder_input(
                  self, prev_input, selection, incr_state_inds
              )
          return dec_input
  
      ####################################
      # Model Generation Overrides       #
      ####################################
  
      ##### Standard TGA Generation Overrides #####
  
      def _get_context(self, batch: Batch, batch_idx: int) -> torch.LongTensor:
          """
          Override TGA._get_context for rag-sequence/turn models.
  
          Reason: the batchsize is artificially higher (n_docs * batchsize)
          """
          if hasattr(self._rag_model_interface, 'get_ctxt_index'):
              batch_idx = self._rag_model_interface.get_ctxt_index(  # type: ignore
                  batch, batch_idx
              )
          return self._generation_agent._get_context(self, batch, batch_idx)
  
      def _generate(
          self,
          batch: Batch,
          beam_size: int,
          max_ts: int,
          prefix_tokens: Optional[torch.LongTensor] = None,
      ) -> Tuple[
          List[Tuple[torch.LongTensor, torch.Tensor, Optional[Dict]]], List[TreeSearch]
      ]:
          """
          Override TGA._generate to potentially call ReGReT.
  
          TGA._generate is implemented in _rag_generate
          """
          if self.regret:
              beam_preds_scores, _ = self._regret_generate(
                  batch, beam_size, self.regret_intermediate_maxlen, prefix_tokens
              )
              preds, _, _ = zip(*beam_preds_scores)
              new_batch = self._regret_rebatchify(batch, preds)  # type: ignore
              gen_outs = self._rag_generate(new_batch, beam_size, max_ts, prefix_tokens)
          else:
              gen_outs = self._rag_generate(batch, beam_size, max_ts, prefix_tokens)
  
          return gen_outs
  
      def _rerank_beams(
          self,
          batch: Batch,
          n_best_beam_preds_scores: List[
              List[Tuple[torch.LongTensor, torch.Tensor, Optional[Dict]]]
          ],
      ) -> List[List[Tuple[torch.LongTensor, torch.Tensor, Optional[Dict]]]]:
          """
          Optional rerank beams, according to RAG Model type.
  
          :param batch:
              current batch
          :param n_best_beam_preds_scores:
              bsz-length list of Tuples of predictions and scores
  
          :return List((pred, score)):
              return a re-ranked version of the n_best_beam_preds_scores
          """
          return self._rag_model_interface.rerank_beams(
              self.model, batch, n_best_beam_preds_scores
          )
  
      ##### RAG-Specific Generation Functions #####
  
      def _rag_generate(
          self,
          batch: Batch,
          beam_size: int,
          max_ts: int,
          prefix_tokens: Optional[torch.LongTensor] = None,
      ) -> Tuple[
          List[Tuple[torch.LongTensor, torch.Tensor, Optional[Dict]]], List[TreeSearch]
      ]:
          """
          Separate from _generate to handle regret.
          """
          batch = self._rag_model_interface.augment_batch_for_generation(
              batch, self.model
          )
          return self._generation_agent._generate(
              self, batch, beam_size, max_ts, prefix_tokens
          )
  
      def _regret_generate(
          self,
          batch: Batch,
          beam_size: int,
          max_ts: int,
          prefix_tokens: Optional[torch.LongTensor] = None,
      ):
          """
          Swap generation model prior to full generation.
          """
          model, model_type = self.model, self.rag_model_type
          self.model = self.regret_model
          self.rag_model_type = self.regret_model.rag_model_type
  
          outs = self._rag_generate(batch, beam_size, max_ts, prefix_tokens)
  
          self.model = model
          self.rag_model_type = model_type
          return outs
  
      def _regret_rebatchify(
          self, batch: Batch, pred_vecs: List[torch.LongTensor]
      ) -> Batch:
          """
          Create a new batch for the model.
  
          Given a model's prediction/generation, create a new batch with this text substituted in as the query.
  
          :param batch:
              prior model batch
          :param pred_vecs:
              predicted, tokenized text from initial batch
  
          :return batch:
              return a new batch for the model.
          """
          new_queries: List[torch.LongTensor] = []
          query_lens: torch.LongTensor = batch.query_lengths
          query_vec: torch.LongTensor = batch.query_vec
  
          for i in range(batch.batchsize):
              vec_i = pred_vecs[i]
              txt_i = self._v2t(vec_i)
              query_i = torch.LongTensor(self.model.tokenize_query(txt_i)).to(query_vec)
              if self.retriever_query == 'one_turn':
                  new_queries.append(query_i)
              else:
                  query_i = torch.cat([query_vec[i][: query_lens[i]], query_i], dim=0)
                  new_queries.append(query_i)  # type: ignore
  
          self.set_batch_query(batch, new_queries)
  
          if (
              batch.input_turn_cnt_vec is not None
              and self.retriever_query == 'full_history'
          ):
              batch.input_turn_cnt_vec += 1
  
          return batch
  
      ########################################
      # Computing Loss / Recording Metrics   #
      ########################################
  
      def _record_retrieval_metrics(self, batch: Batch, encoder_state: Tuple[Any, ...]):
          """
          Compute retrieval metrics, given retrieved documents.
  
          Only works when `--debug` is set.
  
          If there is knowledge in the Batch, we compute the following metrics:
          A) Doc Level:
          1. recall @ 1 --> is the correct document the first document?
          2. recall @ N --> is the correct document in the first N docs?
  
          B) Passage Level:
          1. recall @ 1 --> is the correct passage in the first document?
          2. recall @ N --> is the correct passage in the first N docs?
  
          Only works in debug mode.
  
          :param batch:
              training/eval batch
          :param encoder_state:
              encoder states from RagEncoder
          """
          if batch.valid_indices is None or batch.observations is None:
              return
          docs: List[List[Document]] = []
          _, _, input_turns_cnt, docs, _ = encoder_state
          if input_turns_cnt is not None:
              new_docs = []
              offset = 0
              for it in input_turns_cnt:
                  docs_it = [dd for d in docs[offset : offset + it] for dd in d]
                  new_docs.append(docs_it)
                  offset += it
              docs = new_docs
          title_key = self.opt['gold_knowledge_title_key']
          passage_key = self.opt['gold_knowledge_passage_key']
          batchsize = len(batch.valid_indices)
          n_docs = self.opt['n_docs']
          metrics = {
              k: [0] * batchsize
              for k in [
                  'doc_r@1',
                  f'doc_r@{n_docs}',
                  'passage_r@1',
                  f'passage_r@{n_docs}',
                  'title@1_f1',
                  'passage@1_f1',
              ]
          }
          for i in range(batchsize):
              ex = batch.observations[i]
              label_title = normalize_answer(ex.get(title_key, ''))
              label_passage = normalize_answer(ex.get(passage_key, ''))
  
              for rank, doc in enumerate(docs[i]):
                  model_title = normalize_answer(doc.get_title())
                  model_passage = normalize_answer(doc.get_text())
  
                  title_exact_match = model_title == label_title
                  passage_match = (
                      model_passage in label_passage or label_passage in model_passage
                  )
  
                  if rank == 0:
                      metrics['doc_r@1'][i] = int(title_exact_match)
                      metrics['passage_r@1'][i] = int(passage_match)
                      metrics['title@1_f1'][i] = F1Metric.compute(
                          guess=model_title, answers=[label_title]
                      ).value()
                      metrics['passage@1_f1'][i] = F1Metric.compute(
                          guess=model_passage, answers=[label_passage]
                      ).value()
                  metrics[f'doc_r@{n_docs}'][i] = int(
                      metrics[f'doc_r@{n_docs}'][i] or title_exact_match
                  )
                  metrics[f'passage_r@{n_docs}'][i] = int(
                      metrics[f'passage_r@{n_docs}'][i] or passage_match
                  )
  
          for m in metrics:
              self.record_local_metric(m, AverageMetric.many(metrics[m], [1] * batchsize))
  
      def get_model_output(self, batch: Batch) -> Tuple[Any, ...]:
          """
          Return model output.
  
          :param batch:
              batch to process
  
          :return model_output:
              return output from model
          """
          if not self.regret:
              model_output = self.model(
                  *self._model_input(batch), ys=batch.label_vec
              )  # type: ignore
              scores, preds, enc_state, *_ = model_output
          else:
              with torch.no_grad():
                  beam_preds_scores, beams = self._regret_generate(
                      batch, self.beam_size, self.regret_intermediate_maxlen
                  )
              regret_preds, _, _ = zip(*beam_preds_scores)
              new_batch = self._regret_rebatchify(batch, regret_preds)  # type: ignore
              regret_model_output = self.model(
                  *self._model_input(new_batch), ys=batch.label_vec
              )  # type: ignore
              regret_scores, preds, enc_state = regret_model_output
              scores = regret_scores
  
          return (scores, preds, enc_state)
  
      def compute_loss(
          self, batch: Batch, return_output: bool = False
      ) -> Union[torch.Tensor, Tuple[torch.Tensor, Any]]:
          """
          Override standard TGA.compute_loss to call relevant RAG Model Interface.
          """
          if batch.label_vec is None:
              raise ValueError('Cannot compute loss without a label.')
  
          model_output = self.get_model_output(batch)
          scores, preds, enc_state, *_ = model_output
  
          self._record_retrieval_metrics(batch, enc_state)
          (
              loss,
              metric_loss,
              metric_correct,
              metric_target_tokens,
          ) = self._rag_model_interface.compute_loss(
              self.criterion, scores, preds, enc_state, batch.label_vec
          )
  
          self.record_local_metric(
              'loss', AverageMetric.many(metric_loss, metric_target_tokens)
          )
          self.record_local_metric(
              'ppl', PPLMetric.many(metric_loss, metric_target_tokens)
          )
          self.record_local_metric(
              'token_acc', AverageMetric.many(metric_correct, metric_target_tokens)
          )
          self.record_local_metric(
              'token_em',
              AverageMetric.many(
                  [x == y for x, y in zip(metric_correct, metric_target_tokens)]
              ),
          )
  
          if return_output:
              return loss, model_output
          else:
              return loss
  
      def _construct_token_losses(self, labels, model_output):
          """
          Override default `_construct_token_loss` so that sequence model does not throw
          an error while being evaled with the `--verbose` flag.
  
          TODO: implement this
          """
          return None
  