Update app.py

app.py

@@ -137,8 +137,17 @@ def graphs_from_smiles(smiles_list):
         tf.ragged.constant(pair_indices_list, dtype=tf.int64),
     )
     
+    
+def MPNNDataset(X, y, batch_size=32, shuffle=False):
+    dataset = tf.data.Dataset.from_tensor_slices((X, (y)))
+    if shuffle:
+        dataset = dataset.shuffle(1024)
+    return dataset.batch(batch_size).map(prepare_batch, -1).prefetch(-1)
+   
+       
 model = from_pretrained_keras("keras-io/MPNN-for-molecular-property-prediction")
 
+
 def predict(smiles, label):
     molecules = [molecule_from_smiles(smiles)]
     input = graphs_from_smiles([smiles])