dev/fail: further test implementation of plm_inference()

Author: niklases

pypef/hybrid/hybrid_model.py

@@ -37,9 +37,9 @@
 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.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_tokenizer, inference
+from pypef.plm.esm_lora_tune import get_esm_models
+from pypef.plm.prosst_lora_tune import get_prosst_models
+from pypef.plm.inference import esm_setup, llm_tokenizer, inference
 from pypef.plm.utils import get_batches
 
 # sklearn/base.py:474: FutureWarning: `BaseEstimator._validate_data` is deprecated in 1.6 and 
@@ -84,7 +84,7 @@ def __init__(
                     self.llm_base_model = llm_model_input['esm1v']['llm_base_model']
                     self.llm_model = llm_model_input['esm1v']['llm_model']
                     self.llm_optimizer = llm_model_input['esm1v']['llm_optimizer']
-                    self.llm_train_function = llm_model_input['esm1v']['llm_train_function']
+                    #self.llm_train_function = llm_model_input['esm1v']['llm_train_function']
                     self.llm_inference_function = llm_model_input['esm1v']['llm_inference_function']
                     self.llm_loss_function = llm_model_input['esm1v']['llm_loss_function']
                     self.x_train_llm = llm_model_input['esm1v']['x_llm']
@@ -94,7 +94,7 @@ def __init__(
                     self.llm_base_model = llm_model_input['prosst']['llm_base_model']
                     self.llm_model = llm_model_input['prosst']['llm_model']
                     self.llm_optimizer = llm_model_input['prosst']['llm_optimizer']
-                    self.llm_train_function = llm_model_input['prosst']['llm_train_function']
+                    #self.llm_train_function = llm_model_input['prosst']['llm_train_function']
                     self.llm_inference_function = llm_model_input['prosst']['llm_inference_function']
                     self.llm_loss_function = llm_model_input['prosst']['llm_loss_function']
                     self.x_train_llm = llm_model_input['prosst']['x_llm']

pypef/plm/esm_lora_tune.py

@@ -191,28 +191,4 @@ def esm_train(
     model.train(False)
 
 
-def esm_setup(wt_seq, 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)
-    wt_tokens, _ = tokenize_sequences(
-            [wt_seq],
-            esm_tokenizer,
-            max_length=len(wt_seq) + 2
-    )
-    x_esm, esm_attention_mask = tokenize_sequences(
-        sequences, esm_tokenizer, max_length=len(wt_seq) + 2, verbose=verbose)
-    llm_dict_esm = {
-        'esm1v': {
-            'llm_base_model': esm_base_model,
-            'llm_model': esm_lora_model,
-            'llm_optimizer': esm_optimizer,
-            'llm_train_function': esm_train,
-            'llm_inference_function': esm_infer,
-            'llm_loss_function': corr_loss,
-            'x_llm' : x_esm,
-            'input_ids': wt_tokens,
-            'llm_attention_mask':  esm_attention_mask,
-            'llm_tokenizer': esm_tokenizer
-        }
-    }
-    return llm_dict_esm
+

pypef/plm/inference.py

@@ -10,9 +10,8 @@
 from tqdm import tqdm
 
 from pypef.utils.helpers import get_device
-from pypef.plm.utils import get_batches
-from pypef.plm.esm_lora_tune import esm_infer, esm_setup, tokenize_sequences
-from pypef.plm.prosst_lora_tune import prosst_setup, prosst_simple_vocab_aa_tokenizer, prosst_infer
+from pypef.plm.utils import corr_loss, get_batches
+from pypef.plm.esm_lora_tune import get_esm_models, tokenize_sequences
 
 import logging
 logger = logging.getLogger('pypef.llm.inference')
@@ -427,3 +426,80 @@ def inference(
     else:
         raise RuntimeError("Unknown LLM option.")
     return y_test_pred
+
+
+
+def esm_setup(wt_seq, 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)
+    wt_tokens, _ = tokenize_sequences(
+            [wt_seq],
+            esm_tokenizer,
+            max_length=len(wt_seq) + 2
+    )
+    x_esm, esm_attention_mask = tokenize_sequences(
+        sequences, esm_tokenizer, max_length=len(wt_seq) + 2, verbose=verbose)
+    llm_dict_esm = {
+        'esm1v': {
+            'llm_base_model': esm_base_model,
+            'llm_model': esm_lora_model,
+            'llm_optimizer': esm_optimizer,
+            #'llm_train_function': esm_train,
+            'llm_inference_function': plm_inference,
+            'llm_loss_function': corr_loss,
+            'x_llm' : x_esm,
+            'input_ids': wt_tokens,
+            'llm_attention_mask':  esm_attention_mask,
+            'llm_tokenizer': esm_tokenizer
+        }
+    }
+    return llm_dict_esm
+
+
+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 "
+            "FASTA file input for embedding sequences! "
+            "Specify a FASTA file with the --wt flag."
+        )
+    if pdb_file is None:
+        raise SystemError(
+            "Running ProSST requires a PDB file input "
+            "for embedding sequences! Specify a PDB file "
+            "with the --pdb flag."
+        )
+
+    pdb_seq = str(list(SeqIO.parse(pdb_file, "pdb-atom"))[0].seq)
+    assert wt_seq == pdb_seq, (
+        f"Wild-type sequence is not matching PDB-extracted sequence:"
+        f"\nWT sequence:\n{wt_seq}\nPDB sequence:\n{pdb_seq}"
+    )
+    prosst_base_model, prosst_lora_model, prosst_tokenizer, prosst_optimizer = get_prosst_models()
+    prosst_vocab = prosst_tokenizer.get_vocab()
+    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, verbose=verbose
+    )
+    x_llm_train_prosst, _attention_mask = tokenize_sequences(
+        sequences=sequences, tokenizer=prosst_tokenizer, 
+        max_length=len(wt_seq) + 2, 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': plm_inference,  # prosst_infer,
+            'llm_loss_function': corr_loss,
+            'x_llm' : x_llm_train_prosst,
+            'llm_attention_mask': prosst_attention_mask,
+            'llm_vocab': prosst_vocab,
+            'input_ids': input_ids,
+            'structure_input_ids': structure_input_ids,
+            'llm_tokenizer': prosst_tokenizer
+        }
+    }
+    return llm_dict_prosst

pypef/plm/prosst_lora_tune.py

@@ -28,6 +28,7 @@
 from pypef.utils.helpers import get_device
 from pypef.plm.esm_lora_tune import tokenize_sequences
 from pypef.plm.utils import load_model_and_tokenizer
+from pypef.plm.inference import plm_inference
 
 
 def prosst_simple_vocab_aa_tokenizer(sequences, vocab, verbose=True):
@@ -272,50 +273,3 @@ def get_structure_quantizied(pdb_file, tokenizer, wt_seq, verbose: bool = True):
     return input_ids, attention_mask, structure_input_ids
 
 
-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 "
-            "FASTA file input for embedding sequences! "
-            "Specify a FASTA file with the --wt flag."
-        )
-    if pdb_file is None:
-        raise SystemError(
-            "Running ProSST requires a PDB file input "
-            "for embedding sequences! Specify a PDB file "
-            "with the --pdb flag."
-        )
-
-    pdb_seq = str(list(SeqIO.parse(pdb_file, "pdb-atom"))[0].seq)
-    assert wt_seq == pdb_seq, (
-        f"Wild-type sequence is not matching PDB-extracted sequence:"
-        f"\nWT sequence:\n{wt_seq}\nPDB sequence:\n{pdb_seq}"
-    )
-    prosst_base_model, prosst_lora_model, prosst_tokenizer, prosst_optimizer = get_prosst_models()
-    prosst_vocab = prosst_tokenizer.get_vocab()
-    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, verbose=verbose
-    )
-    x_llm_train_prosst, _attention_mask = tokenize_sequences(
-        sequences=sequences, tokenizer=prosst_tokenizer, 
-        max_length=len(wt_seq) + 2, 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': prosst_infer,
-            'llm_loss_function': corr_loss,
-            'x_llm' : x_llm_train_prosst,
-            'llm_attention_mask': prosst_attention_mask,
-            'llm_vocab': prosst_vocab,
-            'input_ids': input_ids,
-            'structure_input_ids': structure_input_ids,
-            'llm_tokenizer': prosst_tokenizer
-        }
-    }
-    return llm_dict_prosst

