Update app.py

app.py

@@ -6,6 +6,11 @@ import numpy as np
 import tensorflow as tf
 from tensorflow import keras
 
+# Config 
+NUM_ATOMS = 9  # Maximum number of atoms
+ATOM_DIM = 4 + 1  # Number of atom types
+BOND_DIM = 4 + 1  # Number of bond types
+LATENT_DIM = 64  # Size of the latent space
 RDLogger.DisableLog("rdApp.*")
 
 def graph_to_molecule(graph):
@@ -50,7 +55,7 @@ generator = from_pretrained_keras("keras-io/wgan-molecular-graphs")
 
 def predict(num_mol):
     samples = num_mol*2
-    z = tf.random.normal((samples, 64))
+    z = tf.random.normal((samples, LATENT_DIM))
     graph = generator.predict(z)
     # obtain one-hot encoded adjacency tensor
     adjacency = tf.argmax(graph[0], axis=1)
@@ -59,7 +64,7 @@ def predict(num_mol):
     adjacency = tf.linalg.set_diag(adjacency, tf.zeros(tf.shape(adjacency)[:-1]))
     # obtain one-hot encoded feature tensor
     features = tf.argmax(graph[1], axis=2)
-    features = tf.one_hot(features, depth=5, axis=2)
+    features = tf.one_hot(features, depth=ATOM_DIM, axis=2)
     molecules = [
         graph_to_molecule([adjacency[i].numpy(), features[i].numpy()])
         for i in range(samples)