Add CONCURRENCY_LIMIT; Graph config change -> directed

interactive/pyvis_graph.py

@@ -3,12 +3,12 @@ import json
 import torch
 import gradio as gr
 from pyvis.network import Network
-
 sys.path.append(".")
+import re
 from src.benchmarks import get_semistructured_data
 
-
-TITLE = "STaRK Knowledge Base Explorer"
+CONCURRENCY_LIMIT = 1000
+TITLE = "STaRK Semistructure Knowledge Base Explorer"
 BRAND_NAME = {
     "amazon": "STaRK-Amazon",
     "mag": "STaRK-MAG",
@@ -22,20 +22,16 @@ NODE_COLORS = [
     "#00796B",  # Teal
     "#03A9F4",  # Light Blue
     "#CDDC39",  # Lime
-    "#E91E63",  # Pink
     "#3F51B5",  # Indigo
     "#00BCD4",  # Cyan
     "#FFC107",  # Amber
     "#8BC34A",  # Light Green
-    "#795548",  # Brown
     "#9E9E9E",  # Grey
     "#607D8B",  # Blue Grey
     "#FFEB3B",  # Bright Yellow
     "#E1F5FE",  # Light Blue 50
     "#F1F8E9",  # Light Green 50
     "#FFF3E0",  # Orange 50
-    "#FCE4EC",  # Pink 50
-    "#F3E5F5",  # Purple 50
     "#FFFDE7",  # Yellow 50
     "#E0F7FA",  # Cyan 50
     "#E8F5E9",  # Green 50
@@ -90,10 +86,14 @@ def relabel(x, edge_index, batch, pos=None):
     return x, edge_index, batch, pos
 
 
-def generate_network(kb, node_id, max_nodes=10, num_hops="1"):
+def generate_network(kb, node_id, max_nodes=10, num_hops='2'):
     max_nodes = int(max_nodes)
-
-    net = Network()
+    if 'gene/protein' in kb.node_type_dict.values():
+        indirected = True
+        net = Network(directed=False)
+    else:
+        indirected = False
+        net = Network()
 
     def get_one_hop(kb, node_id, max_nodes):
         edge_index = kb.edge_index
@@ -137,7 +137,6 @@ def generate_network(kb, node_id, max_nodes=10, num_hops="1"):
     node_ids, relabel_edge_index, _, _ = relabel(
         torch.arange(kb.num_nodes()), edge_index, batch=torch.zeros(kb.num_nodes())
     )
-
     for idx, n_id in enumerate(node_ids):
         if node_id == n_id:
             net.add_node(
@@ -158,31 +157,45 @@ def generate_network(kb, node_id, max_nodes=10, num_hops="1"):
                 font={"align": "middle", "size": 10},
             )
     for idx in range(relabel_edge_index.size(-1)):
-        net.add_edge(
-            relabel_edge_index[0][idx].item(),
-            relabel_edge_index[1][idx].item(),
-            color=EDGE_COLORS[edge_types[idx].item()],
-            label=kb.edge_type_dict[edge_types[idx].item()]
-            .replace("___", " ")
-            .replace("_", " "),
-            width=1,
-            font={"align": "middle", "size": 10},
-        )
-
+        if indirected:
+            net.add_edge(
+                relabel_edge_index[0][idx].item(),
+                relabel_edge_index[1][idx].item(),
+                color=EDGE_COLORS[edge_types[idx].item()],
+                label=kb.edge_type_dict[edge_types[idx].item()]
+                .replace('___', " ")
+                .replace('_', " "),
+                width=1,
+                font={"align": "middle", "size": 10})
+        else:
+            net.add_edge(
+                relabel_edge_index[0][idx].item(),
+                relabel_edge_index[1][idx].item(),
+                color=EDGE_COLORS[edge_types[idx].item()],
+                label=kb.edge_type_dict[edge_types[idx].item()]
+                .replace('___', " ")
+                .replace('_', " "),
+                width=1,
+                font={"align": "middle", "size": 10},
+                arrows="to",
+                arrowStrikethrough=False)
     return net.get_network_data()
 
 
 def get_text_html(kb, node_id):
     text = kb.get_doc_info(node_id, add_rel=False, compact=False)
-    # need a text box, figure left, text right
-    text = text.replace("\n", "<br>").replace(" ", "&nbsp;")
     # add a title
+    text = text.replace("\n", "<br>").replace(" ", "&nbsp;")
     text = f"<h3>Textual Info of Entity {node_id}:</h3>{text}"
+    text = re.sub(r"\$([^$]+)\$", r"\\(\1\\)", text)
     # show the text as what it is with empty space and can be scrolled
-    return f"""<div style="width: 100%; height: 600px; overflow-x: hidden; overflow-y: scroll; overflow-wrap: break-word; padding: 10px; margin: 0 auto; border: 1px solid #ccc;">{text}</div>"""
+    return f"""<script src="https://polyfill.io/v3/polyfill.min.js?features=es6"></script>
+        <script id="MathJax-script" async src="https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js"></script>
+        <div style="width: 100%; height: 600px; overflow-x: hidden; overflow-y: scroll; overflow-wrap: break-word; hyphens: auto; padding: 10px; margin: 0 auto; border: 1px solid #ccc; line-height: 1.5;
+        font-family: SF Pro Text, SF Pro Icons, Helvetica Neue, Helvetica, Arial, sans-serif;">{text}</div>"""
 
 
-def get_subgraph_html(kb, kb_name, node_id, max_nodes=10, num_hops="1"):
+def get_subgraph_html(kb, kb_name, node_id, max_nodes=10, num_hops=1):
     network = generate_network(kb, node_id, max_nodes, num_hops)
 
     nodes = network[0]
@@ -200,7 +213,7 @@ def get_subgraph_html(kb, kb_name, node_id, max_nodes=10, num_hops="1"):
 
 def main():
     # kb = get_semistructured_data(DATASET_NAME)
-    kbs = {k: get_semistructured_data(k) for k in BRAND_NAME.keys()}
+    kbs = {k: get_semistructured_data(k, indirected=False) for k in BRAND_NAME.keys()}
 
     with gr.Blocks(head=VISJS_HEAD, title=TITLE) as demo:
         gr.Markdown(f"# {TITLE}")
@@ -208,11 +221,14 @@ def main():
             with gr.Tab(BRAND_NAME[name]):
                 with gr.Row():
                     entity_id = gr.Number(
-                        label="Entity ID", elem_id=f"{name}-entity-id-input"
+                        label="Entity ID", 
+                        elem_id=f"{name}-entity-id-input"
+                    )
+                    max_paths = gr.Slider(
+                        1, 200, 10, step=1, label="Max Number of Paths"
                     )
-                    max_paths = gr.Slider(1, 200, 10, step=1, label="Max Paths")
                     num_hops = gr.Dropdown(
-                        ["1", "2", "inf"], value="1", label="Number of Hops"
+                        ["1", "2", "inf"], value="2", label="Number of Hops"
                     )
                     query_btn = gr.Button(
                         value="Show Graph",
@@ -232,7 +248,7 @@ def main():
                     ),
                     inputs=[entity_id, max_paths, num_hops],
                     outputs=[graph_area, text_area],
-                    api_name=f"{name}-fetch-graph",
+                    api_name=f"{name}-fetch-graph"
                 )
 
                 # Hidden inputs for fetch just text
@@ -248,11 +264,12 @@ def main():
                     lambda e, kb=kb: get_text_html(kb, e),
                     inputs=[entity_for_text],
                     outputs=text_area,
-                    api_name=f"{name}-fetch-text",
+                    api_name=f"{name}-fetch-text"
                 )
