import transformers
from transformers import GraphormerForGraphClassification
from transformers.models.graphormer.collating_graphormer import preprocess_item
import gradio as gr
import json
import torch

import os 
try:
    import toml
except ImportError:
    os.system('pip install toml')
    import toml
print('todo en orden')
model = GraphormerForGraphClassification.from_pretrained("PedroLancharesSanchez/graph-regression")

def predict(instancia):
    with open(instancia, "r") as archivo:
        datos=json.load(archivo)
    instancia_preprocesada=preprocess_item(datos)
    inputs={}
    inputs['input_nodes'] = torch.tensor([instancia_preprocesada['input_nodes']])
    inputs['input_edges'] = torch.tensor([instancia_preprocesada['input_edges']])
    inputs['attn_bias'] = torch.tensor([instancia_preprocesada['attn_bias']])
    inputs['in_degree'] = torch.tensor([instancia_preprocesada['in_degree']])
    inputs['out_degree'] = torch.tensor([instancia_preprocesada['out_degree']])
    inputs['spatial_pos'] = torch.tensor([instancia_preprocesada['spatial_pos']])
    inputs['attn_edge_type'] = torch.tensor([instancia_preprocesada['attn_edge_type']])
    print(inputs)
    with torch.no_grad():
        logits = model(**inputs).logits
    return str(logits.item())

ejemplos = '''{"node_feat": [[0],[0],[0],[0],[0],[0],[0],[0],[1],[0],[0],[0],[0],[1],[2],[0],[0],[0],[0],[0],[0],[3],[0],[0]],
     "edge_index": [[0,1,1,1,1,2,3,4,4,4,5,5,6,6,7,7,7,8,8,9,9,10,10,10,11,11,12,12,12,13,14,14,15,15,15,16,16,17,17,18,18,19,19,20,20,20,21,22,22,22,23,23]],
     "edge_attr": [[1],[1],[1],[1],[1],[1],[1],[1],[2],[1],[2],[1],[1],[2],[2],[1],[1],[1],[1],[1],[2],[2],[1],[1],[1],[1],[1],[2],[1],[2],[1],[1],[1],[2],[1],[2],[1],[1],[2],[2],[1],[1],[2],[1],[2],[1],[1],[1],[1],[2],[1],[2]],
     "y": [3.1381945610046387],
     "num_nodes": 24}'''

# Crear la interfaz Gradio
interfaz = gr.Interface(fn=predict, inputs="file", outputs='text', examples=['grafo1.json','grafo2.json','grafo3.json'])
interfaz.launch(share=False)