First QProgressBars

Author: niklases

pypef/__init__.py

@@ -1,4 +1,4 @@
 # PyPEF - Pythonic Protein Engineering Framework
 # https://github.com/niklases/PyPEF
 
-__version__ = '0.4.3'
+__version__ = '0.4.4-dev'

pypef/dca/gremlin_inference.py

@@ -49,7 +49,7 @@
 from tqdm import tqdm
 import torch
 
-from pypef.llm.utils import get_batches
+from pypef.plm.utils import get_batches
 from pypef.utils.variant_data import get_mismatches
 
 

pypef/gui/qt_window.py

@@ -37,10 +37,10 @@
 from pypef.utils.plot import plot_y_true_vs_y_pred
 import pypef.dca.gremlin_inference
 from pypef.dca.gremlin_inference import GREMLIN, get_delta_e_statistical_model
-from pypef.llm.esm_lora_tune import esm_setup, get_esm_models
-from pypef.llm.prosst_lora_tune import get_prosst_models, prosst_setup
-from pypef.llm.inference import llm_embedder, inference
-from pypef.llm.utils import get_batches
+from pypef.plm.esm_lora_tune import esm_setup, get_esm_models
+from pypef.plm.prosst_lora_tune import get_prosst_models, prosst_setup
+from pypef.plm.inference import llm_embedder, inference
+from pypef.plm.utils import get_batches
 
 # sklearn/base.py:474: FutureWarning: `BaseEstimator._validate_data` is deprecated in 1.6 and 
 # will be removed in 1.7. Use `sklearn.utils.validation.validate_data` instead. This function 
@@ -70,7 +70,9 @@ def __init__(
             llm_train: bool = True,
             device: str | None = None,
             seed: int | None = None,
-            verbose: bool = True
+            verbose: bool = True,
+            progress_cb=None, 
+            abort_cb=None
     ):
         if llm_model_input is not None:
             if type(llm_model_input) is not dict:
@@ -141,6 +143,8 @@ def __init__(
             self.y_llm_ttest,
             self.y_llm_lora_ttest
         ) = None, None, None, None, None, None, None, None, None
+        self.progress_cb = progress_cb
+        self.abort_cb = abort_cb
         self.train_and_optimize()
 
     @staticmethod
@@ -465,7 +469,9 @@ def train_llm(self):
                 n_epochs=50,
                 device=self.device,
                 verbose=self.verbose,
