refactor the app to cache kg and models, download kgs from hf hub and add prediction progress

run_prothgt_app.py

@@ -60,6 +60,68 @@ from huggingface_hub import hf_hub_download
 import copy
 import json
 import gzip
+from functools import lru_cache
+
+MODELS_DIR = os.path.join("data", "models")
+KG_PATH = os.path.join("data", "prothgt-kg.pt")
+NAME_INFO_PATH = os.path.join("data", "name_info.json.gz")
+
+MODEL_FILES = {
+    'GO_term_F': ("prothgt-config-molecular-function.yaml", "prothgt-model-molecular-function.pt"),
+    'GO_term_P': ("prothgt-config-biological-process.yaml", "prothgt-model-biological-process.pt"),
+    'GO_term_C': ("prothgt-config-cellular-component.yaml", "prothgt-model-cellular-component.pt"),
+}
+
+
+@lru_cache(maxsize=1)
+def load_name_info(name_file=NAME_INFO_PATH):
+    """Load node display names once per process (the file is ~24 MB uncompressed)."""
+    with gzip.open(name_file, 'rt', encoding='utf-8') as file:
+        return json.load(file)
+
+
+def load_kg(output=KG_PATH):
+    """Download (if needed) and load the ProtHGT knowledge graph. Callers should cache the result."""
+    if not os.path.exists(output):
+        try:
+            cached_path = hf_hub_download(
+                repo_id="HUBioDataLab/ProtHGT",
+                filename="knowledge_graphs/prothgt-kg.pt",
+                repo_type="dataset",
+            )
+            os.makedirs(os.path.dirname(output), exist_ok=True)
+            shutil.copy(cached_path, output)
+            print(f"File downloaded to {output}")
+        except Exception as e:
+            print(f"Error downloading file: {e}")
+            raise
+    print('Loading knowledge graph...')
+    return torch.load(output)
+
+
+def load_model(go_cat, heterodata, model_config_path=None, model_path=None):
+    """Build a ProtHGT model for one GO category and load its weights. Callers should cache the result."""
+    if model_config_path is None or model_path is None:
+        cfg_name, weights_name = MODEL_FILES[go_cat]
+        model_config_path = model_config_path or os.path.join(MODELS_DIR, cfg_name)
+        model_path = model_path or os.path.join(MODELS_DIR, weights_name)
+
+    with open(model_config_path, 'r') as file:
+        model_config = yaml.safe_load(file)
+
+    model = ProtHGT(
+        heterodata,
+        hidden_channels=model_config['hidden_channels'][0],
+        num_heads=model_config['num_heads'],
+        num_layers=model_config['num_layers'],
+        mlp_hidden_layers=model_config['hidden_channels'][1],
+        mlp_dropout=model_config['mlp_dropout']
+    )
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    model.load_state_dict(torch.load(model_path, map_location=device))
+    model.eval()
+    print(f'Loaded model weights from {model_path}')
+    return model
 
 class ProtHGT(torch.nn.Module):
     def __init__(self, data,hidden_channels, num_heads, num_layers, mlp_hidden_layers, mlp_dropout):
@@ -117,10 +179,8 @@ def _build_edge_label_index(heterodata, protein_ids, go_category):
     col = term_indices.repeat(len(protein_indices))
     return torch.stack([row, col], dim=0)
 
-def get_available_proteins(name_file='data/name_info.json.gz'):
-    with gzip.open(name_file, 'rt', encoding='utf-8') as file:
-        name_info = json.load(file)
-    return list(name_info['Protein'].keys())
+def get_available_proteins(name_file=NAME_INFO_PATH):
+    return list(load_name_info(name_file)['Protein'].keys())
 
 def _generate_predictions(heterodata, model, edge_label_index, target_type):
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
@@ -143,9 +203,7 @@ def _create_prediction_df(predictions, heterodata, protein_ids, go_category, thr
         'GO_term_C': 'Cellular Component'
     }
 
-    # Load name information from gzipped file
-    with gzip.open('data/name_info.json.gz', 'rt', encoding='utf-8') as file:
-        name_info = json.load(file)
+    name_info = load_name_info()
 
     id_mapping = heterodata[go_category]['id_mapping']  # dict: GO_id -> index
     n_go_terms = len(id_mapping)