scripts/ProteinGym_runs/official/benchmark_runs/pgym_cv_benchmark.py

@@ -18,9 +18,10 @@
 
 from pypef.utils.variant_data import get_mismatches
 from pypef.plm.prosst_lora_tune import prosst_setup, prosst_simple_vocab_aa_tokenizer
-from pypef.plm.esm_lora_tune import esm_setup, tokenize_sequences
+from pypef.plm.esm_lora_tune import tokenize_sequences
 from pypef.dca.gremlin_inference import GREMLIN, get_delta_e_statistical_model
 from pypef.hybrid.hybrid_model import DCALLMHybridModel
+from pypef.plm.inference import esm_setup
 
 
 @hydra.main(version_base=None, config_path="../configs", config_name="proteingym_data_setup")

tests/test_api_functions.py

@@ -17,9 +17,7 @@
 from pypef.dca.gremlin_inference import GREMLIN
 from pypef.utils.variant_data import get_sequences_from_file, get_wt_sequence
 from pypef.plm.inference import plm_inference
-from pypef.plm.esm_lora_tune import esm_setup
-from pypef.plm.prosst_lora_tune import prosst_setup
-from pypef.plm.inference import inference, llm_tokenizer
+from pypef.plm.inference import esm_setup, prosst_setup, llm_tokenizer
 from pypef.hybrid.hybrid_model import DCALLMHybridModel
 from pypef.plm.esm_lora_tune import (
     get_esm_models, tokenize_sequences,
@@ -115,16 +113,49 @@ def test_hybrid_model_dca_llm():
     )
     assert len(train_seqs_aneh[0]) == len(g.wt_seq)
     aneh_wt_seq = get_wt_sequence(wt_seq_file_aneh)
