Add real-time inference

.gitignore

@@ -5,6 +5,10 @@
 # Ignore python cache files
 __pycache__/
 
+# Ignore model files
+data/*.pt
+
 # Ignore secrets
 .streamlit/secrets.toml
-.streamlit/gravity-user-db.json
+.streamlit/gravity-user-db.json
+test.ipynb

README.md

@@ -1,5 +1,5 @@
 ---
-title: Gravity
+title: GRAVITY
 emoji: 💻
 colorFrom: red
 colorTo: purple

data/kg_edge_types.csv

@@ -0,0 +1,75 @@
+x_type,relation,display_relation,y_type,direction,N
+anatomy,anatomy_protein_present,expression present,gene/protein,forward,3831782
+gene/protein,rev_anatomy_protein_present,expression present,anatomy,reverse,3831782
+drug,drug_drug,synergistic interaction,drug,forward,1433261
+drug,rev_drug_drug,synergistic interaction,drug,reverse,1433261
+anatomy,anatomy_protein_absent,expression absent,gene/protein,forward,324186
+gene/protein,rev_anatomy_protein_absent,expression absent,anatomy,reverse,324186
+gene/protein,protein_protein,ppi,gene/protein,forward,321090
+gene/protein,rev_protein_protein,ppi,gene/protein,reverse,321090
+disease,disease_phenotype_positive,phenotype present,effect/phenotype,forward,200354
+effect/phenotype,rev_disease_phenotype_positive,phenotype present,disease,reverse,200354
+disease,disease_protein,associated with,gene/protein,forward,147984
+gene/protein,rev_disease_protein,associated with,disease,reverse,147984
+biological_process,bioprocess_protein,interacts with,gene/protein,forward,138297
+gene/protein,rev_bioprocess_protein,interacts with,biological_process,reverse,138297
+cellular_component,cellcomp_protein,interacts with,gene/protein,forward,83089
+gene/protein,rev_cellcomp_protein,interacts with,cellular_component,reverse,83089
+disease,disease_protein_negative,expression downregulated,gene/protein,forward,71135
+gene/protein,rev_disease_protein_negative,expression downregulated,disease,reverse,71135
+gene/protein,molfunc_protein,interacts with,molecular_function,forward,70291
+molecular_function,rev_molfunc_protein,interacts with,gene/protein,reverse,70291
+disease,disease_protein_positive,expression upregulated,gene/protein,forward,69488
+gene/protein,rev_disease_protein_positive,expression upregulated,disease,reverse,69488
+drug,drug_effect,side effect,effect/phenotype,forward,64249
+effect/phenotype,rev_drug_effect,side effect,drug,reverse,64249
+biological_process,bioprocess_bioprocess,parent-child,biological_process,forward,50232
+biological_process,rev_bioprocess_bioprocess,parent-child,biological_process,reverse,50232
+gene/protein,pathway_protein,interacts with,pathway,forward,44116
+pathway,rev_pathway_protein,interacts with,gene/protein,reverse,44116
+disease,disease_disease,parent-child,disease,forward,37808
+disease,rev_disease_disease,parent-child,disease,reverse,37808
+disease,contraindication,contraindication,drug,forward,26899
+drug,rev_contraindication,contraindication,disease,reverse,26899
+effect/phenotype,phenotype_phenotype,parent-child,effect/phenotype,forward,20183
+effect/phenotype,rev_phenotype_phenotype,parent-child,effect/phenotype,reverse,20183
+drug,drug_protein,target,gene/protein,forward,18513
+gene/protein,rev_drug_protein,target,drug,reverse,18513
+disease,weak_clinical_evidence,clinical candidate,drug,forward,16111
+drug,rev_weak_clinical_evidence,clinical candidate,disease,reverse,16111
+anatomy,anatomy_anatomy,parent-child,anatomy,forward,14383
+anatomy,rev_anatomy_anatomy,parent-child,anatomy,reverse,14383
+molecular_function,molfunc_molfunc,parent-child,molecular_function,forward,13735
+molecular_function,rev_molfunc_molfunc,parent-child,molecular_function,reverse,13735
+disease,indication,indication,drug,forward,12608
+drug,rev_indication,indication,disease,reverse,12608
+drug,drug_protein,enzyme,gene/protein,forward,5919
+gene/protein,rev_drug_protein,enzyme,drug,reverse,5919
+disease,strong_clinical_evidence,clinical candidate,drug,forward,5352
+drug,rev_strong_clinical_evidence,clinical candidate,disease,reverse,5352
+cellular_component,cellcomp_cellcomp,parent-child,cellular_component,forward,4683
+cellular_component,rev_cellcomp_cellcomp,parent-child,cellular_component,reverse,4683
+effect/phenotype,phenotype_protein,associated with,gene/protein,forward,4437
+gene/protein,rev_phenotype_protein,associated with,effect/phenotype,reverse,4437
+drug,drug_protein,transporter,gene/protein,forward,3349
+gene/protein,rev_drug_protein,transporter,drug,reverse,3349
+pathway,pathway_pathway,parent-child,pathway,forward,2647
+pathway,rev_pathway_pathway,parent-child,pathway,reverse,2647
+disease,exposure_disease,linked to,exposure,forward,2421
+exposure,rev_exposure_disease,linked to,disease,reverse,2421
+disease,off_label_use,off-label use,drug,forward,2370
+drug,rev_off_label_use,off-label use,disease,reverse,2370
+exposure,exposure_exposure,parent-child,exposure,forward,2263
+exposure,rev_exposure_exposure,parent-child,exposure,reverse,2263
+exposure,exposure_protein,interacts with,gene/protein,forward,2012
+gene/protein,rev_exposure_protein,interacts with,exposure,reverse,2012
+biological_process,exposure_bioprocess,interacts with,exposure,forward,1990
+exposure,rev_exposure_bioprocess,interacts with,biological_process,reverse,1990
+drug,drug_protein,carrier,gene/protein,forward,993
+gene/protein,rev_drug_protein,carrier,drug,reverse,993
+disease,disease_phenotype_negative,phenotype absent,effect/phenotype,forward,508
+effect/phenotype,rev_disease_phenotype_negative,phenotype absent,disease,reverse,508
+exposure,exposure_molfunc,interacts with,molecular_function,forward,45
+molecular_function,rev_exposure_molfunc,interacts with,exposure,reverse,45
+cellular_component,exposure_cellcomp,interacts with,exposure,forward,12
+exposure,rev_exposure_cellcomp,interacts with,cellular_component,reverse,12

data/kg_node_types.csv

@@ -0,0 +1,11 @@
+node_type,N
+gene/protein,35198
+biological_process,27668
+disease,22201
+effect/phenotype,16711
+anatomy,14384
+molecular_function,11228
+drug,8160
+cellular_component,4054
+pathway,2629
+exposure,860

pages/about.py

@@ -14,4 +14,7 @@ menu_with_redirect()
 st.image(str(project_config.MEDIA_DIR / 'about_header.svg'), use_column_width=True)
 
 # Main content
-st.markdown(f"Hello, {st.session_state.name}! Welcome to GRAVITY, a GRaph AI VISualization Tool to query and visualize knowledge graph-grounded biomedical AI models.")
+st.markdown(f"Hello, {st.session_state.name}! Welcome to GRAVITY, a **GR**aph **A**I **VI**sualization **T**ool to query and visualize knowledge graph-grounded biomedical AI models.")
+
+# Subheader
+st.subheader("About GRAVITY", divider = "grey")

pages/explore.py

@@ -14,4 +14,7 @@ menu_with_redirect()
 st.image(str(project_config.MEDIA_DIR / 'explore_header.svg'), use_column_width=True)
 
 # Main content
-st.markdown(f"Hello, {st.session_state.name}!")
+# st.markdown(f"Hello, {st.session_state.name}!")
+
+# Coming soon
+st.write("Coming soon...")

pages/input.py

@@ -1,6 +1,10 @@
 import streamlit as st
 from menu import menu_with_redirect
 
+# Standard imports
+import numpy as np
+import pandas as pd
+
 # Path manipulation
 from pathlib import Path
 
@@ -14,4 +18,53 @@ menu_with_redirect()
 st.image(str(project_config.MEDIA_DIR / 'input_header.svg'), use_column_width=True)
 
 # Main content
-st.markdown(f"Hello, {st.session_state.name}!")
+# st.markdown(f"Hello, {st.session_state.name}!")
+
+st.subheader("Construct Query", divider = "red")
+
+# Checkbox to allow reverse edges
+allow_reverse_edges = st.checkbox("Reverse Edges", value = False)
+
+with st.spinner('Loading knowledge graph...'):
+    kg_nodes = nodes = pd.read_csv(project_config.DATA_DIR / 'kg_nodes.csv', dtype = {'node_index': int}, low_memory = False)
+    node_types = pd.read_csv(project_config.DATA_DIR / 'kg_node_types.csv')
+    edge_types = pd.read_csv(project_config.DATA_DIR / 'kg_edge_types.csv')
+
+    if not allow_reverse_edges:
+       edge_types = edge_types[edge_types.direction == 'forward']
+
+# Select source node type
+source_node_type = st.selectbox("Source Node Type", node_types['node_type'],
+                                format_func = lambda x: x.replace("_", " "))
+
+# Select source node
+source_node = st.selectbox("Source Node", kg_nodes[kg_nodes['node_type'] == source_node_type]['node_name'])
+
+# Select target node type
+target_node_type = st.selectbox("Target Node Type", edge_types[edge_types.x_type == source_node_type].y_type.unique(),
+                                format_func = lambda x: x.replace("_", " "))
+
+# Select relation
+relation = st.selectbox("Edge Type", edge_types[(edge_types.x_type == source_node_type) & (edge_types.y_type == target_node_type)].relation.unique(),
+                        format_func = lambda x: x.replace("_", "-"))
+
+# Button to submit query
+if st.button("Submit Query"):
+    
+    # Save query to session state
+    st.session_state.query = {
+        "source_node_type": source_node_type,
+        "source_node": source_node,
+        "target_node_type": target_node_type,
+        "relation": relation
+    }
+
+    # # Write query to console
+    # st.write("Current Query:")
+    # st.write(st.session_state.query)
+    st.write("Query submitted.")
+
+st.subheader("Knowledge Graph", divider = "red")    
+display_data = kg_nodes[['node_id', 'node_type', 'node_name', 'node_source']].copy()
+display_data = display_data.rename(columns = {'node_id': 'ID', 'node_type': 'Type', 'node_name': 'Name', 'node_source': 'Database'})
+st.dataframe(display_data, use_container_width = True)

pages/validate.py

@@ -1,8 +1,16 @@
 import streamlit as st
 from menu import menu_with_redirect
 
+# Standard imports
+import numpy as np
+import pandas as pd
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
 # Path manipulation
 from pathlib import Path
+from huggingface_hub import hf_hub_download
 
 # Custom and other imports
 import project_config
@@ -14,4 +22,93 @@ menu_with_redirect()
 st.image(str(project_config.MEDIA_DIR / 'validate_header.svg'), use_column_width=True)
 
 # Main content
-st.markdown(f"Hello, {st.session_state.name}!")
+# st.markdown(f"Hello, {st.session_state.name}!")
+
+st.subheader("Model Predictions", divider = "green")
+
+# Print current query
+st.markdown(f"**Query:** {st.session_state.query['source_node']} ➡️ {st.session_state.query['relation']} ➡️ {st.session_state.query['target_node_type']}")
+
+with st.spinner('Loading knowledge graph...'):
+    kg_nodes = nodes = pd.read_csv(project_config.DATA_DIR / 'kg_nodes.csv', dtype = {'node_index': int}, low_memory = False)
+
+# Get paths to embeddings, relation weights, and edge types
+with st.spinner('Downloading AI model...'):
+    embed_path = hf_hub_download(repo_id="ayushnoori/galaxy",
+                                filename="2024_03_29_04_12_52_epoch=3-step=54291_embeddings.pt",
+                                token=st.secrets["HF_TOKEN"])
+    relation_weights_path = hf_hub_download(repo_id="ayushnoori/galaxy",
+                                            filename="2024_03_29_04_12_52_epoch=3-step=54291_relation_weights.pt",
+                                            token=st.secrets["HF_TOKEN"])
+    edge_types_path = hf_hub_download(repo_id="ayushnoori/galaxy",
+                                        filename="2024_03_29_04_12_52_epoch=3-step=54291_edge_types.pt",
+                                        token=st.secrets["HF_TOKEN"])
+
+# Load embeddings, relation weights, and edge types
+with st.spinner('Loading AI model...'):
+    embeddings = torch.load(embed_path)
+    relation_weights = torch.load(relation_weights_path)
+    edge_types = torch.load(edge_types_path)
+
+# # Print source node type
+# st.write(f"Source Node Type: {st.session_state.query['source_node_type']}")
+
+# # Print source node
+# st.write(f"Source Node: {st.session_state.query['source_node']}")
+
+# # Print relation
+# st.write(f"Edge Type: {st.session_state.query['relation']}")
+
+# # Print target node type
+# st.write(f"Target Node Type: {st.session_state.query['target_node_type']}")
+
+# Compute predictions
+with st.spinner('Computing predictions...'):
+
+    source_node_type = st.session_state.query['source_node_type']
+    source_node = st.session_state.query['source_node']
+    relation = st.session_state.query['relation']
+    target_node_type = st.session_state.query['target_node_type']
+
+    # Get source node index
+    src_index = kg_nodes[(kg_nodes.node_type == source_node_type) & (kg_nodes.node_name == source_node)].node_index.values[0]
+
+    # Get relation index
+    edge_type_index = [i for i, etype in enumerate(edge_types) if etype == (source_node_type, relation, target_node_type)][0]
+
+    # Get target nodes indices
+    target_nodes = kg_nodes[kg_nodes.node_type == target_node_type]
+    dst_indices = target_nodes.node_index.values
+    src_indices = np.repeat(src_index, len(dst_indices))
+
+    # Retrieve cached embeddings
+    src_embeddings = embeddings[src_indices]
+    dst_embeddings = embeddings[dst_indices]
+
+    # Apply activation function
+    src_embeddings = F.leaky_relu(src_embeddings)
+    dst_embeddings = F.leaky_relu(dst_embeddings)
+
+    # Get relation weights
+    rel_weights = relation_weights[edge_type_index]
+
+    # Compute weighted dot product
+    scores = torch.sum(src_embeddings * rel_weights * dst_embeddings, dim = 1)
+    scores = torch.sigmoid(scores)
+
+    # Add scores to dataframe
+    target_nodes['score'] = scores.detach().numpy()
+
+    # Rank target nodes by score
+    target_nodes = target_nodes.sort_values(by = 'score', ascending = False)
+
+    # Add rank to dataframe
+    target_nodes['rank'] = np.arange(1, target_nodes.shape[0] + 1)
+
+    # Show top ranked nodes
+    top_k = st.slider('Select number of top ranked nodes to show.', 1, target_nodes.shape[0], 50) 
+
+    # Rename columns
+    display_data = target_nodes[['rank', 'node_id', 'node_name', 'node_source', 'score']].iloc[:top_k].copy()
+    display_data = display_data.rename(columns = {'rank': 'Rank', 'node_id': 'ID', 'node_name': 'Name', 'node_source': 'Database', 'score': 'Score'})
+    st.dataframe(display_data, use_container_width = True)

requirements.txt

@@ -7,4 +7,5 @@ pathlib
 torch
 altair<5
 gspread
-oauth2client
+oauth2client
+huggingface_hub