| import html |
| import gradio as gr |
| import spaces |
| from rdkit import Chem |
| from rdkit.Chem.Draw import rdMolDraw2D |
|
|
| from inference import VQAtomTokenizer |
|
|
| TOKENIZER = VQAtomTokenizer("data/model_epoch_3.pt", device="cpu") |
|
|
|
|
| def draw_svg(smiles, tokens): |
| mol = Chem.MolFromSmiles(smiles) |
| if mol is None: |
| return "" |
| for atom, token in zip(mol.GetAtoms(), tokens): |
| atom.SetProp("atomNote", f"VQ{token}") |
| drawer = rdMolDraw2D.MolDraw2DSVG(760, 330) |
| opts = drawer.drawOptions() |
| opts.addAtomIndices = True |
| drawer.DrawMolecule(mol) |
| drawer.FinishDrawing() |
| return drawer.GetDrawingText().replace("svg:", "") |
|
|
|
|
| @spaces.GPU(duration=1) |
| def zerogpu_marker(): |
| """Required by ZeroGPU hosting; live inference is intentionally CPU-side.""" |
| return True |
|
|
|
|
| def encode_molecule(smiles): |
| smiles = (smiles or "").strip() |
| if not smiles: |
| raise gr.Error("Enter a SMILES string.") |
| try: |
| result = TOKENIZER.encode(smiles) |
| except Exception as exc: |
| raise gr.Error(str(exc)) |
| mol = Chem.MolFromSmiles(smiles) |
| atoms = [a.GetSymbol() for a in mol.GetAtoms()] |
| tokens = result["tokens"] |
| chips = "".join( |
| f'<div class="atom-token"><span>{i}</span><b>{html.escape(atom)}</b><code>VQ{token}</code></div>' |
| for i, (atom, token) in enumerate(zip(atoms, tokens)) |
| ) |
| return f'''<div class="result-shell"><div class="result-head"><span>VQATOM ENCODING</span><small>{html.escape(smiles)}</small></div> |
| <div class="mol-svg">{draw_svg(smiles, tokens)}</div> |
| <div class="sequence"><span>PER-ATOM TOKEN SEQUENCE</span><div>{chips}</div></div></div>''' |
|
|
|
|
| CSS = """ |
| @import url('https://fonts.googleapis.com/css2?family=DM+Sans:wght@400;500;600;700&display=swap'); |
| :root{--ink:#17231d;--paper:#f7f6f1;--line:#d9ddd7;--green:#bfff3f;--muted:#56625c} |
| .gradio-container{max-width:1080px!important;margin:0 auto!important;padding:0 24px!important;background:var(--paper)!important;color:var(--ink)!important;font-family:'DM Sans',sans-serif!important} |
| .contain{max-width:none!important;padding:0!important}footer{display:none!important}.main{padding:0!important} |
| .shell{width:100%;margin:auto}.topbar{height:70px;display:flex;align-items:center;justify-content:space-between;border-bottom:1px solid var(--line)} |
| .brand{font-size:20px;font-weight:700}.brand i{font-style:normal;background:var(--ink);color:var(--green);border-radius:50%;font-size:10px;padding:9px 6px;margin-right:8px} |
| .paper-link{color:var(--ink)!important;text-decoration:none;border:1px solid var(--line);border-radius:7px;padding:10px 14px;font-size:15px;font-weight:600} |
| .hero{padding:54px 0 32px;max-width:850px}.hero h1{color:var(--ink)!important;font-size:clamp(46px,6vw,68px);line-height:1.05;letter-spacing:-2.5px;margin:0 0 20px}.hero p{font-size:19px;line-height:1.65;color:var(--muted)!important;max-width:760px;margin:0} |
| .demo-intro{background:white;border:1px solid var(--line);border-radius:12px 12px 0 0;padding:24px 26px 20px}.demo-intro h2{color:var(--ink)!important;font-size:25px;margin:0 0 7px}.demo-intro p{color:var(--muted)!important;font-size:16px;margin:0}.example{display:inline-block;margin-top:13px!important;padding:9px 12px;background:#f0f4ed;border-radius:7px;color:var(--ink)!important;font-family:monospace;font-size:15px!important} |
| .input-card{background:white!important;border:1px solid var(--line)!important;border-top:0!important;border-radius:0 0 12px 12px!important;padding:20px 24px 24px!important;margin:0!important;gap:12px!important}.input-card label span{color:var(--ink)!important;font-size:14px!important;letter-spacing:.5px;font-weight:700!important}.input-card textarea{font-family:monospace!important;font-size:17px!important;background:#fbfcfa!important;color:var(--ink)!important;border-color:var(--line)!important} |
| .encode-btn{background:var(--ink)!important;color:white!important;border:0!important;font-size:16px!important;font-weight:700!important;min-width:190px!important} |
| .output-html{background:transparent!important;border:0!important;padding:0!important;margin:16px 0 0!important} |
| .result-shell{border:1px solid var(--line);border-radius:9px;overflow:hidden;background:white}.result-head{display:flex;justify-content:space-between;padding:14px 16px;border-bottom:1px solid var(--line)} |
| .result-head span,.sequence>span{font-size:9px;letter-spacing:1.4px;color:var(--muted)}.result-head small{font:10px monospace;color:var(--muted)} |
| .mol-svg{padding:12px;background:white;text-align:center}.mol-svg svg{max-width:760px;width:100%;height:auto}.sequence{border-top:1px solid var(--line);padding:16px} |
| .sequence>div{display:flex;gap:7px;flex-wrap:wrap;margin-top:10px}.atom-token{display:grid;grid-template-columns:auto auto;gap:2px 7px;align-items:center;border:1px solid var(--line);background:#f7f9f6;padding:7px 9px;border-radius:6px} |
| .atom-token span{font-size:8px;color:#7d8983}.atom-token b{font-size:12px}.atom-token code{grid-column:1/3;color:#356600;font-size:10px} |
| .method{padding:62px 0 72px}.method h2{color:var(--ink)!important;font-size:34px;margin:0 0 26px}.steps{display:flex;flex-direction:column;gap:0}.steps article{display:grid;grid-template-columns:120px 1fr;gap:28px;align-items:center;border-top:1px solid var(--line);padding:28px 4px;background:transparent}.steps article:last-child{border-bottom:1px solid var(--line)}.step-visual{width:110px;height:82px;color:var(--ink)}.step-copy b{color:#5f8700;font-size:13px}.step-copy h3{color:var(--ink)!important;font-size:23px;margin:7px 0}.step-copy p{font-size:17px;line-height:1.65;color:var(--muted)!important;margin:0;max-width:720px}.citation{font-size:16px;color:var(--muted)!important;margin-top:26px;line-height:1.6}.citation a{color:var(--ink)!important;font-weight:700} |
| @media(max-width:700px){.gradio-container{padding:0 16px!important}.hero{padding:38px 0 25px}.hero h1{font-size:43px}.hero p{font-size:17px}.demo-intro{padding:20px}.input-card{display:block!important;padding:18px!important}.encode-btn{margin-top:10px!important;width:100%!important}.steps article{grid-template-columns:82px 1fr;gap:16px}.step-visual{width:76px}.step-copy p{font-size:15px}.result-head{display:block}.result-head small{display:block;margin-top:6px;overflow-wrap:anywhere}} |
| """ |
|
|
| with gr.Blocks(css=CSS, title="VQAtom Demo", fill_width=True) as demo: |
| gr.HTML('''<div class="shell"><header class="topbar"><div class="brand"><i>VQ</i>Atom</div><a class="paper-link" href="https://arxiv.org/abs/2605.16823" target="_blank">Read the paper ↗</a></header> |
| <section class="hero"><h1>Turn molecular environments into discrete tokens.</h1><p>Enter a SMILES string to assign a learned, context-dependent VQAtom ID to each heavy atom.</p></section> |
| <section class="demo-intro"><h2>Try VQAtom</h2><p>Use the default example or replace it with any valid SMILES string.</p><p class="example">Acetanilide · CC(=O)NC1=CC=CC=C1</p></section>''') |
| with gr.Row(elem_classes="input-card"): |
| smiles = gr.Textbox(value="CC(=O)NC1=CC=CC=C1", label="SMILES", lines=1, scale=5) |
| button = gr.Button("Generate IDs", elem_classes="encode-btn", scale=1) |
| output = gr.HTML(elem_classes="output-html") |
| gr.HTML('''<section class="method"><h2>How the IDs are created</h2><div class="steps"> |
| <article><svg class="step-visual" viewBox="0 0 120 84" aria-hidden="true"><g fill="none" stroke="#17231d" stroke-width="2.5"><path d="M18 42L45 18L78 30L101 13M45 18L55 61L92 69M78 30L92 69M18 42L55 61"/><circle cx="18" cy="42" r="7" fill="#bfff3f"/><circle cx="45" cy="18" r="7" fill="#d9e0dc"/><circle cx="78" cy="30" r="7" fill="#d9e0dc"/><circle cx="55" cy="61" r="7" fill="#d9e0dc"/><circle cx="92" cy="69" r="7" fill="#d9e0dc"/><circle cx="101" cy="13" r="7" fill="#d9e0dc"/></g></svg><div class="step-copy"><b>01 · GRAPH CONTEXT</b><h3>Aggregate the neighborhood</h3><p>A graph neural network collects the structural environment surrounding each atom over three hops.</p></div></article> |
| <article><svg class="step-visual" viewBox="0 0 120 84" aria-hidden="true"><g fill="#d9e0dc" stroke="#17231d" stroke-width="2.5"><circle cx="22" cy="22" r="5"/><circle cx="32" cy="59" r="5"/><circle cx="52" cy="36" r="5"/><path d="M58 42L77 42" fill="none"/><path d="M70 35L77 42L70 49" fill="none"/><rect x="86" y="18" width="11" height="11" fill="#bfff3f"/><rect x="99" y="35" width="11" height="11"/><rect x="82" y="54" width="11" height="11"/></g></svg><div class="step-copy"><b>02 · VECTOR QUANTIZATION</b><h3>Find the nearest code</h3><p>The continuous representation is matched to the nearest entry in an element-partitioned codebook.</p></div></article> |
| <article><svg class="step-visual" viewBox="0 0 120 84" aria-hidden="true"><rect x="15" y="18" width="90" height="48" rx="8" fill="#17231d"/><text x="60" y="49" text-anchor="middle" fill="#bfff3f" font-family="monospace" font-size="18" font-weight="700">VQ 1842</text></svg><div class="step-copy"><b>03 · DISCRETE TOKEN</b><h3>Return a global VQAtom ID</h3><p>The selected codebook entry becomes a reusable discrete token for that atom.</p></div></article> |
| </div><p class="citation">Method and evaluation: <a href="https://arxiv.org/abs/2605.16823" target="_blank">VQ-Atom: Semantic Discretization of Local Atomic Environments for Molecular Representation Learning ↗</a></p></section></div>''') |
| button.click(encode_molecule, smiles, output) |
| smiles.submit(encode_molecule, smiles, output) |
|
|
| if __name__ == "__main__": |
| demo.launch(ssr_mode=False) |
|
|
