yrshi commited on
Commit
f37ea64
1 Parent(s): 5c814f7

modified TOKENIZERS_PARALLELISM to false

Browse files
Files changed (2) hide show
  1. app.py +4 -2
  2. demo.py +192 -97
app.py CHANGED
@@ -18,6 +18,8 @@ from data_provider.data_utils import smiles2data, reformat_smiles
18
  import gradio as gr
19
  from datetime import datetime
20
 
 
 
21
  ## for pyg bug
22
  warnings.filterwarnings('ignore', category=UserWarning, message='TypedStorage is deprecated')
23
  ## for A5000 gpus
@@ -131,7 +133,7 @@ class InferenceRunner:
131
  solvent = smiles_split(solvent) if solvent else []
132
  assert reactant and product
133
  except:
134
- raise KeyError('Please input a valid reaction string')
135
 
136
  extracted_molecules = {product[0]: "$-1$"}
137
  for mol in reactant+solvent:
@@ -304,7 +306,7 @@ def main(args):
304
  btn.click(fn=online_chat, inputs=[reaction_string, temperature], outputs=[out])
305
  clear_btn.click(fn=lambda:("", ""), inputs=[], outputs=[reaction_string, out])
306
 
307
- demo.launch(share=True)
308
 
309
 
310
 
 
18
  import gradio as gr
19
  from datetime import datetime
20
 
21
+ ## disable online tokenizers parallelism to avoid deadlocks
22
+ os.environ["TOKENIZERS_PARALLELISM"] = "false"
23
  ## for pyg bug
24
  warnings.filterwarnings('ignore', category=UserWarning, message='TypedStorage is deprecated')
25
  ## for A5000 gpus
 
133
  solvent = smiles_split(solvent) if solvent else []
134
  assert reactant and product
135
  except:
136
+ raise gr.Error('Please input a valid reaction string')
137
 
138
  extracted_molecules = {product[0]: "$-1$"}
139
  for mol in reactant+solvent:
 
306
  btn.click(fn=online_chat, inputs=[reaction_string, temperature], outputs=[out])
307
  clear_btn.click(fn=lambda:("", ""), inputs=[], outputs=[reaction_string, out])
308
 
309
+ demo.launch()
310
 
311
 
312
 
demo.py CHANGED
@@ -1,30 +1,114 @@
 
 
 
 
1
  import os
2
  import torch
3
  import argparse
4
  import warnings
5
- import pytorch_lightning as pl
6
- from pytorch_lightning import Trainer, strategies
7
- import pytorch_lightning.callbacks as plc
8
- from pytorch_lightning.loggers import CSVLogger
9
- from pytorch_lightning.callbacks import TQDMProgressBar
10
  from data_provider.pretrain_dm import PretrainDM
11
  from data_provider.tune_dm import *
12
  from model.opt_flash_attention import replace_opt_attn_with_flash_attn
13
  from model.blip2_model import Blip2Model
14
- from model.dist_funs import MyDeepSpeedStrategy
15
- from data_provider.reaction_action_dataset import ActionDataset
16
  from data_provider.data_utils import json_read, json_write
17
  from data_provider.data_utils import smiles2data, reformat_smiles
 
 
18
 
 
 
19
  ## for pyg bug
20
  warnings.filterwarnings('ignore', category=UserWarning, message='TypedStorage is deprecated')
21
  ## for A5000 gpus
22
  torch.set_float32_matmul_precision('medium') # can be medium (bfloat16), high (tensorfloat32), highest (float32)
23
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
24
 
25
  class InferenceRunner:
26
  def __init__(self, model, tokenizer, rxn_max_len, smi_max_len,
27
- smiles_type='default', device='cuda', predict_rxn_condition=True, args=None):
28
  self.model = model
29
  self.rxn_max_len = rxn_max_len
30
  self.smi_max_len = smi_max_len
@@ -36,11 +120,42 @@ class InferenceRunner:
36
  self.collater = Collater([], [])
