Initial upload of AllerTrans app

app.py

@@ -0,0 +1,52 @@
+import torch
+import gradio as gr
+import numpy as np
+from transformers import T5Tokenizer, T5EncoderModel
+import esm
+from inference import load_models, predict_ensemble
+
+# Load trained models
+model_protT5, model_cat = load_models()
+
+# Load ProtT5 model
+tokenizer_t5 = T5Tokenizer.from_pretrained("Rostlab/prot_t5_xl_uniref50", do_lower_case=False)
+model_t5 = T5EncoderModel.from_pretrained("Rostlab/prot_t5_xl_uniref50")
+model_t5 = model_t5.eval()
+
+# Load ESM model
+esm_model, alphabet = esm.pretrained.esm2_t33_650M_UR50D()
+batch_converter = alphabet.get_batch_converter()
+esm_model.eval()
+
+
+def extract_prott5_embedding(sequence):
+    sequence = sequence.replace(" ", "")
+    seq = " ".join(list(sequence))
+    ids = tokenizer_t5(seq, return_tensors="pt", padding=True)
+    with torch.no_grad():
+        embedding = model_t5(**ids).last_hidden_state
+    return torch.mean(embedding, dim=1)
+
+
+def extract_esm_embedding(sequence):
+    batch_labels, batch_strs, batch_tokens = batch_converter([("protein1", sequence)])
+    with torch.no_grad():
+        results = esm_model(batch_tokens, repr_layers=[33], return_contacts=False)
+    token_representations = results["representations"][33]
+    return torch.mean(token_representations[0, 1:len(sequence)+1], dim=0).unsqueeze(0)
+
+
+def classify(sequence):
+    protT5_emb = extract_prott5_embedding(sequence)
+    esm_emb = extract_esm_embedding(sequence)
+    concat = torch.cat((esm_emb, protT5_emb), dim=1)
+    pred = predict_ensemble(protT5_emb, concat, model_protT5, model_cat)
+    return "Allergen" if pred.item() == 1 else "Non-Allergen"
+
+
+demo = gr.Interface(fn=classify,
+                    inputs=gr.Textbox(lines=3, placeholder="Enter protein sequence..."),
+                    outputs=gr.Label(label="Prediction"))
+
+if __name__ == "__main__":
+    demo.launch()

inference.py

@@ -0,0 +1,46 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class NeuralNet(nn.Module):
+    def __init__(self, input_size, hidden_size1, hidden_size2, hidden_size3, num_classes):
+        super(NeuralNet, self).__init__()
+        self.fc1 = nn.Linear(input_size, hidden_size1)
+        self.dropout = nn.Dropout(0.1)
+        self.fc2 = nn.Linear(hidden_size1, hidden_size2)
+        self.dropout = nn.Dropout(0.1)
+        self.fc3 = nn.Linear(hidden_size2, hidden_size3)
+        self.dropout = nn.Dropout(0.1)
+        self.fc4 = nn.Linear(hidden_size3, num_classes)
+
+    def forward(self, x):
+        out = F.relu(self.fc1(x))
+        out = F.relu(self.fc2(out))
+        out = F.relu(self.fc3(out))
+        out = self.fc4(out)
+        return out
+
+
+def load_models():
+    model_protT5 = NeuralNet(1024, 200, 100, 50, 2)
+    model_protT5.load_state_dict(torch.load("checkpoints/model17-protT5.pt", map_location=torch.device("cpu")))
+    model_protT5.eval()
+
+    model_cat = NeuralNet(2304, 200, 100, 100, 2)
+    model_cat.load_state_dict(torch.load("checkpoints/model-esm-protT5-5.pt", map_location=torch.device("cpu")))
+    model_cat.eval()
+
+    return model_protT5, model_cat
+
+
+def predict_ensemble(X_protT5, X_concat, model_protT5, model_cat, weight1=0.60, weight2=0.30):
+    with torch.no_grad():
+        outputs1 = model_cat(X_concat)
+        outputs2 = model_protT5(X_protT5)
+        ensemble_outputs = weight1 * outputs1 + weight2 * outputs2
+        _, predicted = torch.max(ensemble_outputs.data, 1)
+    return predicted
+
+
+

requirements.txt

@@ -0,0 +1,9 @@
+torch
+gradio
+transformers
+esm
+fair-esm  # if esm isn't installed via pip
+sentencepiece
+h5py
+git+https://github.com/facebookresearch/esm.git
+git+https://github.com/agemagician/ProtTrans.git