Update app.py

app.py

@@ -1,31 +1,57 @@
-import gradio as gr
 import torch
-from torch_geometric.data import Data
+import torch.nn.functional as F
+from torch_geometric.nn import GCNConv, global_mean_pool
 from torch_geometric.utils import from_networkx
 import networkx as nx
+import gradio as gr
+
+# --- Mismo modelo que en Colab ---
+class GCN(torch.nn.Module):
+    def __init__(self, hidden_channels=64):
+        super().__init__()
+        self.conv1 = GCNConv(7, hidden_channels)
+        self.conv2 = GCNConv(hidden_channels, hidden_channels)
+        self.lin = torch.nn.Linear(hidden_channels, 2)
 
-# Usamos el modelo GCN previamente entrenado (podrías cambiar por GAT si lo prefieres)
+    def forward(self, x, edge_index, batch):
+        x = self.conv1(x, edge_index)
+        x = F.relu(x)
+        x = self.conv2(x, edge_index)
+        x = F.relu(x)
+        x = global_mean_pool(x, batch)
+        return self.lin(x)
+
+# --- Carga de modelo entrenado ---
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+model = GCN().to(device)
+model.load_state_dict(torch.load("model_gcn.pth", map_location=device))
 model.eval()
 
-def predict_mutagenicity():
-    # Creamos un grafo de prueba simple (3 nodos conectados)
+# --- Función de predicción sobre un grafo ejemplo ---
+def demo_predict():
     G = nx.Graph()
     G.add_edges_from([(0, 1), (1, 2)])
-    nx.set_node_attributes(G, {i: [1, 0, 0, 1, 0, 1, 0] for i in G.nodes}, "x")  # vector ficticio
+    nx.set_node_attributes(G, {i: [1, 0, 0, 1, 0, 1, 0] for i in G.nodes}, "x")  # vector de 7 dimensiones
 
-    # Convertimos a objeto PyG
-    pyg_data = from_networkx(G)
-    pyg_data.x = torch.tensor(list(nx.get_node_attributes(G, 'x').values()), dtype=torch.float)
-    pyg_data.edge_index = pyg_data.edge_index
-    pyg_data.batch = torch.tensor([0] * pyg_data.num_nodes)
+    data = from_networkx(G)
+    data.x = torch.tensor(list(nx.get_node_attributes(G, "x").values()), dtype=torch.float)
+    data.edge_index = data.edge_index
+    data.batch = torch.tensor([0] * data.num_nodes)
 
-    pyg_data = pyg_data.to(device)
+    data = data.to(device)
     with torch.no_grad():
-        out = model(pyg_data.x, pyg_data.edge_index, pyg_data.batch)
+        out = model(data.x, data.edge_index, data.batch)
         pred = out.argmax(dim=1).item()
 
-    return "Mutagénico" if pred == 1 else "No mutagénico"
+    return "Mutagénico ✅" if pred == 1 else "No mutagénico ❌"
+
+# --- Interfaz Gradio ---
+demo = gr.Interface(
+    fn=demo_predict,
+    inputs=[],
+    outputs="text",
+    title="Clasificador de Moléculas con GCN",
+    description="Este demo usa una red neuronal en grafo entrenada sobre MUTAG para clasificar moléculas como mutagénicas o no."
+)
 
-gr.Interface(fn=predict_mutagenicity, inputs=[], outputs="text",
-             title="Clasificador de Moléculas con GNN",
-             description="Demo simple de GCN sobre grafos moleculares (MUTAG)").launch()
+demo.launch()