chembfn-webui 1.2.3__tar.gz → 2.1.3__tar.gz

{chembfn_webui-1.2.3 → chembfn_webui-2.1.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: chembfn_webui
-Version: 1.2.3
+Version: 2.1.3
 Summary: WebUI for ChemBFN
 Home-page: https://github.com/Augus1999/ChemBFN-WebUI
 Author: Nianze A. Tao
@@ -15,15 +15,16 @@ Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
 Classifier: Topic :: Scientific/Engineering :: Chemistry
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: bayesianflow_for_chem>=2.2.2
-Requires-Dist: mol2chemfigPy3>=1.5.11
-Requires-Dist: gradio>=5.32.1
-Requires-Dist: torch>=2.7.0
+Requires-Dist: bayesianflow_for_chem>=2.4.0
+Requires-Dist: mol2chemfigPy3>=1.5.12
+Requires-Dist: gradio<7.0.0,>=6.0.0
+Requires-Dist: torch>=2.9.0
 Requires-Dist: selfies>=2.2.0
 Dynamic: author
 Dynamic: author-email
@@ -39,10 +40,20 @@ Dynamic: requires-dist
 Dynamic: requires-python
 Dynamic: summary
 
-## ChemBFN WebUI: WebUI to generate and visualise molecules for ChemBFN method
+## Web-based UI visualisation tool for ChemBFN method
 
+[![PyPI](https://img.shields.io/pypi/v/chembfn-webui?color=green)](https://pypi.org/project/chembfn-webui/)
+![CI](https://github.com/Augus1999/ChemBFN-WebUI/actions/workflows/pytest.yml/badge.svg)
+![black](https://img.shields.io/badge/code%20style-black-black)
+[![Stand With Ukraine](https://raw.githubusercontent.com/vshymanskyy/StandWithUkraine/main/badges/StandWithUkraine.svg)](https://stand-with-ukraine.pp.ua)
 
 
+
+> Important:
+>
+> For the security concerning, it is not recommended to use this application as a public service.
+> When deploying on a local host as a shared application, it is better to install this application in a container or VM, to prevent this application from accessing the Internet, and to limit the premissions of read, create, and delete loacal files and directories.
+
 ### 1. Install
 
 ```bash
@@ -97,6 +108,7 @@ For example,
 > >    "name": JOB_NAME
 > >}
 > >```
+> >The configureation file for base models can be downloaded [here](https://huggingface.co/suenoomozawa/ChemBFN/resolve/main/config.json).
 
 If placed correctly, all these files can be seen in the "model explorer" tab.
 
@@ -146,6 +158,7 @@ Under "advanced control" tab
 
 * You can control semi-autoregressive behaviours by key in `F` for switching off SAR, `T` for switching on SAR, and prompt like `F,F,T,...` to individually control the SAR in an ensemble model.
 * You can add unwanted tokens, e.g., `[Cu],p,[Si]`.
+* You can customise the result preprocessing function, e.g., the model output  a reaction SMILES "CCI.C[O-]>>COCC" which couldn't be recognised by RDKit; you can pass `lambda x: x.split(">>")[-1]` to force the program only looking at the products.
 
 ### 6. Generate molecules
 

{chembfn_webui-1.2.3 → chembfn_webui-2.1.3}/README.md

@@ -1,13 +1,23 @@
-## ChemBFN WebUI: WebUI to generate and visualise molecules for ChemBFN method
+## Web-based UI visualisation tool for ChemBFN method
+
+[![PyPI](https://img.shields.io/pypi/v/chembfn-webui?color=green)](https://pypi.org/project/chembfn-webui/)
+![CI](https://github.com/Augus1999/ChemBFN-WebUI/actions/workflows/pytest.yml/badge.svg)
+![black](https://img.shields.io/badge/code%20style-black-black)
+[![Stand With Ukraine](https://raw.githubusercontent.com/vshymanskyy/StandWithUkraine/main/badges/StandWithUkraine.svg)](https://stand-with-ukraine.pp.ua)
 
 <p align="left">
-<img src="image/screenshot_0.jpeg" alt="screenshot 0" width="400" height="auto">
-<img src="image/screenshot_1.jpeg" alt="screenshot 1" width="400" height="auto">
-<img src="image/screenshot_2.jpeg" alt="screenshot 2" width="400" height="auto">
-<img src="image/screenshot_3.jpeg" alt="screenshot 3" width="400" height="auto">
-<img src="image/screenshot_4.jpeg" alt="screenshot 4" width="400" height="auto">
+<img src="image/screenshot_0.jpeg" alt="screenshot 0" width="360" height="auto">
+<img src="image/screenshot_1.jpeg" alt="screenshot 1" width="360" height="auto">
+<img src="image/screenshot_2.jpeg" alt="screenshot 2" width="360" height="auto">
+<img src="image/screenshot_3.jpeg" alt="screenshot 3" width="360" height="auto">
+<img src="image/screenshot_4.jpeg" alt="screenshot 4" width="360" height="auto">
 </p>
 
+> [!IMPORTANT]
+>
+> For the security concerning, it is not recommended to use this application as a public service.
+> When deploying on a local host as a shared application, it is better to install this application in a container or VM, to prevent this application from accessing the Internet, and to limit the premissions of read, create, and delete loacal files and directories.
+
 ### 1. Install
 
 ```bash
@@ -62,6 +72,7 @@ For example,
 > >    "name": JOB_NAME
 > >}
 > >```
+> >The configureation file for base models can be downloaded [here](https://huggingface.co/suenoomozawa/ChemBFN/resolve/main/config.json).
 
 If placed correctly, all these files can be seen in the "model explorer" tab.
 
@@ -111,6 +122,7 @@ Under "advanced control" tab
 
 * You can control semi-autoregressive behaviours by key in `F` for switching off SAR, `T` for switching on SAR, and prompt like `F,F,T,...` to individually control the SAR in an ensemble model.
 * You can add unwanted tokens, e.g., `[Cu],p,[Si]`.
+* You can customise the result preprocessing function, e.g., the model output  a reaction SMILES "CCI.C[O-]>>COCC" which couldn't be recognised by RDKit; you can pass `lambda x: x.split(">>")[-1]` to force the program only looking at the products.
 
 ### 6. Generate molecules
 

{chembfn_webui-1.2.3 → chembfn_webui-2.1.3}/chembfn_webui/bin/app.py

@@ -6,8 +6,9 @@ Define application behaviours.
 import sys
 import argparse
 from pathlib import Path
+from copy import deepcopy
 from functools import partial
-from typing import Tuple, List, Dict, Union, Literal
+from typing import Tuple, List, Dict, Optional, Union, Literal
 
 sys.path.append(str(Path(__file__).parent.parent))
 from rdkit.Chem import Draw, MolFromSmiles  # type: ignore
@@ -38,6 +39,7 @@ from lib.utilities import (
     parse_prompt,
     parse_exclude_token,
     parse_sar_control,
+    build_result_prep_fn,
 )
 from lib.version import __version__
 
@@ -45,7 +47,7 @@ vocabs = find_vocab()
 models = find_model()
 cache_dir = Path(__file__).parent.parent / "cache"
 favicon_dir = Path(__file__).parent / "favicon.png"
-_result_count = 0
+_RESULT_COUNT = 0
 
 HTML_STYLE = gr.InputHTMLAttributes(
     autocapitalize="off",
@@ -76,7 +78,7 @@ def selfies2vec(sel: str, vocab_dict: Dict[str, int]) -> List[int]:
     return [vocab_dict.get(i, unknown_id) for i in s]
 
 
-def refresh(
+def _refresh(
     model_selected: str, vocab_selected: str, tokeniser_selected: str
 ) -> Tuple[
     List[str], List[str], List[List[str]], List[List[str]], gr.Dropdown, gr.Dropdown
@@ -121,7 +123,7 @@ def refresh(
     return a, b, c, d, e, f
 
 
-def select_lora(evt: gr.SelectData, prompt: str) -> str:
+def _select_lora(evt: gr.SelectData, prompt: str) -> str:
     """
     Select LoRA model name from Dataframe object.
 
@@ -141,7 +143,7 @@ def select_lora(evt: gr.SelectData, prompt: str) -> str:
     return f"{prompt};\n<{selected_lora}:1>"
 
 
-def token_name_change_evt(
+def _token_name_change_evt(
     token_name: str, vocab_fn: str
 ) -> Tuple[gr.Dropdown, gr.Tab, gr.Tab]:
     """
@@ -178,14 +180,15 @@ def run(
     guidance_strength: float,
     method: Literal["BFN", "ODE"],
     temperature: float,
-    prompt: str,
-    scaffold: str,
-    template: str,
-    sar_control: str,
-    exclude_token: str,
+    prompt: Optional[str],
+    scaffold: Optional[str],
+    template: Optional[str],
+    sar_control: Optional[str],
+    exclude_token: Optional[str],
     quantise: Literal["on", "off"],
     jited: Literal["on", "off"],
     sorted_: Literal["on", "off"],
+    result_prep_fn: Optional[str],
 ) -> Tuple[Union[List, None], List[str], str, gr.TextArea, str]:
     """
     Run generation or inpainting.
@@ -207,6 +210,7 @@ def run(
     :param quantise: `"on"` or `"off"`
     :param jited: `"on"` or `"off"`
     :param sorted\\_: whether to sort the reulst; `"on"` or `"off"`
+    :param result_prep_fn: a string form result preprocessing function
     :type model_name: str
     :type token_name: str
     :type vocab_fn: str
@@ -216,14 +220,15 @@ def run(
     :type guidance_strength: float
     :type method: str
     :type temperature: float
-    :type prompt: str
-    :type scaffold: str
-    :type template: str
-    :type sar_control: str
-    :type exclude_token: str
+    :type prompt: str | None
+    :type scaffold: str | None
+    :type template: str | None
+    :type sar_control: str | None
+    :type exclude_token: str | None
     :type quantise: str
     :type jited: str
     :type sorted\\_: str
+    :type result_prep_fn: str | None
     :return: list of images \n
              list of generated molecules \n
              Chemfig code \n
@@ -239,6 +244,8 @@ def run(
     lora_label_dict = dict([[i[0], i[2] != []] for i in models["lora"]])
     standalone_lmax_dict = dict([[i[0], i[3]] for i in models["standalone"]])
     lora_lmax_dict = dict([[i[0], i[3]] for i in models["lora"]])
+    # ------- build result preprocessing function -------
+    _result_prep_fn = build_result_prep_fn(result_prep_fn)
     # ------- build tokeniser -------
     if token_name == "SMILES & SAFE":
         vocab_keys = VOCAB_KEYS
@@ -246,13 +253,13 @@ def run(
         trans_fn = lambda x: [i for i in x if (MolFromSmiles(i) and i)]
         img_fn = lambda x: [Draw.MolToImage(MolFromSmiles(i), (500, 500)) for i in x]
         chemfig_fn = lambda x: [mol2chemfig(i, "-r", inline=True) for i in x]
-    if token_name == "FASTA":
+    elif token_name == "FASTA":
         vocab_keys = FASTA_VOCAB_KEYS
         tokeniser = fasta2vec
         trans_fn = lambda x: [i for i in x if i]
         img_fn = lambda _: None  # senseless to provide dumb 2D images
         chemfig_fn = lambda _: [""]  # senseless to provide very long Chemfig code
-    if token_name == "SELFIES":
+    elif token_name == "SELFIES":
         vocab_data = load_vocab(vocabs[vocab_fn])
         vocab_keys = vocab_data["vocab_keys"]
         vocab_dict = vocab_data["vocab_dict"]
@@ -262,11 +269,15 @@ def run(
             Draw.MolToImage(MolFromSmiles(decoder(i)), (500, 500)) for i in x
         ]
         chemfig_fn = lambda x: [mol2chemfig(decoder(i), "-r", inline=True) for i in x]
+    else:
+        raise RuntimeError("Oops, maybe something wrong with Gradio.")
     _method = "bfn" if method == "BFN" else f"ode:{temperature}"
     # ------- build model -------
     prompt_info = parse_prompt(prompt)
     sar_flag = parse_sar_control(sar_control)
-    print("Prompt summary:", prompt_info, "semi-autoregression:", sar_flag)  # prompt
+    _info = deepcopy(prompt_info)
+    _info["semi-autoregression"] = deepcopy(sar_flag)
+    print("Prompt summary:", _info)  # prompt
     if not prompt_info["lora"]:
         if model_name in base_model_dict:
             lmax = sequence_size
@@ -287,7 +298,7 @@ def run(
                     mlp = MLP.from_checkpoint(
                         standalone_model_dict[model_name] / "mlp.pt"
                     )
-                    y = torch.tensor([prompt_info["objective"]], dtype=torch.float32)
+                    y = torch.tensor([prompt_info["objective"][0]], dtype=torch.float32)
                     y = mlp.forward(y)
             else:
                 y = None
@@ -317,7 +328,7 @@ def run(
                 mlp = MLP.from_checkpoint(
                     lora_model_dict[prompt_info["lora"][0]] / "mlp.pt"
                 )
-                y = torch.tensor([prompt_info["objective"]], dtype=torch.float32)
+                y = torch.tensor([prompt_info["objective"][0]], dtype=torch.float32)
                 y = mlp.forward(y)
         else:
             y = None
@@ -330,7 +341,7 @@ def run(
         if jited == "on":
             bfn.compile()
     else:
-        lmax = max([lora_lmax_dict[i] for i in prompt_info["lora"]])
+        lmax = max(lora_lmax_dict[i] for i in prompt_info["lora"])
         if model_name in base_model_dict:
             base_model_dir = base_model_dict[model_name]
         else:
@@ -355,10 +366,15 @@ def run(
         if jited == "on":
             bfn.compile()
         _message.append(f"Sequence length set to {lmax} from model metadata.")
+    result_prep_fn_ = lambda x: [_result_prep_fn(i) for i in x]
     # ------- inference -------
     allowed_tokens = parse_exclude_token(exclude_token, vocab_keys)
     if not allowed_tokens:
         allowed_tokens = "all"
+    if scaffold is None:
+        scaffold = ""
+    if template is None:
+        template = ""
     scaffold = scaffold.strip()
     template = template.strip()
     if scaffold:
@@ -376,7 +392,7 @@ def run(
             allowed_tokens=allowed_tokens,
             sort=sorted_ == "on",
         )
-        mols = trans_fn(mols)
+        mols = trans_fn(result_prep_fn_(mols))
         imgs = img_fn(mols)
         chemfigs = chemfig_fn(mols)
         if template:
@@ -396,7 +412,7 @@ def run(
             allowed_tokens=allowed_tokens,
             sort=sorted_ == "on",
         )
-        mols = trans_fn(mols)
+        mols = trans_fn(result_prep_fn_(mols))
         imgs = img_fn(mols)
         chemfigs = chemfig_fn(mols)
     else:
@@ -412,17 +428,17 @@ def run(
             allowed_tokens=allowed_tokens,
             sort=sorted_ == "on",
         )
-        mols = trans_fn(mols)
+        mols = trans_fn(result_prep_fn_(mols))
         imgs = img_fn(mols)
         chemfigs = chemfig_fn(mols)
-    n_mol = len(mols)
     with open(cache_dir / "results.csv", "w", encoding="utf-8", newline="") as rf:
         rf.write("\n".join(mols))
     _message.append(
-        f"{n_mol} smaples generated and saved to cache that can be downloaded."
+        f"{(n_mol := len(mols))} {'smaple' if n_mol in (0, 1) else 'samples'} "
+        "generated and saved to cache that can be downloaded."
     )
-    global _result_count
-    _result_count = n_mol
+    global _RESULT_COUNT
+    _RESULT_COUNT = n_mol
     return (
         imgs,
         mols,
@@ -432,14 +448,11 @@ def run(
     )
 
 
-with gr.Blocks(
-    title="ChemBFN WebUI",
-    css="footer {display: none !important} .custom_footer {text-align:center;bottom:0;}",
-    analytics_enabled=False,
-) as app:
+with gr.Blocks(title="ChemBFN WebUI", analytics_enabled=False) as app:
     with gr.Row():
         with gr.Column(scale=1):
             btn = gr.Button("RUN", variant="primary")
+            stop = gr.Button("\u23f9", variant="stop", visible=False)
             model_name = gr.Dropdown(
                 [i[0] for i in models["base"]] + [i[0] for i in models["standalone"]],
                 label="model",
@@ -484,14 +497,15 @@ with gr.Blocks(
                 message = gr.TextArea(label="message", lines=2)
             with gr.Tab(label="result viewer"):
                 with gr.Tab(label="result"):
-                    btn_download = gr.File(label="download", visible=False)
+                    btn_download = gr.File(
+                        str(cache_dir / "results.csv"), label="download", visible=False
+                    )
                     result = gr.Dataframe(
                         headers=["molecule"],
-                        col_count=(1, "fixed"),
+                        column_count=(1, "fixed"),
                         label="",
-                        show_fullscreen_button=True,
+                        interactive=False,
                         show_row_numbers=True,
-                        show_copy_button=True,
                     )
                 with gr.Tab(
                     label="LATEX Chemfig", visible=token_name.value != "FASTA"
@@ -509,7 +523,7 @@ with gr.Blocks(
                     vocab_table = gr.Dataframe(
                         list(vocabs.keys()),
                         headers=["name"],
-                        col_count=(1, "fixed"),
+                        column_count=(1, "fixed"),
                         label="",
                         interactive=False,
                         show_row_numbers=True,
@@ -518,7 +532,7 @@ with gr.Blocks(
                     base_table = gr.Dataframe(
                         [i[0] for i in models["base"]],
                         headers=["name"],
-                        col_count=(1, "fixed"),
+                        column_count=(1, "fixed"),
                         label="",
                         interactive=False,
                         show_row_numbers=True,
@@ -527,7 +541,7 @@ with gr.Blocks(
                     standalone_table = gr.Dataframe(
                         [[i[0], i[2]] for i in models["standalone"]],
                         headers=["name", "objective"],
-                        col_count=(2, "fixed"),
+                        column_count=(2, "fixed"),
                         label="",
                         interactive=False,
                         show_row_numbers=True,
@@ -536,7 +550,7 @@ with gr.Blocks(
                     lora_tabel = gr.Dataframe(
                         [[i[0], i[2]] for i in models["lora"]],
                         headers=["name", "objective"],
-                        col_count=(2, "fixed"),
+                        column_count=(2, "fixed"),
                         label="",
                         interactive=False,
                         show_row_numbers=True,
@@ -553,17 +567,33 @@ with gr.Blocks(
                     placeholder="key in unwanted tokens separated by comma.",
                     html_attributes=HTML_STYLE,
                 )
-                quantise = gr.Radio(["on", "off"], value="off", label="quantisation")
-                jited = gr.Radio(["on", "off"], value="off", label="JIT")
-                sorted_ = gr.Radio(
-                    ["on", "off"],
-                    value="off",
-                    label="sort result",
-                    info="sorting based on entropy",
+                result_prep_fn = gr.Textbox(
+                    "lambda x: x",
+                    label="result preprocessing function",
+                    placeholder="lambda x: x",
+                    html_attributes=HTML_STYLE,
                 )
+                with gr.Row(scale=1):
+                    quantise = gr.Radio(
+                        ["on", "off"], value="off", label="quantisation"
+                    )
+                    jited = gr.Radio(["on", "off"], value="off", label="JIT")
+                    sorted_ = gr.Radio(
+                        ["on", "off"], value="off", label="sort result based on entropy"
+                    )
     gr.HTML(sys_info(), elem_classes="custom_footer", elem_id="footer")
     # ------ user interaction events -------
-    btn.click(
+    gen = btn.click(
+        fn=lambda: (
+            gr.Button("RUN", variant="primary", visible=False),
+            gr.Button("\u23f9", variant="stop", visible=True),
+        ),
+        inputs=None,
+        outputs=[btn, stop],
+        api_name="switch_to_stop_mode",
+        api_description="Switch to STOP.",
+        api_visibility="private",
+    ).then(
         fn=run,
         inputs=[
             model_name,
@@ -583,11 +613,36 @@ with gr.Blocks(
             quantise,
             jited,
             sorted_,
+            result_prep_fn,
         ],
         outputs=[img, result, chemfig, message, btn_download],
+        api_name="run",
+        api_description="Run ChemBFN model.",
+    )
+    gen.then(
+        fn=lambda: (
+            gr.Button("RUN", variant="primary", visible=True),
+            gr.Button("\u23f9", variant="stop", visible=False),
+        ),
+        inputs=None,
+        outputs=[btn, stop],
+        api_name="switch_back_to_run_mode",
+        api_description="Swtch back to RUN.",
+        api_visibility="private",
+    )
+    stop.click(
+        fn=lambda: (
+            gr.Button("RUN", variant="primary", visible=True),
+            gr.Button("\u23f9", variant="stop", visible=False),
+        ),
+        inputs=None,
+        outputs=[btn, stop],
+        cancels=[gen],
+        api_name="stop",
+        api_description="Stop the model.",
     )
     btn_refresh.click(
-        fn=refresh,
+        fn=_refresh,
         inputs=[model_name, vocab_fn, token_name],
         outputs=[
             vocab_table,
@@ -597,11 +652,14 @@ with gr.Blocks(
             model_name,
             vocab_fn,
         ],
+        api_name="refresh_model_list",
+        api_description="Refresh the model list.",
     )
     token_name.input(
-        fn=token_name_change_evt,
+        fn=_token_name_change_evt,
         inputs=[token_name, vocab_fn],
         outputs=[vocab_fn, code, gallery],
+        api_visibility="private",
     )
     method.input(
         fn=lambda x, y: gr.Slider(
@@ -614,12 +672,25 @@ with gr.Blocks(
         ),
         inputs=[method, temperature],
         outputs=temperature,
+        api_name="select_sampling_method",
+        api_description="Select sampling method between 'BFN' and 'ODE'.",
+        api_visibility="private",
+    )
+    lora_tabel.select(
+        fn=_select_lora,
+        inputs=prompt,
+        outputs=prompt,
+        api_name="select_lora",
+        api_description="Select LoRA model from the model list.",
+        api_visibility="private",
     )
-    lora_tabel.select(fn=select_lora, inputs=prompt, outputs=prompt)
     result.change(
-        fn=lambda x: gr.File(x, label="download", visible=_result_count > 0),
+        fn=lambda x: gr.File(x, label="download", visible=_RESULT_COUNT > 0),
         inputs=btn_download,
         outputs=btn_download,
+        api_name="change_download_state",
+        api_description="Hide or show the file downloading item.",
+        api_visibility="private",
     )
 
 
@@ -630,9 +701,13 @@ def main() -> None:
     :return:
     :rtype: None
     """
+    from rdkit import RDLogger
+
+    RDLogger.DisableLog("rdApp.*")  # type: ignore
     parser = argparse.ArgumentParser(
         description="A web-based visualisation tool for ChemBFN method.",
-        epilog=f"ChemBFN WebUI {__version__}, developed in Hiroshima University by chemists for chemists. "
+        epilog=f"ChemBFN WebUI {__version__}, "
+        "developed in Hiroshima University by chemists for chemists. "
         "Visit https://augus1999.github.io/bayesian-flow-network-for-chemistry/ for more details.",
         formatter_class=argparse.ArgumentDefaultsHelpFormatter,
     )
@@ -644,8 +719,10 @@ def main() -> None:
     print(f"This is ChemBFN WebUI version {__version__}")
     app.launch(
         share=args.public,
+        footer_links=["api"],
         allowed_paths=[cache_dir.absolute().__str__()],
         favicon_path=favicon_dir.absolute().__str__(),
+        css=".custom_footer {text-align:center;bottom:0;}",
     )
 
 

chembfn_webui-2.1.3/chembfn_webui/cache/results.csv

@@ -0,0 +1 @@
+c1ccccc1