adding download from huggingface when the finetune_output_dir is empty

BiomedSciAI · Nov 14, 2024 · 08dfa48 · 08dfa48
1 parent d8d5893
commit 08dfa48
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 24 deletions.
diff --git a/mammal/examples/dti_bindingdb_kd/main_infer.py b/mammal/examples/dti_bindingdb_kd/main_infer.py
@@ -8,7 +8,7 @@
 
 
 @click.command()
-@click.argument("finetune_output_dir")
+@click.argument("finetune_output_dir", default="")
 @click.argument(
     "target_seq",
     default="NLMKRCTRGFRKLGKCTTLEEEKCKTLYPRGQCTCSDSKMNTHSCDCKSC",
@@ -17,8 +17,8 @@
     "drug_seq",
     default="CC(=O)NCCC1=CNc2c1cc(OC)cc2",
 )
-@click.argument("norm_y_mean", type=float)
-@click.argument("norm_y_std", type=float)
+@click.argument("norm_y_mean", default=5.79384684128215, type=float)
+@click.argument("norm_y_std", default=1.33808027428196, type=float)
 @click.option(
     "--device", default="cpu", help="Specify the device to use (default: 'cpu')."
 )
@@ -55,17 +55,24 @@ def dti_bindingdb_kd_infer(
     :param norm_y_mean: specify the mean and std values used in fine-tuning
     :param norm_y_std: specify the mean and std values used in fine-tuning
     """
-    # load tokenizer
-    tokenizer_op = ModularTokenizerOp.from_pretrained(
-        os.path.join(finetune_output_dir, "tokenizer")
-    )
-
-    # Load model
-    nn_model = Mammal.from_pretrained(
-        pretrained_model_name_or_path=os.path.join(
-            finetune_output_dir, "best_epoch.ckpt"
+    if finetune_output_dir:
+        # load tokenizer and model from finetune_output_dir
+        tokenizer_op = ModularTokenizerOp.from_pretrained(
+            os.path.join(finetune_output_dir, "tokenizer")
+        )
+        nn_model = Mammal.from_pretrained(
+            pretrained_model_name_or_path=os.path.join(
+                finetune_output_dir, "best_epoch.ckpt"
+            )
+        )
+    else:
+        # load tokenizer and model from huggingface
+        tokenizer_op = ModularTokenizerOp.from_pretrained(
+            "ibm/biomed.omics.bl.sm.ma-ted-458m.dti_bindingdb_pkd"
+        )
+        nn_model = Mammal.from_pretrained(
+            "ibm/biomed.omics.bl.sm.ma-ted-458m.dti_bindingdb_pkd"
         )
-    )
     nn_model.eval()
     nn_model.to(device=device)
 

diff --git a/mammal/examples/protein_solubility/main_infer.py b/mammal/examples/protein_solubility/main_infer.py
@@ -9,7 +9,7 @@
 
 
 @click.command()
-@click.argument("finetune_output_dir")
+@click.argument("finetune_output_dir", default="")
 @click.argument(
     "protein_seq",
     default="NLMKRCTRGFRKLGKCTTLEEEKCKTLYPRGQCTCSDSKMNTHSCDCKSC",
@@ -28,17 +28,24 @@ def protein_solubility_infer(finetune_output_dir: str, protein_seq: str, device:
     :param finetune_output_dir: model_dir argument in finetuning
     :param protein_seq: amino acid sequence of a protein
     """
-    # load tokenizer
-    tokenizer_op = ModularTokenizerOp.from_pretrained(
-        os.path.join(finetune_output_dir, "tokenizer")
-    )
-
-    # Load model
-    nn_model = Mammal.from_pretrained(
-        pretrained_model_name_or_path=os.path.join(
-            finetune_output_dir, "best_epoch.ckpt"
+    if finetune_output_dir:
+        # load tokenizer and model from finetune_output_dir
+        tokenizer_op = ModularTokenizerOp.from_pretrained(
+            os.path.join(finetune_output_dir, "tokenizer")
+        )
+        nn_model = Mammal.from_pretrained(
+            pretrained_model_name_or_path=os.path.join(
+                finetune_output_dir, "best_epoch.ckpt"
+            )
+        )
+    else:
+        # load tokenizer and model from huggingface
+        tokenizer_op = ModularTokenizerOp.from_pretrained(
+            "ibm/biomed.omics.bl.sm.ma-ted-458m.protein_solubility"
+        )
+        nn_model = Mammal.from_pretrained(
+            "ibm/biomed.omics.bl.sm.ma-ted-458m.protein_solubility"
         )
-    )
     nn_model.eval()
     nn_model.to(device=device)