Fix model is None and add verbosity parameters to many functions

Author: niklases

pypef/hybrid/hybrid_model.py

@@ -399,7 +399,6 @@ def train_llm(self):
                 input_ids=self.input_ids,
                 attention_mask=self.llm_attention_mask,
                 structure_input_ids=self.structure_input_ids,
-                train=False,
                 device=self.device
             )
             y_llm_ttrain = self.llm_inference_function(
@@ -408,7 +407,6 @@ def train_llm(self):
                 input_ids=self.input_ids,
                 attention_mask=self.llm_attention_mask,
                 structure_input_ids=self.structure_input_ids,
-                train=False,
                 device=self.device
             )
         elif self.llm_key == 'esm1v':
@@ -598,7 +596,6 @@ def hybrid_prediction(
                     self.input_ids,
                     self.llm_attention_mask, 
                     self.structure_input_ids,
-                    train=False,
                     verbose=verbose,
                     device=self.device).detach().cpu().numpy()
                 y_llm_lora = self.llm_inference_function(
@@ -607,7 +604,6 @@ def hybrid_prediction(
                     self.input_ids,
                     self.llm_attention_mask, 
                     self.structure_input_ids,
-                    train=False,
                     verbose=verbose,
                     device=self.device).detach().cpu().numpy()
             elif self.llm_key == 'esm1v':
@@ -1004,7 +1000,8 @@ def performance_ls_ts(
             x_train_dca=np.array(x_train),
             y_train=np.array(y_train),
             llm_model_input=llm_dict,
-            x_wt=x_wt
+            x_wt=x_wt,
+            device=device
         )
         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/llm/esm_lora_tune.py

@@ -81,7 +81,8 @@ def get_y_pred_scores(encoded_sequences, attention_masks,
     return log_probs
 
 
-def esm_test(xs, attention_mask, scores, loss_fn, model, device: str | None = None, verbose: bool = True):
+def esm_test(xs, attention_mask, scores, loss_fn, model, 
+             device: str | None = None, verbose: bool = True):
     if device is None:
         device = get_device()
     attention_masks = torch.Tensor(np.full(
@@ -148,15 +149,18 @@ def esm_train(xs, attention_mask, scores, loss_fn, model, optimizer, n_epochs=3,
     attention_masks = torch.Tensor(np.full(
         shape=np.shape(xs), fill_value=attention_mask)).to(torch.int64)
     xs, attention_masks, scores = xs.to(device), attention_masks.to(device), scores.to(device) 
-    pbar_epochs = tqdm(range(1, n_epochs + 1))
+    pbar_epochs = tqdm(range(1, n_epochs + 1), disable=not verbose)
     loss = np.nan
     for epoch in pbar_epochs:
         try:
             pbar_epochs.set_description(f'Epoch: {epoch}/{n_epochs}. Loss: {loss.detach():>1f}')
         except AttributeError:
             pbar_epochs.set_description(f'Epoch: {epoch}/{n_epochs}')
         model.train()
-        pbar_batches = tqdm(zip(xs, attention_masks, scores), total=len(xs), leave=False, disable=not verbose)
+        pbar_batches = tqdm(
+            zip(xs, attention_masks, scores), 
+            total=len(xs), leave=False, disable=not verbose
+        )
         for batch, (xs_b, attns_b, scores_b) in enumerate(pbar_batches):
             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)
@@ -173,11 +177,11 @@ def esm_train(xs, attention_mask, scores, loss_fn, model, optimizer, n_epochs=3,
     model.train(False)
 
 
-def esm_setup(sequences, device: str | None = None):
+def esm_setup(sequences, device: str | None = None, verbose: bool = True):
     esm_base_model, esm_lora_model, esm_tokenizer, esm_optimizer = get_esm_models()
     esm_base_model = esm_base_model.to(device)
     x_esm, esm_attention_mask = esm_tokenize_sequences(
-        sequences, esm_tokenizer, max_length=len(sequences[0]))
+        sequences, esm_tokenizer, max_length=len(sequences[0]), verbose=verbose)
     llm_dict_esm = {
         'esm1v': {
             'llm_base_model': esm_base_model,

pypef/llm/inference.py

@@ -8,8 +8,8 @@
 
 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
-from pypef.llm.prosst_lora_tune import prosst_setup, prosst_tokenize_sequences
+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
 
 import logging
 logger = logging.getLogger('pypef.llm.inference')
@@ -40,38 +40,42 @@ def inference(
         pdb_file: str | None = None,
         wt_seq: str | None = None,
         device: str| None = None,
-        model = None
+        model = None,
+        verbose: bool = True
 ):
     """
     Inference of input or base model.
     """
     if device is None:
         device = get_device()
+    if llm == 'esm':
         logger.info("Zero-shot LLM inference on test set using ESM1v...")
-        llm_dict = esm_setup(sequences)
-        if llm == 'esm':
-            if model is None:
-                model = llm_dict['esm1v']['llm_base_model']
-        x_llm_test = llm_embedder(llm_dict, sequences)
-        y_test_pred = llm_dict['esm1v']['llm_inference_function'](
+        llm_dict = esm_setup(sequences, verbose=verbose)
+        if model is None:
+            model = llm_dict['esm1v']['llm_base_model']
+        x_llm_test = llm_embedder(llm_dict, sequences, verbose)
+        y_test_pred = esm_infer(#llm_dict['esm1v']['llm_inference_function'](
             xs=get_batches(x_llm_test, batch_size=1, dtype=int), 
             attention_mask=llm_dict['esm1v']['llm_attention_mask'], 
             model=model, 
-            device=device
+            device=device,
+            verbose=verbose
         ).cpu()
     elif llm == 'prosst':
-        if model is None:
-            model = llm_dict['prosst']['llm_base_model']
         logger.info("Zero-shot LLM inference on test set using ProSST...")
         llm_dict = prosst_setup(
-            wt_seq, pdb_file, sequences=sequences)
-        x_llm_test = llm_embedder(llm_dict, sequences)
-        y_test_pred = llm_dict['prosst']['llm_inference_function'](
+            wt_seq, pdb_file, sequences=sequences, verbose=verbose
+        )
+        if model is None:
+            model = llm_dict['prosst']['llm_base_model']
+        x_llm_test = llm_embedder(llm_dict, sequences, verbose)
+        y_test_pred = prosst_infer(#llm_dict['prosst']['llm_inference_function'](
             xs=x_llm_test, 
             model=model, 
             input_ids=llm_dict['prosst']['input_ids'], 
             attention_mask=llm_dict['prosst']['llm_attention_mask'], 
             structure_input_ids=llm_dict['prosst']['structure_input_ids'],
+            verbose=verbose,
             device=device
         ).cpu()
     else:

pypef/llm/prosst_lora_tune.py

@@ -94,6 +94,27 @@ def get_logits_from_full_seqs(
     return log_probs
 
 
+def prosst_infer(
+        xs, 
+        model, 
+        input_ids, 
+        attention_mask, 
+        structure_input_ids,
+        verbose: bool = False,
+        device: str | None = None
+):
+    return get_logits_from_full_seqs(
+        xs, 
+        model, 
+        input_ids, 
+        attention_mask, 
+        structure_input_ids,
+        train = False,
+        verbose = verbose,
+        device = device
+    )
+
+
 def checkpoint(model, filename):
     torch.save(model.state_dict(), filename)
 
@@ -205,8 +226,8 @@ def get_prosst_models():
     return prosst_base_model, prosst_lora_model, tokenizer, optimizer
 
 
-def get_structure_quantizied(pdb_file, tokenizer, wt_seq):
-    structure_sequence = PdbQuantizer()(pdb_file=pdb_file)
+def get_structure_quantizied(pdb_file, tokenizer, wt_seq, verbose: bool = True):
+    structure_sequence = PdbQuantizer(verbose=verbose)(pdb_file=pdb_file)
     structure_sequence_offset = [i + 3 for i in structure_sequence]
     tokenized_res = tokenizer([wt_seq], return_tensors='pt')
     input_ids = tokenized_res['input_ids']
@@ -216,7 +237,7 @@ def get_structure_quantizied(pdb_file, tokenizer, wt_seq):
     return input_ids, attention_mask, structure_input_ids
 
 
-def prosst_setup(wt_seq, pdb_file, sequences, device: str | None = None):
+def prosst_setup(wt_seq, pdb_file, sequences, device: str | None = None, verbose: bool = True):
     if wt_seq is None:
         raise SystemError(
             "Running ProSST requires a wild-type sequence "
@@ -240,16 +261,18 @@ def prosst_setup(wt_seq, pdb_file, sequences, device: str | None = None):
     prosst_base_model = prosst_base_model.to(device)
     prosst_optimizer = torch.optim.Adam(prosst_lora_model.parameters(), lr=0.0001)
     input_ids, prosst_attention_mask, structure_input_ids = get_structure_quantizied(
-        pdb_file, prosst_tokenizer, wt_seq)
+        pdb_file, prosst_tokenizer, wt_seq, verbose=verbose
+    )
     x_llm_train_prosst = prosst_tokenize_sequences(
-        sequences=sequences, vocab=prosst_vocab)
+        sequences=sequences, vocab=prosst_vocab, verbose=verbose
+    )
     llm_dict_prosst = {
         'prosst': {
             'llm_base_model': prosst_base_model,
             'llm_model': prosst_lora_model,
             'llm_optimizer': prosst_optimizer,
             'llm_train_function': prosst_train,
-            'llm_inference_function': get_logits_from_full_seqs,
+            'llm_inference_function': prosst_infer,
             'llm_loss_function': corr_loss,
             'x_llm' : x_llm_train_prosst,
             'llm_attention_mask': prosst_attention_mask,

pypef/llm/prosst_structure/quantizer.py

@@ -408,7 +408,8 @@ def process_pdb_file(
     pdb_file,
     subgraph_depth,
     subgraph_interval,
-    max_distance
+    max_distance,
+    verbose: bool = True
 ):
     result_dict, subgraph_dict = {}, {}
     result_dict["name"] = Path(pdb_file).name
@@ -436,7 +437,7 @@ def process_subgraph(anchor_node):
         subgraph = convert_graph(subgraph)
         return anchor_node, subgraph
 
-    for anchor_node in tqdm(anchor_nodes, desc='Getting ProSST structure embeddings'):
+    for anchor_node in tqdm(anchor_nodes, desc='Getting ProSST structure embeddings', disable=not verbose):
          anchor, subgraph = process_subgraph(anchor_node)
          subgraph_dict[anchor] = subgraph
 
@@ -449,7 +450,8 @@ def pdb_conventer(
     pdb_files,
     subgraph_depth,
     subgraph_interval,
-    max_distance
+    max_distance,
+    verbose: bool = True
 ):
     error_proteins, error_messages = [], []
     dataset, results, node_counts = [], [], []
@@ -460,6 +462,7 @@ def pdb_conventer(
             subgraph_depth,
             subgraph_interval,
             max_distance,
+            verbose=verbose
         )
 
         if pdb_subgraphs is None:
@@ -502,7 +505,8 @@ def __init__(
         model_path=None,
         cluster_dir=None,
         cluster_model=None,
-        device=None
+        device=None,
+        verbose: bool = True
     ) -> None:
         self.max_distance = max_distance
         self.subgraph_depth = subgraph_depth
@@ -512,6 +516,7 @@ def __init__(
             self.device = get_device()
         else:
             self.device = device
+        self.verbose = verbose
         if model_path is None:
             if self.device == 'cpu':
                 self.model_path = str(Path(__file__).parent / "static" / "AE_CPU.pt")
@@ -554,7 +559,8 @@ def __call__(self, pdb_file, return_residue_seq=False):
             ],
             self.subgraph_depth,
             self.subgraph_interval,
-            self.max_distance
+            self.max_distance,
+            verbose=self.verbose
         )
         sturctures = predict_structure(
             self.model, self.cluster_models, data_loader, self.device

scripts/ProteinGym_runs/protgym_hybrid_perf_test_crossval.py

@@ -140,13 +140,13 @@ def compute_performances(mut_data, mut_sep=':', start_i: int = 0, already_tested
             gremlin = GREMLIN(alignment=msa_path, opt_iter=100, optimize=True)
             sequences_batched = get_batches(sequences, batch_size=1000, 
                                             dtype=str, keep_remaining=True, verbose=True)
-            x_dca = []
-            for seq_b in tqdm(sequences_batched, desc="Getting GREMLIN sequence encodings"):
+            x_dca = []  # required later on also
+            for seq_b in tqdm(sequences_batched, desc="Getting GREMLIN sequence encodings", disable=True):
                 for x in gremlin.collect_encoded_sequences(seq_b):
                     x_dca.append(x)
             x_wt = gremlin.x_wt
             y_pred_dca = get_delta_e_statistical_model(x_dca, x_wt)
-            print(f'DCA (unsupervised performance): {spearmanr(fitnesses, y_pred_dca)[0]:.3f}') 
+            print(f'DCA (unsupervised performance): {spearmanr(fitnesses, y_pred_dca)[0]:.3f}')
             dca_unopt_perf = spearmanr(fitnesses, y_pred_dca)[0]
             # ESM unsupervised
             try:
@@ -158,7 +158,7 @@ def compute_performances(mut_data, mut_sep=':', start_i: int = 0, already_tested
                 #    esm_attention_mask, 
                 #    esm_base_model
                 #)
-                y_esm = inference(sequences, 'esm', model=esm_base_model)
+                y_esm = inference(sequences, 'esm', model=esm_base_model, verbose=False)
                 print(f'ESM1v (unsupervised performance): '
                       f'{spearmanr(fitnesses, y_esm.cpu())[0]:.3f}')
                 esm_unopt_perf = spearmanr(fitnesses, y_esm.cpu())[0]
@@ -167,13 +167,14 @@ def compute_performances(mut_data, mut_sep=':', start_i: int = 0, already_tested
             # ProSST unsupervised
             try:
                 input_ids, prosst_attention_mask, structure_input_ids = get_structure_quantizied(
-                    pdb, prosst_tokenizer, wt_seq)
+                    pdb, prosst_tokenizer, wt_seq, verbose=False
+                    )
                 x_prosst = prosst_tokenize_sequences(sequences=sequences, vocab=prosst_vocab, verbose=False)
                 #y_prosst = get_logits_from_full_seqs(
                 #        x_prosst, prosst_base_model, input_ids, prosst_attention_mask, 
                 #        structure_input_ids, train=False
                 #)
-                y_prosst = inference(sequences, 'prosst', pdb_file=pdb, wt_seq=wt_seq, model=prosst_base_model)
+                y_prosst = inference(sequences, 'prosst', pdb_file=pdb, wt_seq=wt_seq, model=prosst_base_model, verbose=False)
                 print(f'ProSST (unsupervised performance): '
                       f'{spearmanr(fitnesses, y_prosst.cpu())[0]:.3f}')
                 prosst_unopt_perf = spearmanr(fitnesses, y_prosst.cpu())[0]
@@ -195,9 +196,10 @@ def compute_performances(mut_data, mut_sep=':', start_i: int = 0, already_tested
                 for i_split, (train_i, test_i) in enumerate(zip(
                     train_indices, test_indices
                 )):
-                    print(f'Split: {i_split + 1}')
+                    print(f'    Split: {i_split + 1}')
                     temp_results[category].update({f'Split {i_split}': {}})
                     try:
+                        _train_sequences, test_sequences = np.asarray(sequences)[train_i], np.asarray(sequences)[test_i]
                         x_dca_train, x_dca_test = np.asarray(x_dca)[train_i], np.asarray(x_dca)[test_i]
                         x_llm_train_prosst, x_llm_test_prosst = np.asarray(x_prosst)[train_i], np.asarray(x_prosst)[test_i]
                         x_llm_train_esm, x_llm_test_esm = np.asarray(x_esm)[train_i], np.asarray(x_esm)[test_i]
@@ -253,7 +255,7 @@ def compute_performances(mut_data, mut_sep=':', start_i: int = 0, already_tested
                             'structure_input_ids': structure_input_ids
                         }
                     }
-                    print(f'Train: {len(np.array(y_train))} --> Test: {len(np.array(y_test))}')
+                    print(f'        Train: {len(np.array(y_train))} --> Test: {len(np.array(y_test))}')
                     if len(y_test) <= 20: # TODO: 50
                         print(f"Only {len(fitnesses)} in total, splitting the data "
                               f"in N_Train = {len(y_train)} and N_Test = {len(y_test)} "
@@ -264,6 +266,17 @@ def compute_performances(mut_data, mut_sep=':', start_i: int = 0, already_tested
                         ns_y_test.append(np.nan)
                         continue
                     #get_vram()
+
+                    y_test_pred_dca = get_delta_e_statistical_model(x_dca_test, x_wt)
+                    temp_results[category][f'Split {i_split}'].update({'DCA': spearmanr(y_test, y_test_pred_dca)[0]})
+                    print(f'        DCA ZeroShot (split {i_split + 1}) performance: {spearmanr(y_test, y_test_pred_dca)[0]:.3f}')
+                    y_test_pred_esm = inference(test_sequences, 'esm', model=esm_base_model, verbose=False)
+                    temp_results[category][f'Split {i_split}'].update({'ESM1v': spearmanr(y_test, y_test_pred_esm)[0]})
+                    print(f'        ESM1v ZeroShot (split {i_split + 1}) performance: {spearmanr(y_test, y_test_pred_esm)[0]:.3f}')
+                    y_test_pred_prosst = inference(test_sequences, 'prosst', model=prosst_base_model, pdb_file=pdb, wt_seq=wt_seq, verbose=False)
+                    temp_results[category][f'Split {i_split}'].update({'ProSST': spearmanr(y_test, y_test_pred_prosst)[0]})
+                    print(f'        ProSST ZeroShot (split {i_split + 1}) performance: {spearmanr(y_test, y_test_pred_prosst)[0]:.3f}')
+
                     for i_m, method in enumerate([None, llm_dict_esm, llm_dict_prosst]):
                         m_str = ['DCA hybrid', 'DCA+ESM1v hybrid', 'DCA+ProSST hybrid'][i_m]
                         #print('\n~~~ ' + m_str + ' ~~~')
@@ -284,7 +297,7 @@ def compute_performances(mut_data, mut_sep=':', start_i: int = 0, already_tested
                                 ][i_m],
                                 verbose=False
                             )
-                            print(f'{m_str} (split {i_split + 1}) performance: {spearmanr(y_test, y_test_pred)[0]:.3f} '
+                            print(f'        {m_str} (split {i_split + 1}) performance: {spearmanr(y_test, y_test_pred)[0]:.3f} '
                                   f'(train size={train_size}, test_size={test_size})')
                             temp_results[category][f'Split {i_split}'].update({m_str: spearmanr(y_test, y_test_pred)[0]})
                         except RuntimeError as e:  # modeling_prosst.py, line 920, in forward