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 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 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', args=None): self.model = model self.rxn_max_len = rxn_max_len self.smi_max_len = smi_max_len self.tokenizer = tokenizer self.collater = Collater([], []) self.mol_ph = '' * args.num_query_token self.mol_token_id = tokenizer.mol_token_id self.is_gal = args.opt_model.find('galactica') >= 0 self.collater = Collater([], []) self.device = device self.smiles_type = smiles_type 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)}$" 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: ' smiles_wrapper = lambda x: reformat_smiles(x, smiles_type=self.smiles_type)[:smi_max_len] for smi in param_dict['REACTANT']: prompt += f'{param_dict["extracted_molecules"][smi]}: [START_SMILES]{smiles_wrapper(smi)}[END_SMILES] ' smiles_list.append(smi) prompt += 'Product: ' for smi in param_dict['PRODUCT']: prompt += f'{param_dict["extracted_molecules"][smi]}: [START_SMILES]{smiles_wrapper(smi)}[END_SMILES] ' smiles_list.append(smi) if param_dict['CATALYST']: prompt += 'Catalysts: ' for smi in param_dict['CATALYST']: if smi in param_dict["extracted_molecules"]: prompt += f'{param_dict["extracted_molecules"][smi]}: [START_SMILES]{smiles_wrapper(smi)}[END_SMILES] ' else: prompt += f'[START_SMILES]{smiles_wrapper(smi)}[END_SMILES] ' smiles_list.append(smi) if param_dict['SOLVENT']: prompt += 'Solvents: ' for smi in param_dict['SOLVENT']: if smi in param_dict["extracted_molecules"]: prompt += f'{param_dict["extracted_molecules"][smi]}: [START_SMILES]{smiles_wrapper(smi)}[END_SMILES] ' else: prompt += f'[START_SMILES]{smiles_wrapper(smi)}[END_SMILES] ' smiles_list.append(smi) prompt += 'Action Squence: ' return prompt, smiles_list def get_action_elements(self, rxn_dict): 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 graph_list, input_text @spaces.GPU @torch.no_grad() 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} 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): 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] ## deal with prompt self.tokenizer.padding_side = 'left' input_prompt_tokens = self.tokenizer(input_prompt, truncation=True, padding='max_length', add_special_tokens=True, max_length=self.rxn_max_len, return_tensors='pt', 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 graphs, input_prompt_tokens def main(args): device = torch.device('cuda') # model if args.init_checkpoint: model = Blip2Model(args).to(device) ckpt = torch.load(args.init_checkpoint, map_location='cpu') model.load_state_dict(ckpt['state_dict'], strict=False) print(f"loaded model from {args.init_checkpoint}") else: model = Blip2Model(args).to(device) model.eval() print('total params:', sum(p.numel() for p in model.parameters())) if args.opt_model.find('galactica') >= 0 or args.opt_model.find('t5') >= 0: tokenizer = model.blip2opt.opt_tokenizer elif args.opt_model.find('llama') >= 0 or args.opt_model.find('vicuna') >= 0: tokenizer = model.blip2opt.llm_tokenizer else: raise NotImplementedError infer_runner = InferenceRunner( model=model, tokenizer=tokenizer, rxn_max_len=args.rxn_max_len, smi_max_len=args.smi_max_len, device=device, args=args ) example_inputs = json_read('demo.json') example_inputs = [[e] for e in example_inputs] 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( """

ReactXT

This is the demo page of our ACL 2024 paper ReactXT: Understanding Molecular “Reaction-ship” via Reaction-Contextualized Molecule-Text Pretraining.

""") with gr.Row(elem_classes="center"): gr.Image(value="./figures/frameworks.jpg", elem_classes="center", width=800, label="Framework of ReactXT") gr.HTML( """

Please input one chemical reaction below, and we will generate the predicted experimental procedure.

The reaction should be in form of Reactants>Reagents>Product.

""") 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]) demo.launch() if __name__ == '__main__': args = get_args() vars(args).update(app_config) main(args)