0.5.0 senseppi.ckpt added to the package as a default model

senseppi/__init__.py

-__version__ = "0.4.1"
+__version__ = "0.5.0"
 __author__ = "Konstantin Volzhenin"
 
 from . import model, commands, esm2_model, dataset, utils, network_utils

senseppi/commands/predict.py

@@ -65,11 +65,12 @@ def add_args(parser):
     parser = add_general_args(parser)
 
     predict_args = parser.add_argument_group(title="Predict args")
-    parser._action_groups[0].add_argument("model_path", type=str,
-                                          help="A path to .ckpt file that contains weights to a pretrained model.")
     parser._action_groups[0].add_argument("fasta_file", type=pathlib.Path,
                                           help="FASTA file on which to extract the ESM2 representations and then test.",
                                           )
+    predict_args.add_argument("--model_path", type=str, default=os.path.join(os.path.dirname(__file__), "..", "senseppi.ckpt"),
+                                          help="A path to .ckpt file that contains weights to a pretrained model. If "
+                                               "None, the senseppi trained version is used.")
     predict_args.add_argument("--pairs_file", type=str, default=None,
                               help="A path to a .tsv file with pairs of proteins to test (Optional). If not provided, "
                                    "all-to-all pairs will be generated.")

senseppi/commands/predict_string.py

@@ -184,11 +184,12 @@ def add_args(parser):
     parser = add_general_args(parser)
 
     string_pred_args = parser.add_argument_group(title="General options")
-    parser._action_groups[0].add_argument("model_path", type=str,
-                                          help="A path to .ckpt file that contains weights to a pretrained model.")
     parser._action_groups[0].add_argument("genes", type=str, nargs="+",
                                           help="Name of gene to fetch from STRING database. Several names can be "
                                                "typed (separated by whitespaces).")
+    string_pred_args.add_argument("--model_path", type=str, default=os.path.join(os.path.dirname(__file__), "..", "senseppi.ckpt"),
+                                  help="A path to .ckpt file that contains weights to a pretrained model. If "
+                                       "None, the senseppi trained version is used.")
     string_pred_args.add_argument("-s", "--species", type=int, default=9606,
                                   help="Species from STRING database. Default: 9606 (H. Sapiens)")
     string_pred_args.add_argument("-n", "--nodes", type=int, default=10,

senseppi/commands/test.py

@@ -37,8 +37,6 @@ def add_args(parser):
     parser = add_general_args(parser)
 
     test_args = parser.add_argument_group(title="Predict args")
-    parser._action_groups[0].add_argument("model_path", type=str,
-                                          help="A path to .ckpt file that contains weights to a pretrained model.")
     parser._action_groups[0].add_argument("pairs_file", type=str, default=None,
                                           help="A path to a .tsv file with pairs of proteins to test.")
     parser._action_groups[0].add_argument("fasta_file",
@@ -46,12 +44,15 @@ def add_args(parser):
                                           help="FASTA file on which to extract the ESM2 "
                                                "representations and then evaluate.",
                                           )
+    test_args.add_argument("--model_path", type=str, default=os.path.join(os.path.dirname(__file__), "..", "senseppi.ckpt"),
+                           help="A path to .ckpt file that contains weights to a pretrained model. If "
+                                "None, the senseppi trained version is used.")
     test_args.add_argument("-o", "--output", type=str, default="test_metrics",
-                              help="A path to a file where the test metrics will be saved. "
-                                   "(.tsv format will be added automatically)")
+                           help="A path to a file where the test metrics will be saved. "
+                                "(.tsv format will be added automatically)")
     test_args.add_argument("--crop_data_to_model_lims", action="store_true",
-                              help="If set, the data will be cropped to the limits of the model: "
-                                   "evaluations will be done only for proteins >50aa and <800aa.")
+                           help="If set, the data will be cropped to the limits of the model: "
+                                "evaluations will be done only for proteins >50aa and <800aa.")
 
     parser = SensePPIModel.add_model_specific_args(parser)
     remove_argument(parser, "--lr")