Spaces:
Runtime error
Runtime error
| import torch | |
| import torch.nn.functional as F | |
| import gradio as gr | |
| import selfies as sf | |
| import numpy as np | |
| import pandas as pd | |
| import os | |
| from rdkit import Chem | |
| from rdkit.Chem import QED | |
| CKPT_PATH = "diffusion_cfg_mtr_stage2_DEPLOY.pt" | |
| checkpoint = torch.load(CKPT_PATH, map_location="cpu") | |
| vocab = checkpoint["vocab"] | |
| inv_vocab = checkpoint["inv_vocab"] | |
| max_len = checkpoint["max_len"] | |
| raw_config = checkpoint["config"] | |
| allowed_keys = {"dim", "depth", "heads", "dropout", "num_mtr_targets"} | |
| model_config = {k: v for k, v in raw_config.items() if k in allowed_keys} | |
| class PositionalEmbedding(torch.nn.Module): | |
| def __init__(self, max_len, dim): | |
| super().__init__() | |
| self.emb = torch.nn.Embedding(max_len, dim) | |
| def forward(self, x): | |
| positions = torch.arange(x.size(1), device=x.device) | |
| return self.emb(positions)[None, :, :] | |
| class DiffusionTransformer(torch.nn.Module): | |
| def __init__( | |
| self, | |
| vocab_size, | |
| max_len, | |
| dim, | |
| depth, | |
| heads, | |
| dropout=0.1, | |
| num_mtr_targets=209, | |
| ): | |
| super().__init__() | |
| self.token_emb = torch.nn.Embedding(vocab_size, dim) | |
| self.pos_emb = PositionalEmbedding(max_len, dim) | |
| layer = torch.nn.TransformerEncoderLayer( | |
| d_model=dim, | |
| nhead=heads, | |
| dim_feedforward=4 * dim, | |
| dropout=dropout, | |
| batch_first=True, | |
| activation="gelu", | |
| ) | |
| self.encoder = torch.nn.TransformerEncoder(layer, depth) | |
| self.norm = torch.nn.LayerNorm(dim) | |
| self.head = torch.nn.Linear(dim, vocab_size) | |
| self.mtr_head = torch.nn.Sequential( | |
| torch.nn.Linear(dim, 512), | |
| torch.nn.ReLU(), | |
| torch.nn.Linear(512, num_mtr_targets), | |
| ) | |
| def forward(self, x, return_mtr=False): | |
| h = self.token_emb(x) + self.pos_emb(x) | |
| h = self.encoder(h) | |
| h = self.norm(h) | |
| logits = self.head(h) | |
| if return_mtr: | |
| mtr = self.mtr_head(h[:, 0, :]) | |
| return logits, mtr | |
| return logits | |
| model = DiffusionTransformer( | |
| vocab_size=len(vocab), | |
| max_len=max_len, | |
| **model_config, | |
| ) | |
| state_dict = checkpoint["model_state_dict"] | |
| new_state_dict = {} | |
| for k, v in state_dict.items(): | |
| if k.startswith("pos_emb.pos_emb."): | |
| new_k = k.replace("pos_emb.pos_emb.", "pos_emb.emb.") | |
| new_state_dict[new_k] = v | |
| else: | |
| new_state_dict[k] = v | |
| model.load_state_dict(new_state_dict) | |
| model.eval() | |
| def finalize_unmask(x): | |
| logits = model(x) | |
| return logits.argmax(dim=-1) | |
| def sample_selfies_cfg(qed_bin, gamma, steps=30): | |
| device = "cpu" | |
| BOS = vocab["[BOS]"] | |
| EOS = vocab["[EOS]"] | |
| MASK = vocab["[MASK]"] | |
| QED_TOKEN = vocab[f"[QED_{qed_bin}]"] | |
| QED_NULL = vocab["[QED_NULL]"] | |
| x = torch.full((1, max_len), MASK,) | |
| x[0, 0] = BOS | |
| x[0, 1] = QED_TOKEN | |
| x[0, -1] = EOS | |
| x_uncond = x.clone() | |
| x_uncond[0, 1] = QED_NULL | |
| protected = (x == BOS) | (x == EOS) | (x == QED_TOKEN) | |
| for step in range(steps): | |
| logits_u = model(x_uncond) | |
| logits_c = model(x) | |
| if step >= steps // 2: | |
| logits = logits_u + gamma * (logits_c - logits_u) | |
| else: | |
| logits = logits_u | |
| probs = F.softmax(logits, dim=-1) | |
| for tok in ["[PAD]", "[BOS]", "[EOS]"]: | |
| probs[..., vocab[tok]] = 0 | |
| probs = probs / probs.sum(dim=-1, keepdim=True) | |
| sampled = torch.multinomial( | |
| probs.view(-1, probs.size(-1)), 1 | |
| ).view(1, max_len) | |
| conf = probs.max(dim=-1).values | |
| threshold = 0.9 * (1 - step / steps) | |
| remask = (conf < threshold) & (~protected) | |
| x = torch.where(remask, MASK, sampled) | |
| x_uncond = torch.where(remask, MASK, sampled) | |
| x[0, 0], x[0, 1], x[0, -1] = BOS, QED_TOKEN, EOS | |
| x_uncond[0, 0], x_uncond[0, 1], x_uncond[0, -1] = BOS, QED_NULL, EOS | |
| ids = finalize_unmask(x) | |
| selfie = "".join( | |
| inv_vocab[i.item()] | |
| for i in ids[0] | |
| if inv_vocab[i.item()] not in ["[PAD]", "[BOS]", "[EOS]", "[QED_NULL]"] | |
| and not inv_vocab[i.item()].startswith("[QED_") | |
| ) | |
| try: | |
| smiles = sf.decoder(selfie) | |
| mol = Chem.MolFromSmiles(smiles) | |
| qed = QED.qed(mol) if mol else None | |
| except: | |
| smiles, qed = None, None | |
| return smiles, qed | |
| def generate(qed_bin, gamma, n_samples): | |
| rows = [] | |
| for _ in range(n_samples): | |
| smiles, qed = sample_selfies_cfg(qed_bin, gamma) | |
| if smiles is not None and qed is not None: | |
| rows.append([smiles, round(qed, 4)]) | |
| if len(rows) == 0: | |
| return [["No valid molecules generated", None]] | |
| df = pd.DataFrame(rows, columns=["SMILES", "QED"]) | |
| return df | |
| def download_csv(df): | |
| if df is None or len(df) == 0: | |
| return None | |
| path = "/tmp/apexoracle_results.csv" | |
| df.to_csv(path, index=False) | |
| return path | |
| with gr.Blocks(title="ApexOracle – QED-Controlled Molecular Generation") as demo: | |
| gr.Markdown( | |
| """ | |
| # 🧪 ApexOracle – QED-Controlled Molecular Generation | |
| **Discrete diffusion model with CFG + Multi-Target Regression** | |
| """ | |
| ) | |
| # 🔹 Controls (TOP) | |
| with gr.Column(): | |
| qed_bin = gr.Slider( | |
| 0, 9, step=1, value=8, label="Target QED Bin" | |
| ) | |
| gamma = gr.Slider( | |
| 0.0, 10.0, step=0.5, value=6.0, label="CFG Gamma" | |
| ) | |
| n_samples = gr.Slider( | |
| 1, 20, step=1, value=5, label="Number of Molecules" | |
| ) | |
| submit_btn = gr.Button("Generate", variant="primary") | |
| clear_btn = gr.Button("Clear") | |
| gr.Markdown("### Generated Molecules") | |
| output_df = gr.Dataframe( | |
| headers=["SMILES", "QED"], | |
| datatype=["str", "number"], | |
| row_count=(5, "dynamic"), | |
| ) | |
| download_btn = gr.Button("Download CSV") | |
| csv_file = gr.File(label="Download CSV") | |
| submit_btn.click( | |
| fn=generate, | |
| inputs=[qed_bin, gamma, n_samples], | |
| outputs=output_df, | |
| ) | |
| clear_btn.click( | |
| fn=lambda: None, | |
| inputs=None, | |
| outputs=output_df, | |
| ) | |
| download_btn.click( | |
| fn=download_csv, | |
| inputs=output_df, | |
| outputs=csv_file, | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch( | |
| auth=(os.environ.get("app_user"), os.environ.get("app_pass")), | |
| ssr_mode=False | |
| ) | |