fix app batch

NeuroPredPLM/model.py

@@ -44,10 +44,9 @@ class EsmModel(nn.Module):
         representation_cov =  F.relu(self.cov_2(representation_cov))
         representations = rearrange(representation_cov, '(b h) d l -> b h l d', h=self.head)
         att = torch.einsum('bhld,hd->bhl', representations, self.W)
-        mask = length_to_mask(torch.tensor(peptide_length)).to(device)
+        mask = length_to_mask(torch.tensor(peptide_length)).to(device).int()
         att = att.masked_fill(mask.unsqueeze(1)==0, -np.inf)
         att= F.softmax(att, dim=-1)
-        # print(att)
         representations = rearrange(representations * att.unsqueeze(-1), 'b h l d -> b l (h d)')
         representations = torch.sum(representations, dim=1)
         return self.fcn(representations), att

NeuroPredPLM/predict.py

@@ -24,8 +24,11 @@ def batch_predict(peptide_list, cutoff, model_path, device='cpu'):
         state_dict = torch.load(model_path, map_location="cpu")
         neuroPred_model.load_state_dict(state_dict)
         neuroPred_model = neuroPred_model.to(device)
-        prob, att = neuroPred_model(peptide_list, device)
-        pred = torch.softmax(prob, dim=-1).cpu().tolist()
-        att = att.cpu().numpy()
-        out = [[i[0], i[1], f"{j[1]:.3f}", 'Neuropeptide' if j[1] >cutoff else 'Non-neuropeptide'] for i, j in zip(peptide_list, pred)]
+        out = []
+        for item in peptide_list:
+            prob, att = neuroPred_model([item], device)
+            pred = torch.softmax(prob, dim=-1).cpu().tolist()
+            att = att.cpu().numpy()
+            temp = [[i[0], i[1], f"{j[1]:.3f}", 'Neuropeptide' if j[1] >cutoff else 'Non-neuropeptide'] for i, j in zip([item], pred)]
+            out.append(temp[0])
     return out

app.py

@@ -9,7 +9,7 @@ def classifier(peptide_seq):
     data = []
     for record in SeqIO.parse(handle, 'fasta'):
         data.append((record.id, str(record.seq)))
-    device = "cuda" if torch.cuda.is_available() else "cpu" 
+    device = "cuda" if torch.cuda.is_available() else "cpu"
     neuropeptide_pred = predict(data, './model.pth', device)
     return neuropeptide_pred
 # {peptide_id:[Type:int(1->neuropeptide,0->non-neuropeptide), attention score:nd.array]}
@@ -19,7 +19,7 @@ def batch_classifier(file, cutoff):
     data = []
     for record in SeqIO.parse(file.name, 'fasta'):
         data.append((record.id, str(record.seq)))
-    device = "cuda" if torch.cuda.is_available() else "cpu" 
+    device = "cuda" if torch.cuda.is_available() else "cpu"
     neuropeptide_pred = batch_predict(data, cutoff, './model.pth', device)
     return neuropeptide_pred
 
@@ -49,7 +49,7 @@ with gr.Blocks() as demo:
             with gr.Column(scale=2):
                 batch_cutoff = gr.Slider(0, 1, step=0.1, value=0.5, interactive=True, label="Threshold")
                 gr.Markdown("### Note")
-                gr.Markdown("- Limit the number of input sequences to less than 100")
+                gr.Markdown("- Limit the number of input sequences to less than 30")
                 gr.Markdown("- The file should be the Fasta format")
                 gr.Markdown("- The input peptide sequence length should be between 5-100")
                 image_button = gr.Button("Submit")