Upload app.py

app.py

@@ -6,8 +6,26 @@ import gradio as gr
 from glycowork.ml.processing import dataset_to_dataloader
 import numpy as np
 import torch
+import torch.nn as nn
 
-
+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']