Update app.py

app.py

@@ -136,8 +136,36 @@ def graphs_from_smiles(smiles_list):
         tf.ragged.constant(bond_features_list, dtype=tf.float32),
         tf.ragged.constant(pair_indices_list, dtype=tf.int64),
     )
-    
-    
+  
+      
+def prepare_batch(x_batch, y_batch):
+    """Merges (sub)graphs of batch into a single global (disconnected) graph
+    """
+
+    atom_features, bond_features, pair_indices = x_batch
+
+    # Obtain number of atoms and bonds for each graph (molecule)
+    num_atoms = atom_features.row_lengths()
+    num_bonds = bond_features.row_lengths()
+
+    # Obtain partition indices (molecule_indicator), which will be used to
+    # gather (sub)graphs from global graph in model later on
+    molecule_indices = tf.range(len(num_atoms))
+    molecule_indicator = tf.repeat(molecule_indices, num_atoms)
+
+    # Merge (sub)graphs into a global (disconnected) graph. Adding 'increment' to
+    # 'pair_indices' (and merging ragged tensors) actualizes the global graph
+    gather_indices = tf.repeat(molecule_indices[:-1], num_bonds[1:])
+    increment = tf.cumsum(num_atoms[:-1])
+    increment = tf.pad(tf.gather(increment, gather_indices), [(num_bonds[0], 0)])
+    pair_indices = pair_indices.merge_dims(outer_axis=0, inner_axis=1).to_tensor()
+    pair_indices = pair_indices + increment[:, tf.newaxis]
+    atom_features = atom_features.merge_dims(outer_axis=0, inner_axis=1).to_tensor()
+    bond_features = bond_features.merge_dims(outer_axis=0, inner_axis=1).to_tensor()
+
+    return (atom_features, bond_features, pair_indices, molecule_indicator), y_batch
+
+        
 def MPNNDataset(X, y, batch_size=32, shuffle=False):
     dataset = tf.data.Dataset.from_tensor_slices((X, (y)))
     if shuffle: