Add test for new plm_inference() function

Author: niklases

README.md

@@ -48,9 +48,9 @@ When incorporating DCA and PLM features, both models are fine-tuned via few-shot
 A quick installation of the PyPEF command line framework using PyPI for Linux and Windows and Python >= 3.10 can be performed with:
 
 ```bash
-pip install -U pypef
-# optionally, for GPU support (see requirements section below):
+# For GPU support (e.g., using CUDA 12.8, see requirements section below):
 # pip install torch --index-url https://download.pytorch.org/whl/cu128
+pip install -U pypef
 ```
 
 After successful installation, PyPEF should work by calling `pypef` in the shell:

pypef/plm/esm_lora_tune.py

@@ -143,13 +143,14 @@ def esm_infer(xs, attention_mask, model, device: str | None = None, verbose=Fals
     return torch.flatten(y_preds_total)
 
 
-def esm_unmasked_wt_score(
+def unmasked_wt_score(
         tokenized_sequences, 
         attention_mask, 
         wt_input_ids,
         model, 
         train: bool = False,
-        device=None, 
+        cut_special_tokens: bool = True,  # assumption: cut first and last token
+        device=None,
         **kwargs
     ):
     if device is None:
@@ -189,27 +190,29 @@ def esm_unmasked_wt_score(
 
     logits = outputs.logits
     logits = logits.squeeze(0)   # remove batch dim
-    #print('logits.shape:', logits.shape)
     # Better make sure that special tokens are always removed / masked 
     # and only pure amino acid sequence tokens are present / unmasked
-    #logits = logits[1:-1]        # drop CLS/EOS
+    tokenized_seq_len = tokenized_sequences.shape[1]
+    if cut_special_tokens:
+        logits = logits[1:-1]        # drop CLS/EOS
+        tokenized_seq_len -= 2
     token_probs = torch.log_softmax(logits, dim=-1)
-    assert len(tokenized_sequences[0]) == token_probs.shape[0], f"{len(tokenized_sequences[0])} != {token_probs.shape[0]}"
-    #print('token_probs.shape:', token_probs.shape)
-
-    for i_s, tokenized_seq in enumerate(tokenized_sequences):
-        for i_aa, aa in enumerate(tokenized_seq):
-            # alternative: use Tensor.index_select() function
-            if i_aa == 0:
-                seq_log_probs = token_probs[i_aa, aa].reshape(1)
-            else:
-                seq_log_probs = torch.cat(
-                    (seq_log_probs, token_probs[i_aa, aa].reshape(1)), 0)
-        if i_s == 0:
-            log_probs = torch.sum(torch.Tensor(seq_log_probs)).reshape(1)
-        else:
-            log_probs = torch.cat(
-                (log_probs, torch.sum(torch.Tensor(seq_log_probs)).reshape(1)), 0)
+    assert tokenized_seq_len == token_probs.shape[0], (
+        f"{tokenized_seq_len} != {token_probs.shape[0]}")
+
+    log_probs = []
+    for tokenized_seq in tokenized_sequences:
+        if cut_special_tokens:
+            tokenized_seq = tokenized_seq[1:-1]
+    
+        seq_lp = token_probs[
+            torch.arange(tokenized_seq.shape[0], device=tokenized_seq.device),
+            tokenized_seq
+        ].sum(dtype=torch.float64)
+
+        log_probs.append(seq_lp)
+    
+    log_probs = torch.stack(log_probs)
     return log_probs
 
 
@@ -285,7 +288,7 @@ def esm_mutation_all_pos_masked_pll(
     verbose: bool = False,
 ):
     """
-    Correct mutation-only pseudo-log-likelihood for ONE sequence.
+    Correct mutation-only pseudo-log-likelihood for sequences.
     """
     model.eval()
 
@@ -332,13 +335,14 @@ def esm_mutation_all_pos_masked_pll(
     return plls
 
 
-def esm_infer_pll(
+def plm_inference(
     xs,
     wt_input_ids,
     attention_mask,
     model,
     mask_token_id,
     inference_type='unmasked',
+    wt_structure_input_ids=None,
     batch_size=5,
     train=False,
     device=None,
@@ -354,23 +358,19 @@ def esm_infer_pll(
 
     if not isinstance(attention_mask, torch.Tensor):
         attention_mask = torch.tensor(attention_mask, dtype=torch.long)
-    wt_structure_input_ids = None
     if inference_type == 'mutation-masking':
         inference_function = esm_mutation_only_mutation_masked_pll
     elif inference_type in ['full-masking', 'all-pos-masking']:
         inference_function = esm_mutation_all_pos_masked_pll
     elif inference_type in ['unmasked', 'wt-marginals']:
-        inference_function = esm_unmasked_wt_score
-    elif inference_type == 'prosst':
-        wt_input_ids, wt_structure_input_ids = wt_input_ids
-        inference_function = esm_unmasked_wt_score
+        inference_function = unmasked_wt_score
     else:
-        raise SystemError("Choose between 'mutation_masking', 'unmasked', and 'full_masking'")
+        raise SystemError("Choose between 'mutation-masking', 'unmasked', and 'full-masking'")
 
     scores = []
 
     xs_b = get_batches(xs, dtype=int, batch_size=batch_size, keep_remaining=True, verbose=True)
-    desc = f"ESM inference: {inference_type} batch (size={batch_size}) processing ({device.upper()})'"
+    desc = f"Inference: {inference_type} batch (size={batch_size}) processing ({device.upper()})'"
 
     pbar = tqdm(
         range(len(xs_b)),

pypef/plm/inference.py

@@ -9,7 +9,7 @@
 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_tokenize_sequences, prosst_infer
+from pypef.plm.prosst_lora_tune import prosst_setup, prosst_simple_vocab_aa_tokenizer, prosst_infer
 
 import logging
 logger = logging.getLogger('pypef.llm.inference')
@@ -26,7 +26,7 @@ def llm_tokenizer(llm_dict, seqs, verbose=True):
             max_length=len(seqs[0]), verbose=verbose
         )
     elif list(llm_dict.keys())[0] == 'prosst':
-        x_llm_seqs = prosst_tokenize_sequences(
+        x_llm_seqs = prosst_simple_vocab_aa_tokenizer(
             seqs, vocab=llm_dict['prosst']['llm_vocab'], verbose=verbose
         )
     else:
@@ -80,4 +80,4 @@ def inference(
         ).cpu()
     else:
         raise RuntimeError("Unknown LLM option.")
-    return y_test_pred
+    return y_test_pred

pypef/plm/prosst_lora_tune.py

@@ -29,21 +29,22 @@
 from pypef.plm.utils import load_model_and_tokenizer
 
 
-def prosst_tokenize_sequences(sequences, vocab, verbose=True):
+def prosst_simple_vocab_aa_tokenizer(sequences, vocab, verbose=True):
     print(vocab)
     sequences = np.atleast_1d(sequences).tolist()
     x_sequences = []
     for sequence in tqdm(
         sequences, desc='Tokenizing sequences for ProSST modeling', 
         disable=not verbose
     ):
-        x_sequence = [vocab['<cls>']]
+        #x_sequence = [vocab['<cls>']]
+        x_sequence = []
         for aa in sequence:
             try:
                 x_sequence.append(vocab[aa])
             except KeyError:
                 x_sequence.append(vocab['<unk>'])
-        x_sequence.append(vocab['<eos>'])
+        #x_sequence.append(vocab['<eos>'])
         x_sequences.append(x_sequence)
     return torch.Tensor(x_sequences).to(torch.int)
 
@@ -296,7 +297,7 @@ def prosst_setup(wt_seq, pdb_file, sequences, device: str | None = None, verbose
     input_ids, prosst_attention_mask, structure_input_ids = get_structure_quantizied(
         pdb_file, prosst_tokenizer, wt_seq, verbose=verbose
     )
-    x_llm_train_prosst = prosst_tokenize_sequences(
+    x_llm_train_prosst = prosst_simple_vocab_aa_tokenizer(
         sequences=sequences, vocab=prosst_vocab, verbose=verbose
     )
     llm_dict_prosst = {

scripts/ProteinGym_runs/official/benchmark_runs/pgym_cv_benchmark.py

@@ -17,7 +17,7 @@
 
 
 from pypef.utils.variant_data import get_mismatches
-from pypef.plm.prosst_lora_tune import prosst_setup, prosst_tokenize_sequences
+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.dca.gremlin_inference import GREMLIN, get_delta_e_statistical_model
 from pypef.hybrid.hybrid_model import DCALLMHybridModel
@@ -177,7 +177,7 @@ def main(cfg: DictConfig) -> None:
                 device='cuda'
             )
             vocab = llm_kwargs['prosst']['llm_vocab']
-            x_llm_test = np.asarray(prosst_tokenize_sequences(
+            x_llm_test = np.asarray(prosst_simple_vocab_aa_tokenizer(
                 sequences=s_test, vocab=vocab, verbose=False))
         elif llm == "esm1v":
             llm_kwargs = esm_setup(sequences=s_train)

scripts/ProteinGym_runs/protgym_hybrid_perf_test_crossval.py

@@ -27,7 +27,7 @@
 )
 from pypef.plm.prosst_lora_tune import (
     get_logits_from_full_seqs, get_prosst_models, get_structure_quantizied, 
-    prosst_tokenize_sequences, prosst_train
+    prosst_simple_vocab_aa_tokenizer, prosst_train
 )
 from pypef.plm.inference import inference
 from pypef.utils.variant_data import get_seqs_from_var_name
@@ -165,7 +165,7 @@ def compute_performances(mut_data, mut_sep=':', start_i: int = 0, already_tested
                 input_ids, prosst_attention_mask, structure_input_ids = get_structure_quantizied(
                     pdb, prosst_tokenizer, wt_seq, verbose=False
                     )
-                x_prosst = prosst_tokenize_sequences(sequences=sequences, vocab=prosst_vocab, verbose=False)
+                x_prosst = prosst_simple_vocab_aa_tokenizer(sequences=sequences, vocab=prosst_vocab, verbose=False)
                 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}')

scripts/ProteinGym_runs/protgym_hybrid_perf_test_low_n.py

@@ -29,7 +29,7 @@
 )
 from pypef.plm.prosst_lora_tune import (
     get_logits_from_full_seqs, get_prosst_models, get_structure_quantizied, 
-    prosst_tokenize_sequences, prosst_train
+    prosst_simple_vocab_aa_tokenizer, prosst_train
 )
 from pypef.utils.variant_data import get_seqs_from_var_name
 from pypef.utils.helpers import get_vram, get_device
@@ -159,7 +159,7 @@ def compute_performances(mut_data, mut_sep=':', start_i: int = 0, already_tested
             try:
                 input_ids, prosst_attention_mask, structure_input_ids = get_structure_quantizied(
                     pdb, prosst_tokenizer, wt_seq)
-                x_prosst = prosst_tokenize_sequences(sequences=sequences, vocab=prosst_vocab)
+                x_prosst = prosst_simple_vocab_aa_tokenizer(sequences=sequences, vocab=prosst_vocab)
                 y_prosst = get_logits_from_full_seqs(
                         x_prosst, prosst_base_model, input_ids, prosst_attention_mask, 
                         structure_input_ids, train=False

tests/test_api_functions.py

@@ -17,7 +17,7 @@
 from pypef.ml.regression import AAIndexEncoding, full_aaidx_txt_path, get_regressor_performances
 from pypef.dca.gremlin_inference import GREMLIN
 from pypef.utils.variant_data import get_sequences_from_file, get_wt_sequence
-from pypef.plm.esm_lora_tune import esm_infer, esm_infer_pll, esm_setup, esm_train
+from pypef.plm.esm_lora_tune import esm_infer, plm_inference, esm_setup, esm_train
 from pypef.plm.prosst_lora_tune import prosst_setup
 from pypef.plm.inference import inference, llm_tokenizer
 from pypef.hybrid.hybrid_model import DCALLMHybridModel
@@ -26,7 +26,7 @@
 )
 from pypef.plm.prosst_lora_tune import (
     get_logits_from_full_seqs, get_prosst_models, get_structure_quantizied, 
-    prosst_tokenize_sequences
+    prosst_simple_vocab_aa_tokenizer
 )
 from pypef.utils.helpers import get_device
 
@@ -258,10 +258,6 @@ def test_plm_corr_blat_ecolx():
     prosst_base_model = prosst_base_model.to(device)
     df = pd.read_csv(csv_blat_ecolx_stiffler2015)
     sequences = df['mutated_sequence'].to_list()
-    print(sequences[0][23])
-    print(sequences[1][23])
-    print('len(sequences[0]):', len(sequences[0]))
-    print('len(blat_ecolx_wt_seq):', len(blat_ecolx_wt_seq))
     y_true = df['DMS_score'].to_list()
     for x in ['facebook/esm1v_t33_650M_UR90S_3']:
         esm_base_model, _esm_lora_model, esm_tokenizer, esm_optimizer = get_esm_models(model=x)
@@ -275,7 +271,7 @@ def test_plm_corr_blat_ecolx():
             max_length=len(blat_ecolx_wt_seq) + 2
         )
         wt_tokens = torch.tensor(wt_tokens[0], dtype=torch.long)  # shape (L,)
-        y_esm = esm_infer_pll(
+        y_esm = plm_inference(
             xs=x_esm,
             wt_input_ids=wt_tokens,
             attention_mask=esm_attention_mask,
@@ -289,7 +285,7 @@ def test_plm_corr_blat_ecolx():
         print(f'{x}: ESM1v (unsupervised performance): '  
               f'{spearmanr(y_true, y_esm.cpu())[0]}')
         np.testing.assert_almost_equal(spearmanr(y_true, y_esm.cpu())[0], 0.6367826285982324, decimal=6)
-        y_esm = esm_infer_pll(
+        y_esm = plm_inference(
             xs=x_esm,
             wt_input_ids=wt_tokens,
             attention_mask=esm_attention_mask,
@@ -303,7 +299,7 @@ def test_plm_corr_blat_ecolx():
         print(f'{x}: ESM1v (unsupervised performance): '  
               f'{spearmanr(y_true, y_esm.cpu())[0]}')
         np.testing.assert_almost_equal(spearmanr(y_true, y_esm.cpu())[0], 0.6498987261125897, decimal=6)
-        #y_esm = esm_infer_pll(
+        #y_esm = plm_inference(
         #    xs=x_esm,
         #    wt_input_ids=wt_tokens,
         #    attention_mask=esm_attention_mask,
@@ -317,31 +313,34 @@ def test_plm_corr_blat_ecolx():
         #print(f'{x}: ESM1v (unsupervised performance): '  
         #      f'{spearmanr(y_true, y_esm.cpu())[0]}')
         #np.testing.assert_almost_equal(spearmanr(y_true, y_esm.cpu())[0], 0.666666666666666, decimal=6)
-
     wt_input_ids, prosst_attention_mask, wt_structure_input_ids = get_structure_quantizied(
         pdb_blat_ecolx, prosst_tokenizer, blat_ecolx_wt_seq)
-    x_prosst = tokenize_sequences(sequences=sequences, tokenizer=prosst_tokenizer)
-    y_prosst = get_logits_from_full_seqs(
-            x_prosst, prosst_base_model, wt_input_ids, prosst_attention_mask, 
-            wt_structure_input_ids, train=False, verbose=True
+    x_prosst2 = prosst_simple_vocab_aa_tokenizer(sequences, prosst_vocab)
+    x_prosst, prosst_attention_mask_ = tokenize_sequences(
+        sequences=sequences, 
+        tokenizer=prosst_tokenizer, 
+        max_length=len(blat_ecolx_wt_seq) + 2
     )
-    print(f'ProSST (unsupervised performance): '  # ProteinGym: ProSST: 0.760
-          f'{spearmanr(y_true, y_prosst.cpu())[0]:.3f}')
+    assert x_prosst[0][1:-1] == x_prosst2.tolist()[0], (
+        f"{x_prosst[0][1:-1]} != {x_prosst2.tolist()[0]}")
+    assert prosst_attention_mask.tolist()[0] == prosst_attention_mask_, (
+        f"{prosst_attention_mask.tolist()[0]} != {prosst_attention_mask_}")
 
-    y_prosst = esm_infer_pll(
+    y_prosst = plm_inference(
             xs=x_prosst,
-            wt_input_ids=(wt_input_ids, wt_structure_input_ids), ## TODO
+            wt_input_ids=wt_input_ids,
             attention_mask=prosst_attention_mask,
             model=prosst_base_model,
             mask_token_id=prosst_tokenizer.mask_token_id,
-            inference_type='prosst',  ## TODO
+            inference_type='unmasked',
+            wt_structure_input_ids=wt_structure_input_ids,
             batch_size=5,
             train=False,
             verbose=True        
     )
     print(f'ProSST (unsupervised performance): '  # ProteinGym: ProSST: 0.760
-          f'{spearmanr(y_true, y_prosst.cpu())[0]:.3f}')
-    # ACTUAL OLD VERSION: 0.743
+          f'{spearmanr(y_true, y_prosst.cpu())[0]}')
+    np.testing.assert_almost_equal(spearmanr(y_true, y_prosst.cpu())[0], 0.7430279087189432, decimal=6)