Spaces:
Paused
Paused
| """EC number prediction backend for the ProtHGT Space. | |
| Unlike GO term prediction (one flat decoder per sub-ontology), EC prediction uses | |
| a single joint model with a shared HGT backbone (plus a normalized skip | |
| connection) and four level-specific decoders, run as a top-down cascade so the | |
| predicted EC levels form a consistent hierarchy. | |
| """ | |
| import os | |
| import gzip | |
| import json | |
| import shutil | |
| from functools import lru_cache | |
| import torch | |
| from torch.nn import Linear, LayerNorm | |
| from torch_geometric.nn import HGTConv, MLP | |
| import pandas as pd | |
| import yaml | |
| from huggingface_hub import hf_hub_download | |
| from run_prothgt_app import load_name_info | |
| EC_MODELS_DIR = os.path.join("data", "models", "ec") | |
| EC_CONFIG_PATH = os.path.join(EC_MODELS_DIR, "prothgt-ec-config.yaml") | |
| EC_MODEL_PATH = os.path.join(EC_MODELS_DIR, "prothgt-ec-model.pt") | |
| EC_KG_PATH = os.path.join("data", "prothgt-ec-kg.pt") | |
| # EC class/subclass/sub-subclass (L1-L3) names from the ExPASy ENZYME nomenclature; | |
| # L4 (full EC number) names come from name_info instead. | |
| EC_CLASS_NAMES_PATH = os.path.join("data", "ec_class_names.json") | |
| LEVELS = ['l1', 'l2', 'l3', 'l4'] | |
| LEVEL_TARGET_TYPE = {'l1': 'EC_class_1', 'l2': 'EC_class_2', 'l3': 'EC_class_3', 'l4': 'EC_number'} | |
| LEVEL_LABEL = {'l1': 'L1', 'l2': 'L2', 'l3': 'L3', 'l4': 'L4'} | |
| CHILDREN_KEY = {'l2': 'l1_children', 'l3': 'l2_children', 'l4': 'l3_children'} | |
| # EC taxonomy edges used to find each node's children for the cascade. | |
| CHILD_EDGE_TYPES = { | |
| 'l1_children': ('EC_class_1', 'rev_ec_parent_1', 'EC_class_2'), | |
| 'l2_children': ('EC_class_2', 'rev_ec_parent_2', 'EC_class_3'), | |
| 'l3_children': ('EC_class_3', 'rev_ec_parent_3', 'EC_number'), | |
| } | |
| def load_ec_class_names(path=EC_CLASS_NAMES_PATH): | |
| """{EC class prefix (L1-L3) -> name}, e.g. '2.7.11' -> 'Protein-serine/threonine kinases'.""" | |
| with open(path, encoding='utf-8') as f: | |
| return json.load(f) | |
| # --------------------------------------------------------------------------- | |
| # Model (ported from the ProtHGT src/ec package) | |
| # --------------------------------------------------------------------------- | |
| class ProtHGT(torch.nn.Module): | |
| """HGT backbone with a normalized pre/post-GNN skip connection.""" | |
| def __init__(self, data, hidden_channels, num_heads, num_layers, mlp_hidden_layers, mlp_dropout): | |
| super().__init__() | |
| self.lin_dict = torch.nn.ModuleDict() | |
| for node_type in data.node_types: | |
| self.lin_dict[node_type] = Linear(data[node_type].x.size(1), hidden_channels) | |
| self.convs = torch.nn.ModuleList() | |
| for _ in range(num_layers): | |
| self.convs.append(HGTConv(hidden_channels, hidden_channels, data.metadata(), num_heads, group='sum')) | |
| self.initial_norm = LayerNorm(hidden_channels) | |
| self.gnn_norm = LayerNorm(hidden_channels) | |
| self.mlp = MLP(mlp_hidden_layers, dropout=mlp_dropout, norm=None) | |
| def generate_embeddings(self, x_dict, edge_index_dict): | |
| initial_dict = {nt: self.lin_dict[nt](x).relu_() for nt, x in x_dict.items()} | |
| gnn_dict = initial_dict | |
| for conv in self.convs: | |
| out = conv(gnn_dict, edge_index_dict) | |
| gnn_dict = {nt: (out[nt] if out.get(nt) is not None else prev) | |
| for nt, prev in gnn_dict.items()} | |
| return {nt: torch.cat([self.initial_norm(initial_dict[nt]), self.gnn_norm(gnn_dict[nt])], dim=-1) | |
| for nt in initial_dict} | |
| class ProtHGTHier(torch.nn.Module): | |
| """Shared ProtHGT backbone + one MLP decoder per EC level.""" | |
| def __init__(self, data, hidden_channels, num_heads, num_layers, mlp_hidden_layers, mlp_dropout): | |
| super().__init__() | |
| self.backbone = ProtHGT(data, hidden_channels, num_heads, num_layers, mlp_hidden_layers, mlp_dropout) | |
| self.backbone.mlp = torch.nn.Identity() | |
| self.decoders = torch.nn.ModuleDict({ | |
| level: MLP(mlp_hidden_layers, dropout=mlp_dropout, norm=None) for level in LEVELS | |
| }) | |
| def generate_embeddings(self, x_dict, edge_index_dict): | |
| return self.backbone.generate_embeddings(x_dict, edge_index_dict) | |
| def load_hierarchy_masks(data): | |
| """{'l1_children': {ec_class_1_idx: [ec_class_2_idx, ...]}, ...} from the EC taxonomy edges.""" | |
| from collections import defaultdict | |
| masks = {} | |
| for name, et in CHILD_EDGE_TYPES.items(): | |
| row, col = data[et].edge_index | |
| children = defaultdict(list) | |
| for r, c in zip(row.tolist(), col.tolist()): | |
| children[r].append(c) | |
| masks[name] = dict(children) | |
| return masks | |
| # --------------------------------------------------------------------------- | |
| # Loading | |
| # --------------------------------------------------------------------------- | |
| def load_ec_kg(output=EC_KG_PATH): | |
| """Download (if needed) and load the EC knowledge graph. Callers should cache the result.""" | |
| if not os.path.exists(output): | |
| cached_path = hf_hub_download( | |
| repo_id="HUBioDataLab/ProtHGT", | |
| filename="knowledge_graphs/ec/prothgt-ec-kg.pt", | |
| repo_type="dataset", | |
| ) | |
| os.makedirs(os.path.dirname(output), exist_ok=True) | |
| shutil.copy(cached_path, output) | |
| print(f"EC KG downloaded to {output}") | |
| print('Loading EC knowledge graph...') | |
| return torch.load(output) | |
| def load_ec_model(heterodata, config_path=EC_CONFIG_PATH, model_path=EC_MODEL_PATH): | |
| """Build ProtHGTHier and load its weights. Callers should cache the result.""" | |
| with open(config_path) as f: | |
| cfg = yaml.safe_load(f) | |
| h = cfg['hidden_channels'] | |
| model = ProtHGTHier( | |
| heterodata, hidden_channels=h, num_heads=cfg['num_heads'], num_layers=cfg['num_layers'], | |
| mlp_hidden_layers=[4 * h, h, 1], mlp_dropout=cfg['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 EC model weights from {model_path}') | |
| return model | |
| # --------------------------------------------------------------------------- | |
| # Cascade prediction | |
| # --------------------------------------------------------------------------- | |
| def _cascade_protein(model, x_dict, masks, protein_idx, inv_maps, threshold, device): | |
| """Top-down cascade for one protein: L1 over all classes, each deeper level only | |
| over the children of the kept nodes above. Returns (level_key, ec_string, prob) rows.""" | |
| rows = [] | |
| kept_prev = [] | |
| for level in LEVELS: | |
| target_type = LEVEL_TARGET_TYPE[level] | |
| if level == 'l1': | |
| candidates = list(range(x_dict[target_type].size(0))) | |
| else: | |
| ckey = CHILDREN_KEY[level] | |
| candidates = sorted({c for p in kept_prev for c in masks[ckey].get(p, [])}) | |
| if not candidates: | |
| break | |
| cand = torch.tensor(candidates, dtype=torch.long, device=device) | |
| p_emb = x_dict['Protein'][protein_idx].unsqueeze(0).expand(len(candidates), -1) | |
| z = torch.cat([p_emb, x_dict[target_type][cand]], dim=-1) | |
| probs = torch.sigmoid(model.decoders[level](z).view(-1)) | |
| keep = (probs > threshold).nonzero(as_tuple=False).view(-1) | |
| kept_prev = [candidates[i] for i in keep.tolist()] | |
| for i in keep.tolist(): | |
| rows.append((level, inv_maps[target_type][candidates[i]], float(probs[i]))) | |
| return rows | |
| def generate_ec_prediction_df(protein_ids, heterodata, model, threshold=0.5, progress_callback=None): | |
| """Cascade EC predictions for the given proteins. The backbone embeddings are | |
| computed once over the whole graph; the decoders then run per protein.""" | |
| if isinstance(protein_ids, str): | |
| protein_ids = [protein_ids] | |
| # The full-graph forward over the whole KG is memory-heavy for this model | |
| # (hidden=256); fall back to CPU if a GPU runs out of memory. | |
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
| def _embed(dev): | |
| model.to(dev) | |
| hd = heterodata.to(dev) | |
| with torch.no_grad(): | |
| return hd, model.generate_embeddings(hd.x_dict, hd.edge_index_dict) | |
| try: | |
| heterodata, x_dict = _embed(device) | |
| except RuntimeError as e: | |
| if device.type == 'cuda' and 'out of memory' in str(e).lower(): | |
| print('GPU out of memory during the full-graph forward; falling back to CPU...') | |
| torch.cuda.empty_cache() | |
| device = torch.device('cpu') | |
| heterodata, x_dict = _embed(device) | |
| else: | |
| raise | |
| masks = load_hierarchy_masks(heterodata) | |
| inv_maps = { | |
| LEVEL_TARGET_TYPE[level]: {idx: tok for tok, idx in heterodata[LEVEL_TARGET_TYPE[level]]['id_mapping'].items()} | |
| for level in LEVELS | |
| } | |
| id_mapping = heterodata['Protein']['id_mapping'] | |
| name_info = load_name_info() | |
| ec_l4_names = name_info['EC_number'] # full EC number (L4) -> enzyme name | |
| ec_class_names = load_ec_class_names() # L1-L3 class prefix -> class name | |
| rows = [] | |
| for i, pid in enumerate(protein_ids): | |
| if progress_callback is not None: | |
| progress_callback(i, len(protein_ids)) | |
| protein_name = name_info['Protein'].get(pid, pid) | |
| for level, token, prob in _cascade_protein( | |
| model, x_dict, masks, id_mapping[pid], inv_maps, threshold, device): | |
| ec_name = ec_l4_names.get(token, '') if level == 'l4' else ec_class_names.get(token, '') | |
| rows.append({ | |
| 'UniProt_ID': pid, | |
| 'Protein': protein_name, | |
| 'EC': token, | |
| 'EC_name': ec_name, | |
| 'EC_level': LEVEL_LABEL[level], | |
| 'Probability': prob, | |
| }) | |
| df = pd.DataFrame(rows, columns=['UniProt_ID', 'Protein', 'EC', 'EC_name', 'EC_level', 'Probability']) | |
| return heterodata, df | |