ProTrek_650M_UniRef50_protein_encoder
This model is a SentenceTransformer-compatible version of ProTrek_650M_UniRef50_protein_encoder. It has been converted for use with the sentence-transformers library, enabling easy integration into semantic similarity tasks, such as semantic search, clustering, and feature extraction.
Github repo: https://github.com/westlake-repl/ProTrek
Hugging Face repo https://huggingface.co/westlake-repl/ProTrek_650M_UniRef50
Model Details
Model Description
- Model Type: Sentence Transformer
- Maximum Sequence Length: 1026 tokens
- Output Dimensionality: 1024 dimensions
- Similarity Function: Cosine Similarity
Full Model Architecture
SentenceTransformer(
(0): Transformer({'max_seq_length': 1026, 'do_lower_case': False, 'architecture': 'EsmModel'})
(1): Pooling({'word_embedding_dimension': 1280, 'pooling_mode_cls_token': True, 'pooling_mode_mean_tokens': False, 'pooling_mode_max_tokens': False, 'pooling_mode_mean_sqrt_len_tokens': False, 'pooling_mode_weightedmean_tokens': False, 'pooling_mode_lasttoken': False, 'include_prompt': True})
(2): Dense({'in_features': 1280, 'out_features': 1024, 'bias': True, 'activation_function': 'torch.nn.modules.linear.Identity'})
(3): Normalize()
)
Usage
Direct Usage (Sentence Transformers)
First install the Sentence Transformers library:
pip install -U sentence-transformers
Then you can load this model and run inference.
from sentence_transformers import SentenceTransformer
# Download from the 🤗 Hub
protein_encoder = SentenceTransformer("yosshstd/ProTrek_650M_UniRef50_protein_encoder")
text_encoder = SentenceTransformer("yosshstd/ProTrek_650M_UniRef50_text_encoder")
structure_encoder = SentenceTransformer("yosshstd/ProTrek_650M_UniRef50_structure_encoder")
# Run inference
ProTrek_650M_UniRef50_temperature = 0.0186767578
def sim(a, b): return (a @ b.T / ProTrek_650M_UniRef50_temperature).item()
aa_seq = "MALWMRLLPLLALLALWGPDPAAAFVNQHLCGSHLVEALYLVCGERGFFYTPKTRREAEDLQVGQVELGGGPGAGSLQPLALEGSLQKRGIVEQCCTSICSLYQLENYCN"
text = 'Insulin decreases blood glucose concentration. It increases cell permeability to monosaccharides, amino acids and fatty acids. It accelerates glycolysis, the pentose phosphate cycle, and glycogen synthesis in liver.'
foldseek_seq = 'DVVVVVVVVVVVVCVVPPDDPVPPFDFDFDADVVLVVLLCVLLVPLAFDDDDPDPVVVVVVVVDDDPPDDDPPDPDPDPPVVVVVVVVDDCSVVRRVGIDGSVSSNVRGD'.lower()
seq_emb = protein_encoder.encode([aa_seq], convert_to_tensor=True)
text_emb = text_encoder.encode([text], convert_to_tensor=True)
struc_emb = structure_encoder.encode([foldseek_seq], convert_to_tensor=True)
print("Seq-Text similarity:", sim(seq_emb, text_emb))
print("Seq-Structure similarity:", sim(seq_emb, struc_emb))
print("Text-Structure similarity:", sim(text_emb, struc_emb))
Overview
ProTrek is a multimodal model that integrates protein sequence, protein structure, and text information for better protein understanding. It adopts contrastive learning to learn the representations of protein sequence and structure. During the pre-training phase, we calculate the InfoNCE loss for each two modalities as CLIP does.
Model architecture
Protein sequence encoder: esm2_t33_650M_UR50D
Protein structure encoder: foldseek_t30_150M (identical architecture with esm2 except that the vocabulary only contains 3Di tokens)
Text encoder: BiomedNLP-PubMedBERT-base-uncased-abstract-fulltext
Training Details
Framework Versions
- Python: 3.11.11
- Sentence Transformers: 5.0.0
- Transformers: 4.53.2
- PyTorch: 2.2.1+cu121
- Accelerate:
- Datasets:
- Tokenizers: 0.21.2
Citation
@article{su2024protrek,
title={ProTrek: Navigating the Protein Universe through Tri-Modal Contrastive Learning},
author={Su, Jin and Zhou, Xibin and Zhang, Xuting and Yuan, Fajie},
journal={bioRxiv},
pages={2024--05},
year={2024},
publisher={Cold Spring Harbor Laboratory}
}