Spaces:
Runtime error
Runtime error
import os | |
os.system("pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cpu") | |
os.system("pip install torch-scatter torch-sparse torch-cluster torch-spline-conv torch-geometric -f https://data.pyg.org/whl/torch-1.12.0+cpu.html") | |
import gradio as gr | |
from glycowork.ml.processing import dataset_to_dataloader | |
import numpy as np | |
import torch | |
import torch.nn as nn | |
from glycowork.motif.graph import glycan_to_nxGraph | |
import networkx as nx | |
import matplotlib.pyplot as plt | |
from networkx.drawing.nx_agraph import write_dot | |
# import pygraphviz as pgv | |
class EnsembleModel(nn.Module): | |
def __init__(self, models): | |
super().__init__() | |
self.models = models | |
def forward(self, data): | |
# Check if GPU available | |
device = "cpu" | |
if torch.cuda.is_available(): | |
device = "cuda:0" | |
# Prepare data | |
x = data.labels.to(device) | |
edge_index = data.edge_index.to(device) | |
batch = data.batch.to(device) | |
y_pred = [model(x,edge_index, batch).cpu().detach().numpy() for model in self.models] | |
y_pred = np.mean(y_pred,axis=0)[0] | |
return y_pred | |
class_list=['Amoebozoa', 'Animalia', 'Bacteria', 'Bamfordvirae', 'Chromista', 'Euryarchaeota', 'Excavata', 'Fungi', 'Heunggongvirae', | |
'Orthornavirae', 'Pararnavirae', 'Plantae', 'Proteoarchaeota', 'Protista', 'Riboviria'] | |
model1 = torch.load("model1.pt", map_location=torch.device('cpu')) | |
model2 = torch.load("model2.pt", map_location=torch.device('cpu')) | |
model3 = torch.load("model3.pt", map_location=torch.device('cpu')) | |
def fn(glycan, model): | |
# Draw graph | |
graph = glycan_to_nxGraph(glycan) | |
node_labels = nx.get_node_attributes(graph, 'string_labels') | |
labels = {i:node_labels[i] for i in range(len(graph.nodes))} | |
graph = nx.relabel_nodes(graph, labels) | |
write_dot(graph, "graph.dot") | |
graph=pgv.AGraph("graph.dot") | |
graph.layout(prog='dot') | |
graph.draw("graph.png") | |
# Perform inference | |
if model == "No data augmentation": | |
model_pred = model1 | |
model_pred.eval() | |
elif model == "Ensemble": | |
model_pred = model3 | |
model_pred.eval() | |
else: | |
model_pred = model2 | |
model_pred.eval() | |
glycan = [glycan] | |
label = [0] | |
data = next(iter(dataset_to_dataloader(glycan, label, batch_size=1))) | |
if model == "Ensemble": | |
pred = model_pred(data) | |
else: | |
device = "cpu" | |
x = data.labels | |
edge_index = data.edge_index | |
batch = data.batch | |
x = x.to(device) | |
edge_index = edge_index.to(device) | |
batch = batch.to(device) | |
pred = model_pred(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, "graph.png" | |
demo = gr.Interface( | |
fn=fn, | |
inputs=[gr.Textbox(label="Glycan sequence"), gr.Radio(label="Model",choices=["No data augmentation", "Random node deletion", "Ensemble"])], | |
outputs=[gr.Label(num_top_classes=15, label="Prediction"), gr.Image(label="Glycan graph")], | |
allow_flagging=False, | |
title="SweetNet demo", | |
examples=[["GlcOSN(a1-4)GlcA(b1-4)GlcOSN(a1-4)GlcAOS(b1-4)GlcOSN(a1-4)GlcOSN", "No data augmentation"], | |
["Man(a1-2)Man(a1-3)[Man(a1-3)Man(a1-6)]Man(b1-4)GlcNAc(b1-4)GlcNAc", "Random node deletion"], | |
["Man(a1-2)Man(a1-3)[Man(a1-6)]Man(a1-6)[Man(a1-2)Man(a1-2)Man(a1-3)]Man(b1-4)GlcNAc", "Ensemble"]] | |
) | |
demo.launch(debug=True) |