class_docstr: A helper class for exporting simple greedy-search models via TorchScript.
  
  Models with extra inputs will need to override to include more variables.
code: def __init__(self, agent: TorchAgent):
      super().__init__()
  
      self.is_bart = isinstance(agent, BartAgent)
      self.device = agent.model.encoder.embeddings.weight.device
      # Dictionary/tokenization setup
      for key, val in self.CAIRAOKE_DICT_PARAMS.items():
          if type(val) == list and len(val) > 0:
              assert (
                  agent.opt.get(key, val[0]) in val
              ), f'The only currently supported values of "{key}" are {", ".join(val)}!'
          else:
              assert (
                  agent.opt.get(key, val) == val
              ), f'The only currently supported value of "{key}" is {val}!'
      orig_dict: DictionaryAgent = agent.dict
      orig_bpe: Gpt2BpeHelper = orig_dict.bpe
      assert all(len(key) == 2 for key in orig_bpe.bpe_ranks.keys())
      if any(
          i
          for key in orig_bpe.bpe_ranks.keys()
          for i in key
          if "\n" in i
      ):
          raise AssertionError(
              "We need to temporarily merge the bpe_ranks dict's keys with a newline "
              "character in order to use it as a TorchScript arg, but at least one of the "
              "dict's keys contains a newline character already!"
          )
      fused_key_bpe_ranks = {
          "\n".join(key): float(val)
          for key, val in orig_bpe.bpe_ranks.items()
      }
      # Cast the values as floats to be able to compare to float('inf') when doing BPE
      # splitting
      self.dict = ScriptableDictionaryAgent(
          null_token=orig_dict.null_token,
          end_token=orig_dict.end_token,
          unk_token=orig_dict.unk_token,
          start_token=orig_dict.start_token,
          freq=orig_dict.freq,
          tok2ind=orig_dict.tok2ind,
          ind2tok=orig_dict.ind2tok,
          bpe_add_prefix_space=agent.opt["bpe_add_prefix_space"],
          bpe_encoder=orig_bpe.encoder,
          bpe_byte_encoder=orig_bpe.byte_encoder,
          fused_key_bpe_ranks=fused_key_bpe_ranks,
          special_tokens=agent._get_special_tokens(),
          subword_bpe_version=(0, 0),
          fused_bpe_codes={},
          subword_bpe_separator="",
      )
  
      # History tracking and start/end tokens
      self.delimiter_tok = agent.history.delimiter_tok
      self.history_size = agent.opt["history_size"]
      if (
          agent.opt.get("history_add_global_end_token", None)
          is not None
      ):
          self.global_end_token = agent.dict[agent.dict.end_token]
      else:
          self.global_end_token = None
      self.text_truncate = (
          agent.opt.get("text_truncate") or agent.opt["truncate"]
      )
      self.text_truncate = (
          self.text_truncate if self.text_truncate >= 0 else None
      )
  
      self.start_idx = agent.model.START_IDX
      self.end_idx = agent.model.END_IDX
      self.null_idx = agent.model.NULL_IDX
      if self.is_bart:
          self.initial_decoder_input = [self.end_idx, self.start_idx]
      else:
          self.initial_decoder_input = [self.start_idx]
  
      agent.model.eval()
  
      # Create versions of the model and decoder that will flatten the incremental
      # state dict, as required by TorchScript
      wrapped_decoder = DecoderIncrStateFlattener(agent.model.decoder)
      wrapped_model = ModelIncrStateFlattener(agent.model)
  
      # Create sample inputs for tracing
      sample_tokens = torch.tensor(
          [[1, 2, 3, 4, 5]], dtype=torch.long, device=self.device
      )
      sample_tokens = sample_tokens.to(self.device)
      encoder_states = agent.model.encoder(sample_tokens)
      initial_generations = self._get_initial_decoder_input(
          sample_tokens
      )
      latent, initial_incr_state = wrapped_decoder(
          initial_generations, encoder_states
      )
      logits = agent.model.output(latent[:, -1:, :])
      _, preds = logits.max(dim=2)
      incr_state = {
          k: torch.clone(v) for k, v in initial_incr_state.items()
      }
      # Copy the initial incremental state, used when tracing the
      # .reorder_decoder_incremental_state() method below, to avoid having it be
      # mutated by the following line
      incr_state = wrapped_model.reorder_decoder_incremental_state(
          incr_state,
          torch.tensor(
              [0], dtype=torch.long, device=sample_tokens.device
          ),
      )
      generations = torch.cat([initial_generations, preds], dim=1)
  
      # Do tracing
      self.encoder = torch.jit.trace(agent.model.encoder, sample_tokens)
      self.decoder_first_pass = torch.jit.trace(
          wrapped_decoder,
          (initial_generations, encoder_states),
          strict=False,
      )
      # We do strict=False to avoid an error when passing a Dict out of
      # decoder.forward()
      self.partially_traced_model = torch.jit.trace_module(
          wrapped_model,
          {
              "output": (latent[:, -1:, :]),
              "reorder_decoder_incremental_state": (
                  initial_incr_state,
                  torch.tensor(
                      [0], dtype=torch.long, device=sample_tokens.device
                  ),
              ),
          },
          strict=False,
      )
      self.decoder_later_pass = torch.jit.trace(
          wrapped_decoder,
          (generations, encoder_states, incr_state),
          strict=False,
      )
  
imports: from collections import defaultdict
  from typing import Dict, List, Optional, Tuple
  import torch.jit
  from parlai.agents.bart.bart import BartAgent
  from parlai.core.dict import DictionaryAgent
  from parlai.core.torch_agent import TorchAgent
  from parlai.utils.bpe import Gpt2BpeHelper, SubwordBPEHelper
  from parlai.torchscript.tokenizer import ScriptableDictionaryAgent
  from torch import nn
label_desc: Use image node_id 1 to create a new TorchScriptGreedySearch object from inherited base classes: nn with example: obj = TorchScriptGreedySearch(agent)