Upload 2 files

FAPM_inference.py

@@ -0,0 +1,86 @@
+import torch
+import torch.nn as nn
+import pandas as pd
+import torch.nn.functional as F
+from lavis.models.protein_models.protein_function_opt import Blip2ProteinMistral
+# from lavis.models.base_model import FAPMConfig
+# from lavis.models.blip2_models.blip2_opt import Blip2ProteinOPT
+import random
+from lavis.models.base_model import FAPMConfig
+import argparse
+
+prop = True
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='FAPM')
+    parser.add_argument('--model_path', type=str, help='Dataset path')
+    parser.add_argument('--example_path', type=str, help='Example protein path')
+    parser.add_argument('--device', type=str, default='cuda', help='Which gpu to use if any (default: cuda)')
+    parser.add_argument('--prompt', type=str, default='none', help='Input prompt for protein function prediction')
+    parser.add_argument('--ground_truth', type=str, default='none', help='ground truth function')
+    args = parser.parse_args()
+    test_sdf_paths = args.model_path
+
+    # model = Blip2ProteinOPT(config=FAPMConfig(), esm_size='3b')
+    # model.load_checkpoint('/cluster/home/wenkai/LAVIS/lavis/output/BLIP2/Pretrain_stage2/20240327081/checkpoint_2.pth')
+    model = Blip2ProteinMistral(config=FAPMConfig(), esm_size='3b')
+    model.load_checkpoint(args.model_path)
+    model.to(args.device)
+
+    # esm_emb = torch.load('/cluster/home/wenkai/LAVIS/data/pretrain/ipr_domain_emb_esm2_3b/Gp49.pt')['representations'][36]
+    esm_emb = torch.load(args.example_path)['representations'][36]
+    esm_emb = F.pad(esm_emb.t(), (0, 1024 - len(esm_emb))).t().to('cuda')
+    samples = {'name': ['P18281'],
+               'image': torch.unsqueeze(esm_emb, dim=0),
+               'text_input': [args.ground_truth],
+               'prompt': [args.prompt]}
+    prediction = model.generate(samples, length_penalty=0., num_beams=15, num_captions=10, temperature=1., repetition_penalty=1.0)
+    print(f"Text Prediction: {prediction}")
+
+
+    if prop == True:
+        from data.evaluate_data.utils import Ontology
+        import difflib
+        import re
+
+        # godb = Ontology(f'/cluster/home/wenkai/LAVIS/data/go1.4-basic.obo', with_rels=True)
+        godb = Ontology(f'data/go1.4-basic.obo', with_rels=True)
+
+        go_des = pd.read_csv('data/go_descriptions1.4.txt', sep='|', header=None)
+        go_des.columns = ['id', 'text']
+        go_des = go_des.dropna()
+        go_des['id'] = go_des['id'].apply(lambda x: re.sub('_', ':', x))
+        go_obo_set = set(go_des['id'].tolist())
+        go_des['text'] = go_des['text'].apply(lambda x: x.lower())
+        GO_dict = dict(zip(go_des['text'], go_des['id']))
+        Func_dict = dict(zip(go_des['id'], go_des['text']))
+
+        # terms_mf = pd.read_pickle('/cluster/home/wenkai/deepgo2/data/mf/terms.pkl')
+        terms_mf = pd.read_pickle('data/terms/mf_terms.pkl')
+        choices_mf = [Func_dict[i] for i in list(set(terms_mf['gos']))]
+        choices = {x.lower(): x for x in choices_mf}
+
+        pred_terms_list = []
+        pred_go_list = []
+        prop_annotations = []
+        for x in prediction:
+            x = [eval(i) for i in x.split('; ')]
+            pred_terms = []
+            pred_go = []
+            annot_set = set()
+            for i in x:
+                txt = i[0]
+                prob = i[1]
+                sim_list = difflib.get_close_matches(txt.lower(), choices, n=1, cutoff=0.9)
+                if len(sim_list) > 0:
+                    pred_terms.append((sim_list[0], prob))
+                    pred_go.append((GO_dict[sim_list[0]], prob))
+                    annot_set |= godb.get_anchestors(GO_dict[sim_list[0]])
+            pred_terms_list.append(pred_terms)
+            pred_go_list.append(pred_go)
+            annots = list(annot_set)
+            prop_annotations.append(annots)
+
+        print(f"Predictions of GO terms: \n{pred_terms_list} \nPredictions of GO id: \n{pred_go_list} \nPredictions of GO id propgated: \n{prop_annotations}")
+
+

README.md

@@ -0,0 +1,75 @@
+## Introduction
+<p align="center">
+    <br>
+    <img src="assets/FAPM.png"/>
+    <br>
+<p>
+
+## Installation
+
+1. (Optional) Creating conda environment
+
+```bash
+conda create -n lavis python=3.8
+conda activate lavis
+```
+ 
+2. for development, you may build from source
+
+```bash
+git clone https://github.com/xiangwenkai/FAPM.git
+cd FAPM
+pip install -e .
+
+pip install Biopython
+pip install fair-esm
+```
+
+### Datasets
+#### 1.raw dataset
+Raw data are avaliable at *https://ftp.uniprot.org/pub/databases/uniprot/previous_releases/release-2023_04/knowledgebase/*, this file is very large and need to be processed to get its name, sequence, GO label, function description and prompt.  
+The domain level protein dataset we used are avaliable at *https://ftp.ebi.ac.uk/pub/databases/interpro/releases/95.0/protein2ipr.dat.gz*  
+In this respository, We provide the experimental train/val/test sets of Swiss-Prot, which are avaliable at data/swissprot_exp  
+#### 2.ESM2 embeddings  
+Source code for ESM2 embeddings generation: *https://github.com/facebookresearch/esm*  
+The generation command:  
+```bash
+python esm_scripts/extract.py esm2_t33_3B_UR50D you_path/protein.fasta you_path_to_save_embedding_files --repr_layers 36 --truncation_seq_length 1024 --include per_tok
+```
+The default path to save embedding files in this respository is **data/emb_esm2_3b**
+
+## Pretraining language models  
+Source: *https://huggingface.co/teknium/OpenHermes-2.5-Mistral-7B*
+
+## Training
+data config: lavis/configs/datasets/protein/GO_defaults_cap.yaml  
+stage1 config: lavis/projects/blip2/train/protein_pretrain_stage1.yaml  
+stage1 training command: run_scripts/blip2/train/protein_pretrain_domain_stage1.sh  
+stage2 config: lavis/projects/blip2/train/protein_pretrain_stage2.yaml  
+stage2 training/finetuning command: run_scripts/blip2/train/protein_pretrain_domain_stage2.sh  
+
+## Trained models
+You can download our trained models from drive: *https://drive.google.com/drive/folders/1aA0eSYxNw3DvrU5GU1Cu-4q2kIxxAGSE?usp=drive_link*  
+
+## Testing
+config: lavis/projects/blip2/eval/caption_protein_eval.yaml  
+command: run_scripts/blip2/eval/eval_cap_protein.sh  
+
+## Inference example
+```
+python FAPM_inference.py \
+--model_path model/checkpoint_mf2.pth \
+--example_path data/emb_esm2_3b/P18281.pt \
+--device cuda \
+--prompt Acanthamoeba
+```  
+
+
+
+
+
+
+
+
+
+