-
+    demo.queue(max_size=2*CONCURRENCY_LIMIT, default_concurrency_limit=CONCURRENCY_LIMIT)
     demo.launch(share=True)
 
 
 if __name__ == "__main__":
-    main()
+
+    main()

src/benchmarks/get_semistruct.py

@@ -2,21 +2,22 @@ import os.path as osp
 from src.benchmarks.semistruct import AmazonSemiStruct, PrimeKGSemiStruct, MagSemiStruct
 
 
-def get_semistructured_data(name, root='data/', download_processed=True):
+def get_semistructured_data(name, root='data/', download_processed=True, **kwargs):
     data_root = osp.join(root, name)
     if name == 'amazon':
         categories = ['Sports_and_Outdoors']
         kb = AmazonSemiStruct(root=data_root,
-                                    categories=categories,
-                                    meta_link_types=['brand'],
-                                    indirected=True,
-                                    download_processed=download_processed
-                                    )
+                                categories=categories,
+                                meta_link_types=['brand'],
+                                download_processed=download_processed,
+                                **kwargs
+                                )
     if name == 'primekg':
         kb = PrimeKGSemiStruct(root=data_root, 
-                                     download_processed=download_processed)
+                               download_processed=download_processed,
+                               **kwargs)
     
     if name == 'mag':
         kb = MagSemiStruct(root=data_root, 
-                                 download_processed=download_processed)
+                           download_processed=download_processed)
     return kb