37
  self.device = device
38
  self.smiles_type = smiles_type
39
- self.predict_rxn_condition = predict_rxn_condition
40
  self.args = args
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
41
 
42
- def make_prompt(self, param_dict, smi_max_len=128, predict_rxn_condition=False):
43
- action_sequence = param_dict['actions']
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
44
  smiles_list = []
45
  prompt = ''
46
  prompt += 'Reactants: '
@@ -72,58 +187,44 @@ class InferenceRunner:
72
  prompt += f'[START_SMILES]{smiles_wrapper(smi)}[END_SMILES] '
73
  smiles_list.append(smi)
74
 
75
- if predict_rxn_condition:
76
- for value, token in param_dict['extracted_duration'].items():
77
- action_sequence = action_sequence.replace(token, value)
78
- for value, token in param_dict['extracted_temperature'].items():
79
- action_sequence = action_sequence.replace(token, value)
80
- else:
81
- prompt += 'Temperatures: '
82
- for value, token in param_dict['extracted_temperature'].items():
83
- prompt += f'{token}: {value} '
84
-
85
- prompt += 'Durations: '
86
- for value, token in param_dict['extracted_duration'].items():
87
- prompt += f'{token}: {value} '
88
-
89
  prompt += 'Action Squence: '
90
- return prompt, smiles_list, action_sequence
91
 
92
  def get_action_elements(self, rxn_dict):
93
- rxn_id = rxn_dict['index']
94
- input_text, smiles_list, output_text = self.make_prompt(rxn_dict, self.smi_max_len, self.predict_rxn_condition)
95
- output_text = output_text.strip() + '\n'
96
 
97
  graph_list = []
98
  for smiles in smiles_list:
99
  graph_item = smiles2data(smiles)
100
  graph_list.append(graph_item)
101
- return rxn_id, graph_list, output_text, input_text
102
-
 
103
  @torch.no_grad()
104
- def predict(self, rxn_dict):
105
- rxn_id, graphs, prompt_tokens, output_text, input_text = self.tokenize(rxn_dict)
106
- result_dict = {
107
- 'raw': rxn_dict,
108
- 'index': rxn_id,
109
- 'input': input_text,
110
- 'target': output_text
111
- }
112
  samples = {'graphs': graphs, 'prompt_tokens': prompt_tokens}
113
- with torch.no_grad():
114
- result_dict['prediction'] = self.model.blip2opt.generate(
115
- samples,
116
- do_sample=self.args.do_sample,
117
- num_beams=self.args.num_beams,
118
- max_length=self.args.max_inference_len,
119
- min_length=self.args.min_inference_len,
120
- num_captions=self.args.num_generate_captions,
121
- use_graph=True
122
- )
 
 
 
 
 
123
  return result_dict
124
 
 
125
  def tokenize(self, rxn_dict):
126
- rxn_id, graph_list, output_text, input_text = self.get_action_elements(rxn_dict)
127
  if graph_list:
128
  graphs = self.collater(graph_list).to(self.device)
129
  input_prompt = smiles_handler(input_text, self.mol_ph, self.is_gal)[0]
@@ -139,13 +240,10 @@ class InferenceRunner:
139
  return_attention_mask=True).to(self.device)
140
  is_mol_token = input_prompt_tokens.input_ids == self.mol_token_id
141
  input_prompt_tokens['is_mol_token'] = is_mol_token
142
- return rxn_id, graphs, input_prompt_tokens, output_text, input_text
143
-
144
 
145
  def main(args):
146
  device = torch.device('cuda')
147
- data_list = json_read('demo.json')
148
- pl.seed_everything(args.seed)
149
  # model
150
  if args.init_checkpoint:
151
  model = Blip2Model(args).to(device)
@@ -171,54 +269,51 @@ def main(args):
171
  rxn_max_len=args.rxn_max_len,
172
  smi_max_len=args.smi_max_len,
