demo.py

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 = '<mol>' * 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())