GNN-MUTAG / app.py
AdrianRevi's picture
Create app.py
d5f2d19 verified
raw
history blame
1.2 kB
import gradio as gr
import torch
from torch_geometric.data import Data
from torch_geometric.utils import from_networkx
import networkx as nx
# Usamos el modelo GCN previamente entrenado (podrías cambiar por GAT si lo prefieres)
model.eval()
def predict_mutagenicity():
# Creamos un grafo de prueba simple (3 nodos conectados)
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
# 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)
pyg_data = pyg_data.to(device)
with torch.no_grad():
out = model(pyg_data.x, pyg_data.edge_index, pyg_data.batch)
pred = out.argmax(dim=1).item()
return "Mutagénico" if pred == 1 else "No mutagénico"
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()