modified TOKENIZERS_PARALLELISM to false

app.py

@@ -18,6 +18,8 @@ from data_provider.data_utils import smiles2data, reformat_smiles
 import gradio as gr
 from datetime import datetime
 
+## disable online tokenizers parallelism to avoid deadlocks
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
 ## for pyg bug
 warnings.filterwarnings('ignore', category=UserWarning, message='TypedStorage is deprecated')
 ## for A5000 gpus
@@ -131,7 +133,7 @@ class InferenceRunner:
             solvent = smiles_split(solvent) if solvent else []
             assert reactant and product
         except:
-            raise KeyError('Please input a valid reaction string')
+            raise gr.Error('Please input a valid reaction string')
 
         extracted_molecules = {product[0]: "$-1$"}
         for mol in reactant+solvent:
@@ -304,7 +306,7 @@ def main(args):
         btn.click(fn=online_chat, inputs=[reaction_string, temperature], outputs=[out])
         clear_btn.click(fn=lambda:("", ""), inputs=[], outputs=[reaction_string, out])
 
-    demo.launch(share=True)
+    demo.launch()
 
 
 

demo.py

@@ -1,30 +1,114 @@
+import subprocess
+subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
+subprocess.run('pip install -U timm', shell=True)
+import spaces
 import os
 import torch
 import argparse
 import warnings
-import pytorch_lightning as pl
-from pytorch_lightning import Trainer, strategies
-import pytorch_lightning.callbacks as plc
-from pytorch_lightning.loggers import CSVLogger
-from pytorch_lightning.callbacks import TQDMProgressBar
+from rdkit import Chem
+from rdkit.Chem import CanonSmiles
+from rdkit.Chem import MolFromSmiles, MolToSmiles
 from data_provider.pretrain_dm import PretrainDM
 from data_provider.tune_dm import *
 from model.opt_flash_attention import replace_opt_attn_with_flash_attn
 from model.blip2_model import Blip2Model
-from model.dist_funs import MyDeepSpeedStrategy
-from data_provider.reaction_action_dataset import ActionDataset
 from data_provider.data_utils import json_read, json_write
 from data_provider.data_utils import smiles2data, reformat_smiles
+import gradio as gr
+from datetime import datetime
 
+## disable online tokenizers parallelism to avoid deadlocks
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
 ## for pyg bug
 warnings.filterwarnings('ignore', category=UserWarning, message='TypedStorage is deprecated')
 ## for A5000 gpus
 torch.set_float32_matmul_precision('medium') # can be medium (bfloat16), high (tensorfloat32), highest (float32)
 
+def smiles_split(string, separator='.'):
+    string = str(string)
+    mols = []
+    for smi in string.split(separator):
+        mol = MolFromSmiles(smi)
+        if mol is None:
+            continue  # Skip invalid SMILES strings
+        mols.append(mol)
+
+    parts = []
+    current_part = []
+    charge_count = 0
+
+    for mol in mols:
+        charge = Chem.GetFormalCharge(mol)
+        if charge==0:
+            if current_part:
+                smiles = '.'.join([MolToSmiles(m) for m in current_part])
+                smiles = CanonSmiles(smiles)
+                parts.append(smiles)
+                current_part = []
+                charge_count = 0
+            parts.append(MolToSmiles(mol))
+        else:
+            charge_count += charge
+            current_part.append(mol)
+            if charge_count == 0:
+                smiles = '.'.join([MolToSmiles(m) for m in current_part])
+                smiles = CanonSmiles(smiles)
+                parts.append(smiles)
+                current_part = []
+                charge_count = 0
+    if current_part:
+        smiles = '.'.join([MolToSmiles(m) for m in current_part])
+        smiles = CanonSmiles(smiles)
+        parts.append(smiles)
+
+    return parts
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--filename', type=str, default="main")
+    parser.add_argument('--seed', type=int, default=42, help='random seed')
+    # MM settings
+    parser.add_argument('--mode', type=str, default='pretrain', choices=['pretrain', 'ft', 'eval', 'pretrain_eval'])
+    parser.add_argument('--strategy_name', type=str, default='mydeepspeed')
+    parser.add_argument('--iupac_prediction', action='store_true', default=False)
+    parser.add_argument('--ckpt_path', type=str, default=None)
+    # parser = Trainer.add_argparse_args(parser)
+    parser = Blip2Model.add_model_specific_args(parser)  # add model args
+    parser = PretrainDM.add_model_specific_args(parser)
+    parser.add_argument('--accelerator', type=str, default='gpu')
+    parser.add_argument('--devices', type=str, default='0,1,2,3')
+    parser.add_argument('--precision', type=str, default='bf16-mixed')
+    parser.add_argument('--downstream_task', type=str, default='action', choices=['action', 'synthesis', 'caption', 'chebi'])
+    parser.add_argument('--max_epochs', type=int, default=10)
+    parser.add_argument('--enable_flash', action='store_true', default=False)
+    parser.add_argument('--disable_graph_cache', action='store_true', default=False)
+    parser.add_argument('--generate_restrict_tokens', action='store_true', default=False)
+    parser.add_argument('--train_restrict_tokens', action='store_true', default=False)
+    parser.add_argument('--smiles_type', type=str, default='default', choices=['default', 'canonical', 'restricted', 'unrestricted', 'r_smiles'])
+    parser.add_argument('--accumulate_grad_batches', type=int, default=1)
+    parser.add_argument('--tqdm_interval', type=int, default=50)
+    parser.add_argument('--check_val_every_n_epoch', type=int, default=1)
+    args = parser.parse_args()
+
+    if args.enable_flash:
+        replace_opt_attn_with_flash_attn()
+    return args
+
+app_config = {
+    "init_checkpoint": "all_checkpoints/ckpt_tune_hybridFeb11_May31/last_converted.ckpt",
+    "filename": "app",
+    "opt_model": "facebook/galactica-1.3b",
+    "num_workers": 4,
+    "rxn_max_len": 512,
+    "text_max_len": 512,
+    "precision": "bf16-mixed",
+    "max_inference_len": 512,
+}
 
 class InferenceRunner:
     def __init__(self, model, tokenizer, rxn_max_len, smi_max_len,
-                 smiles_type='default', device='cuda', predict_rxn_condition=True, args=None):
+                 smiles_type='default', device='cuda', args=None):
         self.model = model
         self.rxn_max_len = rxn_max_len
         self.smi_max_len = smi_max_len
@@ -36,11 +120,42 @@ class InferenceRunner:
         self.collater = Collater([], [])
         self.device = device
         self.smiles_type = smiles_type
-        self.predict_rxn_condition = predict_rxn_condition
         self.args = args
+        time_stamp = datetime.now().strftime("%Y.%m.%d-%H:%M")
+        self.cache_dir = f'results/{self.args.filename}/{time_stamp}'
+        os.makedirs(self.cache_dir, exist_ok=True)
+    
+    def make_query_dict(self, rxn_string):
+        try:
+            reactant, solvent, product = rxn_string.split('>')
+            reactant = smiles_split(reactant)
+            product = smiles_split(product)
+            solvent = smiles_split(solvent) if solvent else []
+            assert reactant and product
+        except:
+            raise gr.Error('Please input a valid reaction string')
+
+        extracted_molecules = {product[0]: "$-1$"}
+        for mol in reactant+solvent:
+            extracted_molecules[mol] = f"${len(extracted_molecules)}$"
 
-    def make_prompt(self, param_dict, smi_max_len=128, predict_rxn_condition=False):
-        action_sequence = param_dict['actions']
+        result_dict = {}
+        result_dict['time_stamp'] = datetime.now().strftime("%Y.%m.%d %H:%M:%S.%f")[:-3]
+        result_dict['reaction_string'] = rxn_string
+        result_dict['REACTANT'] = reactant
+        result_dict['SOLVENT'] = solvent
+        result_dict['CATALYST'] = []
+        result_dict['PRODUCT'] = product
+        result_dict['extracted_molecules'] = extracted_molecules
+        return result_dict
+    
+    def save_prediction(self, result_dict):
+        os.makedirs(self.cache_dir, exist_ok=True)
+        result_id = result_dict['time_stamp']
+        result_path = os.path.join(self.cache_dir, f'{result_id}.json')
+        json_write(result_path, result_dict)
+
+    def make_prompt(self, param_dict, smi_max_len=128):
         smiles_list = []
         prompt = ''
         prompt += 'Reactants: '
@@ -72,58 +187,44 @@ class InferenceRunner:
                     prompt += f'[START_SMILES]{smiles_wrapper(smi)}[END_SMILES] '
                 smiles_list.append(smi)
 
