jannisborn commited on
Commit
1f4fe6e
1 Parent(s): a0a4d74
Files changed (1) hide show
  1. app.py +12 -4
app.py CHANGED
@@ -12,8 +12,8 @@ logger = logging.getLogger(__name__)
12
  logger.addHandler(logging.NullHandler())
13
 
14
  MOLFORMER_VERSIONS = {
15
- "Molformer_classification": ["bace", "bbbp", "hiv"],
16
- "Molformer_regression": [
17
  "alpha",
18
  "cv",
19
  "g298",
@@ -26,7 +26,7 @@ MOLFORMER_VERSIONS = {
26
  "r2",
27
  "u0",
28
  ],
29
- "Molformer_multitask_classification": ["clintox", "sider", "tox21"],
30
  }
31
 
32
  REMOVE = ["docking", "docking_tdc", "molecule_one", "askcos", "plogp"]
@@ -42,10 +42,16 @@ MODEL_PROP_DESCRIPTION = {
42
 
43
  def main(property: str, smiles: str, smiles_file: str):
44
 
 
 
 
45
  algo, config = MOLECULE_PROPERTY_PREDICTOR_FACTORY[property.lower()]
46
  kwargs = (
47
  {"algorithm_version": "v0"} if property in MODEL_PROP_DESCRIPTION.keys() else {}
48
  )
 
 
 
49
  model = algo(config(**kwargs))
50
  if smiles is not None and smiles_file is not None:
51
  raise ValueError("Pass either smiles or smiles_file, not both.")
@@ -80,7 +86,9 @@ if __name__ == "__main__":
80
 
81
  # MolFormer options
82
  for key in MOLFORMER_VERSIONS.keys():
83
- properties.extend([f"{key} ({version})" for version in MOLFORMER_VERSIONS[key]])
 
 
84
 
85
  # Load metadata
86
  metadata_root = pathlib.Path(__file__).parent.joinpath("model_cards")
 
12
  logger.addHandler(logging.NullHandler())
13
 
14
  MOLFORMER_VERSIONS = {
15
+ "molformer_classification": ["bace", "bbbp", "hiv"],
16
+ "molformer_regression": [
17
  "alpha",
18
  "cv",
19
  "g298",
 
26
  "r2",
27
  "u0",
28
  ],
29
+ "molformer_multitask_classification": ["clintox", "sider", "tox21"],
30
  }
31
 
32
  REMOVE = ["docking", "docking_tdc", "molecule_one", "askcos", "plogp"]
 
42
 
43
  def main(property: str, smiles: str, smiles_file: str):
44
 
45
+ if "Molformer" in property:
46
+ version = property.split(" ")[-1].split("(")[-1].split(")")[0]
47
+ property = property.split(" ")[0]
48
  algo, config = MOLECULE_PROPERTY_PREDICTOR_FACTORY[property.lower()]
49
  kwargs = (
50
  {"algorithm_version": "v0"} if property in MODEL_PROP_DESCRIPTION.keys() else {}
51
  )
52
+ if property in MOLFORMER_VERSIONS.keys():
53
+ kwargs["algorithm_version"] = version
54
+
55
  model = algo(config(**kwargs))
56
  if smiles is not None and smiles_file is not None:
57
  raise ValueError("Pass either smiles or smiles_file, not both.")
 
86
 
87
  # MolFormer options
88
  for key in MOLFORMER_VERSIONS.keys():
89
+ properties.extend(
90
+ [f"{key.capitalize()} ({version})" for version in MOLFORMER_VERSIONS[key]]
91
+ )
92
 
93
  # Load metadata
94
  metadata_root = pathlib.Path(__file__).parent.joinpath("model_cards")