173
  device=device,
174
- predict_rxn_condition=args.predict_rxn_condition,
175
  args=args
176
  )
 
 
177
 
178
- import time
179
- for data_item in data_list:
180
- t1 = time.time()
181
- result = infer_runner.predict(data_item)
182
- print(result)
183
- print(f"Time: {time.time() - t1:.2f}s")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
184
 
 
 
 
 
 
 
 
 
 
 
185
 
186
- def get_args():
187
- parser = argparse.ArgumentParser()
188
- parser.add_argument('--filename', type=str, default="main")
189
- parser.add_argument('--seed', type=int, default=42, help='random seed')
190
- # MM settings
191
- parser.add_argument('--mode', type=str, default='pretrain', choices=['pretrain', 'ft', 'eval', 'pretrain_eval'])
192
- parser.add_argument('--strategy_name', type=str, default='mydeepspeed')
193
- parser.add_argument('--iupac_prediction', action='store_true', default=False)
194
- parser.add_argument('--ckpt_path', type=str, default=None)
195
- # parser = Trainer.add_argparse_args(parser)
196
- parser = Blip2Model.add_model_specific_args(parser) # add model args
197
- parser = PretrainDM.add_model_specific_args(parser)
198
- parser.add_argument('--accelerator', type=str, default='gpu')
199
- parser.add_argument('--devices', type=str, default='0,1,2,3')
200
- parser.add_argument('--precision', type=str, default='bf16-mixed')
201
- parser.add_argument('--downstream_task', type=str, default='action', choices=['action', 'synthesis', 'caption', 'chebi'])
202
- parser.add_argument('--max_epochs', type=int, default=10)
203
- parser.add_argument('--enable_flash', action='store_true', default=False)
204
- parser.add_argument('--disable_graph_cache', action='store_true', default=False)
205
- parser.add_argument('--predict_rxn_condition', action='store_true', default=False)
206
- parser.add_argument('--generate_restrict_tokens', action='store_true', default=False)
207
- parser.add_argument('--train_restrict_tokens', action='store_true', default=False)
208
- parser.add_argument('--smiles_type', type=str, default='default', choices=['default', 'canonical', 'restricted', 'unrestricted', 'r_smiles'])
209
- parser.add_argument('--accumulate_grad_batches', type=int, default=1)
210
- parser.add_argument('--tqdm_interval', type=int, default=50)
211
- parser.add_argument('--check_val_every_n_epoch', type=int, default=1)
212
- args = parser.parse_args()
213
 
214
- if args.enable_flash:
215
- replace_opt_attn_with_flash_attn()
216
- print("=========================================")
217
- for k, v in sorted(vars(args).items()):
218
- print(k, '=', v)
219
- print("=========================================")
220
- return args
221
 
222
- if __name__ == '__main__':
223
- main(get_args())
224
 
 
 
 
 
 
1
+ import subprocess
2
+ subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
3
+ subprocess.run('pip install -U timm', shell=True)
4
+ import spaces
5
  import os
6
  import torch
7
  import argparse
8
  import warnings
9
+ from rdkit import Chem
10
+ from rdkit.Chem import CanonSmiles
11
+ from rdkit.Chem import MolFromSmiles, MolToSmiles
 
 
12
  from data_provider.pretrain_dm import PretrainDM
13
  from data_provider.tune_dm import *
14
  from model.opt_flash_attention import replace_opt_attn_with_flash_attn
15
  from model.blip2_model import Blip2Model
 
 
16
  from data_provider.data_utils import json_read, json_write
17
  from data_provider.data_utils import smiles2data, reformat_smiles
18
+ import gradio as gr
19
+ from datetime import datetime
20
 
21
+ ## disable online tokenizers parallelism to avoid deadlocks
22
+ os.environ["TOKENIZERS_PARALLELISM"] = "false"
23
  ## for pyg bug
