Merge pull request #34 from IBM/code-dev

🛠️🪲📝: Improvements to `prediction` mode

README.md

@@ -218,6 +218,13 @@ autopeptideml dataset.csv
 autopeptideml dataset.csv --val_method graph-part --val_threshold 0.3 --val_alignment needle
 ```
 
+For running predictions
+
+```bash
+
+autopeptideml-predict peptides.csv --ensemble AB_1/ensemble
+```
+
 
 
 License <a name="license"></a>

autopeptideml/main.py

@@ -65,6 +65,8 @@ def parse_cli_predict():
     parser.add_argument('--plm', type=str, default='esm2-8m',
                         help='PLM for computing peptide representations. Check GitHub Repository for available options.')
     parser.add_argument('--plm_batch_size', type=int, default=12)
+    parser.add_argument('--device', type=str, default=None,
+                        help='Device where the representations should be computed.')
     return parser.parse_args()
 
 
@@ -154,6 +156,8 @@ def predict():
     args = parse_cli_predict()
 
     re = RepresentationEngine(args.plm, args.plm_batch_size)
+    if args.device is not None:
+        re.move_to_device(args.device)
     apml = AutoPeptideML(args.verbose, args.threads, 1)
     df = apml.curate_dataset(args.dataset, args.outputdir)
     apml.predict(df, re, args.ensemble, args.outputdir)

setup.py

@@ -37,7 +37,7 @@
     entry_points={
         'console_scripts': [
             'autopeptideml=autopeptideml.main:main',
-            'autopeptdeml-predict=autopeptideml.main:predict',
+            'autopeptideml-predict=autopeptideml.main:predict',
             'autopeptideml-setup=autopeptideml.data.preprocess_db:main'
         ],
     },