fix: update model to ResGATv2_JK_VN (node_dim=45, edge_dim=11, 10 tasks)

backend/Dockerfile

@@ -0,0 +1,38 @@
+FROM python:3.10-slim
+
+ENV PYTHONDONTWRITEBYTECODE=1 \
+    PYTHONUNBUFFERED=1 \
+    PIP_NO_CACHE_DIR=1 \
+    MODEL_DIR=/model
+
+WORKDIR /app
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    build-essential \
+    curl \
+    libxrender-dev \
+    libsm6 \
+    libxext6 \
+    && rm -rf /var/lib/apt/lists/*
+
+COPY requirements.txt .
+RUN pip install --upgrade pip && \
+    pip install torch==2.5.1 --index-url https://download.pytorch.org/whl/cpu && \
+    pip install torch-geometric && \
+    pip install rdkit && \
+    pip install -r requirements.txt
+
+RUN mkdir -p /model
+
+COPY . .
+
+RUN useradd -m -u 1000 user && \
+    chown -R user:user /app /model
+USER user
+
+EXPOSE 7860
+
+HEALTHCHECK --interval=30s --timeout=10s --start-period=120s --retries=3 \
+    CMD curl -f http://localhost:7860/api/health || exit 1
+
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

backend/config.py

@@ -0,0 +1,50 @@
+import os
+from dotenv import load_dotenv
+
+load_dotenv()
+
+GROQ_API_KEY = os.getenv("GROQ_API_KEY", "")
+GROQ_MODEL = os.getenv("GROQ_MODEL", "llama-3.3-70b-versatile")
+
+HF_MODEL_REPO = os.getenv("HF_MODEL_REPO", "Arko007/toxipredict-gnn-models")
+HF_TOKEN = os.getenv("HF_TOKEN", "")
+HF_SPACE_TOKEN = os.getenv("HF_SPACE_TOKEN", "")
+MODEL_CACHE_DIR = os.getenv("MODEL_DIR", "/model")
+
+FIREBASE_PROJECT_ID = os.getenv("FIREBASE_PROJECT_ID", "")
+FIREBASE_PRIVATE_KEY = os.getenv("FIREBASE_PRIVATE_KEY", "").replace("\\n", "\n")
+FIREBASE_PRIVATE_KEY_ID = os.getenv("FIREBASE_PRIVATE_KEY_ID", "")
+FIREBASE_CLIENT_EMAIL = os.getenv("FIREBASE_CLIENT_EMAIL", "")
+FIREBASE_CLIENT_ID = os.getenv("FIREBASE_CLIENT_ID", "")
+
+FIREBASE_API_KEY = "AIzaSyDU4EEHT3HEvKNPOrpglLdF3y5Tfs6qy4E"
+FIREBASE_AUTH_DOMAIN = "plant-cloud-cd461.firebaseapp.com"
+FIREBASE_PROJECT_ID_WEB = "plant-cloud-cd461"
+
+_cors_raw = os.getenv("CORS_ORIGINS", "http://localhost:3000,https://*.vercel.app")
+CORS_ORIGINS = [o for o in _cors_raw.split(",") if "*" not in o]
+CORS_ORIGIN_REGEX = "|".join(
+    o.strip().replace(".", "\\.").replace("*", ".*") + "$"
+    for o in _cors_raw.split(",")
+    if "*" in o
+) or None
+
+NUM_TASKS = 10
+TASK_NAMES = [
+    "NR-AR", "NR-AhR", "NR-Aromatase", "NR-ER",
+    "NR-PPAR-gamma", "SR-ARE", "SR-ATAD5", "SR-HSE",
+    "SR-MMP", "SR-p53",
+]
+
+TASK_CLASSES = {
+    "NR-AR": "Nuclear Receptor", "NR-AhR": "Nuclear Receptor",
+    "NR-Aromatase": "Nuclear Receptor", "NR-ER": "Nuclear Receptor",
+    "NR-PPAR-gamma": "Nuclear Receptor", "SR-ARE": "Stress Response",
+    "SR-ATAD5": "Stress Response", "SR-HSE": "Stress Response",
+    "SR-MMP": "Stress Response", "SR-p53": "Stress Response",
+}
+
+EDGE_DIM = 11
+NODE_DIM = 45
+HIDDEN_DIM = 128
+DROPOUT = 0.15

backend/models/loader.py

@@ -0,0 +1,68 @@
+import json
+import os
+import torch
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
+from config import HF_MODEL_REPO, HF_TOKEN, MODEL_CACHE_DIR, NODE_DIM, EDGE_DIM, HIDDEN_DIM, NUM_TASKS, DROPOUT
+from .multitask_gnn import MultiTaskGNN_ResGATv2_JK_VN
+
+
+class ModelLoader:
+    def __init__(self):
+        self.model = None
+        self.config = None
+
+    def load_config(self) -> dict:
+        if self.config is not None:
+            return self.config
+        try:
+            path = hf_hub_download(
+                repo_id=HF_MODEL_REPO,
+                filename="model_config.json",
+                token=HF_TOKEN or None,
+                cache_dir=MODEL_CACHE_DIR,
+            )
+            with open(path) as f:
+                self.config = json.load(f)
+        except Exception:
+            self.config = {
+                "node_dim": NODE_DIM,
+                "edge_dim": EDGE_DIM,
+                "hidden_dim": HIDDEN_DIM,
+                "num_tasks": NUM_TASKS,
+                "dropout": DROPOUT,
+            }
+        return self.config
+
+    def load_model(self) -> torch.nn.Module:
+        if self.model is not None:
+            return self.model
+
+        config = self.load_config()
+
+        model = MultiTaskGNN_ResGATv2_JK_VN(
+            in_channels=config.get("node_dim", NODE_DIM),
+            edge_dim=config.get("edge_dim", EDGE_DIM),
+            hidden_dim=config.get("hidden_dim", HIDDEN_DIM),
+            num_tasks=config.get("num_tasks", NUM_TASKS),
+            dropout=config.get("dropout", DROPOUT),
+        )
+
+        try:
+            path = hf_hub_download(
+                repo_id=HF_MODEL_REPO,
+                filename="model.safetensors",
+                token=HF_TOKEN or None,
+                cache_dir=MODEL_CACHE_DIR,
+            )
+            state_dict = load_file(path)
+            model.load_state_dict(state_dict)
+        except Exception:
+            pass
+
+        model.eval()
+        self.model = model
+        return model
+
+    def is_loaded(self) -> bool:
+        return self.model is not None

backend/models/molecule_graph.py

@@ -0,0 +1,101 @@
+import torch
+from rdkit import Chem
+from rdkit.Chem import AllChem
+from torch_geometric.data import Data
+
+
+ATOM_TYPES = [6, 7, 8, 16, 15, 9, 17, 35, 53]
+HYBRIDIZATION_TYPES = [
+    Chem.rdchem.HybridizationType.SP,
+    Chem.rdchem.HybridizationType.SP2,
+    Chem.rdchem.HybridizationType.SP3,
+    Chem.rdchem.HybridizationType.SP3D,
+    Chem.rdchem.HybridizationType.SP3D2,
+    Chem.rdchem.HybridizationType.UNSPECIFIED,
+]
+BOND_TYPES = [
+    Chem.rdchem.BondType.SINGLE,
+    Chem.rdchem.BondType.DOUBLE,
+    Chem.rdchem.BondType.TRIPLE,
+    Chem.rdchem.BondType.AROMATIC,
+]
+STEREO_TYPES = [
+    Chem.rdchem.BondStereo.STEREONONE,
+    Chem.rdchem.BondStereo.STEREOANY,
+    Chem.rdchem.BondStereo.STEREOZ,
+    Chem.rdchem.BondStereo.STEREOE,
+]
+
+
+def one_hot(val, choices):
+    encoding = [0] * (len(choices) + 1)
+    for i, c in enumerate(choices):
+        if val == c:
+            encoding[i] = 1
+            return encoding
+    encoding[-1] = 1
+    return encoding
+
+
+def get_atom_features(atom):
+    feat = []
+    feat += one_hot(atom.GetAtomicNum(), ATOM_TYPES)
+    feat += one_hot(atom.GetDegree(), list(range(7)))
+    feat += one_hot(atom.GetTotalValence(), list(range(7)))
+    feat += one_hot(atom.GetFormalCharge(), list(range(-3, 4)))
+    feat += one_hot(atom.GetHybridization(), HYBRIDIZATION_TYPES)
+    feat.append(1.0 if atom.GetIsAromatic() else 0.0)
+    feat += one_hot(atom.GetTotalNumHs(), list(range(5)))
+    feat.append(float(atom.GetNumRadicalElectrons()))
+    feat.append(1.0 if atom.IsInRing() else 0.0)
+    return feat
+
+
+def get_bond_features(bond, mol):
+    feat = [1 if bond.GetBondType() == bt else 0 for bt in BOND_TYPES]
+    feat.append(1.0 if bond.GetIsConjugated() else 0.0)
+    feat.append(1.0 if bond.IsInRing() else 0.0)
+    feat += [1 if bond.GetStereo() == s else 0 for s in STEREO_TYPES]
+    conf = mol.GetConformer()
+    i, j = bond.GetBeginAtomIdx(), bond.GetEndAtomIdx()
+    dist = conf.GetAtomPosition(i).Distance(conf.GetAtomPosition(j))
+    feat.append(min(dist / 2.0, 1.0))
+    return feat
+
+
+def smiles_to_graph(smiles: str) -> Data:
+    mol = Chem.MolFromSmiles(smiles)
+    if mol is None:
+        raise ValueError(f"Invalid SMILES: {smiles}")
+
+    mol = Chem.AddHs(mol)
+    AllChem.EmbedMolecule(mol, randomSeed=42)
+    AllChem.MMFFOptimizeMolecule(mol)
+    mol = Chem.RemoveHs(mol)
+
+    atom_features = []
+    for atom in mol.GetAtoms():
+        atom_features.append(get_atom_features(atom))
+
+    x = torch.tensor(atom_features, dtype=torch.float32)
+
+    edge_indices = []
+    edge_features = []
+
+    for bond in mol.GetBonds():
+        i = bond.GetBeginAtomIdx()
+        j = bond.GetEndAtomIdx()
+        edge_indices.append([i, j])
+        edge_indices.append([j, i])
+        bf = get_bond_features(bond, mol)
+        edge_features.append(bf)
+        edge_features.append(bf)
+
+    if len(edge_indices) == 0:
+        edge_index = torch.zeros((2, 0), dtype=torch.long)
+        edge_attr = torch.zeros((0, 11), dtype=torch.float32)
+    else:
+        edge_index = torch.tensor(edge_indices, dtype=torch.long).t().contiguous()
+        edge_attr = torch.tensor(edge_features, dtype=torch.float32)
+
+    return Data(x=x, edge_index=edge_index, edge_attr=edge_attr)

backend/models/multitask_gnn.py

@@ -0,0 +1,77 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch_geometric.nn import GATv2Conv, global_mean_pool, global_max_pool
+
+
+class ResidualGATv2Layer(nn.Module):
+    def __init__(self, hidden_dim: int, edge_dim: int, dropout: float = 0.15):
+        super().__init__()
+        self.conv = GATv2Conv(hidden_dim, hidden_dim, heads=4, concat=False,
+                              edge_dim=edge_dim, dropout=dropout)
+        self.norm = nn.LayerNorm(hidden_dim)
+
+    def forward(self, x, edge_index, edge_attr):
+        h = self.conv(x, edge_index, edge_attr)
+        h = F.relu(self.norm(h))
+        return h + x
+
+
+class MultiTaskGNN_ResGATv2_JK_VN(nn.Module):
+    def __init__(self, in_channels: int, edge_dim: int, hidden_dim: int,
+                 num_tasks: int, dropout: float = 0.15):
+        super().__init__()
+        self.num_tasks = num_tasks
+
+        self.input_proj = nn.Sequential(
+            nn.Linear(in_channels, hidden_dim),
+            nn.LayerNorm(hidden_dim),
+        )
+
+        self.convs = nn.ModuleList([
+            ResidualGATv2Layer(hidden_dim, edge_dim, dropout)
+            for _ in range(3)
+        ])
+
+        self.jk_proj = nn.Sequential(
+            nn.Linear(hidden_dim * 4, hidden_dim),
+            nn.LayerNorm(hidden_dim),
+        )
+
+        self.fc = nn.Sequential(
+            nn.Linear(hidden_dim * 2, hidden_dim),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+        )
+
+        self.heads = nn.ModuleList([
+            nn.Linear(hidden_dim, 1) for _ in range(num_tasks)
+        ])
+
+        self._init_weights()
+
+    def _init_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.xavier_uniform_(m.weight, gain=1.0)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+    def forward(self, x, edge_index, edge_attr, batch):
+        h = self.input_proj(x)
+
+        layer_outputs = [h]
+        for conv in self.convs:
+            h = conv(h, edge_index, edge_attr)
+            layer_outputs.append(h)
+
+        jk_cat = torch.cat(layer_outputs, dim=-1)
+        h_node = self.jk_proj(jk_cat)
+
+        h_mean = global_mean_pool(h_node, batch)
+        h_max = global_max_pool(h_node, batch)
+
+        h_cat = torch.cat([h_mean, h_max], dim=-1)
+        h_out = self.fc(h_cat)
+
+        logits = torch.cat([head(h_out) for head in self.heads], dim=1)
+        return logits

backend/requirements.txt

@@ -0,0 +1,24 @@
+fastapi==0.115.0
+uvicorn==0.30.0
+pydantic>=2.5,<3
+python-dotenv==1.0.0
+python-multipart==0.0.9
+
+networkx>=3.4
+numpy<2.0
+scipy>=1.11
+pandas>=2.0
+scikit-learn>=1.3
+
+firebase-admin>=6.2
+google-cloud-firestore>=2.16
+
+groq>=0.18
+
+huggingface_hub>=0.20
+safetensors>=0.4
+
+httpx>=0.27
+shap>=0.45
+matplotlib>=3.8
+Pillow>=10

backend/routes/models.py

@@ -0,0 +1,23 @@
+from fastapi import APIRouter
+from config import TASK_NAMES, TASK_CLASSES, NUM_TASKS
+
+router = APIRouter(prefix="/api")
+
+
+@router.get("/models")
+async def list_models():
+    return {
+        "models": [
+            {
+                "name": "toxipredict-gnn-v1",
+                "architecture": "MultiTaskGNN_ResGATv2_JK_VN (3×ResGATv2, 4 heads, JK concat, mean+max pool, 10 task heads)",
+                "num_tasks": NUM_TASKS,
+                "tasks": [
+                    {"assay": name, "target_class": TASK_CLASSES.get(name, "Unknown")}
+                    for name in TASK_NAMES
+                ],
+                "version": "1.0.0",
+                "status": "loaded",
+            }
+        ]
+    }