import os
import streamlit as st
from rapidfuzz import process
import random
# with st.spinner("Initializing the environment... This may take up to 10 minutes at the start of each session."):
# # Create a temporary placeholder for the message
# loading_placeholder = st.empty()
# # Show the info message only while the spinner is active
# loading_placeholder.info("""
# **Note:** This initialization is required at the start of each session.
# Once the app is ready, you can run multiple predictions without re-initializing by clicking the **Reset** button in the sidebar.
# """)
# # Run setup script if not already executed
# if not os.path.exists(".setup_done"):
# start_time = time.time()
# os.system("bash setup.sh")
# end_time = time.time()
# print(f"Environment prepared in {end_time - start_time:.2f} seconds")
# with open(".setup_done", "w") as f:
# f.write("done")
# # ❌ Remove the info message after initialization is complete
# loading_placeholder.empty()
from run_prothgt_app import *
from visualize_kg import *
def convert_df(df):
return df.to_csv(index=False).encode('utf-8')
# Initialize session state variables
if 'predictions_df' not in st.session_state:
st.session_state.predictions_df = None
if 'heterodata' not in st.session_state:
st.session_state.heterodata = None
if 'submitted' not in st.session_state:
st.session_state.submitted = False
if 'previous_inputs' not in st.session_state:
st.session_state.previous_inputs = None
if 'generating_predictions' not in st.session_state:
st.session_state.generating_predictions = False
if 'protein_visualizations' not in st.session_state:
st.session_state.protein_visualizations = {}
def reset_prediction_state():
st.session_state.generating_predictions = False
st.session_state.submitted = False
st.session_state.predictions_df = None
st.session_state.previous_inputs = None
# Clean up visualization files
if 'protein_visualizations' in st.session_state:
for viz_info in st.session_state.protein_visualizations.values():
try:
os.unlink(viz_info['path'])
except:
pass
st.session_state.protein_visualizations = {}
def set_generating_predictions():
st.session_state.generating_predictions = True
st.session_state.submitted = True
with st.expander("🚀 Upcoming Features"):
st.info("""
We are actively working on enhancing ProtHGT application with new capabilities:
- **Real-time data retrieval for new proteins**: Currently, ProtHGT can only generate predictions for proteins that already exist in our knowledge graph. We are developing a new feature that will allow users to **predict functions for entirely new proteins starting from their sequences**. This will work by **retrieving relevant relationship data in real time from external source databases** (e.g., UniProt, STRING, and other biological repositories). The system will dynamically construct a knowledge graph for the query protein, incorporating its interactions, domains, pathways, and other biological associations before running function prediction. This approach will enable ProtHGT to analyze newly discovered or less-studied proteins even if they are not pre-annotated in our dataset.
- **Expanded embedding options**: Currently, this application represents proteins using **TAPE embeddings**, which serve as the initial numerical representations of protein sequences before being processed in the heterogeneous graph model. We are working on integrating **ProtT5** and **ESM-2** as alternative initial embeddings, allowing users to choose different sequence representations that may enhance performance for specific tasks. A detailed comparison of how these embeddings influence function prediction accuracy will be included in our upcoming publication.
- **Knowledge graph visualization for interpretability**: To improve model explainability, we are developing an interactive **knowledge graph visualization** feature. This will allow users to explore the biological relationships that contributed to ProtHGT’s predictions for a given protein. Users will be able to inspect **protein interactions, GO annotations, domains, pathways, and other key connections** in a structured graphical format, making it easier to interpret and validate predictions.
Stay tuned for updates and future publications!
""")
with st.sidebar:
disabled = st.session_state.generating_predictions
st.markdown("""
ProtHGT
Heterogeneous Graph Transformers for Automated Protein Function Prediction Using Knowledge Graphs and Language Models
""", unsafe_allow_html=True)
available_proteins = get_available_proteins()
if 'example_proteins' not in st.session_state:
st.session_state.example_proteins = random.sample(available_proteins, 5)
selected_proteins = []
# Add protein selection methods
selection_method = st.radio(
"Choose input method:",
["Use example query", "Search proteins", "Upload protein ID file"],
disabled=disabled
)
if selection_method == "Use example query":
selected_proteins = st.session_state.example_proteins
st.write(f"Selected proteins:")
st.markdown(
f"""
{'
'.join(selected_proteins)}
""",
unsafe_allow_html=True
)
elif selection_method == "Search proteins":
# User enters search term
search_query = st.text_input(
"1\\. Start typing a protein ID (at least 3 characters) and press Enter to see search results in the dropdown menu below (2)",
"",
disabled=disabled
)
# Initialize selected_proteins in session state if not exists
if 'selected_proteins_search' not in st.session_state:
st.session_state.selected_proteins_search = []
# Apply fuzzy search only if query length is >= 3
filtered_proteins = []
if len(search_query) >= 3:
# Case-insensitive search by converting query and proteins to lowercase
matches = process.extract(
search_query.upper(),
{p: p.upper() for p in available_proteins},
limit=50
)
filtered_proteins = [match[0] for match in matches] # Show top 50 matches
with st.container():
# Include previously selected proteins in options
all_options = list(set(filtered_proteins + st.session_state.selected_proteins_search))
selected_proteins = st.multiselect(
"2\\. Select proteins from search results",
options=all_options,
default=st.session_state.selected_proteins_search,
placeholder="Start typing a protein ID above (1) to see search results...",
max_selections=100,
disabled=disabled,
key="protein_selector"
)
# Update session state with current selection
st.session_state.selected_proteins_search = selected_proteins
# Apply custom CSS to make container scrollable
st.markdown("""
""", unsafe_allow_html=True)
else: # Upload file option
uploaded_file = st.file_uploader(
"Upload a text file with UniProt IDs (one per line, max 100)*",
type=['txt'],
disabled=disabled
)
if uploaded_file:
protein_list = [line.strip() for line in uploaded_file.read().decode('utf-8').splitlines()]
# Remove empty lines and duplicates
protein_list = list(filter(None, protein_list))
protein_list = list(dict.fromkeys(protein_list))
# Check for proteins not in available_proteins
proteins_not_found = [p for p in protein_list if p not in available_proteins]
# Filter to keep only available proteins
protein_list = [p for p in protein_list if p in available_proteins]
if len(protein_list) > 100:
st.error("Please upload a file with maximum 100 protein IDs.")
selected_proteins = []
else:
selected_proteins = protein_list
st.write(f"Loaded {len(selected_proteins)} proteins")
if proteins_not_found:
st.warning(f"""
The following proteins were not found in our input knowledge graph and have been discarded:
""")
with st.expander("View Discarded Proteins"):
# Create scrollable container with fixed height
st.markdown(
f"""
{'
'.join(proteins_not_found)}
""",
unsafe_allow_html=True
)
st.warning(f"""
Currently, our system can only generate predictions for proteins that are already included in our knowledge graph. **Real-time retrieval of relationship data from external source databases is not yet supported.**
We are actively working on integrating this capability in future updates. Stay tuned!
""")
if selected_proteins:
st.write(f"Total proteins selected: {len(selected_proteins)}")
# Add download button
proteins_text = '\n'.join(selected_proteins)
st.download_button(
label="Download Selected Proteins List",
data=proteins_text,
file_name="selected_proteins.txt",
mime="text/plain",
key="download_selected_proteins"
)
# Add GO category selection
go_category_options = {
'All Categories': None,
'Molecular Function': 'GO_term_F',
'Biological Process': 'GO_term_P',
'Cellular Component': 'GO_term_C'
}
selected_go_category = st.selectbox(
"Select GO Category for predictions",
options=list(go_category_options.keys()),
help="Choose which GO category to generate predictions for. Selecting 'All Categories' will generate predictions for all three categories.",
disabled=disabled
)
if selected_proteins and selected_go_category:
button_disabled = st.session_state.submitted
# Add custom CSS for red button
st.markdown("""
""", unsafe_allow_html=True)
if st.button("Generate Predictions",
disabled=button_disabled,
key="generate_predictions",
on_click=set_generating_predictions):
pass
# Create a tuple of current inputs to track changes
current_inputs = (tuple(selected_proteins), selected_go_category)
# Check if inputs have changed
if st.session_state.previous_inputs != current_inputs:
st.session_state.predictions_df = None
st.session_state.submitted = False
st.session_state.previous_inputs = current_inputs
st.warning("⚠️ Due to memory and computational constraints, the maximum number of proteins that can be processed at once is limited to 100 proteins. For larger datasets, please consider running the model locally using our [GitHub repository](https://github.com/HUBioDataLab/ProtHGT).")
if st.session_state.submitted:
with st.spinner("Generating predictions..."):
# Generate predictions only if not already in session state
if st.session_state.predictions_df is None:
# Load model config from JSON file
import json
import os
# Define data directory path
data_dir = "data"
models_dir = os.path.join(data_dir, "models")
# Load model configuration
model_config_paths = {
'GO_term_F': os.path.join(models_dir, "prothgt-config-molecular-function.yaml"),
'GO_term_P': os.path.join(models_dir, "prothgt-config-biological-process.yaml"),
'GO_term_C': os.path.join(models_dir, "prothgt-config-cellular-component.yaml")
}
# Paths for model and data
model_paths = {
'GO_term_F': os.path.join(models_dir, "prothgt-model-molecular-function.pt"),
'GO_term_P': os.path.join(models_dir, "prothgt-model-biological-process.pt"),
'GO_term_C': os.path.join(models_dir, "prothgt-model-cellular-component.pt")
}
# Get the selected GO category
go_category = go_category_options[selected_go_category]
# If a specific category is selected, use that model path
if go_category:
model_config_paths = [model_config_paths[go_category]]
model_paths = [model_paths[go_category]]
go_categories = [go_category]
else:
model_config_paths = [model_config_paths[cat] for cat in ['GO_term_F', 'GO_term_P', 'GO_term_C']]
model_paths = [model_paths[cat] for cat in ['GO_term_F', 'GO_term_P', 'GO_term_C']]
go_categories = ['GO_term_F', 'GO_term_P', 'GO_term_C']
# Generate predictions
heterodata, predictions_df = generate_prediction_df(
protein_ids=selected_proteins,
model_paths=model_paths,
model_config_paths=model_config_paths,
go_category=go_categories
)
st.session_state.heterodata = heterodata
st.session_state.predictions_df = predictions_df
# Reset only the generating_predictions flag to release the sidebar
st.session_state.generating_predictions = False
st.rerun()
# Display and filter predictions
st.success("Predictions generated successfully!")
# tabs for predictions and visualizations
predictions_tab, kg_viz_tab = st.tabs(["View Predictions", "View Knowledge Graphs"])
with predictions_tab:
st.markdown("### Filter and View Predictions")
# Create filters
col1, col2, col3, col4 = st.columns(4)
with col1:
# Extract UniProt IDs from URLs for the selectbox
uniprot_ids = st.session_state.predictions_df['UniProt_ID'].unique().tolist()
# Protein filter
selected_protein = st.selectbox(
"Filter by Protein",
options=['All'] + sorted(uniprot_ids)
)
with col2:
# GO category filter
selected_category = st.selectbox(
"Filter by GO Category",
options=['All'] + sorted(st.session_state.predictions_df['GO_category'].unique().tolist())
)
with col3:
# GO term filter
go_term_filter = st.text_input(
"Filter by GO Term ID",
placeholder="e.g., GO:0003674",
help="Enter a GO term ID to filter results"
).strip()
with col4:
# Probability threshold range slider
probability_range = st.slider(
"Probability Range",
min_value=0.0,
max_value=1.0,
value=(0.5, 1.0), # (min, max) default values
step=0.05
)
min_probability_threshold, max_probability_threshold = probability_range
# Filter the dataframe using session state data
filtered_df = st.session_state.predictions_df.copy()
if selected_protein != 'All':
filtered_df = filtered_df[filtered_df['UniProt_ID'].str.contains(selected_protein)]
if selected_category != 'All':
filtered_df = filtered_df[filtered_df['GO_category'] == selected_category]
if go_term_filter:
filtered_df = filtered_df[filtered_df['GO_ID'] == go_term_filter]
filtered_df = filtered_df[(filtered_df['Probability'] >= min_probability_threshold) &
(filtered_df['Probability'] <= max_probability_threshold)]
filtered_df['UniProt_ID'] = [f"https://www.uniprot.org/uniprotkb/{pid}/entry" for pid in filtered_df['UniProt_ID']]
filtered_df['GO_ID'] = [f"https://www.ebi.ac.uk/QuickGO/term/{go_id}" for go_id in filtered_df['GO_ID']]
# Custom CSS to increase table width and improve layout
st.markdown("""
""", unsafe_allow_html=True)
# Add pagination controls
col1, col2, col3 = st.columns([2, 1, 2])
with col2:
rows_per_page = st.selectbox("Rows per page", [50, 100, 200, 500], index=1)
total_rows = len(filtered_df)
total_pages = (total_rows + rows_per_page - 1) // rows_per_page
# Initialize page number in session state
if "page_number" not in st.session_state:
st.session_state.page_number = 0
# Calculate start and end indices for current page
start_idx = st.session_state.page_number * rows_per_page
end_idx = min(start_idx + rows_per_page, total_rows)
st.dataframe(
filtered_df.iloc[start_idx:end_idx],
hide_index=True,
use_container_width=True,
column_config={
"UniProt_ID": st.column_config.LinkColumn(
"UniProt ID",
help="Click to view protein in UniProt",
validate="^https://www\\.uniprot\\.org/uniprotkb/[A-Z0-9]+/entry$",
display_text="^https://www\\.uniprot\\.org/uniprotkb/([A-Z0-9]+)/entry$"
),
"GO_ID": st.column_config.LinkColumn(
"GO ID",
help="Click to view GO term in QuickGO",
validate="^https://www\\.ebi\\.ac\\.uk/QuickGO/term/GO:[0-9]+$",
display_text="^https://www\\.ebi\\.ac\\.uk/QuickGO/term/(GO:[0-9]+)$"
),
"Probability": st.column_config.ProgressColumn(
"Probability",
format="%.2f",
min_value=0,
max_value=1,
),
"Protein": st.column_config.TextColumn(
"Protein",
help="Protein Name",
),
"GO_category": st.column_config.TextColumn(
"GO Category",
help="Gene Ontology Category",
),
"GO_term": st.column_config.TextColumn(
"GO Term",
help="Gene Ontology Term Name",
),
}
)
# Pagination controls with better layout
col1, col2, col3 = st.columns([1, 3, 1])
with col1:
if st.button("Previous", disabled=st.session_state.page_number == 0):
st.session_state.page_number -= 1
st.rerun()
with col2:
st.markdown(f"""
""", unsafe_allow_html=True)
with col3:
if st.button("Next", disabled=st.session_state.page_number >= total_pages - 1):
st.session_state.page_number += 1
st.rerun()
downloadable_df = filtered_df.copy()
downloadable_df['UniProt_ID'] = downloadable_df['UniProt_ID'].apply(
lambda x: x.split('/')[-2] # Gets the ID part from the URL
)
downloadable_df['GO_ID'] = downloadable_df['GO_ID'].apply(
lambda x: x.split('/')[-1] # Gets the ID part from the URL
)
# Download filtered results
st.download_button(
label="Download Filtered Results",
data=convert_df(downloadable_df),
file_name="filtered_predictions.csv",
mime="text/csv",
key="download_filtered_predictions"
)
with kg_viz_tab:
st.markdown("### Knowledge Graph Visualization")
if not selected_proteins:
st.info("Please select proteins from the sidebar to visualize their knowledge graphs.")
elif len(selected_proteins) <= 10:
st.text("Visualize the knowledge graph for each protein to understand the biological relationships that contributed to the predictions.")
protein_tabs = st.tabs([f"{protein_id}" for protein_id in selected_proteins])
# Create visualizations in each tab
for idx, protein_id in enumerate(selected_proteins):
with protein_tabs[idx]:
max_node_count = st.slider(
"Maximum neighbors per edge type",
min_value=5,
max_value=50,
value=10,
step=5,
help="Control the maximum number of neighboring nodes shown for each relationship type",
key=f"slider_{protein_id}"
)
# Check if visualization exists for this protein
viz_exists = (protein_id in st.session_state.protein_visualizations and
os.path.exists(st.session_state.protein_visualizations[protein_id]['path']))
if not viz_exists:
if st.button(f"Generate Visualization", key=f"viz_{protein_id}"):
# Generate visualization with selected max_node_count
html_path, visualized_edges = visualize_protein_subgraph(
st.session_state.heterodata,
protein_id,
st.session_state.predictions_df,
limit=max_node_count
)
# Store visualization info in session state
st.session_state.protein_visualizations[protein_id] = {
'path': html_path,
'edges': visualized_edges
}
st.rerun()
# If visualization exists, display it
if viz_exists:
viz_info = st.session_state.protein_visualizations[protein_id]
# Add download button for edges
formatted_edges = {}
for edge_type, edges in viz_info['edges'].items():
edge_type_str = f"{edge_type[0]}_{edge_type[1]}_{edge_type[2]}"
formatted_edges[edge_type_str] = [
{"source": edge[0][0], "target": edge[0][1], "probability": edge[1]}
for edge in edges
]
kg_viz_button_columns = st.columns([1, 1, 1])
with kg_viz_button_columns[0]:
st.download_button(
label='Download Visualized Edges',
data=json.dumps(formatted_edges, indent=2),
file_name=f'{protein_id}_visualized_edges.json',
mime='application/json'
)
with kg_viz_button_columns[1]:
if st.button("Regenerate Visualization", key=f"regenerate_{protein_id}"):
# Clean up old file
try:
os.unlink(viz_info['path'])
except FileNotFoundError:
pass
# Remove from session state
del st.session_state.protein_visualizations[protein_id]
st.rerun()
with open(viz_info['path'], 'r', encoding='utf-8') as f:
html_content = f.read()
st.components.v1.html(html_content, height=1200)
else:
st.warning("Knowledge graph visualization is only available when 10 or fewer proteins are selected.") 