code: import os
  import time
  import math
  import numpy as np
  import torch
  
  # torch.multiprocessing.set_start_method('spawn')
  torch.multiprocessing.set_start_method('forkserver', force=True)
  import torch.nn as nn
  from torch.utils.data import DataLoader
  from tqdm import tqdm
  from argparse import Namespace
  from typing import List
  from dglt.data.dataset.molecular import MoleculeDataset
  from dglt.data.transformer.scaler import StandardScaler
  from dglt.data.transformer.collator import MolCollator
  from dglt.data.dataset.utils import get_data, get_data_from_smiles
  from dglt.utils import load_args, load_checkpoint, load_scalers
  from deploy import get_newest_train_args
  from third_party.dimorphite_dl.acid_base import mol_cls
  
  class MoleProp(object):
      """Molecular Properties Prediction Service"""
  
      def __init__(self, checkpoint_dir, debug=print):
          self.debug_ = debug
          self.checkpoint_paths_ = []
          for root, _, files in os.walk(checkpoint_dir):
              for fname in files:
                  if fname.endswith('.pt'):
                      self.checkpoint_paths_.append(os.path.join(root, fname))
  
      def load_model(self, args: Namespace):
          """
          Load checkpoints
  
          :param args: Arguments.
          :return:
          """
          self.scaler_, self.features_scaler_ = load_scalers(self.checkpoint_paths_[0])
          self.train_args = load_args(self.checkpoint_paths_[0])
          self.args_ = args
          for key, value in vars(self.train_args).items():
              if not hasattr(self.args_, key):
                  setattr(self.args_, key, value)
  
          # update args with newest training args
          newest_train_args = get_newest_train_args()
          for key, value in vars(newest_train_args).items():
              if not hasattr(args, key):
                  setattr(args, key, value)
          if args.features_path:
              args.features_path = None
              args.features_generator = ['rdkit_2d_normalized']
          self.models_ = []
          for checkpoint_path in tqdm(self.checkpoint_paths_, total=len(self.checkpoint_paths_)):
              self.models_.append(load_checkpoint(checkpoint_path, cuda=self.args_.cuda, current_args=self.args_))
  
      def inference(self,
                    model: nn.Module,
                    data: MoleculeDataset,
                    args,
                    batch_size: int,
                    shared_dict,
                    scaler: StandardScaler = None
                    ) -> List[List[float]]:
          """
          Do inference
          :param model: model.
          :param data: input data.
          :param args: Arguments.
          :param batch_size: batch size.
          :param shared_dict: shared_dict of model.
          :param scaler: scaler of input data.
          :return: prediction of molecular properties.
          """
          # model.share_memory()
          model.eval()
          args.bond_drop_rate = 0
          preds = []
          iter_count = 0
          mol_collator = MolCollator(args=args, shared_dict=shared_dict)
          mol_loader = DataLoader(data, batch_size=batch_size, shuffle=False, num_workers=0, collate_fn=mol_collator)
          for i, item in enumerate(mol_loader):
              smiles_batch, batch, features_batch, mask, _ = item
  
              with torch.no_grad():
                  batch_preds = model(batch, features_batch)
                  iter_count += args.batch_size
              batch_preds = batch_preds.data.cpu().numpy()
              if scaler is not None:
                  batch_preds = scaler.inverse_transform(batch_preds)
              batch_preds = batch_preds.tolist()
              preds.extend(batch_preds)
  
          return preds
  
      def postprocessing(self, task: str = None, smiles: List[str] = None, preds: np.ndarray = None):
          if task == 'caco2':
              for i in range(preds.shape[0]):
                  if preds[i] is not None:
                      for j in range(len(preds[i])):
                          preds[i][j] = (math.pow(10, preds[i][j]) - 1) / 10
          elif task == 'pka':
              acid_base = mol_cls(smiles)
              preds[acid_base == None] = np.nan
              preds = np.column_stack((preds, np.array(acid_base, dtype=np.float)))
          elif task == 'ppb':
              preds[preds > 1] = 1
              preds[preds < 0] = 0
          return preds
  
      def predict(self, task: str = None, smiles: List[str] = None):
          """
          Predict molecular properties.
          :param smiles: input data.
          :return: molecular properties.
          """
          self.debug_('Loading data')
          tic = time.time()
          self.args_.max_workers = 30
          if smiles is not None:
              test_data = get_data_from_smiles(smiles=smiles, skip_invalid_smiles=True, args=self.args_)
          else:
              test_data = get_data(path=self.args_.input_file, args=self.args_,
                                   use_compound_names=self.args_.use_compound_names,
                                   skip_invalid_smiles=True)
          toc = time.time()
          self.debug_('loading data: {}s'.format(toc - tic))
          self.debug_('Validating SMILES')
          tic = time.time()
          valid_indices = [i for i in range(len(test_data)) if test_data[i].mol is not None]
          full_data = test_data
          test_data = MoleculeDataset([test_data[i] for i in valid_indices])
  
          # Edge case if empty list of smiles is provided
          if len(test_data) == 0:
              return [None] * len(full_data)
  
          # Normalize features
          if self.train_args.features_scaling:
              test_data.normalize_features(self.features_scaler)
  
          sum_preds = np.zeros((len(test_data), self.args_.num_tasks))
          toc = time.time()
          self.debug_('validating smiles: {}s'.format(toc - tic))
          self.debug_(f'Predicting...')
          tic = time.time()
          shared_dict = {}
          for model in self.models_:
              model_preds = self.inference(
                  model=model,
                  data=test_data,
                  batch_size=self.args_.batch_size,
                  scaler=self.scaler_,
                  shared_dict=shared_dict,
                  args=self.args_
              )
              sum_preds += np.array(model_preds)
          toc = time.time()
          self.debug_('predicting: {}s'.format(toc - tic))
          avg_preds = sum_preds / len(self.checkpoint_paths_)
          avg_preds = self.postprocessing(task=task, smiles=smiles, preds=avg_preds)
          avg_preds = avg_preds.tolist()
          assert len(test_data) == len(avg_preds)
          test_smiles = test_data.smiles()
          res = {}
          for i in range(len(avg_preds)):
              res[test_smiles[i]] = avg_preds[i]
  
          return {'task': task, 'task_score': res}
   import torch
  import torch.nn as nn
  from torch.distributions.gumbel import Gumbel
  
  from dglt.contrib.moses.moses.model.gvae.utils import Flatten
  
  class CVAE(nn.Module):
      """Character Variational Autoencoder."""
  
      def __init__(self, vocab, config):
          """Constructor function.
  
          Args:
          * vocab: model's vocabulary
          * config: model's configuration
          """
  
          super().__init__()
  
          # get the vocabulary and configurations
          self.eos_symbol = '&'  # symbol '&' is not used in SMILES
          self.vocab = [self.eos_symbol] + vocab
          self.config = config
          self.vocab_len = len(self.vocab)
  
          # build encoder & decoder
          self.__build_encoder()
          self.__build_decoder()
  
          # pack all modules into one
          self.vae = nn.ModuleList([
              self.encoder,
              self.q_mu,
              self.q_logvar,
              self.mapper,
              self.decoder,
              self.decoder_fc
          ])
  
      def forward(self, x):
          """Perform the forward passing and compute losses.
  
          Args:
          * x: training samples (as torch.tensor)
  
          Returns:
          * kl_loss: KL-divergence loss
          * recon_loss: reconstruction loss
          """
  
          # send the data to model's device
          x = x.to(self.device)
  
          # encode into latent vectors
          x_trans = torch.transpose(x, 1, 2)
          x_encoded = self.encoder(x_trans)
          mu = self.q_mu(x_encoded)
          logvar = self.q_logvar(x_encoded)
  
          # decode from latent vectors
          z = mu + (logvar / 2).exp() * torch.randn_like(mu)
          z_mapped = self.mapper(z)
          z_tiled = z_mapped.unsqueeze(1).repeat(1, self.config.smiles_maxlen, 1)
          z_decoded, __ = self.decoder(z_tiled, None)
          y = self.decoder_fc(z_decoded)
  
          # compute KL-divergence & re-construction losses
          kl_loss = 0.5 * (logvar.exp() + mu ** 2 - 1 - logvar).sum(1).mean()
          recon_loss = -(x * y.log()).sum([1, 2]).mean()
  
          return kl_loss, recon_loss
  
      def sample(self, n_batch, max_len=100):
          """Sample SMILES strings from the prior distribution.
  
          Args:
          * n_batch: # of SMILES strings
          * max_len: maximal length of a SMILES string
  
          Returns:
          * string_list: list of SMILES strings
          """
  
          with torch.no_grad():
              # sample latent vectors from the prior distribution
              z = torch.randn(n_batch, self.config.d_z)
              z = z.to(self.device)
  
              # decode from latent vectors
              z_mapped = self.mapper(z)
              z_tiled = z_mapped.unsqueeze(1).repeat(1, max_len, 1)
              z_decoded, __ = self.decoder(z_tiled, None)
              y = self.decoder_fc(z_decoded)
  
              # convert tensors into SMILES strings
              m = Gumbel(torch.tensor([0.0]), torch.tensor([0.1]))
              noise = torch.squeeze(m.sample(y.size()))
              noise = noise.to(self.device)
              y_idxs = torch.argmax(y.log() + noise, dim=-1)
              tensor_list = torch.split(y_idxs, 1, dim=0)
              string_list = [self.tensor2string(tensor) for tensor in tensor_list]
  
          return string_list
  
      def string2tensor(self, string, device='model'):
          """Convert a SMILES string to torch.tensor.
  
          Args:
          * string: SMILES string
          * device: where to place the torch.tensor
  
          Returns:
          * tensor: torch.tensor consists of one-hot vectors
          """
  
          # obtain a list of non-zero entries' indices
          string += self.eos_symbol * (self.config.smiles_maxlen - len(string))
          ids = list(map(lambda x: self.vocab.index(x), string))
  
          # convert into a 2-D tensor consists of one-hot vectors
          tensor = torch.zeros(self.config.smiles_maxlen, self.vocab_len)
          tensor.scatter_(1, torch.tensor(ids).view([-1, 1]), 1)
          tensor.to(self.device if device == 'model' else device)
  
          return tensor
  
      def tensor2string(self, tensor):
          """Convert a torch.tensor to SMILES string.
  
          Args:
          * tensor: torch.tensor consists of non-zero entries' indices
  
          Returns:
          * string: SMILES string
          """
  
          # convert into a SMILES string with end-of-sequence characters removed
          ids = tensor.view(-1).tolist()
          string = ''.join([self.vocab[id] for id in ids])
          string = string.replace(self.eos_symbol, '')
  
          return string
  
      @property
      def device(self):
          """The model's device."""
  
          return next(self.parameters()).device
  
      def __build_encoder(self):
          """Build the encoder sub-network.
  
          NOTE: encoder's input must be of size <batch_size * vocab_len * smiles_maxlen>.
          """
  
          # configure parameters for convolutional and linear layers
          conv1_param = (self.vocab_len, 9, 9)  # in_channels / out_channels / kernel_size
          conv2_param = (9, 9, 9)
          conv3_param = (9, 10, 11)
          nb_idims_fc = conv3_param[1] * (self.config.smiles_maxlen -
              (conv1_param[2] - 1) - (conv2_param[2] - 1) - (conv3_param[2] - 1))
          nb_odims_fc = 435
  
          # encoder sub-network
          self.encoder = nn.Sequential(
              nn.Conv1d(conv1_param[0], conv1_param[1], conv1_param[2]),
              nn.ReLU(),
              nn.Conv1d(conv2_param[0], conv2_param[1], conv2_param[2]),
              nn.ReLU(),
              nn.Conv1d(conv3_param[0], conv3_param[1], conv3_param[2]),
              nn.ReLU(),
              Flatten(),
              nn.Linear(nb_idims_fc, nb_odims_fc),
              nn.ReLU()
          )
  
          # latent vector's mean & variance
          self.q_mu = nn.Linear(nb_odims_fc, self.config.d_z)
          self.q_logvar = nn.Linear(nb_odims_fc, self.config.d_z)
  
      def __build_decoder(self):
          """Build the decoder sub-network."""
  
          # map the latent vector for decoding
          # <self.mapper>'s output should be repeated before feeding into <self.decoder>
          self.mapper = nn.Sequential(
              nn.Linear(self.config.d_z, self.config.d_z),
              nn.ReLU()
          )
  
          # decoder sub-network
          self.decoder = nn.GRU(
              self.config.d_z,
              self.config.d_d_h,
              num_layers=3,
              batch_first=True,
              dropout=self.config.d_dropout if self.config.d_n_layers > 1 else 0
          )
  
          # probabilistic outputs
          self.decoder_fc = nn.Sequential(
              nn.Linear(self.config.d_d_h, self.vocab_len),
              nn.Softmax(dim=-1),
          )
  
apis: numpy.array torch.utils.data.DataLoader torch.multiprocessing.set_start_method torch.no_grad
apis: torch.nn.Linear torch.randn_like torch.split torch.randn torch.nn.Softmax torch.no_grad torch.tensor torch.nn.GRU torch.nn.Conv1d torch.nn.ModuleList torch.zeros torch.nn.ReLU torch.transpose