-    y_pred_esm = inference(train_seqs_aneh, 'esm', wt_seq=aneh_wt_seq)
+    #y_pred_esm = inference(train_seqs_aneh, 'esm', wt_seq=aneh_wt_seq)
+
+    esm_base_model, _esm_lora_model, esm_tokenizer, _esm_optimizer = get_esm_models(
+        model='facebook/esm1v_t33_650M_UR90S_3')
+    esm_base_model = esm_base_model.to(get_device())
+    x_esm, esm_attention_mask = tokenize_sequences(
+        train_seqs_aneh, esm_tokenizer, max_length=len(wt_seq_file_aneh) + 2)
+    # Tokenize WT sequence once
+    wt_tokens, _ = tokenize_sequences(
+            [aneh_wt_seq],
+            esm_tokenizer,
+            max_length=len(aneh_wt_seq) + 2
+    )
+    wt_tokens = torch.tensor(wt_tokens[0], dtype=torch.long)  # shape (L,)
+    print(wt_tokens.shape)
+    print(esm_attention_mask.shape)
+    print(x_esm.shape)
+
+    y_pred_esm = plm_inference(xs=x_esm, wt_input_ids=wt_tokens, 
+                               attention_mask=esm_attention_mask, model=esm_base_model)
     np.testing.assert_almost_equal(
         spearmanr(train_ys_aneh, y_pred_esm)[0], 
          -0.713214007088901, 
         decimal=7
     )
-    y_pred_prosst = inference(
-        train_seqs_aneh, 'prosst', 
-        pdb_file=pdb_file_aneh, wt_seq=aneh_wt_seq
+
+    #y_pred_prosst = inference(
+    #    train_seqs_aneh, 'prosst', 
+    #    pdb_file=pdb_file_aneh, wt_seq=aneh_wt_seq
+    #)
+    prosst_base_model, prosst_lora_model, prosst_tokenizer, prosst_optimizer = get_prosst_models()
+    prosst_vocab = prosst_tokenizer.get_vocab()
+    prosst_base_model = prosst_base_model.to(get_device())
+    wt_input_ids, prosst_attention_mask, wt_structure_input_ids = get_structure_quantizied(
+        pdb_blat_ecolx, prosst_tokenizer, aneh_wt_seq)
+    x_prosst, prosst_attention_mask_ = tokenize_sequences(
+        sequences=train_seqs_aneh, 
+        tokenizer=prosst_tokenizer, 
+        max_length=len(wt_seq_file_aneh) + 2
     )
+    y_pred_prosst = plm_inference(xs=x_prosst, wt_input_ids=wt_input_ids, 
+                                  attention_mask=prosst_attention_mask, model=prosst_base_model, 
+                                  wt_structure_input_ids=wt_structure_input_ids)
     np.testing.assert_almost_equal(
         spearmanr(train_ys_aneh, y_pred_prosst)[0], 
         -0.7394433335146882, 
@@ -383,8 +414,8 @@ def test_plm_corr_blat_ecolx():
 
 
 if __name__ == "__main__":
-    #test_gremlin_avgfp()
-    #test_hybrid_model_dca_llm()
-    #test_dataset_b_results()
+    test_gremlin_avgfp()
+    test_hybrid_model_dca_llm()
+    test_dataset_b_results()
     test_plm_corr_blat_ecolx()