src/benchmarks/semistruct/amazon.py

@@ -63,8 +63,8 @@ class AmazonSemiStruct(SemiStructureKB):
                  categories: list, 
                  meta_link_types=['brand'],
                  max_entries=25,
-                 indirected=True,
-                 download_processed=True):
+                 download_processed=True,
+                 **kwargs):
         '''
             Args: 
                 root (str): root directory to store the data
@@ -108,7 +108,7 @@ class AmazonSemiStruct(SemiStructureKB):
             if meta_link_types: 
                 # customize the graph by adding meta links
                 processed_data = self.post_process(processed_data, meta_link_types=meta_link_types, cache_path=cache_path)
-        super(AmazonSemiStruct, self).__init__(**processed_data, indirected=indirected)
+        super(AmazonSemiStruct, self).__init__(**processed_data, **kwargs)
     
     def __getitem__(self, idx):
         idx = int(idx)

src/benchmarks/semistruct/mag.py

@@ -40,7 +40,7 @@ class MagSemiStruct(SemiStructureKB):
     ogbn_papers100M_url = 'https://snap.stanford.edu/ogb/data/misc/ogbn_papers100M/paperinfo.zip'
     mag_mapping_url = 'https://zenodo.org/records/2628216/files'
     
-    def __init__(self, root, download_processed=True):
+    def __init__(self, root, download_processed=True, **kwargs):
         '''
         Args:
             root (str): root directory to store the dataset folder
@@ -88,7 +88,7 @@ class MagSemiStruct(SemiStructureKB):
             processed_data = self._process_raw()
         processed_data.update({'node_type_dict': self.node_type_dict, 
                                'edge_type_dict': self.edge_type_dict})
-        super(MagSemiStruct, self).__init__(**processed_data)
+        super(MagSemiStruct, self).__init__(**processed_data, **kwargs)
 
     def load_edge(self, edge_type):
         edge_dir = osp.join(self.graph_data_root, f"raw/relations/{edge_type}/edge.csv.gz")

src/benchmarks/semistruct/primekg.py

@@ -30,7 +30,7 @@ class PrimeKGSemiStruct(SemiStructureKB):
     candidate_types = NODE_TYPES 
     raw_data_url = 'https://drive.google.com/uc?id=1d__3yP6YZYjKWR2F9fGg-y1rW7-HJPpr'
 
-    def __init__(self, root, download_processed=True):
+    def __init__(self, root, download_processed=True, **kwargs):
         '''
         Args:
             root (str): root directory to store the dataset folder
@@ -61,7 +61,7 @@ class PrimeKGSemiStruct(SemiStructureKB):
             print(f'Loaded from {self.processed_data_dir}!')
         else:
             processed_data = self._process_raw()
-        super(PrimeKGSemiStruct, self).__init__(**processed_data)
+        super(PrimeKGSemiStruct, self).__init__(**processed_data, **kwargs)
 
         self.node_info = clean_dict(self.node_info)
         self.node_attr_dict = {}