add app

.gitattributes

@@ -0,0 +1,36 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,3 @@
+trash/
+__pycache__/
+scripts/

.streamlit/config.toml

@@ -0,0 +1,7 @@
+[browser]
+gatherUsageStats = false
+
+[server]
+headless = true
+enableCORS = false
+enableXsrfProtection = false

Dockerfile

@@ -0,0 +1,29 @@
+FROM python:3.11
+
+ENV DEBIAN_FRONTEND=noninteractive
+SHELL ["/bin/bash", "-c"]
+RUN apt-get update -y \
+    && apt-get install -y build-essential python3-dev r-base make apt-utils unzip gpg doxygen git curl aria2 vim screen rsync wget locales gfortran mafft libglew-dev libeigen3-dev \
+    libpng-dev libfreetype6-dev libxml2-dev \
+    libmsgpack-dev python3-pyqt5.qtopengl libglm-dev libnetcdf-dev \
+    && locale-gen en_US.UTF-8 \
+    && rm -rf /var/lib/apt/lists/*
+
+WORKDIR /app
+
+COPY . .
+
+RUN pip3 install torch torchvision --index-url https://download.pytorch.org/whl/cpu
+RUN pip3 install -r requirements.txt
+
+# Ensure Hugging Face cache directory exists and is writable
+RUN mkdir -p /app/.cache && chmod 777 /app/.cache
+RUN mkdir -p /app/.cache/offload && chmod 777 /app/.cache/offload
+
+EXPOSE 8501
+
+HEALTHCHECK CMD curl --fail http://localhost:8501/_stcore/health
+
+ENV HF_HOME=/app/.cache
+
+ENTRYPOINT ["streamlit", "run", "src/streamlit_app.py", "--server.port=8501", "--server.address=0.0.0.0"]

README.md

@@ -0,0 +1,15 @@
+---
+title: DELM
+emoji: 🧬
+colorFrom: red
+colorTo: red
+sdk: docker
+app_port: 8501
+tags:
+- streamlit
+pinned: false
+short_description: Prediction of enzymatic degradation
+license: cc-by-nc-sa-4.0
+---
+
+

requirements.txt

@@ -0,0 +1,25 @@
+tqdm
+transformers
+fair-esm
+lightning
+cupy-cuda11x
+scikit-learn
+line_profiler
+sentence-transformers
+pandas
+openpyxl
+timm
+wandb
+accelerate
+ipykernel
+einops
+SentencePiece
+seaborn
+streamlit
+biotite
+matplotlib
+git+https://github.com/Ramprasad-Group/psmiles.git
+py3Dmol 
+stmol
+ipython_genutils
+cryptography

src/checkpoints/weights.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:42319d679b0fac05d5e21979771295e9c97b333bc3e13a1ef0f4a8189e47dbc7
+size 122999708

src/configs/train.yml

@@ -0,0 +1,13 @@
+plm: esm2_t33_650M_UR50D
+train_csv: data/train_addition.csv
+test_csv: data/test.csv
+batch_size: 32
+epochs: 100
+lr: 0.001
+wd: 0
+num_workers: 4
+amp: true
+seed: 42
+nfolds: 5
+ckpt_dir: checkpoints
+patience: 20

src/configs/tune.yml

@@ -0,0 +1,17 @@
+program: train.py
+method: grid
+metric:
+  name: val_acc
+  goal: maximize
+parameters:
+  plm:
+    values:
+      - esm2_t33_650M_UR50D
+      - esm2_t48_15B_UR50D
+      - esm2_t36_3B_UR50D
+      - esm2_t12_35M_UR50D
+      - esm2_t30_150M_UR50D
+      - esm2_t6_8M_UR50D
+      - esm1b_t33_650M_UR50S
+      - prot_t5_xl_half_uniref50-enc
+      - prot_t5_xl_bfd

src/data/polymer2tok.csv

@@ -0,0 +1,38 @@
+polymer_id,polymer
+0,PE
+1,PET
+2,PCL
+3,PHB
+4,PEF
+5,PBS
+6,PBSA
+7,PLA
+8,PHV
+9,PU
+10,PES
+11,PHA
+12,PHO
+13,PVA
+14,PPL
+15,P3HP
+16,P4HB
+17,PEA
+18,O-PVA
+19,P(3HB-co-3MP)
+20,PEG
+21,PHBV
+22,PHPV
+23,Nylon
+24,PBS-Blend
+25,PBSA-Blend
+26,P3HV
+27,PBAT
+28,PMCL
+29,LDPE
+30,PS
+31,NR
+32,PBSeT
+33,Ecovio-FT
+34,PHBH
+35,PHBVH
+36,Impranil

src/models/dataset.py

@@ -0,0 +1,102 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+import lightning as L
+from pathlib import Path
+import pandas as pd
+from models.plm import get_model
+from models.polybert import PolyEncoder, polymer2psmiles
+from argparse import Namespace as Args
+from sklearn.model_selection import KFold
+from tqdm import tqdm
+from torch.utils.data import WeightedRandomSampler
+
+
+class EnzymeDataset(Dataset):
+    def __init__(self, csv_file: str, plm: str):
+        self.data_list = []
+        for i, row in pd.read_csv(csv_file).iterrows():
+            self.data_list.append(
+                (row['category'], row['sequence'].upper(), row['degradation'], row['sequence_id'], row['polymer_id']))
+        (cache_dir := Path('cache')).mkdir(parents=True, exist_ok=True)
+        Path(cache_dir, 'protein').mkdir(parents=True, exist_ok=True)
+        Path(cache_dir, 'protein', plm).mkdir(parents=True, exist_ok=True)
+        Path(cache_dir, 'polymer').mkdir(parents=True, exist_ok=True)
+        if not all(Path(cache_dir, 'protein', plm, f"{seqid}.pt").exists() for _, _, _, seqid, _ in self.data_list):
+            plm_func = get_model(plm, 'cuda')
+            for _, seq, _, seqid, _ in tqdm(self.data_list, desc='Encoding enzyme sequences'):
+                seq_path = Path(cache_dir, 'protein', plm, f'{seqid}.pt')
+                if not seq_path.exists():
+                    seq_tensor = plm_func([seq])
+                    torch.save(seq_tensor, seq_path)
+
+    def __len__(self):
+        return len(self.data_list)
+
+    def __getitem__(self, idx):
+        return self.data_list[idx]
+
+
+class EnzymeDataModule(L.LightningDataModule):
+    def __init__(self, args: Args):
+        super().__init__()
+        self.args = args
+        self.train_csv = args.train_csv
+        self.test_csv = args.test_csv
+        self.batch_size = args.batch_size
+        self.num_workers = args.num_workers
+        self.plm = args.plm
+
+        self.train_val_set = EnzymeDataset(self.train_csv, self.plm)
+        self.test_set = EnzymeDataset(self.test_csv, self.plm)
+
+        self.kfold = KFold(
+            n_splits=args.nfolds, shuffle=True,
+            random_state=self.args.seed)
+        self.indices = list(range(len(self.train_val_set)))
+        self.splits = list(self.kfold.split(self.indices))
+
+    def setup_k_fold(self, fold_idx):
+        train_idx, val_idx = self.splits[fold_idx]
+
+        self.train_set = torch.utils.data.Subset(
+            self.train_val_set, train_idx)
+        self.val_set = torch.utils.data.Subset(
+            self.train_val_set, val_idx)
+        self.sampler = self.data_sampler()
+
+    def data_sampler(self):
+        # Get labels for train_set
+        if hasattr(self, 'train_set'):
+            # train_set is a Subset, get indices
+            indices = self.train_set.indices if hasattr(
+                self.train_set, 'indices') else range(len(self.train_set))
+            labels = [self.train_val_set[i][2] for i in indices]
+            # Compute class weights
+            label_counts = pd.Series(labels).value_counts()
+            weights = [1.0 / label_counts[label] for label in labels]
+            sampler = WeightedRandomSampler(
+                weights, num_samples=len(weights), replacement=True)
+            return sampler
+        else:
+            raise AttributeError(
+                'train_set not initialized. Call setup_k_fold first.')
+
+    def train_dataloader(self):
+        return DataLoader(
+            self.train_set, batch_size=self.batch_size,
+            # shuffle=True,
+            num_workers=self.num_workers,
+            sampler=self.sampler,
+        )
+
+    def val_dataloader(self):
+        return DataLoader(
+            self.val_set, batch_size=self.batch_size,
+            shuffle=False, num_workers=self.num_workers,)
+
+    def test_dataloader(self):
+        return DataLoader(
+            self.test_set, batch_size=self.batch_size,
+            shuffle=False, num_workers=self.num_workers)

src/models/plm.py

@@ -0,0 +1,239 @@
+import torch
+import torch.nn as nn
+from transformers import AutoModel, AutoTokenizer, T5EncoderModel, T5Tokenizer
+import re
+
+
+def EsmModelInfo(name: str):
+    """Get model info by name:
+    Args:
+        name: str, model name
+    Returns:
+        dict, model info: dim, layers, model
+    """
+    return {
+        "esm2_t48_15B_UR50D": {
+            "dim": 5120,
+            "layers": 48,
+            "model": "facebook/esm2_t48_15B_UR50D",
+        },
+        "esm2_t36_3B_UR50D": {
+            "dim": 2560,
+            "layers": 36,
+            "model": "facebook/esm2_t36_3B_UR50D",
+        },
+        "esm2_t33_650M_UR50D": {
+            "dim": 1280,
+            "layers": 33,
+            "model": "facebook/esm2_t33_650M_UR50D",
+        },
+        "esm2_t30_150M_UR50D": {
+            "dim": 640,
+            "layers": 30,
+            "model": "facebook/esm2_t30_150M_UR50D",
+        },
+        "esm2_t12_35M_UR50D": {
+            "dim": 480,
+            "layers": 12,
+            "model": "facebook/esm2_t12_35M_UR50D",
+        },
+        "esm2_t6_8M_UR50D": {
+            "dim": 320,
+            "layers": 6,
+            "model": "facebook/esm2_t6_8M_UR50D",
+        },
+        "esm1b_t33_650M_UR50S": {
+            "dim": 1280,
+            "layers": 33,
+            "model": "facebook/esm1b_t33_650M_UR50S",
+        },
+        "prot_t5_xl_half_uniref50-enc": {
+            "dim": 1024,
+            "layers": 24,
+            "model": "Rostlab/prot_t5_xl_uniref50",
+        },
+        "prot_t5_xl_bfd": {
+            "dim": 1024,
+            "layers": 24,
+            "model": "Rostlab/prot_t5_xl_bfd",
+        },
+        "esmc-6b-2024-12": {
+            "dim": 2560,
+            "layers": -1,
+            "model": "esmc-6b-2024-12",
+        },
+        "esmc_300m": {
+            "dim": 768,
+            "layers": -1,
+            "model": "esmc_300m",
+        },
+        "esmc_600m": {
+            "dim": 1152,
+            "layers": -1,
+            "model": "esmc_600m",
+        },
+    }[name]
+
+
+plm2abbr = {
+    'esm2_t48_15B_UR50D': 'ESM2_T48',
+    'esm2_t36_3B_UR50D': 'ESM2_T36',
+    'esm2_t33_650M_UR50D': 'ESM2_T33',
+    'esm2_t30_150M_UR50D': 'ESM2_T30',
+    'esm2_t12_35M_UR50D': 'ESM2_T12',
+    'esm2_t6_8M_UR50D': 'ESM2_T6',
+    'esm1b_t33_650M_UR50S': 'ESM1B_T33',
+    'prot_t5_xl_half_uniref50-enc': 'PT_UR',
+    'prot_t5_xl_bfd': 'PT_BFD',
+}
+
+
+class EsmEncoder(nn.Module):
+    def __init__(self, model_name, dev):
+        super().__init__()
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.model = AutoModel.from_pretrained(
+            # auto, balanced_low_0
+            model_name,
+            device_map="balanced",
+            # torch_dtype=torch.float16,
+            torch_dtype=torch.float32,
+            offload_folder=".cache/offload",
+            offload_state_dict=True,
+        )
+        if model_name == "facebook/esm2_t48_15B_UR50D":
+            self.max_len = 512
+        else:
+            self.max_len = 960
+        self.overlap = 31
+        self.model.eval()
+        # self.model.half()
+
+    def forward(self, _seqs):
+        with torch.no_grad():
+            assert len(_seqs) == 1, "currently only support batch size 1"
+            seqs = _seqs[0]
+            # left overlappping, right overlappping
+            seqs = [
+                seqs[max(0, i - self.overlap): (i + self.max_len + self.overlap)]
+                for i in range(0, len(seqs), self.max_len)
+            ]
+            segs = []
+            for seq in seqs:
+                inputs = self.tokenizer(
+                    [seq],
+                    return_tensors="pt",
+                ).to(self.model.device)
+                outputs = (
+                    self.model(
+                        **inputs).last_hidden_state.squeeze(0).detach().cpu()
+                )
+                outputs0 = self.model.embeddings(
+                    **inputs).squeeze(0).detach().cpu()
+                segs.append(torch.stack([outputs0, outputs], dim=-1))
+            t = []
+            for i in range(len(seqs)):
+                if i == 0:
+                    t.append(segs[i][1: (1 + self.max_len)])
+                elif i == len(seqs) - 1:
+                    t.append(segs[i][1 + self.overlap:])
+                else:
+                    t.append(
+                        segs[i][1 + self.overlap: 1 +
+                                self.max_len + self.overlap]
+                    )
+            outputs = torch.cat(t, dim=0)[: len(_seqs[0])]
+            assert outputs.shape[0] == len(_seqs[0])
+            return outputs
+
+
+class T5Encoder(nn.Module):
+    def __init__(self, name: str, dev) -> None:
+        super().__init__()
+        self.dev = dev
+        if name == "Rostlab/prot_t5_xl_uniref50":
+            # Load the tokenizer
+            self.tokenizer = T5Tokenizer.from_pretrained(
+                "Rostlab/prot_t5_xl_half_uniref50-enc",
+                do_lower_case=False,
+                legacy=False,
+            )
+            # Load the model
+            self.model = T5EncoderModel.from_pretrained(
+                "Rostlab/prot_t5_xl_half_uniref50-enc"
+            ).to(dev)
+        elif name == "Rostlab/prot_t5_xl_bfd":
+            # Load the tokenizer
+            self.tokenizer = T5Tokenizer.from_pretrained(
+                "Rostlab/prot_t5_xl_bfd",
+                do_lower_case=False,
+                legacy=False,
+            )
+            # Load the model
+            self.model = T5EncoderModel.from_pretrained("Rostlab/prot_t5_xl_bfd").to(
+                dev
+            )
+
+        self.max_len = 960  # start_token, end_token occupy 2 positions
+        self.overlap = 31
+        self.model.eval()
+        self.model.half()
+
+    def forward(self, _seqs):
+        with torch.no_grad():
+            assert len(_seqs) == 1, "currently only support batch size 1"
+            seqs = _seqs[0]
+            # replace non-amino acids with X
+            seqs = re.sub(r"[^A-Z]", "X", seqs)
+            # left overlappping, right overlappping
+            seqs = [
+                seqs[max(0, i - self.overlap)
+                         : (i + self.max_len + self.overlap)]
+                for i in range(0, len(seqs), self.max_len)
+            ]
+            input_ids = self.tokenizer.batch_encode_plus(
+                [" ".join(list(s)) for s in seqs],
+                add_special_tokens=True,
+                padding="longest",
+            )["input_ids"]
+            input_ids = torch.tensor(input_ids).to(self.dev)
+            outputs = self.model(input_ids=input_ids)
+            outputs0 = self.model.get_input_embeddings()(input_ids)
+            outputs = outputs.last_hidden_state
+            outputs = torch.stack([outputs0, outputs], dim=-1)
+            t = []
+            for i in range(len(seqs)):
+                if i == 0:
+                    t.append(outputs[i, 1: (1 + self.max_len)])
+                elif i == len(seqs) - 1:
+                    t.append(outputs[i, 1 + self.overlap:])
+                else:
+                    t.append(
+                        outputs[i, 1 + self.overlap: 1 +
+                                self.max_len + self.overlap]
+                    )
+
+            outputs = torch.cat(t, dim=0)[: len(_seqs[0])]
+            assert outputs.shape[0] == len(_seqs[0]), \
+                f"outputs shape {outputs.shape} does not match input seqs length {len(_seqs[0])}: {seqs}"
+            return outputs
+
+
+def get_model(name: str, dev):
+    "Get model by name"
+    if name in [
+        "esm2_t48_15B_UR50D",
+        "esm2_t36_3B_UR50D",
+        "esm2_t33_650M_UR50D",
+        "esm2_t30_150M_UR50D",
+        "esm2_t12_35M_UR50D",
+        "esm2_t6_8M_UR50D",
+        "esm1b_t33_650M_UR50S",
+    ]:
+        d = EsmModelInfo(name)
+        return EsmEncoder(d["model"], dev)
+    elif name in ["prot_t5_xl_half_uniref50-enc", "prot_t5_xl_bfd"]:
+        d = EsmModelInfo(name)
+        return T5Encoder(d["model"], dev)
+    else:
+        raise ValueError(f"Unknown model name: {name}")

src/models/polybert.py

@@ -0,0 +1,41 @@
+from transformers import AutoTokenizer, AutoModel
+import torch
+import torch.nn as nn
+
+polymer2psmiles = {
+    'PHB': '[*]OC(C)CC(=O)[*]',  
+    'PCL': '[*]OCCCCCC(=O)[*]', 
+    'PVA': '[*]C(O)C[*]',      
+    'PPL': '[*]CCC(=O)O[*]',  
+    'P3HP': '[*]OCCC(=O)[*]', 
+    'P4HB': '[*]C(=O)CCCO[*]',
+    'PEA': '[*]OCCOC(=O)CCCCC(=O)[*]', 
+    'PES': '[*]OCCOC(=O)CCC(=O)[*]',  
+    'O-PVA': '[*]C(=O)C[*]', 
+    'PBS': '[*]C(=O)CCC(=O)OCCCCO[*]', 
+    'PLA': '[*]C(C)C(=O)O[*]', 
+    'PEG': '[*]CCO[*]', 
+    'PBSA': '[*]C(=O)CCC(=O)OCCCCOC(=O)CCC(=O)[*]', 
+    'PET': '[*]CCOC(=O)c1ccc(C(=O)O[*])cc1', 
+    'PE': '[*]CC[*]',
+    'PMCL': '[*]C(=O)CCC(C)CCO[*]', 
+    'PEF': '[*]OC(=O)c1oc(C(=O)OCC[*])cc1', 
+    'PS': '[*]C(c1ccccc1)C[*]', 
+    'NR': '[*]CC(C)=CC[*]', 
+    'PHV': '[*]OC(CC)CC(=O)[*]',
+}
+
+
+class PolyEncoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.tokenizer = AutoTokenizer.from_pretrained('kuelumbus/polyBERT')
+        self.polyBERT = AutoModel.from_pretrained('kuelumbus/polyBERT')
+
+    def forward(self, psmiles_strings):
+        assert len(psmiles_strings) == 1, "Batch size must be 1 for PolyEncoder"
+        encoded_input = self.tokenizer(
+            psmiles_strings, padding=True, truncation=True, return_tensors='pt')
+        with torch.no_grad():
+            model_output = self.polyBERT(**encoded_input)
+        return model_output[0]

src/models/training.py

@@ -0,0 +1,263 @@
+import json
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from models.dataset import EnzymeDataModule
+import lightning as L
+from timm.scheduler.cosine_lr import CosineLRScheduler
+from argparse import Namespace as Args
+import wandb
+import time
+from models.dataset import polymer2psmiles
+from models.plm import EsmModelInfo
+from models.polybert import PolyEncoder
+from models.utils import is_wandb_running
+from pathlib import Path
+from einops import rearrange
+import numpy as np
+from sklearn.metrics import roc_auc_score, f1_score, matthews_corrcoef, precision_score, recall_score
+
+
+class CrossAttnLayer(nn.Module):
+    def __init__(self, protein_dim, smiles_dim, nheads=8):
+        super().__init__()
+        self.fc_smiles = nn.Linear(smiles_dim, protein_dim)
+        self.fc_protein = nn.Linear(protein_dim, smiles_dim)
+        self.smiles2protein = nn.MultiheadAttention(
+            smiles_dim, nheads, batch_first=True)
+        self.protein2smiles = nn.MultiheadAttention(
+            protein_dim, nheads, batch_first=True)
+
+    def forward(self, protein, smiles):
+        down_protein = self.fc_protein(protein)
+        up_smiles = self.fc_smiles(smiles)
+        l_attn, l_weights = self.smiles2protein(
+            smiles, down_protein, down_protein)
+        p_attn, p_weights = self.protein2smiles(protein, up_smiles, up_smiles)
+        return l_attn, p_attn, l_weights, p_weights
+
+
+class BaseModel(nn.Module):
+    def __init__(self, in_dim1, in_dim2, n_classes):
+        super().__init__()
+        self.attn = CrossAttnLayer(in_dim1, in_dim2)
+        self.fc = nn.Linear(in_dim1 + in_dim2, n_classes)
+
+    def forward(self, x):
+        protein, smiles = x
+        P, L, P_weights, L_weights = self.attn(protein, smiles)
+        x = torch.cat((P.mean(dim=1), L.mean(dim=1)), dim=-1)
+        x = self.fc(x)
+        return x, P_weights, L_weights
+
+
+class PlasticPredictor(L.LightningModule):
+    def __init__(self, args: L.LightningModule):
+        super().__init__()
+        self.args = args
+        info = EsmModelInfo(args.plm)
+        plm_dim = info['dim']*2  # the first and last layers are concatenated
+        pbert_dim = 600
+        self.model = BaseModel(
+            in_dim1=plm_dim, in_dim2=pbert_dim, n_classes=2)
+
+        self.cached_proteins = {}
+        self.cached_smiles = {}
+
+        self.encoder = {}  # trick: use dictionary to exclude modules
+        self.encoder['polybert'] = PolyEncoder()
+
+        self.automatic_optimization = False
+
+    def forward(self, x):
+        pass
+
+    def get_protein_embedding(self, seq_id):
+        if seq_id not in self.cached_proteins:
+            seq_path = f'cache/protein/{self.args.plm}/{seq_id}.pt'
+            if not Path(seq_path).exists():
+                raise FileNotFoundError(
+                    f"Protein embedding for {seq_id} not found.")
+            emb = torch.load(seq_path)
+            emb = rearrange(emb, 'b l d -> b (l d)')
+            self.cached_proteins[seq_id] = emb
+        return self.cached_proteins[seq_id]
+
+    def get_smiles_embedding(self, polymer):
+        smi = polymer2psmiles[polymer]
+        # mol = Chem.MolFromSmiles(smi)
+        # smi = Chem.MolToSmiles(mol, doRandom=True)
+        # # replace * with [*]
+        # smi = smi.replace('*', '[*]')
+
+        if smi not in self.cached_smiles:
+            # first dimension is 1
+            with torch.no_grad(), torch.inference_mode():
+                emb = self.encoder['polybert']([smi])[0, 2:-1, :]
+            self.cached_smiles[smi] = emb
+        return self.cached_smiles[smi]
+
+    def step(self, batch):
+        polymer, seq, deg, seq_id, poly_id = zip(batch)
+        seqs = [self.get_protein_embedding(s.item()) for s in seq_id[0]]
+        polys = [self.get_smiles_embedding(p) for p in polymer[0]]
+        protein_lengths = [len(s) for s in seqs]
+        smiles_lengths = [len(p) for p in polys]
+
+        seqs = nn.utils.rnn.pad_sequence(
+            seqs, batch_first=True).to(self.device)
+        polys = nn.utils.rnn.pad_sequence(
+            polys, batch_first=True).to(self.device)
+        protein_lengths = torch.tensor(
+            protein_lengths, dtype=torch.long).to(self.device)
+        smiles_lengths = torch.tensor(
+            smiles_lengths, dtype=torch.long).to(self.device)
+
+        logits, P_weights, L_weights = self.model((seqs, polys))
+        # Flatten the output for cross-entropy loss
+        logits = logits.view(-1, 2)
+        deg = deg[0].to(self.device)
+
+        loss = F.cross_entropy(logits, deg, reduction='mean')
+
+        y_prob = torch.softmax(logits, dim=-1)[:, 1]
+
+        return deg, y_prob, loss
+
+    def training_step(self, batch, batch_idx):
+        y, y_prob, loss = self.step(batch)
+        self.log_dict({"train/loss": loss, }, prog_bar=True)
+
+        self.manual_backward(loss)
+
+        self.optimizers().step()
+        self.lr_scheduler_step()
+        self.optimizers().zero_grad()
+
+    def validation_step(self, batch, batch_idx):
+        y, y_prob, loss = self.step(batch)
+        self.y.append(y.detach().cpu().numpy())
+        self.y_prob.append(y_prob.detach().cpu().numpy())
+
+    def on_validation_epoch_start(self):
+        self.y, self.y_prob = [], []
+
+    def on_validation_epoch_end(self):
+        y_prob = np.concatenate(self.y_prob, axis=0)
+        y = np.concatenate(self.y, axis=0)
+        auc = roc_auc_score(y, y_prob)
+        f1 = f1_score(y, y_prob > 0.5)
+        mcc = matthews_corrcoef(y, y_prob > 0.5)
+        precision = precision_score(y, y_prob > 0.5)
+        recall = recall_score(y, y_prob > 0.5)
+        self.log_dict({
+            "val_auc": auc,
+            "val_f1": f1,
+            "val_mcc": mcc,
+            "val_pre": precision,
+            "val_rec": recall,
+        }, prog_bar=True)
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.AdamW(
+            self.model.parameters(), lr=self.args.lr, weight_decay=self.args.wd
+        )
+        warmup_steps = round(self.args.t_initial * 0.1)
+        lr_scheduler = CosineLRScheduler(
+            optimizer,
+            t_initial=self.args.t_initial,
+            lr_min=1e-5,
+            warmup_t=warmup_steps,
+            warmup_lr_init=1e-5,
+            warmup_prefix=True,
+        )
+        self.lr_scheduler = lr_scheduler
+        return [optimizer]
+
+    def lr_scheduler_step(self, *args, **kwargs):
+        if self.trainer.global_step < self.trainer.max_steps:
+            self.lr_scheduler.step_update(self.trainer.global_step)
+
+
+def train_plastic(args: Args):
+    L.seed_everything(args.seed)
+    if is_wandb_running():
+        wandb.init(project="plastic-predictor",)
+        args.__dict__.update(dict(wandb.config))
+
+    dm = EnzymeDataModule(args)
+
+    for kfold in range(args.nfolds):
+        print(f"Training fold {kfold + 1}/{args.nfolds}")
+
+        model = PlasticPredictor(args)
+        dm.setup_k_fold(kfold)
+
+        print(
+            f'Data loaded: {len(dm.train_dataloader())} train, {len(dm.val_dataloader())} val, {len(dm.test_dataloader())} test')
+
+        devices = 1
+        logger = None
+        # devices = torch.cuda.device_count()
+        # logger = L.pytorch.loggers.WandbLogger(project="plastic-predictor",)
+
+        strategy = "ddp" if devices > 1 else "auto"
+        steps_per_epoch = len(dm.train_dataloader())
+        args.__dict__.update(
+            {
+                "batch_size": args.batch_size // devices,
+                "dev_count": devices,
+                "t_initial": args.epochs * steps_per_epoch,
+                "steps_per_epoch": steps_per_epoch,
+            }
+        )
+        print(f"Total steps: {args.t_initial}")
+        checkpoint = L.pytorch.callbacks.ModelCheckpoint(
+            dirpath=args.ckpt_dir,
+            filename=f"plastic-{{epoch:02d}}-{{val_auc:.4f}}",
+            # monitor="val_auc",
+            # mode="max",
+        )
+
+        early_stopping = L.pytorch.callbacks.EarlyStopping(
+            monitor="val_auc",
+            patience=args.patience,
+            mode="max",
+            verbose=True,
+        )
+        precision = "16-mixed" if args.amp else "32-true"
+        trainer = L.Trainer(
+            max_epochs=args.epochs,
+            accelerator="gpu",
+            devices=devices,
+            strategy=strategy,
+            precision=precision,
+            log_every_n_steps=1,
+            callbacks=[checkpoint, early_stopping],
+            # callbacks=[checkpoint],
+            # enable_checkpointing=False,
+            logger=logger,
+        )
+
+        trainer.fit(model, dm)
+
+        trainer.validate(model, dm.val_dataloader(),
+                         ckpt_path="best", verbose=True)
+        time.sleep(1)
+        val_test_metrics = trainer.callback_metrics.copy()
+        trainer.validate(model, dm.test_dataloader(),
+                         ckpt_path="best", verbose=True)
+
+        time.sleep(1)
+        val_test_metrics.update(
+            [(k.replace("val_", "test_"), v)
+             for k, v in trainer.callback_metrics.items()]
+        )
+
+        # val_test_metrics['fold'] = kfold + 1
+        # add _fold suffix to each key
+        val_test_metrics = {
+            k + f"_fold{kfold + 1}": v for k, v in val_test_metrics.items()
+        }
+        if is_wandb_running():
+            wandb.log(val_test_metrics)

src/models/utils.py

@@ -0,0 +1,342 @@
+import os
+import signal
+
+import psutil
+import torch
+import yaml
+from functools import wraps
+import errno
+import signal
+import numpy as np
+from scipy.spatial import KDTree
+from math import ceil
+from tqdm import tqdm
+import line_profiler
+import os
+import base64
+import pickle
+from cryptography.hazmat.primitives.asymmetric import padding
+from cryptography.hazmat.primitives import serialization, hashes
+from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
+from cryptography.hazmat.backends import default_backend
+
+import io
+
+
+def num_parameters(model: torch.nn.Module) -> int:
+    """Return the number of parameters in the model"""
+    return sum(p.numel() for p in model.parameters() if p.requires_grad)
+
+
+class Config:
+    """Read configuration from a YAML file and store as attributes"""
+
+    def __init__(self, yaml_file: str):
+        with open(yaml_file, "r") as f:
+            config = yaml.safe_load(f)
+
+        for k, v in config.items():
+            setattr(self, k, v)
+
+    def update(self, new_yaml_file: str):
+        with open(new_yaml_file, "r") as f:
+            config = yaml.safe_load(f)
+
+        for k, v in config.items():
+            setattr(self, k, v)
+
+    def save(self, yaml_file: str):
+        with open(yaml_file, "w") as f:
+            yaml.dump(self.__dict__, f)
+
+
+def memory_usage_psutil():
+    """Return the memory usage in percentage like top"""
+    process = psutil.Process(os.getpid())
+    mem = process.memory_percent()
+    return mem
+
+
+def is_wandb_running():
+    """Check if wandb is running"""
+    return "WANDB_SWEEP_ID" in os.environ
+
+
+class TimeoutError(Exception):
+    pass
+
+
+def timeout(seconds=10, error_message=os.strerror(errno.ETIME)):
+    def decorator(func):
+        def _handle_timeout(signum, frame):
+            raise TimeoutError(error_message)
+
+        def wrapper(*args, **kwargs):
+            signal.signal(signal.SIGALRM, _handle_timeout)
+            signal.alarm(seconds)
+            try:
+                result = func(*args, **kwargs)
+            finally:
+                signal.alarm(0)
+            return result
+
+        return wraps(func)(wrapper)
+
+    return decorator
+
+
+def shorten_path(path: str, max_len: int = 30) -> str:
+    """Shorten the path to max_len characters"""
+    if len(path) > max_len:
+        return path[:max_len // 2] + "..." + path[-max_len // 2:]
+    return path
+
+
+def cluster_points(data: torch.Tensor, d: float) -> torch.Tensor:
+    """
+    Cluster points based on the Euclidean distance.
+
+    :param data: Input data, shape (n_points, n_features), type torch.Tensor.
+    :param d: Distance threshold for clustering.
+    :return: Cluster indices, shape (n_points,), type torch.Tensor.
+    """
+    dist = torch.cdist(data, data)
+    indices = torch.full((data.shape[0],), -1, dtype=torch.long)
+    cluster_id = 0
+    for i in range(data.shape[0]):
+        if indices[i] == -1:
+            indices[dist[i] < d] = cluster_id
+            cluster_id += 1
+    return indices
+
+
+def bron_kerbosch(R, P, X, graph):
+    if not P and not X:
+        yield R
+    while P:
+        v = P.pop()
+        yield from bron_kerbosch(
+            R | {v},
+            P & set(graph[v]),
+            X & set(graph[v]),
+            graph
+        )
+        X.add(v)
+
+
+def find_cliques(graph):
+    """
+    Find all maximal cliques in an undirected graph with the Bron–Kerbosch algorithm.
+
+    :param graph: Input graph as a NetworkX graph
+    :return: List of maximal cliques
+    """
+    return list(bron_kerbosch(set(), set(graph.nodes()), set(), graph))
+
+
+def segment_cmd(cmd_str: str, max_len: int = 1000):
+    cmds = ['']
+    prev = 0
+    for i, c in enumerate(cmd_str):
+        if c == ';':
+            if len(cmds[-1]) + len(cmd_str[prev:i]) > max_len:
+                cmds.append('')
+            cmds[-1] += cmd_str[prev:i + 1]
+            prev = i + 1
+    return cmds
+
+
+def get_color(v):
+    assert 0 <= v <= 1, f'v should be in [0, 1], got {v}'
+    # green to brown
+    color1 = np.array([0, 128, 0])
+    color2 = np.array([165, 42, 42])
+    v = v * (color2 - color1) + color1
+    v /= 255
+    return f'[{v[0]:.2f},{v[1]:.2f},{v[2]:.2f}]'
+
+
+def generate_pymol_script(possible_sites):
+    cmd = ''
+    for i, pos in enumerate(possible_sites):
+        cmd += f"pseudoatom s{i},pos=[{pos[0]:.1f},{pos[1]:.1f},{pos[2]:.1f}];color blue,s{i};"
+    return cmd
+
+
+def remove_close_points_kdtree(points, min_distance):
+    tree = KDTree(points)
+    keep = np.ones(len(points), dtype=bool)
+    for i, point in enumerate(points):
+        if not keep[i]:
+            continue
+        neighbors = tree.query_ball_point(
+            point, min_distance)
+        keep[neighbors] = False
+        keep[i] = True  # Keep the current point
+    return points[keep]
+
+
+@line_profiler.profile
+def pack_bit(x: torch.Tensor):
+    """ Pack the bit tensor to a sequence of bytes.
+    Args:
+        x (torch.Tensor): The input tensor to be packed.
+    Returns:
+        torch.Tensor: The packed tensor.
+    """
+    batch_size, num_bits = x.shape
+    num_bytes = (num_bits + 7) // 8
+    output = torch.zeros(batch_size, num_bytes,
+                         dtype=torch.uint8, device=x.device)
+    for i in range(num_bits):
+        byte_index = i // 8
+        bit_index = i % 8
+        output[:, byte_index] |= (x[:, i] << bit_index).to(torch.uint8)
+    return output
+
+
+@line_profiler.profile
+def unpack_bit(x: torch.Tensor, num_bits: int):
+    """ Unpack the bit tensor from a sequence of bytes.
+    Args:
+        x (torch.Tensor): The input tensor to be unpacked.
+        num_bits (int): The number of bits to unpack.
+    Returns:
+        torch.Tensor: The unpacked tensor.
+    """
+    batch_size, num_bytes = x.shape
+    output = torch.zeros(batch_size, num_bits,
+                         dtype=torch.uint8, device=x.device)
+    for i in range(num_bits):
+        byte_index = i // 8
+        bit_index = i % 8
+        output[:, i] = (x[:, byte_index] >> bit_index) & 1
+    return output
+
+
+def safe_dist(vec1: torch.Tensor, vec2: torch.Tensor, max_size: int = 100_000_000, p: int = 2):
+    """ compute the minimum distance between two vectors:
+
+    vec1: (N, 3), N could be very very large, i.e., all atoms' coordinates in a large protein
+
+    vec2: (M, 3), M are not very large, usually the coordinates of the binding sites
+
+    max_size: the maximum size of the distance matrix to compute at once
+
+    p: the p-norm to use for distance calculation
+
+    return: (M, ) the minimum distance of each binding site to the protein
+    """
+    size1 = vec1.shape
+    size2 = vec2.shape
+    batch_size = ceil(max_size / size1[0])
+    dists = []
+    for i in range(0, size2[0], batch_size):
+        dist = torch.cdist(vec1, vec2[i:i + batch_size], p=p)
+        dists.append(dist.min(dim=0).values)
+    return torch.cat(dists)
+
+
+@line_profiler.profile
+def safe_filter(nos: torch.Tensor, pos: torch.Tensor, thr: torch.Tensor, all: torch.Tensor, lb: float, max_size: int = 100_000_000):
+    """ filter the binding sites based on the distance matrix 
+    nos: (N, 3), N are the coordinates of the binding sites
+    *pos: (M, 3), M are the coordinates of the protein, could be very very large
+    thr: (N, 2), the distance threshold for each binding site
+    all: (P, 3), P are the coordinates of all atoms in the protein
+    lb: the lower bound of the distance
+
+    return: (N, M) available binding sites
+    """
+    N, M, P = nos.shape[0], pos.shape[0], all.shape[0]
+    batch_size = ceil(max_size / N)
+    output = []
+    interests = []
+    for i in tqdm(range(0, M, batch_size), leave=False, desc=f'Filtering (batch_size: {batch_size})'):
+        dist = torch.cdist(pos[i:i + batch_size], nos)
+        dist = (dist <= thr[:, 1].unsqueeze(0)) & \
+            (dist >= thr[:, 0].unsqueeze(0))
+        dist_all = safe_dist(all, pos[i:i + batch_size]) > lb
+        dist = dist & dist_all.unsqueeze(-1)
+
+        mask = dist.any(dim=1)
+        output.append(pack_bit(dist[mask]).T)
+        interests.append(mask)
+    return torch.cat(output, dim=1), torch.cat(interests)
+
+
+def backbone(atoms, chain_id):
+    """ return the atoms of the backbone of a chain """
+    return atoms[
+        (atoms.chain_id == chain_id) &
+        (atoms.atom_name == "CA") &
+        (atoms.element == "C")]
+
+
+def get_color(v):
+    assert 0 <= v <= 1, f'v should be in [0, 1], got {v}'
+    # green to brown
+    color1 = np.array([0, 128, 0])
+    color2 = np.array([165, 42, 42])
+    v = v * (color2 - color1) + color1
+    v /= 255
+    return f'[{v[0]:.2f},{v[1]:.2f},{v[2]:.2f}]'
+
+
+def load_private_key_from_file(private_key_file=None):
+    if private_key_file is None:
+        private_key_b64 = os.environ.get('ModelCheckpointPrivateKey')
+    else:
+        with open(private_key_file, 'r') as f:
+            private_key_b64 = f.read().strip()
+
+    private_pem = base64.b64decode(private_key_b64)
+    private_key = serialization.load_pem_private_key(
+        private_pem,
+        password=None,
+        backend=default_backend()
+    )
+    return private_key
+
+
+def decrypt_checkpoint(encrypted_path, private_key):
+    backend = default_backend()
+
+    with open(encrypted_path, 'rb') as f:
+
+        key_length = int.from_bytes(f.read(4), 'big')
+
+        encrypted_aes_key = f.read(key_length)
+        iv = f.read(16)
+        original_size = int.from_bytes(f.read(8), 'big')
+        encrypted_data = f.read()
+
+    try:
+        aes_key = private_key.decrypt(
+            encrypted_aes_key,
+            padding.OAEP(
+                mgf=padding.MGF1(algorithm=hashes.SHA256()),
+                algorithm=hashes.SHA256(),
+                label=None
+            )
+        )
+
+        cipher = Cipher(algorithms.AES(aes_key),
+                        modes.CBC(iv), backend=backend)
+        decryptor = cipher.decryptor()
+        decrypted_padded = decryptor.update(
+            encrypted_data) + decryptor.finalize()
+
+        decrypted_data = decrypted_padded[:original_size]
+
+        try:
+            buffer = io.BytesIO(decrypted_data)
+            checkpoint_dict = torch.load(buffer, map_location='cpu')
+            return checkpoint_dict
+        except:
+            checkpoint_dict = pickle.loads(decrypted_data)
+            return checkpoint_dict
+
+    except Exception as e:
+        print(f"Error: {e}")
+        raise

src/predict.py

@@ -0,0 +1,87 @@
+from pathlib import Path
+from einops import rearrange
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import time
+from models.polybert import PolyEncoder
+from models.training import BaseModel
+from models.utils import decrypt_checkpoint, load_private_key_from_file
+import argparse
+from tqdm import tqdm
+
+from models.utils import Config
+from models.plm import EsmModelInfo, get_model
+import pandas as pd
+
+
+if __name__ == "__main__":
+    # fmt: off
+    parser = argparse.ArgumentParser(description="Predict plastic degradation")
+    parser.add_argument("--ckpt", type=str, help="Path to the model checkpoint")
+    parser.add_argument("--plm", type=str, help="Protein language model to use", default='esm2_t33_650M_UR50D')
+    parser.add_argument("--csv", type=str, help="Path to the CSV file with test data", default=None)
+    parser.add_argument("--output",'-o', type=str, help="Path to the output file", default='predictions.csv')
+    parser.add_argument("--attn", action='store_true', help="Save attention weights to files")
+    # fmt: on
+    args = parser.parse_args()
+
+    info = EsmModelInfo(args.plm)
+    plm_dim = info['dim']*2
+    pbert_dim = 600
+
+    dev = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    model = BaseModel(plm_dim, pbert_dim, n_classes=2).to(dev)
+
+    # load weights
+    private_key = load_private_key_from_file()
+    state_dict = decrypt_checkpoint(args.ckpt, private_key)
+    state_dict = {
+        k.replace('model.', ''): v for k, v in state_dict['state_dict'].items() if k.startswith('model.')}
+    model.load_state_dict(state_dict)
+    model.eval()
+    print(f'Load predictor from {args.ckpt}')
+
+    plm_func = get_model(args.plm, 'cuda')
+    print(f'Loaded PLM model {args.plm}')
+
+    polybert_func = PolyEncoder()
+    print('Loaded PolyEncoder model')
+
+    outfile = Path(
+        'predictions.csv' if args.output is None else args.output)
+    # get protein embedding
+    with torch.no_grad(), torch.inference_mode():
+        df = pd.read_csv(args.csv)
+        probs = []
+        running_time = []
+        for i, row in tqdm(df.iterrows()):
+            start_time = time.time()
+
+            seq = row['sequence'].upper()
+            poly = row['polymer']
+            seq_emb = plm_func([seq]).to(dev)
+            seq_emb = rearrange(seq_emb, 'b l d -> b (l d)').unsqueeze(0)
+            poly_emb = polybert_func([poly]).to(dev)
+            logits, p_weights, l_weights = model((seq_emb, poly_emb))
+            prob = F.softmax(logits, dim=-1)[:, 1].item()
+            if args.attn:
+                outfile.with_suffix('.attn').mkdir(
+                    parents=True, exist_ok=True)
+                torch.save(
+                    (p_weights, l_weights),
+                    outfile.with_suffix('.attn') / f'{i}.pt')
+            probs.append(prob)
+            running_time.append(time.time() - start_time)
+
+        df['prob'] = probs
+        df['pred'] = df['prob'].apply(lambda x: 'Yes' if x >= 0.5 else 'No')
+        df['time'] = running_time
+
+        # move pred and prob to the front
+        df = df[['pred', 'prob'] +
+                [col for col in df.columns if col not in ['pred', 'prob']]]
+
+        df.to_csv(outfile, index=False)
+        print(f'Predictions saved to {outfile}')
+        print(f'Attention weights saved to current directory as <index>.pt')

src/streamlit_app.py

@@ -0,0 +1,248 @@
+#fmt: off
+import streamlit as st
+import pandas as pd
+import os
+import tempfile
+import subprocess
+import requests
+import csv
+from models.polybert import polymer2psmiles
+import py3Dmol
+
+# Fix for permission error - disable usage stats
+if 'STREAMLIT_CONFIG_DIR' not in os.environ:
+    os.environ['STREAMLIT_CONFIG_DIR'] = '/tmp/.streamlit'
+
+# Create streamlit config directory if it doesn't exist
+streamlit_dir = os.environ.get('STREAMLIT_CONFIG_DIR', '/tmp/.streamlit')
+os.makedirs(streamlit_dir, exist_ok=True)
+
+# Create config.toml to disable usage stats
+config_path = os.path.join(streamlit_dir, 'config.toml')
+if not os.path.exists(config_path):
+    with open(config_path, 'w') as f:
+        f.write("""[browser]
+gatherUsageStats = false
+
+[server]
+headless = true
+enableCORS = false
+enableXsrfProtection = false
+""")
+# fmt: on
+
+aa2resn = {
+    'A': 'ALA',
+    'C': 'CYS',
+    'D': 'ASP',
+    'E': 'GLU',
+    'F': 'PHE',
+    'G': 'GLY',
+    'H': 'HIS',
+    'I': 'ILE',
+    'K': 'LYS',
+    'L': 'LEU',
+    'M': 'MET',
+    'N': 'ASN',
+    'P': 'PRO',
+    'Q': 'GLN',
+    'R': 'ARG',
+    'S': 'SER',
+    'T': 'THR',
+    'V': 'VAL',
+    'W': 'TRP',
+    'Y': 'TYR'
+}
+
+# Fancy header
+st.markdown("""
+<div style='text-align: center;'>
+    <h1 style='color:#377EB9;font-size:2.5em;'>🧬 Plastic Degradation Predictor</h1>
+    <h3 style='color:#4DAE48;'>Predict the degradability of plastics using protein sequences and polymer SMILES</h3>
+</div>
+<hr style='border:1px solid #974F9F;'>
+""", unsafe_allow_html=True)
+
+st.write("Enter a UniProt ID or paste a protein sequence. Select a polymer from the list below.")
+
+
+# Load polymer names and SMILES
+
+# Only show polymers with SMILES in the dropdown
+polymer_csv = os.path.join(os.path.dirname(
+    __file__), 'data/polymer2tok.csv')
+polymer_options = []
+with open(polymer_csv, newline='') as f:
+    reader = csv.DictReader(f)
+    for row in reader:
+        name = row['polymer']
+        smiles = polymer2psmiles.get(name, '')
+        if smiles:  # Only include polymers with SMILES
+            polymer_options.append(f"{name} | {smiles}")
+
+input_type = st.radio("Input type", ["UniProt ID", "Protein Sequence"])
+
+if input_type == "UniProt ID":
+    uniprot_id = st.text_input("Enter UniProt ID", "P69905")
+    sequence = ""
+    if uniprot_id:
+        # Fetch sequence from UniProt
+        url = f"https://rest.uniprot.org/uniprotkb/{uniprot_id}.fasta"
+        resp = requests.get(url)
+        if resp.status_code == 200:
+            fasta = resp.text
+            sequence = "".join(fasta.split("\n")[1:])
+            st.success(f"Fetched sequence for {uniprot_id}")
+            st.code(sequence)
+        else:
+            st.error("Failed to fetch sequence from UniProt.")
+else:
+    sequence = st.text_area("Paste protein sequence",
+                            "MVLSPADKTNVKAAWGKVGAHAGEYGAEALERMFLSFPTTKTYFPHFDLSHGSAQVKGHG")
+
+polymer = st.selectbox("Select polymer", polymer_options)
+selected_polymer = polymer.split('|')[0].strip() if '|' in polymer else polymer
+
+
+ckpt = "src/checkpoints/weights.ckpt"
+plm = "esm2_t33_650M_UR50D"
+
+if st.button("Predict degradation", type="primary"):
+    if not sequence or not selected_polymer:
+        st.error("Please provide both sequence and polymer.")
+    else:
+        # Create temp CSV
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".csv", mode="w") as tmp:
+            tmp.write("sequence,polymer\n")
+            tmp.write(f"{sequence},{selected_polymer}\n")
+            tmp_path = tmp.name
+        output_path = os.path.join(tempfile.gettempdir(), "predictions.csv")
+        st.write("Running prediction...")
+        result = subprocess.run([
+            "python", "src/predict.py",
+            "--ckpt", ckpt,
+            "--plm", plm,
+            "--csv", tmp_path,
+            "--output", output_path,
+            "--attn"
+        ], capture_output=True, text=True)
+        if result.returncode == 0 and os.path.exists(output_path):
+            df = pd.read_csv(output_path)
+            if 'time' in df.columns:
+                df = df.rename(columns={'time': 'running time'})
+            st.markdown(f"""
+<div style='background: linear-gradient(90deg, #377EB9 0%, #4DAE48 100%); padding: 1.5em; border-radius: 12px; color: white; margin-bottom: 1em;'>
+    <h2 style='margin:0;'><span style='font-size:18pt'>✅</span> Prediction Complete!</h2>
+    <p style='font-size:12pt;'>Your input has been processed. See the results below:</p>
+    <p style='font-size:12pt;'>Degradation: {df['pred'].values[0]} (Probability: {df['prob'].values[0]:.4f})</p>
+</div>
+""", unsafe_allow_html=True)
+            st.download_button("⬇️ Download Results", data=df.to_csv(
+                index=False), file_name="predictions.csv", type="primary")
+
+            # Show top-N attention residues if attention file exists
+            attn_dir = os.path.join(os.path.dirname(
+                output_path), "predictions.attn")
+            attn_path = os.path.join(attn_dir, "0.pt")
+            if os.path.exists(attn_path):
+                import torch
+                attn = torch.load(attn_path)
+                # attn[0][0]: shape (num_heads, seq_len, seq_len) or (1, seq_len, seq_len)
+                attn_matrix = attn[0][0] if isinstance(
+                    attn[0], (list, tuple)) else attn[0]
+                # Average over heads if needed
+                if attn_matrix.ndim == 3:
+                    attn_matrix = attn_matrix.mean(0)
+                # For each residue, sum attention weights
+                residue_scores = attn_matrix.sum(0).cpu().numpy()
+                topN = min(10, len(residue_scores))
+                top_idx = residue_scores.argsort()[::-1][:topN]
+                st.markdown(f"**Top {topN} high-attention residues:**")
+                st.write(pd.DataFrame({
+                    "Amino Acid": [sequence[i] for i in top_idx],
+                    "Residue Index": top_idx+1,
+                    "Attention Score": residue_scores[top_idx]
+                }))
+            else:
+                st.info("No attention file found for visualization.")
+        else:
+            st.error("Prediction failed. See details below:")
+            st.text(result.stderr)
+
+        # If UniProt ID, try to download AlphaFold structure
+        structure_path = None
+
+        if input_type == "UniProt ID" and uniprot_id:
+            af_url = f"https://alphafold.ebi.ac.uk/files/AF-{uniprot_id}-F1-model_v4.cif"
+            # If attention available, highlight top residues
+            highlight_residues = None
+            attn_dir = os.path.join(tempfile.gettempdir(), "predictions.attn")
+            attn_path = os.path.join(attn_dir, "0.pt")
+            if os.path.exists(attn_path):
+                import torch
+                attn = torch.load(attn_path)
+                attn_matrix = attn[0][0] if isinstance(
+                    attn[0], (list, tuple)) else attn[0]
+                if attn_matrix.ndim == 3:
+                    attn_matrix = attn_matrix.mean(0)
+                residue_scores = attn_matrix.sum(0).cpu().numpy()
+                topN = min(10, len(residue_scores))
+                top_idx = residue_scores.argsort()[::-1][:topN]
+                # Molstar selection: list of residue numbers (1-based)
+                highlight_residues = [int(i+1) for i in top_idx]
+
+            structure_path = os.path.join(
+                tempfile.gettempdir(), f"AF-{uniprot_id}-F1-model_v4.cif")
+            try:
+                r = requests.get(af_url)
+                if r.status_code == 200:
+                    with open(structure_path, "wb") as f:
+                        f.write(r.content)
+                    st.success(
+                        f"AlphaFold structure downloaded: {structure_path}")
+                else:
+                    st.warning(
+                        "AlphaFoldDB structure not found for this UniProt ID.")
+            except Exception as e:
+                st.warning(f"AlphaFoldDB download error: {e}")
+
+        if input_type == "UniProt ID" and uniprot_id and os.path.exists(attn_path) and os.path.exists(structure_path):
+            st.markdown("### 3D Structure Visualization (stmol)")
+            import torch
+            from stmol import showmol
+            attn = torch.load(attn_path)
+            attn_matrix = attn[0][0] if isinstance(
+                attn[0], (list, tuple)) else attn[0]
+            if attn_matrix.ndim == 3:
+                attn_matrix = attn_matrix.mean(0)
+            residue_scores = attn_matrix.sum(0).cpu().numpy()
+            topN = min(10, len(residue_scores))
+            top_idx = residue_scores.argsort()[::-1][:topN]
+            labels = [
+                f"{sequence[i]}{i+1}: {residue_scores[i]:.4g}" for i in top_idx]
+            with open(structure_path, "r") as cif_file:
+                cif_data = cif_file.read()
+            view = py3Dmol.view(width=600, height=400)
+            view.addModel(cif_data, "cif")
+            view.setStyle({"cartoon": {"color": "lightgray"}})
+            for i, idx in enumerate(top_idx):
+                resi_num = int(idx+1)
+                view.setStyle(
+                    {"resi": resi_num}, {
+                        "cartoon": {"color": "red"}})
+                view.addResLabels(
+                    {"resi": resi_num},
+                    {
+                        "font": 'Arial', "fontColor": 'black',
+                        "showBackground": False, "screenOffset": {"x": 0, "y": 0}})
+            view.zoomTo()
+            showmol(view, height=600, width='100%')
+
+
+# --- Footer: License and References ---
+st.markdown("""
+---
+<h4>License</h4>
+Attribution-NonCommercial-ShareAlike 4.0 International (CC BY-NC-SA 4.0)<br>
+<a href='https://creativecommons.org/licenses/by-nc-sa/4.0/' target='_blank'>View full license details</a><br>
+""", unsafe_allow_html=True)