Update README.md

README.md

@@ -1,4 +1,5 @@
 ---
+license: mit
 tags:
 - sentence-transformers
 - sentence-similarity
@@ -6,11 +7,15 @@ tags:
 - dense
 pipeline_tag: sentence-similarity
 library_name: sentence-transformers
+base_model:
+- westlake-repl/ProTrek_650M_UniRef50
 ---
 
-# SentenceTransformer
+# ProTrek_650M_UniRef50_text_encoder
+This model is a SentenceTransformer-compatible version of ProTrek_650M_UniRef50_text_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.
 
-This is a [sentence-transformers](https://www.SBERT.net) model trained. It maps sentences & paragraphs to a 1024-dimensional dense vector space and can be used for semantic textual similarity, semantic search, paraphrase mining, text classification, clustering, and more.
+**Github repo: https://github.com/westlake-repl/ProTrek**  
+**Hugging Face repo https://huggingface.co/westlake-repl/ProTrek_650M_UniRef50**
 
 ## Model Details
 
@@ -24,12 +29,6 @@ This is a [sentence-transformers](https://www.SBERT.net) model trained. It maps
 <!-- - **Language:** Unknown -->
 <!-- - **License:** Unknown -->
 
-### Model Sources
-
-- **Documentation:** [Sentence Transformers Documentation](https://sbert.net)
-- **Repository:** [Sentence Transformers on GitHub](https://github.com/UKPLab/sentence-transformers)
-- **Hugging Face:** [Sentence Transformers on Hugging Face](https://huggingface.co/models?library=sentence-transformers)
-
 ### Full Model Architecture
 
 ```
@@ -56,25 +55,38 @@ Then you can load this model and run inference.
 from sentence_transformers import SentenceTransformer
 
 # Download from the 🤗 Hub
-model = SentenceTransformer("sentence_transformers_model_id")
+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
-sentences = [
-    'The weather is lovely today.',
-    "It's so sunny outside!",
-    'He drove to the stadium.',
-]
-embeddings = model.encode(sentences)
-print(embeddings.shape)
-# [3, 1024]
-
-# Get the similarity scores for the embeddings
-similarities = model.similarity(embeddings, embeddings)
-print(similarities)
-# tensor([[1.0000, 0.7137, 0.6395],
-#         [0.7137, 1.0000, 0.5949],
-#         [0.6395, 0.5949, 1.0000]])
+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](https://arxiv.org/abs/2103.00020)
+does.
+
+## Model architecture
+**Protein sequence encoder**:  [esm2_t33_650M_UR50D](https://huggingface.co/facebook/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](https://huggingface.co/microsoft/BiomedNLP-BiomedBERT-base-uncased-abstract-fulltext)
+
 <!--
 ### Direct Usage (Transformers)
 
@@ -123,8 +135,16 @@ You can finetune this model on your own dataset.
 - Tokenizers: 0.21.2
 
 ## Citation
-
-### BibTeX
+```
+@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}
+}
+```
 
 <!--
 ## Glossary