text: import numpy as np
  import glob
  import os
  
  from multiprocessing import connection
  from threading import Thread
  from time import time
  
  from config import Config
  import torch
  from torch.autograd import Variable as V
  import torch.multiprocessing as mp
  from torch.multiprocessing import Pipe
  import torch.nn.utils.rnn as rnn
  
  mp = mp.get_context('forkserver')
  
  '''try:
      set_start_method('spawn')
  except RuntimeError:
      pass'''
  
  class CacheDicts:
      def __init__(self, process_num):
          self.dicts = [{} for _ in range(process_num)]
  
      def add(self, keys, values, id):
          for key, value in zip(keys[id], values[id]):
              if key not in self.dicts[id]:
                  self.dicts[id][key] = value
  
      def update(self, keys, id):
          dict = self.dicts[id]
          new_dict = {}
          for key in keys[id]:
              new_dict[key] = dict[key]
          self.dicts[id] = new_dict
  
  class APIClient:
      def __init__(self, conn):
          self.connection = conn
  
      def predict(self, x):
          self.connection.send(x)
          return self.connection.recv()
  
  class APIServer:
      def __init__(self, config: Config):
          """
          :param config:
          """
          self.config = config
          self.connections = []
          self.model = None
          self.latest_file = None
          self.cache_dicts = CacheDicts(config.multi_process_num)
  
      def get_api_client(self):
          me, you = Pipe()
          self.connections.append(me)
          return APIClient(you)
  
      def start_serve(self):
          self.load_model()
          self.prediction_worker = Thread(target=self.prediction_worker, name="prediction_worker")
          self.prediction_worker.daemon = True
          self.prediction_worker.start()
  
      def load_model(self):
          from lstm import Model, initialize
          self.model = Model(self.config)
          if self.config.use_cuda:
              self.model = self.model.cuda()
          list_of_files = glob.glob(self.config.model_path + '/*')
          latest_file = None
          if list_of_files:
              latest_file = max(list_of_files, key=os.path.getctime)
          model_ckpt = latest_file if not list_of_files else None
          print(model_ckpt)
          if model_ckpt is not None:
              self.model.load_state_dict(torch.load(model_ckpt)['state_dict'])
          else:
              self.model.apply(initialize)
  
      # update model whenever a ckpt file is created
      def update_model(self):
  
          try:
              list_of_files = glob.glob(self.config.model_path + '/*')
              if list_of_files:
                  latest_file = max(list_of_files, key=os.path.getctime)
                  if latest_file != self.latest_file:
                      # TODO: if restore interferes with other processes, modify this
                      try:
                          self.model.load_state_dict(torch.load(latest_file)['state_dict'])
                          print(latest_file)
                          print('loaded')
                          self.latest_file = latest_file
                      except:
                          return
                  for file in list_of_files:
                      if int(file.split('.')[-2].split('-')[-1]) % self.config.saving_period != 0\
                              and file != latest_file:
                          try:
                              os.remove(file)
                          except OSError:
                              pass
          except:
              pass
  
      def prediction_worker(self):
          with torch.no_grad():
              while True:
                  self.update_model()
                  ready_conns = connection.wait(self.connections, timeout=0.00001)
                  if not ready_conns:
                      continue
                  action, lengths, cache_tmp, cur_cs, sizes, ids, ids2, rootss = [], [], [], [], [], [], [], []
                  #TODO: think about the boundary cases
                  for conn in ready_conns:
                      a, c, cur_c, id, roots = conn.recv()
                      self.cache_dicts.update(cur_c, id)
                      action += a
                      cache_tmp += c
                      ids += [id]
                      ids2 += [id]*len(a)
                      sizes.append(len(a))
                      rootss += roots
                  # cache = (bs, depth, hid_dim)
                  original_action = action
                  original_cache = cache_tmp
                  cache = []
                  none_id = []
                  for i in range(len(cache_tmp)):
                      if cache_tmp[i] is not None:
                          cache.append(self.cache_dicts.dicts[ids2[i]][cache_tmp[i]])
                          none_id.append(True)
                      else:
                          if self.config.arch == 'Transformer':
                              cache.append(torch.zeros([1, self.config.hidden_dim, self.config.depth, 2], device=self.config.device))
                          elif self.config.arch == 'sru':
                              cache.append(torch.zeros([1, self.config.hidden_dim, self.config.depth], device=self.config.device))
                          none_id.append(False)
                  lengths = [len(c) if n else 0 for c, n in zip(cache, none_id)]
                  action = [[a] for a in action]
                  #TODO: timing signal to be implemented
                  action = torch.tensor(action, device=self.config.device)
                  if self.config.arch == 'Transformer':
                      #TODO: implement individual cache method
                      # padding for decoding unnecessary
                      # cache is inverted in length dimension
                      #cache = [(i, c) for i, c in enumerate(cache)]
                      #this line doesn't work for non-individual method
                      cache = sorted(enumerate(cache), key=lambda x: len(x[1]), reverse=True)
                      order, cache = list(map(list, zip(*cache)))
                      cache = [c.unbind(-1) for c in cache.unbind(-1)]
                      new_tensor = [[[None]*len(order) for i in range(2)] for j in range(self.config.depth)]
                      for i in range(len(order)):
                          for j in range(self.config.depth):
                              for k in range(2):
                                  new_tensor[j][k][i] = cache[i][j][k]
                      cache = [[flip(rnn.pad_sequence(c, batch_first=True), 1) for c in c_depth] for c_depth in cache]
                      bs, max_length, _ = list(cache[0][0].size())
                      mask = action.new_tensor([[1]*(max_length-lengths[i])+[0]*(lengths[i]) for i in range(bs)]).byte()
                  elif self.config.arch == 'sru':
                      # cache = [[for depth] for batch]
                      cache = list(map(list, zip(*cache)))
                      cache = [torch.stack(c) for c in cache]
                      mask = None
                  policy, value, cache = self.model(action, cache, mask)
                  if self.config.arch == 'Transformer':
                      cache = torch.stack([torch.stack(c,-1) for c in cache],-1)
                      cache = [(o, c) for o, c in zip(order, flip(cache, 1))]
                      cache = sorted(enumerate(cache), key=lambda x: x[0])
                      cache = [c[:lengths[i]+1] for i, c in enumerate(torch.unbind(cache,0))]
                  elif self.config.arch == 'sru':
                      cache = [c.unbind() for c in cache]
                      cache = list(map(list, zip(*cache)))
                  policy = policy.data.cpu().numpy()
                  value = value.data.cpu().numpy()
                  def elm_add(x,y):
                      new_tuple = []
                      for i, j in zip(x,y):
                          new_tuple += [tuple(list(i)+[j])]
                      return new_tuple
                  idx = 0
                  for conn, s, id in zip(ready_conns, sizes, ids):
                      conn.send((policy[idx:idx + s], value[idx:idx + s]))
                      self.cache_dicts.add(elm_add(original_cache[idx:idx + s], original_action[idx:idx + s]),
                                           cache[idx:idx + s], id)
                      idx += s
  
  def convert(x, config):
      return torch.tensor(x.astype(np.float32), device=config.device)
  
  def multi_meshgrid(*args):
      """
      Creates a meshgrid from possibly many
      elements (instead of only 2).
      Returns a nd tensor with as many dimensions
      as there are arguments
      """
      args = list(args)
      template = [1 for _ in args]
      for i in range(len(args)):
          n = args[i].shape[0]
          template_copy = template.copy()
          template_copy[i] = n
          args[i] = args[i].view(*template_copy)
          # there will be some broadcast magic going on
      return tuple(args)
  
  def flip(tensor, dims):
      if not isinstance(dims, (tuple, list)):
          dims = [dims]
      indices = [torch.arange(tensor.shape[dim] - 1, -1, -1,
          dtype=torch.int64) for dim in dims]
      multi_indices = multi_meshgrid(*indices)
      final_indices = [slice(i) for i in tensor.shape]
      for i, dim in enumerate(dims):
          final_indices[dim] = multi_indices[i]
      flipped = tensor[final_indices]
      # need to permute the final dimensions
      # if dims is not consecutive, but I'm lazy
      # now :-)
      return flipped