24
  warnings.filterwarnings('ignore', category=UserWarning, message='TypedStorage is deprecated')
25
  ## for A5000 gpus
26
  torch.set_float32_matmul_precision('medium') # can be medium (bfloat16), high (tensorfloat32), highest (float32)
27
 
28
+ def smiles_split(string, separator='.'):
29
+ string = str(string)
30
+ mols = []
31
+ for smi in string.split(separator):
32
+ mol = MolFromSmiles(smi)
33
+ if mol is None:
34
+ continue # Skip invalid SMILES strings
35
+ mols.append(mol)
36
+
37
+ parts = []
38
+ current_part = []
39
+ charge_count = 0
40
+
41
+ for mol in mols:
42
+ charge = Chem.GetFormalCharge(mol)
43
+ if charge==0:
44
+ if current_part:
45
+ smiles = '.'.join([MolToSmiles(m) for m in current_part])
46
+ smiles = CanonSmiles(smiles)
47
+ parts.append(smiles)
48
+ current_part = []
49
+ charge_count = 0
50
+ parts.append(MolToSmiles(mol))
51
+ else:
52
+ charge_count += charge
53
+ current_part.append(mol)
54
+ if charge_count == 0:
55
+ smiles = '.'.join([MolToSmiles(m) for m in current_part])
56
+ smiles = CanonSmiles(smiles)
57
+ parts.append(smiles)
58
+ current_part = []
59
+ charge_count = 0
60
+ if current_part:
61
+ smiles = '.'.join([MolToSmiles(m) for m in current_part])
62
+ smiles = CanonSmiles(smiles)
63
+ parts.append(smiles)
64
+
65
+ return parts
66
+
67
+ def get_args():
68
+ parser = argparse.ArgumentParser()
69
+ parser.add_argument('--filename', type=str, default="main")
70
+ parser.add_argument('--seed', type=int, default=42, help='random seed')
71
+ # MM settings
72
+ parser.add_argument('--mode', type=str, default='pretrain', choices=['pretrain', 'ft', 'eval', 'pretrain_eval'])
73
+ parser.add_argument('--strategy_name', type=str, default='mydeepspeed')
74
+ parser.add_argument('--iupac_prediction', action='store_true', default=False)
75
+ parser.add_argument('--ckpt_path', type=str, default=None)
76
+ # parser = Trainer.add_argparse_args(parser)
77
+ parser = Blip2Model.add_model_specific_args(parser) # add model args
78
+ parser = PretrainDM.add_model_specific_args(parser)
79
+ parser.add_argument('--accelerator', type=str, default='gpu')
80
+ parser.add_argument('--devices', type=str, default='0,1,2,3')
81
+ parser.add_argument('--precision', type=str, default='bf16-mixed')
82
+ parser.add_argument('--downstream_task', type=str, default='action', choices=['action', 'synthesis', 'caption', 'chebi'])
83
+ parser.add_argument('--max_epochs', type=int, default=10)
84
+ parser.add_argument('--enable_flash', action='store_true', default=False)
85
+ parser.add_argument('--disable_graph_cache', action='store_true', default=False)
86
+ parser.add_argument('--generate_restrict_tokens', action='store_true', default=False)
87
+ parser.add_argument('--train_restrict_tokens', action='store_true', default=False)
88
+ parser.add_argument('--smiles_type', type=str, default='default', choices=['default', 'canonical', 'restricted', 'unrestricted', 'r_smiles'])
89
+ parser.add_argument('--accumulate_grad_batches', type=int, default=1)
90
+ parser.add_argument('--tqdm_interval', type=int, default=50)
91
+ parser.add_argument('--check_val_every_n_epoch', type=int, default=1)
92
+ args = parser.parse_args()
93
+
94
+ if args.enable_flash:
95
+ replace_opt_attn_with_flash_attn()
96
+ return args
97
+
98
+ app_config = {
99
+ "init_checkpoint": "all_checkpoints/ckpt_tune_hybridFeb11_May31/last_converted.ckpt",
100
+ "filename": "app",
101
+ "opt_model": "facebook/galactica-1.3b",
102
+ "num_workers": 4,
103
+ "rxn_max_len": 512,
104
+ "text_max_len": 512,
105
+ "precision": "bf16-mixed",
106
+ "max_inference_len": 512,
107
+ }
108
 
