Upload README.md

README.md

@@ -34,9 +34,7 @@ from transformers import AutoModel
 import torch
 
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-print(f"Using {device} device")
 def pad_inner_lists_to_length(outer_list,target_length=16):
     for inner_list in outer_list:
         padding_length = target_length - len(inner_list)
@@ -45,22 +43,20 @@ def pad_inner_lists_to_length(outer_list,target_length=16):
     return outer_list
 
 
-tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D")
-model = AutoModel.from_pretrained("SkywalkerLu/TransHLA_I", trust_remote_code=True)
-
-
-
-model.to(device)
-peptide_examples = ['EDSAIVTPSR','SVWEPAKAKYVFR']
-peptide_encoding = tokenizer(peptide_examples)['input_ids']
-peptide_encoding = pad_inner_lists_to_length(peptide_encoding)
-print(peptide_encoding)
-
-peptide_encoding = torch.tensor(peptide_encoding)
-
-outputs,representations = model(peptide_encoding.to(device))
-print(outputs)
-print(representations)
+if __name__ = "__main__":
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Using {device} device")
+    tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D")
+    model = AutoModel.from_pretrained("SkywalkerLu/TransHLA_I", trust_remote_code=True)
+    model.to(device)
+    peptide_examples = ['EDSAIVTPSR','SVWEPAKAKYVFR']
+    peptide_encoding = tokenizer(peptide_examples)['input_ids']
+    peptide_encoding = pad_inner_lists_to_length(peptide_encoding)
+    print(peptide_encoding)
+    peptide_encoding = torch.tensor(peptide_encoding)
+    outputs,representations = model(peptide_encoding.to(device))
+    print(outputs)
+    print(representations)
 ```
 And here is how to use TransHLA_II model to predict the peptide whether epitope: