init commit

.gitignore

@@ -0,0 +1,96 @@
+data/
+checkpoint/
+
+*.err
+*.out
+
+.DS_Store
+.idea
+.vscode
+.vscode/
+*.bin
+*.pkl
+*.pt
+*.pyc
+*.bk
+*.vim
+tokenizer/temp*.json
+
+replace.sh
+
+Quick*.*
+Temp*.*
+
+*.env
+cov.xml
+
+.ipynb_checkpoints/
+
+*.py[cod]
+
+# C extensions
+*.so
+
+# Packages
+*.egg
+*.egg-info
+dist
+build
+eggs
+parts
+bin
+var
+sdist
+develop-eggs
+.installed.cfg
+lib
+lib64
+
+# Installer logs
+pip-log.txt
+
+# Unit test / coverage reports
+.coverage*
+.tox
+nosetests.xml
+htmlcov
+
+# Translations
+*.mo
+
+# Mr Developer
+.mr.developer.cfg
+.project
+.pydevproject
+
+# Complexity
+output/*.html
+output/*/index.html
+
+# Sphinx
+docs/_build
+
+MANIFEST
+
+*.tif
+
+# Rever
+rever/
+
+# Dev notebooks
+# notebooks/
+
+# MkDocs
+site/
+
+.vscode
+
+.idea/
+
+data/
+output/
+wandb/
+oracle/
+expts/models/
+expts/dev-data/
+expts/notebooks/

README.md

@@ -1,13 +1,15 @@
 ---
-title: FragLlama UI
-emoji: 📈
-colorFrom: green
-colorTo: red
+title: FragLlama Demo
+emoji: 🧪
+colorFrom: yellow
+colorTo: purple
 sdk: gradio
-sdk_version: 5.1.0
+sdk_version: 5.0.1
 app_file: app.py
 pinned: false
-license: mit
+license: other
+short_description: Molecule Generation
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+<h1 align="center"> YDS FragLlama </h1>
+An example chatbot using [Gradio](https://gradio.app), [`huggingface_hub`](https://huggingface.co/docs/huggingface_hub/v0.22.2/en/index), and the [Hugging Face Inference API](https://huggingface.co/docs/api-inference/index).

app.py

@@ -1,7 +1,136 @@
+"""
+For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
+"""
+# from huggingface_hub import InferenceClient
+# client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
+import os
+import argparse
+import warnings
 import gradio as gr
+import datamol as dm
+from rdkit import RDLogger
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+warnings.filterwarnings("ignore", message="DEPRECATION WARNING: please use MorganGenerator")
+RDLogger.DisableLog('rdApp.*')
 
-def greet(name):
-    return "Hello " + name + "!!"
+HF_MODEL = "YDS-Pharmatech/FragLlama-base"
+HF_SPACE = os.getenv('HF_SPACE', False)
+# if HF_SPACE: os.chdir("/data/fraglm")
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+from fraglm.inference import FragLMDesign
+from fraglm.utils import *; from boring_utils.utils import *
+from fraglm.trainer.model import FragLMLlamaModel
+from fraglm.inference.post_processing import PostProcessMode
+device = get_device()
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--model_path', type=str, default="", help='Path to the model')
+parser.add_argument('--tokenizer_path', type=str, default="tokenizer/fraglm_2406_bpe_8k.json", help='Path to the model')
+args = parser.parse_args()
+
+if HF_SPACE:
+    model = FragLMLlamaModel.from_pretrained(HF_MODEL, token=os.getenv('HF_TOKEN')).to(device)
+    designer = FragLMDesign(model=model, tokenizer=HF_MODEL)
+else:
+    model = FragLMLlamaModel.from_pretrained(args.model_path).to(device)
+    designer = FragLMDesign(model=model, tokenizer=args.tokenizer_path)
+
+
+def scaffold_decoration(scaffold, n_samples, n_trials):
+    generated_smiles = designer.scaffold_decoration(
+        scaffold=scaffold,
+        n_samples_per_trial=n_samples,
+        n_trials=n_trials,
+        sanitize=True,
+        do_not_fragment_further=True,
+    )
+    generated_mols = [dm.to_mol(x) for x in generated_smiles]
+
+    img = dm.viz.lasso_highlight_image(
+                generated_mols, 
+                dm.from_smarts(scaffold), 
+                mol_size=(350, 200), 
+                color_list=["#ff80b5"], 
+                scale_padding=0.1,
+                use_svg=False,
+            )
+    # img = dm.to_image(
+    #             generated_mols, 
+    #             mol_size=(350, 200),
+    #             use_svg=False,
+    #         )
+    return img
+
+
+def motif_extension(motif, n_samples, n_trials, min_length, max_length):
+    generated_smiles = designer.motif_extension(
+        motif=motif,
+        n_samples_per_trial=n_samples,
+        n_trials=n_trials,
+        sanitize=True,
+        do_not_fragment_further=False,
+        min_length=min_length,
+        max_length=max_length,
+    )
+    img = dm.to_image(
+                generated_smiles, 
+                mol_size=(350, 200),
+                use_svg=False,
+            )
+    return img
+
+
+def linker_generation(linker1, linker2, n_samples, n_trials):
+    linker_generation = [linker1, linker2]
+    generated_smiles = designer.linker_generation(
+        *linker_generation,
+        n_samples_per_trial=n_samples,
+        n_trials=n_trials,
+        sanitize=True,
+        do_not_fragment_further=False,
+        random_seed=100,
+        post_process_mode=PostProcessMode.DO_NOTHING
+    )
+    img = dm.to_image(
+                generated_smiles, 
+                mol_size=(350, 200),
+                use_svg=False,
+            )
+    return img
+
+
+with gr.Blocks() as demo:
+    gr.Markdown("# FragLlama Demo")
+    
+    with gr.Tab("Scaffold Decoration"):
+        scaffold_input = gr.Textbox(label="Scaffold SMILES")
+        scaffold_n_samples = gr.Slider(1, 20, 12, step=1, label="Samples per Trial")
+        scaffold_n_trials = gr.Number(value=2, label="Number of Trials")
+        scaffold_button = gr.Button("Generate")
+        scaffold_output = gr.Image(type="pil", label="Generated Molecules")
+        scaffold_button.click(scaffold_decoration, inputs=[scaffold_input, scaffold_n_samples, scaffold_n_trials], outputs=scaffold_output)
+    
+    with gr.Tab("Motif Extension"):
+        motif_input = gr.Textbox(label="Motif SMILES")
+        motif_n_samples = gr.Slider(1, 20, 12, step=1, label="Samples per Trial")
+        motif_n_trials = gr.Number(value=1, label="Number of Trials")
+        motif_min_length = gr.Number(value=25, label="Min Length")
+        motif_max_length = gr.Number(value=80, label="Max Length")
+        motif_button = gr.Button("Generate")
+        # motif_output = gr.Image(label="Generated Molecules")
+        motif_output = gr.Image(type="pil", label="Generated Molecules")
+        motif_button.click(motif_extension, inputs=[motif_input, motif_n_samples, motif_n_trials, motif_min_length, motif_max_length], outputs=motif_output)
+    
+    with gr.Tab("Linker Generation"):
+        linker1_input = gr.Textbox(label="Linker 1 SMILES")
+        linker2_input = gr.Textbox(label="Linker 2 SMILES")
+        linker_n_samples = gr.Slider(1, 20, 12, step=1, label="Samples per Trial")
+        linker_n_trials = gr.Number(value=1, label="Number of Trials")
+        linker_button = gr.Button("Generate")
+        # linker_output = gr.Image(label="Generated Molecules")
+        linker_output = gr.Image(type="pil", label="Generated Molecules")
+        linker_button.click(linker_generation, inputs=[linker1_input, linker2_input, linker_n_samples, linker_n_trials], outputs=linker_output)
+
+
+if __name__ == "__main__":
+    demo.launch()

packages.txt

@@ -0,0 +1,3 @@
+tmux
+neovim
+ranger

requirements.txt

@@ -0,0 +1,24 @@
+tqdm
+loguru
+typer
+universal_pathlib
+datamol
+pandas
+numpy
+torch>=2.0 --extra-index-url https://download.pytorch.org/whl/cu113
+transformers
+datasets
+tokenizers
+accelerate>=0.28.0
+evaluate
+wandb
+huggingface_hub
+deepspeed
+ruff
+nbconvert
+jupyterlab
+ipywidgets
+s3cmd
+git+https://github.com/TimS-ml/Boring-Utils.git
+-e /data/fraglm
+--no-deps