109
  class InferenceRunner:
110
  def __init__(self, model, tokenizer, rxn_max_len, smi_max_len,
111
+ smiles_type='default', device='cuda', args=None):
112
  self.model = model
113
  self.rxn_max_len = rxn_max_len
114
  self.smi_max_len = smi_max_len
 
120
  self.collater = Collater([], [])
121
  self.device = device
122
  self.smiles_type = smiles_type
 
123
  self.args = args
124
+ time_stamp = datetime.now().strftime("%Y.%m.%d-%H:%M")
125
+ self.cache_dir = f'results/{self.args.filename}/{time_stamp}'
126
+ os.makedirs(self.cache_dir, exist_ok=True)
127
+
128
+ def make_query_dict(self, rxn_string):
129
+ try:
130
+ reactant, solvent, product = rxn_string.split('>')
131
+ reactant = smiles_split(reactant)
132
+ product = smiles_split(product)
133
+ solvent = smiles_split(solvent) if solvent else []
134
+ assert reactant and product
135
+ except:
136
+ raise gr.Error('Please input a valid reaction string')
137
+
138
+ extracted_molecules = {product[0]: "$-1$"}
139
+ for mol in reactant+solvent:
140
+ extracted_molecules[mol] = f"${len(extracted_molecules)}$"
141
 
142
+ result_dict = {}
143
+ result_dict['time_stamp'] = datetime.now().strftime("%Y.%m.%d %H:%M:%S.%f")[:-3]
144
+ result_dict['reaction_string'] = rxn_string
145
+ result_dict['REACTANT'] = reactant
146
+ result_dict['SOLVENT'] = solvent
147
+ result_dict['CATALYST'] = []
148
+ result_dict['PRODUCT'] = product
149
+ result_dict['extracted_molecules'] = extracted_molecules
150
+ return result_dict
151
+
152
+ def save_prediction(self, result_dict):
153
+ os.makedirs(self.cache_dir, exist_ok=True)
154
+ result_id = result_dict['time_stamp']
155
+ result_path = os.path.join(self.cache_dir, f'{result_id}.json')
156
+ json_write(result_path, result_dict)
157
+
158
+ def make_prompt(self, param_dict, smi_max_len=128):
159
  smiles_list = []
160
  prompt = ''
161
  prompt += 'Reactants: '
 
187
  prompt += f'[START_SMILES]{smiles_wrapper(smi)}[END_SMILES] '
188
  smiles_list.append(smi)
189
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
190
  prompt += 'Action Squence: '
191
+ return prompt, smiles_list
192
 
193
  def get_action_elements(self, rxn_dict):
194
+ input_text, smiles_list = self.make_prompt(rxn_dict, self.smi_max_len)
 
 
195
 
196
  graph_list = []
197
  for smiles in smiles_list:
198
  graph_item = smiles2data(smiles)
199
  graph_list.append(graph_item)
200
+ return graph_list, input_text
201
+
202
+ @spaces.GPU
203
  @torch.no_grad()
204
+ def predict(self, rxn_dict, temperature=1):
205
+ graphs, prompt_tokens = self.tokenize(rxn_dict)
206
+ result_dict = rxn_dict
 
 
 
 
 
207
  samples = {'graphs': graphs, 'prompt_tokens': prompt_tokens}
