Update README.md

README.md

@@ -9,12 +9,13 @@ Here's how to extract PepDoRA embeddings for your input peptide:
 
 ```
 import torch
-from transformers import AutoModel, AutoTokenizer
+from transformers import RobertaTokenizer, RobertaForMaskedLM
 
 # Load the model and tokenizer
 model_name = "ChatterjeeLab/PepDoRA"
-model = AutoModel.from_pretrained(model_name, output_hidden_states=True)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+tokenizer = RobertaTokenizer.from_pretrained(model_name)
+model = RobertaForMaskedLM.from_pretrained(model_name)
 
 # Input peptide sequence
 peptide = "CC(C)C[C@H]1NC(=O)[C@@H](C)NCCCCCCNC(=O)[C@H](CO)NC1=O"
@@ -25,12 +26,13 @@ inputs = tokenizer(peptide, return_tensors="pt")
 # Get the hidden states (embeddings) from the model
 with torch.no_grad():
     outputs = model(**inputs)
-
+   
+    
 # Extract the embeddings from the last hidden layer
-last_hidden_state = outputs.hidden_states[-1]
+embedding = outputs.logits
 
 # Print the embedding shape (or the embedding itself)
-print(last_hidden_state.shape)
+print(embedding.shape)
 ```
 
 ## Repository Authors