Repository: bioinf-jku/FCD (language: python)

Issue:
Port package to pytorch
It's more convenient to work with pytorch as it is used by many of today's generative algorithms. This saves an additional tensorflow install.


Discussion / hints:

Fix patch:
diff --git a/MANIFEST.in b/MANIFEST.in
index a87172a..ee0cd16 100644
--- a/MANIFEST.in
+++ b/MANIFEST.in
@@ -1,2 +1,2 @@
 # MANIFEST.in
-include fcd/ChemNet_v0.13_pretrained.h5
+include fcd/ChemNet_v0.13_pretrained.pt
diff --git a/README.md b/README.md
index d8192d5..65481fd 100644
--- a/README.md
+++ b/README.md
@@ -11,27 +11,21 @@ You can install the FCD using
 pip install fcd
 ```
 
-rdkit is best install via conda
-```
-conda install rdkit -c rdkit
-```
-
 # Requirements
 ```
-python>3<3.8
 numpy
-tensorflow>1.8
-keras>2.1
+torch
 scipy
+rdkit
 ```
-rdkit is causing troubles with python3.8 for me. You might wanna opt for 3.7.
-For the effect of versions on results see `tests/test_results.csv`.
-Using the current versions of tensorflow (2.1.0) and keras (2.3.1) results differ from previous versions but
-are probably negligible.
-
 
+# Updates
 ## Version 1.1 changes
 - Got rid of unneeded imports
 - `load_ref_model` doesn't need an argument any more to load a model.
 - `canonical` and `canonical_smiles` now return `None` for invalid smiles.
-- Added `get_fcd` as a quick way to get a the fcd score from two lists of smiles.
\ No newline at end of file
+- Added `get_fcd` as a quick way to get a the fcd score from two lists of smiles.
+
+## Version 1.2 changes
+- Ported the package to pytorch with the help of https://github.com/insilicomedicine/fcd_torch
+- pytorch allows a lighter package and is more popular than Tensorflow which saves an additional install
\ No newline at end of file
diff --git a/example.ipynb b/example.ipynb
index 2e40743..b0d9471 100644
--- a/example.ipynb
+++ b/example.ipynb
@@ -2,35 +2,20 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 13,
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2020-04-23T09:11:07.345354Z",
-     "start_time": "2020-04-23T09:11:07.321510Z"
-    }
-   },
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "The autoreload extension is already loaded. To reload it, use:\n",
-      "  %reload_ext autoreload\n"
-     ]
-    }
-   ],
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
    "source": [
-    "%load_ext autoreload\n",
-    "%autoreload 2\n",
     "import os\n",
-    "\n",
+    "from rdkit import RDLogger                                                                                                                                                               \n",
     "import numpy as np\n",
     "import pandas as pd\n",
-    "import pickle\n",
     "from fcd import get_fcd, load_ref_model,canonical_smiles, get_predictions, calculate_frechet_distance\n",
     "\n",
+    "RDLogger.DisableLog('rdApp.*')\n",
+    "\n",
     "np.random.seed(0)\n",
-    "os.environ[\"CUDA_VISIBLE_DEVICES\"]= '0' #set gpu "
+    "os.environ[\"CUDA_VISIBLE_DEVICES\"]= '0' #set gpu"
    ]
   },
   {
@@ -46,47 +31,14 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 19,
+   "execution_count": 2,
    "metadata": {
     "ExecuteTime": {
      "end_time": "2020-04-23T09:13:50.403933Z",
      "start_time": "2020-04-23T09:13:47.310624Z"
     }
    },
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "RDKit ERROR: [11:13:48] Explicit valence for atom # 18 N, 4, is greater than permitted\n",
-      "RDKit ERROR: [11:13:48] Explicit valence for atom # 6 N, 4, is greater than permitted\n",
-      "RDKit ERROR: [11:13:48] Explicit valence for atom # 5 N, 4, is greater than permitted\n",
-      "RDKit ERROR: [11:13:48] Explicit valence for atom # 2 N, 5, is greater than permitted\n",
-      "RDKit ERROR: [11:13:48] Explicit valence for atom # 10 N, 4, is greater than permitted\n",
-      "RDKit ERROR: [11:13:48] Explicit valence for atom # 22 N, 4, is greater than permitted\n",
-      "RDKit ERROR: [11:13:48] Explicit valence for atom # 1 N, 4, is greater than permitted\n",
-      "RDKit ERROR: [11:13:48] Explicit valence for atom # 6 N, 4, is greater than permitted\n",
-      "RDKit ERROR: [11:13:48] Explicit valence for atom # 13 N, 4, is greater than permitted\n",
-      "RDKit ERROR: [11:13:48] Explicit valence for atom # 13 N, 4, is greater than permitted\n",
-      "RDKit ERROR: [11:13:48] Explicit valence for atom # 5 N, 5, is greater than permitted\n",
-      "RDKit ERROR: [11:13:48] Explicit valence for atom # 10 N, 4, is greater than permitted\n",
-      "RDKit ERROR: [11:13:48] Explicit valence for atom # 14 N, 4, is greater than permitted\n",
-      "RDKit ERROR: [11:13:48] Explicit valence for atom # 14 N, 4, is greater than permitted\n",
-      "RDKit ERROR: [11:13:48] Explicit valence for atom # 9 N, 4, is greater than permitted\n",
-      "RDKit ERROR: [11:13:49] Explicit valence for atom # 11 N, 4, is greater than permitted\n",
-      "RDKit ERROR: [11:13:49] Explicit valence for atom # 15 N, 4, is greater than permitted\n",
-      "RDKit ERROR: [11:13:49] Explicit valence for atom # 23 N, 4, is greater than permitted\n",
-      "RDKit ERROR: [11:13:49] Explicit valence for atom # 7 N, 5, is greater than permitted\n",
-      "RDKit ERROR: [11:13:49] Explicit valence for atom # 21 N, 4, is greater than permitted\n",
-      "RDKit ERROR: [11:13:49] Explicit valence for atom # 14 N, 4, is greater than permitted\n",
-      "RDKit ERROR: [11:13:49] Explicit valence for atom # 8 N, 4, is greater than permitted\n",
-      "RDKit ERROR: [11:13:49] Explicit valence for atom # 9 N, 5, is greater than permitted\n",
-      "RDKit ERROR: [11:13:49] Explicit valence for atom # 6 N, 4, is greater than permitted\n",
-      "RDKit ERROR: [11:13:50] Explicit valence for atom # 17 N, 5, is greater than permitted\n",
-      "RDKit ERROR: [11:13:50] Explicit valence for atom # 20 N, 5, is greater than permitted\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "# Load chemnet model\n",
     "model = load_ref_model()\n",
@@ -113,7 +65,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 16,
+   "execution_count": 3,
    "metadata": {
     "ExecuteTime": {
      "end_time": "2020-04-23T09:11:27.207953Z",
@@ -125,7 +77,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "FCD:  0.3338613001233881\n"
+      "FCD:  0.333862289051325\n"
      ]
     }
    ],
@@ -151,7 +103,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 17,
+   "execution_count": 4,
    "metadata": {
     "ExecuteTime": {
      "end_time": "2020-04-23T09:11:38.873496Z",
@@ -163,20 +115,20 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "FCD:  0.3338613001233881\n"
+      "FCD:  0.333862289051325\n"
      ]
     }
    ],
    "source": [
     "\"\"\"if you don't need to store the activations you can also take a shortcut.\"\"\"\n",
-    "fcd_score = get_fcd(model, can_sample1, can_sample2)\n",
+    "fcd_score = get_fcd(can_sample1, can_sample2, model)\n",
     "\n",
     "print('FCD: ',fcd_score)"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 18,
+   "execution_count": 5,
    "metadata": {
     "ExecuteTime": {
      "end_time": "2020-04-23T09:11:49.760022Z",
@@ -188,14 +140,13 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "FCD:  25.63927611890624\n"
+      "FCD:  25.635578193222216\n"
      ]
     }
    ],
    "source": [
     "\"\"\"This is what happens if you do not canonicalize the smiles\"\"\"\n",
-    "fcd_score = get_fcd(model, can_sample1, sample2)\n",
-    "\n",
+    "fcd_score = get_fcd(can_sample1, sample2, model)\n",
     "print('FCD: ',fcd_score)"
    ]
   }
@@ -216,7 +167,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.7.4"
+   "version": "3.11.3"
   }
  },
  "nbformat": 4,
diff --git a/fcd/ChemNet_v0.13_pretrained.h5 b/fcd/ChemNet_v0.13_pretrained.h5
deleted file mode 100644
index e41f2a3..0000000
Binary files a/fcd/ChemNet_v0.13_pretrained.h5 and /dev/null differ
diff --git a/fcd/ChemNet_v0.13_pretrained.pt b/fcd/ChemNet_v0.13_pretrained.pt
new file mode 100644
index 0000000..84e622f
Binary files /dev/null and b/fcd/ChemNet_v0.13_pretrained.pt differ
diff --git a/fcd/FCD.py b/fcd/FCD.py
deleted file mode 100644
index 2b30895..0000000
--- a/fcd/FCD.py
+++ /dev/null
@@ -1,231 +0,0 @@
-#!/usr/bin/env python3
-''' Defines the functions necessary for calculating the Frechet ChemNet
-Distance (FCD) to evalulate generative models for molecules.
-
-The FCD metric calculates the distance between two distributions of molecules.
-Typically, we have summary statistics (mean & covariance matrix) of one
-of these distributions, while the 2nd distribution is given by the generative
-model.
-
-The FCD is calculated by assuming that X_1 and X_2 are the activations of
-the preulitmate layer of the CHEMNET for generated samples and real world
-samples respectivly.
-'''
-
-from __future__ import absolute_import, division, print_function
-
-import os
-import pkgutil
-import tempfile
-import warnings
-from functools import lru_cache
-from multiprocessing import Pool
-
-import keras.backend as K
-import numpy as np
-from keras.models import load_model
-from rdkit import Chem
-from scipy import linalg
-
-warnings.filterwarnings('ignore')
-
-
-def calculate_frechet_distance(mu1, sigma1, mu2, sigma2, eps=1e-6):
-    """Numpy implementation of the Frechet Distance.
-    The Frechet distance between two multivariate Gaussians X_1 ~ N(mu_1, C_1)
-    and X_2 ~ N(mu_2, C_2) is
-            d^2 = ||mu_1 - mu_2||^2 + Tr(C_1 + C_2 - 2*sqrt(C_1*C_2)).
-
-    Stable version by Dougal J. Sutherland.
-
-    Params:
-    -- mu1:    The mean of the activations of preultimate layer of the
-               CHEMNET ( like returned by the function 'get_predictions')
-               for generated samples.
-    -- mu2:    The mean of the activations of preultimate layer of the
-               CHEMNET ( like returned by the function 'get_predictions')
-               for real samples.
-    -- sigma1: The covariance matrix of the activations of preultimate layer of the
-               CHEMNET ( like returned by the function 'get_predictions')
-               for generated samples.
-    -- sigma2: The covariance matrix of the activations of preultimate layer of the
-               CHEMNET ( like returned by the function 'get_predictions')
-               for real samples.
-
-    Returns:
-    --   : The Frechet Distance.
-    """
-
-    mu1 = np.atleast_1d(mu1)
-    mu2 = np.atleast_1d(mu2)
-
-    sigma1 = np.atleast_2d(sigma1)
-    sigma2 = np.atleast_2d(sigma2)
-
-    assert mu1.shape == mu2.shape, "Training and test mean vectors have different lengths"
-    assert sigma1.shape == sigma2.shape, "Training and test covariances have different dimensions"
-
-    diff = mu1 - mu2
-
-    # product might be almost singular
-    covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False)
-    if not np.isfinite(covmean).all():
-        offset = np.eye(sigma1.shape[0]) * eps
-        covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset))
-
-    # numerical error might give slight imaginary component
-    if np.iscomplexobj(covmean):
-        if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3):
-            m = np.max(np.abs(covmean.imag))
-            raise ValueError("Imaginary component {}".format(m))
-        covmean = covmean.real
-
-    tr_covmean = np.trace(covmean)
-
-    return diff.dot(diff) + np.trace(sigma1) + np.trace(sigma2) - 2 * tr_covmean
-
-
-def build_masked_loss(loss_function, mask_value):
-    """Builds a loss function that masks based on targets
-
-    Args:
-        loss_function: The loss function to mask
-        mask_value: The value to mask in the targets
-
-    Returns:
-        function: a loss function that acts like loss_function with masked inputs
-    """
-
-    def masked_loss_function(y_true, y_pred):
-        mask = K.cast(K.not_equal(y_true, mask_value), K.floatx())
-        return loss_function(y_true * mask, y_pred * mask)
-
-    return masked_loss_function
-
-
-def masked_accuracy(y_true, y_pred):
-    a = K.sum(K.cast(K.equal(y_true, K.round(y_pred)), K.floatx()))
-    c = K.sum(K.cast(K.not_equal(y_true, 0.5), K.floatx()))
-    acc = (a) / c
-    return acc
-
-
-def get_one_hot(smiles, pad_len=-1):
-    one_hot = [
-        'C', 'N', 'O', 'H', 'F', 'Cl', 'P', 'B', 'Br', 'S', 'I', 'Si',
-        '#', '(', ')', '+', '-', '1', '2', '3', '4', '5', '6', '7', '8', '=', '[', ']', '@',
-        'c', 'n', 'o', 's', 'X', '.']
-    smiles = smiles + '.'
-    if pad_len < 0:
-        vec = np.zeros((len(smiles), len(one_hot)))
-    else:
-        vec = np.zeros((pad_len, len(one_hot)))
-    cont = True
-    j = 0
-    i = 0
-    while cont:
-        if smiles[i + 1] in ['r', 'i', 'l']:
-            sym = smiles[i:i + 2]
-            i += 2
-        else:
-            sym = smiles[i]
-            i += 1
-        if sym in one_hot:
-            vec[j, one_hot.index(sym)] = 1
-        else:
-            vec[j, one_hot.index('X')] = 1
-        j += 1
-        if smiles[i] == '.' or j >= (pad_len - 1) and pad_len > 0:
-            vec[j, one_hot.index('.')] = 1
-            cont = False
-    return (vec)
-
-
-def myGenerator_predict(smilesList, batch_size=128, pad_len=350):
-    while 1:
-        N = len(smilesList)
-        nn = pad_len
-        idxSamples = np.arange(N)
-
-        for j in range(int(np.ceil(N / batch_size))):
-            idx = idxSamples[j * batch_size: min((j + 1) * batch_size, N)]
-
-            x = []
-            for i in range(0, len(idx)):
-                currentSmiles = smilesList[idx[i]]
-                smiEnc = get_one_hot(currentSmiles, pad_len=nn)
-                x.append(smiEnc)
-
-            x = np.asarray(x) / 35
-            yield x
-
-
-@lru_cache(maxsize=4)
-def load_ref_model(model_file=None):
-    """Loads the Chemnet model. If called without arguments it will use the
-    model in the package. In case you want to use a different one provide the path
-
-    Args:
-        model_file: Path to model. (default=None)
-    """
-
-    if model_file is None:
-        chemnet_model_filename = 'ChemNet_v0.13_pretrained.h5'
-        model_bytes = pkgutil.get_data('fcd', chemnet_model_filename)
-
-        tmpdir = tempfile.TemporaryDirectory()
-        model_file = os.path.join(tmpdir.name, chemnet_model_filename)
-
-        with open(model_file, 'wb') as f:
-            f.write(model_bytes)
-
-    masked_loss_function = build_masked_loss(K.binary_crossentropy, 0.5)
-    model = load_model(
-        model_file,
-        custom_objects={
-            'masked_loss_function': masked_loss_function,
-            'masked_accuracy': masked_accuracy})
-    model.pop()
-    model.pop()
-    return model
-
-
-def get_predictions(model, gen_mol):
-    gen_mol_act = model.predict_generator(
-        myGenerator_predict(gen_mol, batch_size=128),
-        steps=np.ceil(len(gen_mol) / 128))
-    return gen_mol_act
-
-
-def canonical(smi):
-    try:
-        return Chem.MolToSmiles(Chem.MolFromSmiles(smi))
-    except:
-        return None
-
-
-def canonical_smiles(smiles, njobs=32):
-    with Pool(njobs) as pool:
-        return pool.map(canonical, smiles)
-
-
-def get_fcd(smiles1, smiles2, model=None):
-    if model is None:
-        model = load_ref_model()
-
-    act1 = get_predictions(model, smiles1)
-    act2 = get_predictions(model, smiles2)
-
-    mu1 = np.mean(act1, axis=0)
-    sigma1 = np.cov(act1.T)
-
-    mu2 = np.mean(act2, axis=0)
-    sigma2 = np.cov(act2.T)
-
-    fcd_score = calculate_frechet_distance(
-        mu1=mu1,
-        mu2=mu2,
-        sigma1=sigma1,
-        sigma2=sigma2)
-
-    return fcd_score
diff --git a/fcd/__init__.py b/fcd/__init__.py
index c388694..4ee7c79 100644
--- a/fcd/__init__.py
+++ b/fcd/__init__.py
@@ -1,4 +1,4 @@
-from .FCD import calculate_frechet_distance, build_masked_loss, masked_accuracy, get_one_hot, myGenerator_predict, \
-    load_ref_model, get_predictions, canonical, canonical_smiles, get_fcd
+from .fcd import get_fcd, get_predictions, load_ref_model
+from .utils import calculate_frechet_distance, canonical_smiles
 
-__version__ = "1.1"
\ No newline at end of file
+__version__ = "1.2"
diff --git a/fcd/fcd.py b/fcd/fcd.py
new file mode 100644
index 0000000..c7e528e
--- /dev/null
+++ b/fcd/fcd.py
@@ -0,0 +1,110 @@
+import os
+import pkgutil
+import tempfile
+from functools import lru_cache
+from typing import List, Optional
+
+import numpy as np
+import torch
+from torch import nn
+from torch.utils.data import DataLoader
+
+from .utils import (
+    SmilesDataset,
+    calculate_frechet_distance,
+    load_imported_model,
+    todevice,
+)
+
+
+@lru_cache(maxsize=1)
+def load_ref_model(model_path: Optional[str] = None):
+    """Loads chemnet model
+
+    Args:
+        model_path (str | None, optional): Path to model file. Defaults to None.
+
+    Returns:
+        Chemnet as torch model
+    """
+
+    if model_path is None:
+        chemnet_model_filename = "ChemNet_v0.13_pretrained.pt"
+        model_bytes = pkgutil.get_data("fcd", chemnet_model_filename)
+
+        tmpdir = tempfile.TemporaryDirectory()
+        model_path = os.path.join(tmpdir.name, chemnet_model_filename)
+        with open(model_path, "wb") as f:
+            f.write(model_bytes)
+
+    model_config = torch.load(model_path)
+    model = load_imported_model(model_config)
+    model.eval()
+    return model
+
+
+def get_predictions(
+    model: nn.Module,
+    smiles_list: List[str],
+    batch_size: int = 128,
+    n_jobs: int = 1,
+    device: str = "cpu",
+) -> np.ndarray:
+    """Calculate Chemnet activations
+
+    Args:
+        model (nn.Module): Chemnet model
+        smiles_list (List[str]): List of smiles to process
+        batch_size (int, optional): Which batch size to use for inference. Defaults to 128.
+        n_jobs (int, optional): How many jobs to use for preprocessing. Defaults to 1.
+        device (str, optional): On which device the chemnet model is run. Defaults to "cpu".
+
+    Returns:
+        np.ndarray: The activation for the input list
+    """
+    if len(smiles_list) == 0:
+        return np.zeros((0, 512))
+
+    dataloader = DataLoader(
+        SmilesDataset(smiles_list), batch_size=batch_size, num_workers=n_jobs
+    )
+    with todevice(model, device), torch.no_grad():
+        chemnet_activations = []
+        for batch in dataloader:
+            chemnet_activations.append(
+                model(batch.transpose(1, 2).float().to(device))
+                .to("cpu")
+                .detach()
+                .numpy()
+            )
+    return np.row_stack(chemnet_activations)
+
+
+def get_fcd(smiles1: List[str], smiles2: List[str], model: nn.Module = None) -> float:
+    """Calculate FCD between two sets of Smiles
+
+    Args:
+        smiles1 (List[str]): First set of smiles
+        smiles2 (List[str]): Second set of smiles
+        model (nn.Module, optional): The model to use. Loads default model if None.
+
+    Returns:
+        float: The FCD score
+    """
+    if model is None:
+        model = load_ref_model()
+
+    act1 = get_predictions(model, smiles1)
+    act2 = get_predictions(model, smiles2)
+
+    mu1 = np.mean(act1, axis=0)
+    sigma1 = np.cov(act1.T)
+
+    mu2 = np.mean(act2, axis=0)
+    sigma2 = np.cov(act2.T)
+
+    fcd_score = calculate_frechet_distance(
+        mu1=mu1, mu2=mu2, sigma1=sigma1, sigma2=sigma2
+    )
+
+    return fcd_score
diff --git a/fcd/ref_combined_stat.p b/fcd/ref_combined_stat.p
deleted file mode 100644
index 9fa2ff0..0000000
Binary files a/fcd/ref_combined_stat.p and /dev/null differ
diff --git a/fcd/torch_layers.py b/fcd/torch_layers.py
new file mode 100644
index 0000000..2421a98
--- /dev/null
+++ b/fcd/torch_layers.py
@@ -0,0 +1,54 @@
+import torch
+from torch import nn
+
+
+class Reverse(nn.Module):
+    def forward(self, x):
+        return torch.flip(x, [1])
+
+
+class IndexTuple(nn.Module):
+    def __init__(self, pos):
+        super().__init__()
+        self.pos = pos
+
+    def forward(self, x):
+        return x[self.pos]
+
+
+class IndexTensor(nn.Module):
+    def __init__(self, pos, dim):
+        super().__init__()
+        self.pos = pos
+        self.dim = dim
+
+    def forward(self, x):
+        return torch.select(x, self.dim, self.pos)
+
+
+class Transpose(nn.Module):
+    def forward(self, x):
+        return x.transpose(1, 2)
+
+
+# https://stackoverflow.com/questions/52975843/comparing-conv2d-with-padding-between-tensorflow-and-pytorch
+class SamePadding1d(nn.Module):
+    def __init__(self, kernel_size, stride):
+        super().__init__()
+        self.kernel_size = kernel_size
+        self.stride = stride
+
+    def forward(self, x):
+        if x.shape[2] % self.stride == 0:
+            pad = max(self.kernel_size - self.stride, 0)
+        else:
+            pad = max(self.kernel_size - (x.shape[2] % self.stride), 0)
+
+        if pad % self.stride == 0:
+            pad_val = pad // self.stride
+            padding = (pad_val, pad_val)
+        else:
+            pad_val_start = pad // self.stride
+            pad_val_end = pad - pad_val_start
+            padding = (pad_val_start, pad_val_end)
+        return torch.nn.functional.pad(x, padding, "constant", 0)
diff --git a/fcd/utils.py b/fcd/utils.py
new file mode 100644
index 0000000..bb7d57a
--- /dev/null
+++ b/fcd/utils.py
@@ -0,0 +1,198 @@
+import re
+from contextlib import contextmanager
+from multiprocessing import Pool
+from typing import List
+
+import numpy as np
+import torch
+from rdkit import Chem
+from scipy import linalg
+from torch import nn
+from torch.utils.data import Dataset
+
+from .torch_layers import IndexTensor, IndexTuple, Reverse, SamePadding1d, Transpose
+
+# fmt: off
+__vocab = ["C","N","O","H","F","Cl","P","B","Br","S","I","Si","#","(",")","+","-","1","2","3","4","5","6","7","8","=","[","]","@","c","n","o","s","X","."]
+# fmt: on
+__vocab_c2i = {k: i for i, k in enumerate(__vocab)}
+__unk = __vocab_c2i["X"]
+
+
+sorted_vocab = __vocab[:]
+sorted_vocab.sort()  # sort alphabetically
+sorted_vocab.sort(key=len, reverse=True)  # sort by length for regex
+FULL_REGEX = "|".join(
+    "(%s)" % re.escape(base_symbol) for base_symbol in sorted_vocab
+)  # Tries to match longer tokens first.
+FULL_REGEX += "|."  # Handle unkown characters
+
+
+def tokenize(smiles: str) -> List[str]:
+    """Tokenizes the given smiles string. Needed for multi-character tokens like 'Cl'
+
+    Args:
+        smiles (str): Input molecule as Smiles
+
+    Returns:
+        List[str]: List of tokens
+    """
+    tok_smile = [mo.group() for mo in re.finditer(FULL_REGEX, smiles)]
+    assert "".join(tok_smile) == smiles
+    return tok_smile
+
+
+def get_one_hot(smiles: str, pad_len: int = -1) -> np.ndarray:
+    """Generate one-hot representation of a Smiles string.
+
+    Args:
+        smiles (str): Input molecule as Smiles
+        pad_len (int, optional): Whether or not to pad to a given size. Defaults to -1.
+
+    Returns:
+        np.ndarray: Array containing the one-hot encoded Smiles
+    """
+    smiles = smiles + "."
+
+    # initialize array
+    array_length = len(smiles) if pad_len < 0 else pad_len
+    vocab_size = len(__vocab)
+    one_hot = np.zeros((array_length, vocab_size))
+
+    tokens = tokenize(smiles)
+    numeric = [__vocab_c2i.get(token, __unk) for token in tokens]
+
+    for pos, num in enumerate(numeric):
+        one_hot[pos, num] = 1
+
+    return one_hot
+
+
+def load_imported_model(keras_config):
+    activations = {
+        "relu": nn.ReLU,
+        "selu": nn.SELU,
+        "tanh": nn.Tanh,
+        "sigmoid": nn.Sigmoid,
+    }
+    layers = []
+    conv = True
+    for layer_type, config in keras_config:
+        state_dict, kwargs, other_info = config
+        if layer_type == "Conv1d":
+            assert conv, "Can't have conv layers after LSTM"
+            if other_info["padding"] == "same":
+                layers.append(SamePadding1d(kwargs["kernel_size"], kwargs["stride"]))
+            layer = nn.Conv1d(**kwargs)
+            layer.load_state_dict(state_dict)
+            layers.append(layer)
+            activation = other_info["activation"]
+            layers.append(activations[activation]())
+        elif layer_type == "LSTM":
+            if conv:
+                conv = False
+                layers.append(Transpose())
+            layer = nn.LSTM(**kwargs)
+            layer.load_state_dict(state_dict)
+            if other_info["reverse"]:
+                layers.append(Reverse())
+            layers.append(layer)
+            layers.append(IndexTuple(0))
+            if other_info["last"]:
+                layers.append(IndexTensor(-1, 1))
+        else:
+            raise ValueError("Unknown layer type")
+    return nn.Sequential(*layers)
+
+
+class SmilesDataset(Dataset):
+    __PAD_LEN = 350
+
+    def __init__(self, smiles_list):
+        super().__init__()
+        self.smiles_list = smiles_list
+
+    def __getitem__(self, idx):
+        smiles = self.smiles_list[idx]
+        features = get_one_hot(smiles, 350)
+        return features / features.shape[1]
+
+    def __len__(self):
+        return len(self.smiles_list)
+
+
+def calculate_frechet_distance(mu1, sigma1, mu2, sigma2, eps=1e-6):
+    """Numpy implementation of the Frechet Distance.
+    The Frechet distance between two multivariate Gaussians X_1 ~ N(mu_1, C_1)
+    and X_2 ~ N(mu_2, C_2) is
+            d^2 = ||mu_1 - mu_2||^2 + Tr(C_1 + C_2 - 2*sqrt(C_1*C_2)).
+
+    Stable version by Dougal J. Sutherland.
+    Params:
+    -- mu1:    The mean of the activations of preultimate layer of the
+               CHEMNET (like returned by the function 'get_predictions')
+               for generated samples.
+    -- mu2:    The mean of the activations of preultimate layer of the
+               CHEMNET (like returned by the function 'get_predictions')
+               for real samples.
+    -- sigma1: The covariance matrix of the activations of preultimate layer
+               of the CHEMNET (like returned by the function 'get_predictions')
+               for generated samples.
+    -- sigma2: The covariance matrix of the activations of preultimate layer
+               of the CHEMNET (like returned by the function 'get_predictions')
+               for real samples.
+    Returns:
+    --   : The Frechet Distance.
+    """
+
+    mu1 = np.atleast_1d(mu1)
+    mu2 = np.atleast_1d(mu2)
+
+    sigma1 = np.atleast_2d(sigma1)
+    sigma2 = np.atleast_2d(sigma2)
+
+    assert (
+        mu1.shape == mu2.shape
+    ), "Training and test mean vectors have different lengths"
+    assert (
+        sigma1.shape == sigma2.shape
+    ), "Training and test covariances have different dimensions"
+
+    diff = mu1 - mu2
+
+    # product might be almost singular
+    covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False)
+    if not np.isfinite(covmean).all():
+        offset = np.eye(sigma1.shape[0]) * eps
+        covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset))
+
+    # numerical error might give slight imaginary component
+    if np.iscomplexobj(covmean):
+        if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3):
+            m = np.max(np.abs(covmean.imag))
+            raise ValueError("Imaginary component {}".format(m))
+        covmean = covmean.real
+
+    tr_covmean = np.trace(covmean)
+
+    return diff.dot(diff) + np.trace(sigma1) + np.trace(sigma2) - 2 * tr_covmean
+
+
+@contextmanager
+def todevice(model, device):
+    model.to(device)
+    yield
+    model.to("cpu")
+    torch.cuda.empty_cache()
+
+
+def canonical(smi):
+    try:
+        return Chem.MolToSmiles(Chem.MolFromSmiles(smi))
+    except:
+        return None
+
+
+def canonical_smiles(smiles, njobs=32):
+    with Pool(njobs) as pool:
+        return pool.map(canonical, smiles)
diff --git a/setup.py b/setup.py
index d020896..e18e275 100644
--- a/setup.py
+++ b/setup.py
@@ -3,27 +3,20 @@
 # from distutils.core import setup
 from setuptools import setup
 
-with open("README.md", "r", encoding='utf8') as fh:
+with open("README.md", "r", encoding="utf8") as fh:
     long_description = fh.read()
 
-setup(name='FCD',
-      version='1.1',
-      author='',
-      author_email='',
-      description='Fréchet ChEMNet Distance',
-      url='https://github.com/bioinf-jku/FCD',
-      packages=['fcd'],
-      license='LGPLv3',
-      long_description=long_description,
-      long_description_content_type="text/markdown",
-      install_requires=[
-          'keras',
-          'numpy',
-          'scipy',
-          'tensorflow'
-      ],
-      extras_require={
-          'rdkit': ['rdkit'],
-      },
-      include_package_data=True,
-      )
+setup(
+    name="FCD",
+    version="1.2",
+    author="Philipp Renz",
+    author_email="renz.ph@gmail.com",
+    description="Fréchet ChEMNet Distance",
+    url="https://github.com/bioinf-jku/FCD",
+    packages=["fcd"],
+    license="LGPLv3",
+    long_description=long_description,
+    long_description_content_type="text/markdown",
+    install_requires=["torch", "numpy", "scipy", "rdkit"],
+    include_package_data=True,
+)


Test patch:
diff --git a/test/test_fcd.py b/test/test_fcd.py
new file mode 100644
index 0000000..92ead4e
--- /dev/null
+++ b/test/test_fcd.py
@@ -0,0 +1,65 @@
+import unittest
+
+import numpy as np
+
+from fcd import get_fcd
+from fcd.utils import get_one_hot
+
+
+class TestFCD(unittest.TestCase):
+    def test_random_smiles(self):
+        smiles_list1 = ["CNOHF", "NOHFCl", "OHFClP", "HFClPB"]
+        smiles_list2 = ["ISi#()", "Si#()+", "#()+-", "()+-1"]
+        target = 8.8086
+        self.assertAlmostEqual(
+            get_fcd(smiles_list1, smiles_list2), target, 3, f"Should be {target}"
+        )
+
+    def test_fcd_torch(self):
+        smiles_list1 = [
+            "COc1cccc(NC(=O)Cc2coc3ccc(OC)cc23)c1",
+            "Cc1noc(C)c1CN(C)C(=O)Nc1cc(F)cc(F)c1",
+        ]
+        smiles_list2 = [
+            "Oc1ccccc1-c1cccc2cnccc12",
+            "Cc1noc(C)c1CN(C)C(=O)Nc1cc(F)cc(F)c1",
+        ]
+        target = 47.773486382119444
+        self.assertAlmostEqual(
+            get_fcd(smiles_list1, smiles_list2), target, 3, f"Should be {target}"
+        )
+
+    def test_one_hot(self):
+        inputs = [
+            "O=C([C@H]1CC[C@@H]2[C@@H]",
+            "CCN(CC)S(=O)(=O)CCC[C@H](",
+            "COc1ccc(\C=C\2/C[N+](C)(C",
+            "COc1cc(CC(=O)N2CCN(CC2)c3",
+            "CN(CC#Cc1cc(ccc1C)ClC2(O)CC",
+            "CYACCyyCLCl",
+        ]
+        # fmt: off
+        outputs = [
+            ((26, 35), [2, 25, 0, 13, 26, 0, 28, 3, 27, 17, 0, 0, 26, 0, 28, 28, 3, 27, 18, 26, 0, 28, 28, 3, 27, 34]),
+            ((26, 35), [0, 0, 1, 13, 0, 0, 14, 9, 13, 25, 2, 14, 13, 25, 2, 14, 0, 0, 0, 26, 0, 28, 3, 27, 13, 34]),
+            ((25, 35), [0, 2, 29, 17, 29, 29, 29, 13, 33, 0, 25, 0, 33, 33, 0, 26, 1, 15, 27, 13, 0, 14, 13, 0, 34]),
+            ((26, 35), [0, 2, 29, 17, 29, 29, 13, 0, 0, 13, 25, 2, 14, 1, 18, 0, 0, 1, 13, 0, 0, 18, 14, 29, 19, 34]),
+            ((28, 35), [0, 1, 13, 0, 0, 12, 0, 29, 17, 29, 29, 13, 29, 29, 29, 17, 0, 14, 5, 0, 18, 13, 2, 14, 0, 0, 34]),
+            ((12, 35), [0, 33, 33, 0, 0, 33, 33, 0, 33, 5, 34]),
+        ]
+        # fmt: on
+
+        for inp, (correct_shape, correct_entries) in zip(inputs, outputs):
+            one_hot = get_one_hot(inp)
+            shape = one_hot.shape
+            entries = np.where(one_hot)[1].tolist()
+            self.assertEqual(shape, correct_shape)
+            self.assertEqual(entries, correct_entries)
+
+            # assert that no duplicate ones and no missing entries. Trailing zero vectors are allowed.
+            non_zero_idx = np.where(one_hot)[0]
+            assert np.all(non_zero_idx == np.arange(len(non_zero_idx)))
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/dummyfcd_results.csv b/tests/dummyfcd_results.csv
deleted file mode 100644
index 65e29c0..0000000
--- a/tests/dummyfcd_results.csv
+++ /dev/null
@@ -1,13 +0,0 @@
-fcd_score, tf version, keras version
-8.808623423861544, 1.8.0, 2.1.0
-8.808623423861544, 1.8.0, 2.2.5
-8.808621430192593, 1.14.0, 2.0.9
-8.808621430192593, 1.14.0, 2.1.5
-8.808621430192593, 1.14.0, 2.1.6
-8.808621430192593, 1.14.0, 2.2.5
-8.808609328184861, 1.14.0, 2.3.1
-8.808627531931002, 2.1.0, 2.3.1
-8.80862709007977, 2.1.0, 2.3.1,
-8.808609328184861, 1.14.0, 2.3.1
-<function dummyfcd at 0x7f5dea9b7170>, 1.14.0, 2.3.1
-<function dummyfcd at 0x7f397779c170>, 1.14.0, 2.3.1
diff --git a/tests/test_dummyfcd.py b/tests/test_dummyfcd.py
deleted file mode 100644
index 9eff82a..0000000
--- a/tests/test_dummyfcd.py
+++ /dev/null
@@ -1,42 +0,0 @@
-import os
-import sys
-
-import keras
-import tensorflow
-from fcd import get_fcd, load_ref_model
-
-
-
-
-
-# Don't use GPU
-os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
-
-
-def dummyfcd():
-    one_hot = [
-        'C', 'N', 'O', 'H', 'F', 'Cl', 'P', 'B', 'Br', 'S', 'I', 'Si',
-        '#', '(', ')', '+', '-', '1', '2', '3', '4', '5', '6', '7', '8', '=', '[', ']', '@',
-        'c', 'n', 'o', 's', 'X', '.']
-
-    # get sum dummy sequences
-    s1 = []
-    s2 = []
-    for i in range(4):
-        s1.append(''.join(one_hot[i:i + 5]))
-        s2.append(''.join(one_hot[i + 10:i + 15]))
-
-    return get_fcd(s1, s2)
-
-
-def test_dummyfcd():
-    fcd_score = dummyfcd()
-    assert abs(fcd_score - 8.808623423861544) < 1e-4
-
-
-if __name__ == '__main__':
-    fcd_score = dummyfcd
-
-    with open(os.path.join(os.path.dirname(__file__), 'dummyfcd_results.csv'), 'a') as f:
-        f.write(', '.join([str(fcd_score), str(
-            tensorflow.__version__), str(keras.__version__)]) + '\n')