--- a
+++ b/demo_heteroencoder.ipynb
@@ -0,0 +1,183 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%load_ext autoreload\n",
+    "%autoreload 2\n",
+    "\n",
+    "\n",
+    "import numpy as np\n",
+    "import rdkit\n",
+    "from rdkit import Chem\n",
+    "\n",
+    "import h5py, ast, pickle\n",
+    "\n",
+    "# Occupy a GPU for the model to be loaded \n",
+    "%env CUDA_DEVICE_ORDER=PCI_BUS_ID\n",
+    "# GPU ID, if occupied change to an available GPU ID listed under !nvidia-smi\n",
+    "%env CUDA_VISIBLE_DEVICES=2 \n",
+    "\n",
+    "from ddc_pub import ddc_v3 as ddc"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Load model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Import existing (trained) model\n",
+    "# Ignore any warning(s) about training configuration or non-seriazable keyword arguments\n",
+    "model_name = \"models/heteroencoder_model\"\n",
+    "model = ddc.DDC(model_name=model_name)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Load data from dataset"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dataset_name = \"datasets/CHEMBL25_TEST.h5\"\n",
+    "npoints = 1000\n",
+    "\n",
+    "dataset = h5py.File(dataset_name, \"r\")\n",
+    "mols    = dataset[\"mols\"][:]\n",
+    "# Select random npoints\n",
+    "mols_in = mols[np.random.choice(len(mols), npoints, replace=False)]\n",
+    "dataset.close()\n",
+    "\n",
+    "# Get the SMILES behind the binary mols\n",
+    "smiles_in = [Chem.MolToSmiles(Chem.Mol(mol)) for mol in mols_in]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Alternatively, use your own SMILES"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Input SMILES to auto-encode\n",
+    "smiles_in = ['Cc1cccn2c(CN(C)C3CCCc4ccccc43)c(C(=O)N3CCOCC3)nc12',\n",
+    "             'COC(=O)NN=C(c1ccc(O)cc1)C1C(=O)N(C)C(=O)N(C)C1=O',\n",
+    "             'CCc1cc(CC)nc(OCCCn2c3c(c4cc(-c5nc(C)no5)ccc42)CC(F)(F)CC3)n1',\n",
+    "             'Cc1ccc2c(C(=O)Nc3ccccc3)c(SSc3c(C(=O)Nc4ccccc4)c4ccc(C)cc4n3C)n(C)c2c1',\n",
+    "             'Cc1cccc(-c2ccccc2)c1Oc1nc(O)nc(NCc2ccc3occc3c2)n1',\n",
+    "             'Cn1nnnc1SCC(=O)NN=Cc1ccc(Cl)cc1',\n",
+    "             'COc1cccc(NS(=O)(=O)c2ccc(OC)c(OC)c2)c1',\n",
+    "             'COc1ccc(OC)c(S(=O)(=O)n2nc(C)cc2C)c1',\n",
+    "             'NCCCn1cc(C2=C(c3ccncc3)C(=O)NC2=O)c2ccccc21',\n",
+    "             'CN(C)C(=O)N1CCN(C(c2ccc(Cl)cc2)c2cccnc2)CC1']\n",
+    "\n",
+    "# MUST convert SMILES to binary mols for the model to accept them (it re-converts them to SMILES internally)\n",
+    "mols_in = [Chem.rdchem.Mol.ToBinary(Chem.MolFromSmiles(smiles)) for smiles in smiles_in]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Encode the binary mols into their latent representations\n",
+    "latent = model.transform(model.vectorize(mols_in))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Convert back to SMILES\n",
+    "smiles_out = []\n",
+    "for lat in latent:   \n",
+    "    smiles, _ = model.predict(lat, temp=0)\n",
+    "    smiles_out.append(smiles)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# To compare the results, convert smiles_out to CANONICAL\n",
+    "for idx, smiles in enumerate(smiles_out):\n",
+    "    mol = Chem.MolFromSmiles(smiles)\n",
+    "    if mol:\n",
+    "        smiles_out[idx] = Chem.MolToSmiles(mol, canonical=True)\n",
+    "    else:\n",
+    "        smiles_out[idx] = \"INVALID\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "scrolled": true
+   },
+   "outputs": [],
+   "source": [
+    "smiles_in"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "scrolled": true
+   },
+   "outputs": [],
+   "source": [
+    "smiles_out"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "ddc",
+   "language": "python",
+   "name": "ddc"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.6.7"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}