YuWang0103/LGGM-Text2Graph

1from omegaconf import OmegaConf2import gradio as gr3 4from dataset import init_dataset, compute_input_output_dims5from extra_features import ExtraFeatures6from demo_model import LGGMText2Graph_Demo7from analysis.spectre_utils import CrossDomainSamplingMetrics8import networkx as nx9import numpy as np10import matplotlib.pyplot as plt11import torch12 13 14cfg = OmegaConf.load('./config.yaml')15hydra_path = '.'16 17 18data_loaders, num_classes, max_n_nodes, nodes_dist, edge_types, node_types, n_nodes, cond_dims, cond_emb = init_dataset(cfg.dataset.name, cfg.train.batch_size, hydra_path, cfg.general.condition, cfg.model.transition)19 20extra_features = ExtraFeatures(cfg.model.extra_features, max_n_nodes)21 22input_dims, output_dims = compute_input_output_dims(data_loaders['train'], extra_features)23 24sampling_metrics = CrossDomainSamplingMetrics(data_loaders)25 26# model = LGGMText2Graph_Demo.load_from_checkpoint('cc-deg.ckpt', map_location=torch.device('cpu'))27model = LGGMText2Graph_Demo.load_from_checkpoint('cc-deg.ckpt', map_location=torch.device("cpu"))28 29model.init_prompt_encoder_pretrained()30 31def calculate_average_degree(graph):32    num_nodes = graph.number_of_nodes()33    num_edges = graph.number_of_edges()34    return (2 * num_edges) / num_nodes if num_nodes > 0 else 035 36 37def predict(text, num_nodes = None):38    # Assuming model.generate and other processes are defined as before39    graphs = model.generate_pretrained(text, int(num_nodes))40    ccs = []41    degs = []42    images = []43 44    for g in graphs:45        ccs.append(nx.average_clustering(g))46        degs.append(calculate_average_degree(g))47 48        fig, ax = plt.subplots()49        nx.draw(g, ax=ax)50        fig.canvas.draw()51        image = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)52        image = image.reshape(fig.canvas.get_width_height()[::-1] + (3,))53        plt.close(fig)54        55        images.append(image)56    57    avg_deg = np.mean(degs)58    avg_cc = np.mean(ccs)59 60    return images[0], images[1], images[2], ccs[0], ccs[1], ccs[2], degs[0], degs[1], degs[2], avg_cc, avg_deg61 62def clear(input_text):63    return None, None, None, None, None, None, None, None, None, None, None64 65 66with gr.Blocks() as demo:67    gr.Markdown("## Text2Graph Generation Demo")68    with gr.Row():69        with gr.Column():70            input_text = gr.Textbox(label="Input your text prompt here", placeholder="Type here...")71        with gr.Column():72            input_num = gr.Slider(5, 100, value=25, step = 1, label="Count", info="Number of nodes in the graph to be generated")73        with gr.Column():74            gr.Markdown("### Suggested Prompts")75            gr.Markdown("1. Create a complex network with high clustering coefficient.\n2. Create a graph with extremely low number of triangles.\n 3. Please give me a Power Network with extremely low number of triangles but with medium level of average degree.")76 77    with gr.Row() as output_row:78        output_images = [gr.Image(label = f"Generated Network #{_}") for _ in range(3)]79    with gr.Row():80        output_texts_cc = [gr.Textbox(label=f"CC #{_}") for _ in range(3)]81    with gr.Row():82        output_texts_deg = [gr.Textbox(label=f"DEG #{_}") for _ in range(3)]83    84    with gr.Row():85        avg_cc_text = gr.Textbox(label="Average Clustering Coefficient")86        avg_deg_text = gr.Textbox(label="Average Degree")87 88    with gr.Row():89        submit_button = gr.Button("Submit")90        clear_button = gr.Button("Clear")91 92    # Change function is linked to the submit button93    submit_button.click(fn=predict, inputs=[input_text, input_num], outputs=output_images + output_texts_cc + output_texts_deg + [avg_cc_text, avg_deg_text])94    input_text.submit(fn=predict, inputs=[input_text, input_num], outputs=output_images + output_texts_cc + output_texts_deg + [avg_cc_text, avg_deg_text])95 96    # Clear function resets the text input and clears the outputs97    clear_button.click(fn=clear, inputs=[input_text], outputs=output_images + output_texts_cc + output_texts_deg + [avg_cc_text, avg_deg_text])98 99demo.launch()