-        if predict_rxn_condition:
-            for value, token in param_dict['extracted_duration'].items():
-                action_sequence = action_sequence.replace(token, value)
-            for value, token in param_dict['extracted_temperature'].items():
-                action_sequence = action_sequence.replace(token, value)
-        else:
-            prompt += 'Temperatures: '
-            for value, token in param_dict['extracted_temperature'].items():
-                prompt += f'{token}: {value} '
-
-            prompt += 'Durations: '
-            for value, token in param_dict['extracted_duration'].items():
-                prompt += f'{token}: {value} '
-
         prompt += 'Action Squence: '
-        return prompt, smiles_list, action_sequence
+        return prompt, smiles_list
 
     def get_action_elements(self, rxn_dict):
-        rxn_id = rxn_dict['index']
-        input_text, smiles_list, output_text = self.make_prompt(rxn_dict, self.smi_max_len, self.predict_rxn_condition)
-        output_text = output_text.strip() + '\n'
+        input_text, smiles_list = self.make_prompt(rxn_dict, self.smi_max_len)
 
         graph_list = []
         for smiles in smiles_list:
             graph_item = smiles2data(smiles)
             graph_list.append(graph_item)
-        return rxn_id, graph_list, output_text, input_text
-    
+        return graph_list, input_text
+
+    @spaces.GPU
     @torch.no_grad()
-    def predict(self, rxn_dict):
-        rxn_id, graphs, prompt_tokens, output_text, input_text = self.tokenize(rxn_dict)
-        result_dict = {
-            'raw': rxn_dict,
-            'index': rxn_id,
-            'input': input_text,
-            'target': output_text
-        }
+    def predict(self, rxn_dict, temperature=1):
+        graphs, prompt_tokens = self.tokenize(rxn_dict)
+        result_dict = rxn_dict
         samples = {'graphs': graphs, 'prompt_tokens': prompt_tokens}
-        with torch.no_grad():
-            result_dict['prediction'] = self.model.blip2opt.generate(
-                samples,
-                do_sample=self.args.do_sample,
-                num_beams=self.args.num_beams,
-                max_length=self.args.max_inference_len,
-                min_length=self.args.min_inference_len,
-                num_captions=self.args.num_generate_captions,
-                use_graph=True
-            )
+        assert prompt_tokens.input_ids.is_cuda
+        assert graphs.is_cuda
+        prediction = self.model.blip2opt.generate(
+            samples,
+            do_sample=self.args.do_sample,
+            num_beams=self.args.num_beams,
+            max_length=self.args.max_inference_len,
+            min_length=self.args.min_inference_len,
+            num_captions=self.args.num_generate_captions,
+            temperature=temperature,
+            use_graph=True
+        )[0]
+        for k, v in result_dict['extracted_molecules'].items():
+            prediction = prediction.replace(v, k)
+        result_dict['prediction'] = prediction
         return result_dict
     
+    @spaces.GPU
     def tokenize(self, rxn_dict):
-        rxn_id, graph_list, output_text, input_text = self.get_action_elements(rxn_dict)
+        graph_list, input_text = self.get_action_elements(rxn_dict)
         if graph_list:
             graphs = self.collater(graph_list).to(self.device)
         input_prompt = smiles_handler(input_text, self.mol_ph, self.is_gal)[0]
@@ -139,13 +240,10 @@ class InferenceRunner:
                                               return_attention_mask=True).to(self.device)
         is_mol_token = input_prompt_tokens.input_ids == self.mol_token_id
         input_prompt_tokens['is_mol_token'] = is_mol_token
-        return rxn_id, graphs, input_prompt_tokens, output_text, input_text
-
+        return graphs, input_prompt_tokens
 
 def main(args):
     device = torch.device('cuda')
-    data_list = json_read('demo.json')
-    pl.seed_everything(args.seed)
     # model
     if args.init_checkpoint:
         model = Blip2Model(args).to(device)
@@ -171,54 +269,51 @@ def main(args):
         rxn_max_len=args.rxn_max_len,
         smi_max_len=args.smi_max_len,
         device=device,
-        predict_rxn_condition=args.predict_rxn_condition,
         args=args
     )
+    example_inputs = json_read('demo.json')
+    example_inputs = [[e] for e in example_inputs]
 
-    import time
-    for data_item in data_list:
-        t1 = time.time()
-        result = infer_runner.predict(data_item)
-        print(result)
-        print(f"Time: {time.time() - t1:.2f}s")
+    def online_chat(reaction_string, temperature=1):
+        data_item = infer_runner.make_query_dict(reaction_string)
+        result = infer_runner.predict(data_item, temperature=temperature)
+        infer_runner.save_prediction(result)
+        prediction = result['prediction'].replace(' ; ', ' ;\n')
+        return prediction
+    
+    with gr.Blocks(css="""
+            .center { display: flex; justify-content: center; }
+        """) as demo:
+        gr.HTML(
+        """
+        <center><h1><b>ReactXT</b></h1></center>
+        <p style="font-size:20px; font-weight:bold;">This is the demo page of our ACL 2024 paper 
+        <i>ReactXT: Understanding Molecular “Reaction-ship” via Reaction-Contextualized Molecule-Text Pretraining.</i></p>
+        """)
+        with gr.Row(elem_classes="center"):
+            gr.Image(value="./figures/frameworks.jpg", elem_classes="center", width=800, label="Framework of ReactXT")
+        gr.HTML(
+        """
+        <p style="font-size:16px;"> Please input one chemical reaction below, and we will generate the predicted experimental procedure.</p>
+        <p style="font-size:16px;"> The reaction should be in form of <b>Reactants>Reagents>Product</b>.</p>
+        """)
 
+        reaction_string = gr.Textbox(placeholder="Input one reaction", label='Input Reaction')
+        gr.Examples(example_inputs, [reaction_string,], fn=online_chat, label='Example Reactions')
+        with gr.Row():
+            btn = gr.Button("Submit")
+            clear_btn = gr.Button("Clear")
+        temperature = gr.Slider(0.1, 1, value=1, label='Temperature')
+        with gr.Row():
+            out = gr.Textbox(label="ReactXT's Output", placeholder="Predicted experimental procedure")
+        btn.click(fn=online_chat, inputs=[reaction_string, temperature], outputs=[out])
+        clear_btn.click(fn=lambda:("", ""), inputs=[], outputs=[reaction_string, out])
 
-def get_args():
-    parser = argparse.ArgumentParser()
-    parser.add_argument('--filename', type=str, default="main")
-    parser.add_argument('--seed', type=int, default=42, help='random seed')
-    # MM settings
-    parser.add_argument('--mode', type=str, default='pretrain', choices=['pretrain', 'ft', 'eval', 'pretrain_eval'])
-    parser.add_argument('--strategy_name', type=str, default='mydeepspeed')
-    parser.add_argument('--iupac_prediction', action='store_true', default=False)
-    parser.add_argument('--ckpt_path', type=str, default=None)
-    # parser = Trainer.add_argparse_args(parser)
-    parser = Blip2Model.add_model_specific_args(parser)  # add model args
-    parser = PretrainDM.add_model_specific_args(parser)
-    parser.add_argument('--accelerator', type=str, default='gpu')
-    parser.add_argument('--devices', type=str, default='0,1,2,3')
-    parser.add_argument('--precision', type=str, default='bf16-mixed')
-    parser.add_argument('--downstream_task', type=str, default='action', choices=['action', 'synthesis', 'caption', 'chebi'])
-    parser.add_argument('--max_epochs', type=int, default=10)
-    parser.add_argument('--enable_flash', action='store_true', default=False)
-    parser.add_argument('--disable_graph_cache', action='store_true', default=False)
-    parser.add_argument('--predict_rxn_condition', action='store_true', default=False)
-    parser.add_argument('--generate_restrict_tokens', action='store_true', default=False)
-    parser.add_argument('--train_restrict_tokens', action='store_true', default=False)
-    parser.add_argument('--smiles_type', type=str, default='default', choices=['default', 'canonical', 'restricted', 'unrestricted', 'r_smiles'])
-    parser.add_argument('--accumulate_grad_batches', type=int, default=1)
-    parser.add_argument('--tqdm_interval', type=int, default=50)
-    parser.add_argument('--check_val_every_n_epoch', type=int, default=1)
-    args = parser.parse_args()
+    demo.launch()
 
-    if args.enable_flash:
-        replace_opt_attn_with_flash_attn()
-    print("=========================================")
-    for k, v in sorted(vars(args).items()):
-        print(k, '=', v)
-    print("=========================================")
-    return args
 
-if __name__ == '__main__':
-    main(get_args())
 
+if __name__ == '__main__':
+    args = get_args()
+    vars(args).update(app_config)
+    main(args)