@@ -206,86 +264,52 @@ def _create_prediction_df(predictions, heterodata, protein_ids, go_category, thr
     
     return prediction_df
 
-def generate_prediction_df(protein_ids, model_paths, model_config_paths, go_category, threshold: float = 0.0):
+def generate_prediction_df(protein_ids, model_paths=None, model_config_paths=None, go_category=None,
+                           threshold: float = 0.0, heterodata=None, models=None, progress_callback=None):
+    """
+    Generate GO term predictions for the given proteins.
+
+    heterodata / models can be passed in pre-loaded (e.g. from st.cache_resource) so the
+    1.6 GB knowledge graph and model weights are not reloaded on every request.
+    progress_callback(i, n, go_cat) is called before each category is scored.
+    """
     all_predictions = []
-    
-    # Convert single protein ID to list if necessary
+
     if isinstance(protein_ids, str):
         protein_ids = [protein_ids]
+    if go_category is None:
+        go_category = list(MODEL_FILES.keys())
+    if model_paths is None:
+        model_paths = [None] * len(go_category)
+    if model_config_paths is None:
+        model_config_paths = [None] * len(go_category)
 
-    # Load dataset once
-    # heterodata = load_dataset('HUBioDataLab/ProtHGT-KG', data_files="prothgt-kg.json.gz")
-    print('Loading data...')
-    # file_id = "18u1o2sm8YjMo9joFw4Ilwvg0-rUU0PXK"
-    output = "data/prothgt-kg.pt"
-
-    if not os.path.exists(output):
-        try:
-            # url = f"https://drive.google.com/uc?id={file_id}"
-            # print(f"Downloading file from {url}...")
-            # gdown.download(url, output, quiet=False)
-            # print(f"File downloaded to {output}")
-            cached_path = hf_hub_download(
-                repo_id="HUBioDataLab/ProtHGT",
-                filename="knowledge_graphs/prothgt-kg.pt",
-                repo_type="dataset",
-            )
-
-            output = "data/prothgt-kg.pt"
-            os.makedirs("data", exist_ok=True)
-            shutil.copy(cached_path, output)
-            print(f"File downloaded to {output}")
-        except Exception as e:
-            print(f"Error downloading file: {e}")
-            raise
-    else:
-        print(f"File already exists at {output}")
-
-    heterodata = torch.load(output)
-    
-    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    if heterodata is None:
+        heterodata = load_kg()
 
-    for go_cat, model_config_path, model_path in zip(go_category, model_config_paths, model_paths):
+    for i, (go_cat, model_config_path, model_path) in enumerate(zip(go_category, model_config_paths, model_paths)):
         print(f'Generating predictions for {go_cat}...')
-        
+        if progress_callback is not None:
+            progress_callback(i, len(go_category), go_cat)
+
         # Build candidate edges for inference (do NOT modify graph edges)
         edge_label_index = _build_edge_label_index(heterodata, protein_ids, go_cat)
-        
-        # Load model config
-        with open(model_config_path, 'r') as file:
-            model_config = yaml.safe_load(file)
-        
-        # Initialize model with configuration
-        model = ProtHGT(
-            heterodata,
-            hidden_channels=model_config['hidden_channels'][0],
-            num_heads=model_config['num_heads'],
-            num_layers=model_config['num_layers'],
-            mlp_hidden_layers=model_config['hidden_channels'][1],
-            mlp_dropout=model_config['mlp_dropout']
-        )
-        
-        # Load model weights
-        model.load_state_dict(torch.load(model_path, map_location=device))
-        print(f'Loaded model weights from {model_path}')
-        
-        # Generate predictions
+
+        if models is not None and go_cat in models:
+            model = models[go_cat]
+        else:
+            model = load_model(go_cat, heterodata, model_config_path, model_path)
+
         predictions = _generate_predictions(heterodata, model, edge_label_index, go_cat)
         prediction_df = _create_prediction_df(predictions, heterodata, protein_ids, go_cat, threshold=threshold)
         all_predictions.append(prediction_df)
-        
-        # Clean up memory
+
         del edge_label_index
-        del model
         del predictions
-        torch.cuda.empty_cache()  # Clear CUDA cache if using GPU
-
+        torch.cuda.empty_cache()
 
-    # Combine all predictions
     final_df = pd.concat(all_predictions, ignore_index=True)
-    
-    # Clean up
     del all_predictions
     torch.cuda.empty_cache()
-    
+
     return heterodata, final_df