Views
No views yet
esm2_t6_8M_UR50D)1from transformers import AutoTokenizer, AutoModelForSequenceClassification
2import torch
3
4# Initialize the tokenizer and model
5model_path_directory = "AmelieSchreiber/esm2_t6_8M_UR50D-finetuned-localization"
6tokenizer = AutoTokenizer.from_pretrained(model_path_directory)
7model = AutoModelForSequenceClassification.from_pretrained(model_path_directory)
8
9# Define a function to predict the category of a protein sequence
10def predict_category(sequence):
11 # Tokenize the sequence and convert it to tensor format
12 inputs = tokenizer(sequence, return_tensors="pt", truncation=True, max_length=512, padding="max_length")
13
14 # Make prediction
15 with torch.no_grad():
16 logits = model(**inputs).logits
17
18 # Determine the category with the highest score
19 predicted_class = torch.argmax(logits, dim=1).item()
20
21 # Return the category: 0 for cytosolic, 1 for membrane
22 return "cytosolic" if predicted_class == 0 else "membrane"
23
24# Example sequence
25new_protein_sequence = "MTQRAGAAMLPSALLLLCVPGCLTVSGPSTVMGAVGESLSVQCRYEEKYKTFNKYWCRQPCLPIWHEMVETGGSEGVVRSDQVIITDHPGDLTFTVTLENLTADDAGKYRCGIATILQEDGLSGFLPDPFFQVQVLVSSASSTENSVKTPASPTRPSQCQGSLPSSTCFLLLPLLKVPLLLSILGAILWVNRPWRTPWTES"
26
27# Predict the category
28category = predict_category(new_protein_sequence)
29print(f"The predicted category for the sequence is: {category}")