update

app.py

@@ -21,6 +21,7 @@ logger.addHandler(logging.NullHandler())
 
 
 MINIMIZATION_FUNCTIONS.pop("callable", None)
+MINIMIZATION_FUNCTIONS.pop("molwt", None)
 
 
 def run_inference(

model_cards/examples.csv

@@ -1 +1,3 @@
-v0|["qed"]||1.2|100|10|4|8|4|1|0.1|3|4|42
+v0|["qed"]||1.2|100|10|4|8|4|1|0.1|3|4|42
+v0|["qed","sa"]||1.2|100|10|4|8|4|1|0.1|3|4|42
+v0|["affinity"]|MVLSPADKTNVKAAWGKVGAHAGEYGAEALERMFLSFPTT|1.2|100|10|4|8|4|1|0.1|3|4|42

utils.py

@@ -3,7 +3,7 @@ from collections import defaultdict
 from typing import List, Callable
 from gt4sd.properties import PropertyPredictorRegistry
 from gt4sd.algorithms.prediction.paccmann.core import PaccMann, AffinityPredictor
-
+import torch
 
 import mols2grid
 import pandas as pd
@@ -13,15 +13,18 @@ logger.addHandler(logging.NullHandler())
 
 
 def get_affinity_function(target: str) -> Callable:
-    return lambda mols: PaccMann(
-        AffinityPredictor(protein_targets=[target] * len(mols), ligands=mols)
-    ).sample(len(mols))
+    return lambda mols: torch.stack(
+        list(
+            PaccMann(
+                AffinityPredictor(protein_targets=[target] * len(mols), ligands=mols)
+            ).sample(len(mols))
+        )
+    ).tolist()
 
 
 EVAL_DICT = {
     "qed": PropertyPredictorRegistry.get_property_predictor("qed"),
-    "sas": PropertyPredictorRegistry.get_property_predictor("sas"),
-    "molwt": PropertyPredictorRegistry.get_property_predictor("molecular_weight"),
+    "sa": PropertyPredictorRegistry.get_property_predictor("sas"),
 }