text: # 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.
  
  import time
  from functools import partial
  from multiprocessing.pool import ThreadPool as Pool
  
  from . import DEFAULT_EOS, GET, SEND
  
  
  class Agent(object):
      "an agent needs to follow this pattern"
  
      def __init__(self, *args, **kwargs):
          pass
  
      def init_states(self, *args, **kwargs):
          raise NotImplementedError
  
      def update_states(self, states, new_state):
          raise NotImplementedError
  
      def finish_eval(self, states, new_state):
          raise NotImplementedError
  
      def policy(self, state):
          raise NotImplementedError
  
      def reset(self):
          raise NotImplementedError
  
      def decode(self, session, low=0, high=100000, num_thread=10):
          corpus_info = session.corpus_info()
          high = min(corpus_info["num_sentences"] - 1, high)
          if low >= high:
              return
  
          t0 = time.time()
          if num_thread > 1:
              with Pool(10) as p:
                  p.map(
                      partial(self._decode_one, session),
                      [sent_id for sent_id in range(low, high + 1)],
                  )
          else:
              for sent_id in range(low, high + 1):
                  self._decode_one(session, sent_id)
  
          print(f"Finished {low} to {high} in {time.time() - t0}s")
  
      def _decode_one(self, session, sent_id):
          action = {}
          self.reset()
          states = self.init_states()
          while action.get("value", None) != DEFAULT_EOS:
              # take an action
              action = self.policy(states)
  
              if action["key"] == GET:
                  new_states = session.get_src(sent_id, action["value"])
                  states = self.update_states(states, new_states)
  
              elif action["key"] == SEND:
                  session.send_hypo(sent_id, action["value"])
          print(" ".join(states["tokens"]["tgt"]))
  