Update README.md

README.md

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