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 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 ## 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) class InferenceRunner: def __init__(self, model, tokenizer, rxn_max_len, smi_max_len, smiles_type='default', device='cuda', predict_rxn_condition=True, 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.predict_rxn_condition = predict_rxn_condition self.args = args def make_prompt(self, param_dict, smi_max_len=128, predict_rxn_condition=False): action_sequence = param_dict['actions'] 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) 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 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' 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 @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 } 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 ) return result_dict def tokenize(self, rxn_dict): rxn_id, graph_list, output_text, 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 rxn_id, graphs, input_prompt_tokens, output_text, input_text 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) 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, predict_rxn_condition=args.predict_rxn_condition, args=args ) 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 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() 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())