Upload app.py

app.py

@@ -8,52 +8,12 @@ import numpy as np
 import torch
 from glycowork.glycan_data.loader import lib
 
-# Update lib
-equivalence_classes = [                     
-["Glc", "Man", "Gal", "Gul", "Alt", "All", "Tal", "Ido" ],
-["GlcNAc", "ManNAc", "GalNAc", "GulNAc", "AltNAc", "AllNAc", "TalNAc", "IdoNAc"],
-["GlcN", "ManN", "GalN", "GulN", "AltN", "AllN", "TalN", "IdoN"],
-["GlcA", "ManA", "GalA", "GulA", "AltA", "AllA", "TalA", "IdoA"],
-["Qui", "Rha", "6dGul", "6dAlt", "6dTal", "Fuc"],
-["QuiNAc", "RhaNAc", "6dAltNAc", "6dTalNAc", "FucNAc"],
-["Oli", "Tyv", "Abe", "Par", "Dig", "Col"],
-["Ara", "Lyx", "Xyl", "Rib"],
-["Kdn", "Neu5Ac", "Neu5Gc", "Neu", "Sia"],
-["Pse", "Leg", "Aci", "4eLeg"],
-["Bac", "LDmanHep", "Kdo", "Dha", "DDmanHep", "MurNAc", "MurNGc", "Mur", "Api", "Fru", "Tag", "Sor", "Psi"]
-]
-
-linkage_classes = [    
-  ["a1-2", "a1-z", "z1-2", "z1-z"],
-  ["a1-3", "a1-z", "z1-3", "z1-z"],
-  ["a1-4", "a1-z", "z1-4", "z1-z"],
-  ["a1-6", "a1-z", "z1-6", "z1-z"],
-  ["b1-2", "b1-z", "z1-2", "z1-z"],
-  ["b1-3", "b1-z", "z1-3", "z1-z"],
-  ["b1-4", "b1-z", "z1-4", "z1-z"],
-  ["b1-6", "b1-z", "z1-6", "z1-z"],
-  ["a2-3", "a2-z", "z2-3", "z2-z"],
-  ["a2-6", "a2-z", "z2-6", "z2-z"],
-  ["a2-8", "a2-z", "z2-8", "z2-z"]
-  ]
-
-# Update lib
-print(len(lib))
-for equivalence_class in equivalence_classes:
-  for target in equivalence_class:
-    if target not in lib:
-      lib.append(target)
-for linkage_class in linkage_classes:
-  for target in linkage_class:
-    if target not in lib:
-      lib.append(target)
-print(len(lib))
 
 def fn(model, class_list):
   def f(glycan):
     glycan = [glycan]
     label = [0]
-    data = next(iter(dataset_to_dataloader(glycan, label, batch_size=1, libr=lib)))
+    data = next(iter(dataset_to_dataloader(glycan, label, batch_size=1)))
     device = "cpu"
     if torch.cuda.is_available():
       device = "cuda:0"
@@ -63,9 +23,10 @@ def fn(model, class_list):
     x = x.to(device)
     edge_index = edge_index.to(device)
     batch = batch.to(device) 
-    pred = model(x,edge_index, batch).cpu().detach().numpy()
-    pred = np.argmax(pred)
-    pred = class_list[pred]
+    pred = model(x,edge_index, batch).cpu().detach().numpy()[0]
+    pred = np.exp(pred)/sum(np.exp(pred)) # Softmax 
+    pred = [float(x) for x in pred]
+    pred = {class_list[i]:pred[i] for i in range(15)}
     return pred
   return f
 
@@ -79,7 +40,7 @@ f = fn(model, class_list)
 demo = gr.Interface(
     fn=f,
     inputs=[gr.Textbox(label="Glycan sequence")],
-    outputs=[gr.Textbox(label="Predicted Class")],
+    outputs=[gr.Label(num_top_classes=15, label="Class prediction")],
     allow_flagging=False,
     title="SweetNet demo",
     examples=["GlcOSN(a1-4)GlcA(b1-4)GlcOSN(a1-4)GlcAOS(b1-4)GlcOSN(a1-4)GlcOSN",