208
+ assert prompt_tokens.input_ids.is_cuda
209
+ assert graphs.is_cuda
210
+ prediction = self.model.blip2opt.generate(
211
+ samples,
212
+ do_sample=self.args.do_sample,
213
+ num_beams=self.args.num_beams,
214
+ max_length=self.args.max_inference_len,
215
+ min_length=self.args.min_inference_len,
216
+ num_captions=self.args.num_generate_captions,
217
+ temperature=temperature,
218
+ use_graph=True
219
+ )[0]
220
+ for k, v in result_dict['extracted_molecules'].items():
221
+ prediction = prediction.replace(v, k)
222
+ result_dict['prediction'] = prediction
223
  return result_dict
224
 
225
+ @spaces.GPU
226
  def tokenize(self, rxn_dict):
227
+ graph_list, input_text = self.get_action_elements(rxn_dict)
228
  if graph_list:
229
  graphs = self.collater(graph_list).to(self.device)
230
  input_prompt = smiles_handler(input_text, self.mol_ph, self.is_gal)[0]
 
240
  return_attention_mask=True).to(self.device)
241
  is_mol_token = input_prompt_tokens.input_ids == self.mol_token_id
242
  input_prompt_tokens['is_mol_token'] = is_mol_token
243
+ return graphs, input_prompt_tokens
 
244
 
245
  def main(args):
246
  device = torch.device('cuda')
 
 
247
  # model
248
  if args.init_checkpoint:
249
  model = Blip2Model(args).to(device)
 
269
  rxn_max_len=args.rxn_max_len,
270
  smi_max_len=args.smi_max_len,
271
  device=device,
 
272
  args=args
273
  )
274
+ example_inputs = json_read('demo.json')
275
+ example_inputs = [[e] for e in example_inputs]
276
 
277
+ def online_chat(reaction_string, temperature=1):
278
+ data_item = infer_runner.make_query_dict(reaction_string)
279
+ result = infer_runner.predict(data_item, temperature=temperature)
280
+ infer_runner.save_prediction(result)
281
+ prediction = result['prediction'].replace(' ; ', ' ;\n')
282
+ return prediction
283
+
284
+ with gr.Blocks(css="""
285
+ .center { display: flex; justify-content: center; }
286
+ """) as demo:
287
+ gr.HTML(
288
+ """
289
+ <center><h1><b>ReactXT</b></h1></center>
290
+ <p style="font-size:20px; font-weight:bold;">This is the demo page of our ACL 2024 paper
291
+ <i>ReactXT: Understanding Molecular “Reaction-ship” via Reaction-Contextualized Molecule-Text Pretraining.</i></p>
292
+ """)
293
+ with gr.Row(elem_classes="center"):
294
+ gr.Image(value="./figures/frameworks.jpg", elem_classes="center", width=800, label="Framework of ReactXT")
295
+ gr.HTML(
296
+ """
297
+ <p style="font-size:16px;"> Please input one chemical reaction below, and we will generate the predicted experimental procedure.</p>
298
+ <p style="font-size:16px;"> The reaction should be in form of <b>Reactants>Reagents>Product</b>.</p>
299
+ """)
300
 
301
+ reaction_string = gr.Textbox(placeholder="Input one reaction", label='Input Reaction')
302
+ gr.Examples(example_inputs, [reaction_string,], fn=online_chat, label='Example Reactions')
303
+ with gr.Row():
304
+ btn = gr.Button("Submit")
305
+ clear_btn = gr.Button("Clear")
306
+ temperature = gr.Slider(0.1, 1, value=1, label='Temperature')
307
+ with gr.Row():
308
+ out = gr.Textbox(label="ReactXT's Output", placeholder="Predicted experimental procedure")
309
+ btn.click(fn=online_chat, inputs=[reaction_string, temperature], outputs=[out])
310
+ clear_btn.click(fn=lambda:("", ""), inputs=[], outputs=[reaction_string, out])
311
 
312
+ demo.launch()
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
313
 
 
 
 
 
 
 
 
314
 
 
 
315
 
316
+ if __name__ == '__main__':
317
+ args = get_args()
318
+ vars(args).update(app_config)
319
+ main(args)