fix a bug in TDC-huggingface

mims-harvard · Aug 2, 2023 · 6af2a41 · 6af2a41
1 parent 0390be5
commit 6af2a41
Showing 1 changed file with 10 additions and 3 deletions.
diff --git a/tdc/tdc_hf.py b/tdc/tdc_hf.py
@@ -3,8 +3,15 @@
 import os
 
 deeppurpose_repo = [
-    'hERG_Karim_Morgan',
-    'BBB_Martins-AttentiveFP'
+    'hERG_Karim-Morgan',
+    'hERG_Karim-CNN', 
+    'hERG_Karim-AttentiveFP', 
+    'BBB_Martins-AttentiveFP',
+    'BBB_Martins-Morgan', 
+    'BBB_Martins-CNN', 
+    'CYP3A4_Veith-Morgan', 
+    'CYP3A4_Veith-CNN', 
+    'CYP3A4_Veith-AttentiveFP', 
 ]
 
 class tdc_hf_interface:
@@ -82,4 +89,4 @@ def predict_deeppurpose(self, model, drugs):
                                 drug_encoding = self.model_name, 
                                 split_method='no_split')
         y_pred = model.predict(X_pred)[0]
-        return y_pred
+        return y_pred