-                raise_error_on_train_fail=False
+                raise_error_on_train_fail=False,
+                progress_cb=self.progress_cb, 
+                abort_cb=self.abort_cb
             )
             y_llm_lora_ttrain = self.llm_inference_function(
                 xs=self.x_llm_ttrain,
@@ -496,7 +502,9 @@ def train_llm(self):
                 self.llm_optimizer,  
                 n_epochs=5, 
                 device=self.device,
-                verbose=self.verbose
+                verbose=self.verbose,
+                progress_cb=self.progress_cb, 
+                abort_cb=self.abort_cb
             )
             y_llm_lora_ttrain = self.llm_inference_function(
                 xs=x_llm_ttrain_b,
@@ -802,6 +810,8 @@ def save_model_to_dict_pickle(
             model.llm_base_model = model.llm_base_model.state_dict()
             model.llm_model_input[model.llm_key]['llm_base_model'] = None
             model.llm_model_input[model.llm_key]['llm_model'] = None
+            model.progress_cb = None
+            model.abort_cb = None
             model_type += model.llm_key.upper()
     pkl_path = os.path.abspath(f'Pickles/{model_type.upper()}')
     pickle.dump(
@@ -986,7 +996,9 @@ def performance_ls_ts(
         wt_seq: str | None = None,
         substitution_sep: str = '/',
         label=False,
-        device: str| None = None
+        device: str| None = None,
+        progress_cb=None, 
+        abort_cb=None
 ):
     test_sequences, test_variants, y_test = get_sequences_from_file(ts_fasta)
 
@@ -1032,7 +1044,9 @@ def performance_ls_ts(
             y_train=np.array(y_train),
             llm_model_input=llm_dict,
             x_wt=x_wt,
-            device=device
+            device=device,
+            progress_cb=progress_cb, 
+            abort_cb=abort_cb
         )
         y_test_pred = hybrid_model.hybrid_prediction(np.array(x_test), x_llm_test)
         logger.info(f'Hybrid performance: {spearmanr(y_test, y_test_pred)[0]:.3f} N={len(y_test)}')

pypef/hybrid/hybrid_model.py

@@ -15,7 +15,7 @@
 from pypef.utils.low_n_mutation_extrapolation import performance_mutation_extrapolation, low_n
 
 
-def run_pypef_hybrid_modeling(arguments):
+def run_pypef_hybrid_modeling(arguments, progress_cb=None, abort_cb=None):
     threads = abs(arguments['--threads']) if arguments['--threads'] is not None else 1
     threads = threads + 1 if threads == 0 else threads
     if arguments['--params'] is not None:
@@ -52,7 +52,9 @@ def run_pypef_hybrid_modeling(arguments):
             pdb_file=arguments['--pdb'],
             wt_seq=get_wt_sequence(arguments['--wt']),
             substitution_sep=arguments['--mutation_sep'],
-            label=arguments['--label']
+            label=arguments['--label'],
+            progress_cb=progress_cb, 
+            abort_cb=abort_cb
         )
 
     elif arguments['--params'] and arguments['--model'] or arguments['--ps']:

pypef/hybrid/hybrid_run.py

@@ -417,7 +417,7 @@ def validate(args):
         exit(e)
 
 
-def run_main(argv=None):
+def run_main(argv=None, progress_cb=None, abort_cb=None):
     """
     Entry point for pip-installed version.
     Arguments are created from Docstring using docopt that 
@@ -434,7 +434,7 @@ def run_main(argv=None):
     elif arguments['ml']:
         run_pypef_pure_ml(arguments)
     elif arguments['hybrid'] or arguments['param_inference'] or arguments['save_msa_info']:
-        run_pypef_hybrid_modeling(arguments)
+        run_pypef_hybrid_modeling(arguments, progress_cb=progress_cb, abort_cb=abort_cb)
     else:
         run_pypef_utils(arguments)
 

pypef/main.py

@@ -30,7 +30,7 @@
 hf_logging.set_verbosity_error()
 
 from pypef.utils.helpers import get_device
-from pypef.llm.utils import corr_loss, load_model_and_tokenizer
+from pypef.plm.utils import corr_loss, load_model_and_tokenizer
 
 
 def get_esm_models():
@@ -143,7 +143,8 @@ def esm_infer(xs, attention_mask, model, device: str | None = None, verbose=Fals
 def esm_train(
         xs, attention_mask, scores, loss_fn, model, optimizer, n_epochs=3, 
         device: str | None = None, seed: int | None = None, 
-        n_batch_grad_accumulations: int = 1, verbose: bool = True
+        n_batch_grad_accumulations: int = 1, verbose: bool = True,
+        progress_cb=None, abort_cb=None
 ):
     if seed is not None:
         torch.manual_seed(seed)
@@ -157,6 +158,7 @@ def esm_train(
     xs, attention_masks, scores = xs.to(device), attention_masks.to(device), scores.to(device) 
     pbar_epochs = tqdm(range(1, n_epochs + 1), disable=not verbose)
     loss = np.nan
+    logger.info(progress_cb)  # TODO: delete
     for epoch in pbar_epochs:
         try:
             pbar_epochs.set_description(f'Epoch: {epoch}/{n_epochs}. Loss: {loss.detach():>1f}')
@@ -171,6 +173,8 @@ def esm_train(
             xs_b, attns_b = xs_b.to(torch.int64), attns_b.to(torch.int64)
             y_preds_b = get_y_pred_scores(xs_b, attns_b, model, device=device)
             loss = loss_fn(scores_b, y_preds_b) / n_batch_grad_accumulations
+            if progress_cb:
+                progress_cb(epoch - 1, batch + 1, len(pbar_epochs), len(pbar_batches), loss)
             loss.backward()
             if (batch + 1) % n_batch_grad_accumulations == 0 or (batch + 1) == len(pbar_batches):
                 optimizer.step()
@@ -180,6 +184,8 @@ def esm_train(
                 f"[batch: {batch+1}/{len(xs)} | sequence: "
                 f"{(batch + 1) * len(xs_b):>5d}/{len(xs) * len(xs_b)}] ({device.upper()})"
             )
+    if progress_cb:
+        progress_cb(epoch, batch + 1, len(pbar_epochs), len(pbar_batches), loss)
     y_preds_b = y_preds_b.detach()
     model.train(False)
 

pypef/llm/__init__.py pypef/plm/__init__.pypypef/llm/__init__.py renamed to pypef/plm/__init__.py

@@ -7,9 +7,9 @@
 import numpy as np
 
 from pypef.utils.helpers import get_device
-from pypef.llm.utils import get_batches
-from pypef.llm.esm_lora_tune import esm_setup, esm_tokenize_sequences, esm_infer
-from pypef.llm.prosst_lora_tune import prosst_setup, prosst_tokenize_sequences, prosst_infer
+from pypef.plm.utils import get_batches
+from pypef.plm.esm_lora_tune import esm_setup, esm_tokenize_sequences, esm_infer
+from pypef.plm.prosst_lora_tune import prosst_setup, prosst_tokenize_sequences, prosst_infer
 
 import logging
 logger = logging.getLogger('pypef.llm.inference')

pypef/llm/esm_lora_tune.py pypef/plm/esm_lora_tune.pypypef/llm/esm_lora_tune.py renamed to pypef/plm/esm_lora_tune.py

@@ -11,7 +11,7 @@
 
 import logging
 
-from pypef.llm.utils import load_model_and_tokenizer
+from pypef.plm.utils import load_model_and_tokenizer
 logger = logging.getLogger('pypef.llm.prosst_lora_tune')
 
 import os
@@ -25,8 +25,8 @@
 from Bio import SeqIO, BiopythonParserWarning
 warnings.filterwarnings(action='ignore', category=BiopythonParserWarning)
 
-from pypef.llm.esm_lora_tune import corr_loss
-from pypef.llm.prosst_structure.quantizer import PdbQuantizer
+from pypef.plm.esm_lora_tune import corr_loss
+from pypef.plm.prosst_structure.quantizer import PdbQuantizer
 from pypef.utils.helpers import get_device
 
 
@@ -139,7 +139,9 @@ def prosst_train(
         input_ids, attention_mask, structure_input_ids,
         n_epochs=50, device: str | None = None, seed: int | None = None,
         early_stop: int = 50, verbose: bool = True, 
-        n_batch_grad_accumulations: int = 1, raise_error_on_train_fail: bool = True):
+        n_batch_grad_accumulations: int = 1, raise_error_on_train_fail: bool = True,
+        progress_cb=None, abort_cb=None
+):
     if seed is not None:
         torch.manual_seed(seed)
     if device is None:
@@ -154,6 +156,7 @@ def prosst_train(
     best_model = None
     best_model_epoch = np.nan
     best_model_perf = np.nan
+    loss = np.nan
     os.makedirs('model_saves', exist_ok=True)
     for epoch in pbar_epochs:
         if epoch == 0:
@@ -171,6 +174,8 @@ def prosst_train(
             )
             y_preds_detached.append(y_preds_b.detach().cpu().numpy().flatten())
             loss = loss_fn(scores_b, y_preds_b) / n_batch_grad_accumulations
+            if progress_cb:
+                progress_cb(epoch - 1, batch + 1, len(pbar_epochs), len(pbar_batches), loss)
             loss.backward()
             if (batch + 1) % n_batch_grad_accumulations == 0 or (batch + 1) == len(pbar_batches):
                 optimizer.step()
@@ -215,6 +220,8 @@ def prosst_train(
         pbar_epochs.set_description(
             f'Epoch {epoch}/{n_epochs} [SpearCorr: {epoch_spearman_2:.3f}, Loss: {loss_total:.3f}] '
             f'(Best epoch: {best_model_epoch}: {best_model_perf:.3f})')
+    if progress_cb:
+        progress_cb(epoch, batch + 1, len(pbar_epochs), len(pbar_batches), loss)
     if best_model is None:
         msg = ("Failed to train a model (probably due to the input "
                "data characteristics and loss/correlation being NaN).")