Update app.py

app.py

@@ -7,6 +7,16 @@ import tensorflow as tf
 from tensorflow import keras
 
 # Config 
+atom_mapping = {
+    "C": 0,
+    0: "C",
+    "N": 1,
+    1: "N",
+    "O": 2,
+    2: "O",
+    "F": 3,
+    3: "F",
+}
 NUM_ATOMS = 9  # Maximum number of atoms
 ATOM_DIM = 4 + 1  # Number of atom types
 BOND_DIM = 4 + 1  # Number of bond types