infer_one.py

import sys
import os
import pickle
import numpy as np
import rdkit
import random
import torch
import torch.nn as nn
import pandas as pd
from rdchiral.main import rdchiralReaction, rdchiralReactants, rdchiralRun

from datetime import datetime
from pathlib import Path
from typing import Dict, List, Optional, Tuple, Union
from scipy import sparse
from tqdm import tqdm
from rdkit import RDLogger
from rdkit import Chem

from model import TemplateNN_Highway
from prepare_data import mol_smi_to_count_fp
from infer_config import infer_config

DATA_FOLDER = Path(__file__).resolve().parent / 'data'
CHECKPOINT_FOLDER = Path(__file__).resolve().parent / 'checkpoint'

class Proposer:
    def __init__(self, infer_config: Dict) -> None:
        super().__init__()
        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

        print(f"Loading templates from file: {infer_config['templates_file']}")
        with open(DATA_FOLDER / infer_config['templates_file'], 'r') as f:
            templates = f.readlines()
        self.templates_filtered = []
        for p in templates:
            pa, cnt = p.strip().split(': ')
            if int(cnt) >= infer_config['min_freq']:
                self.templates_filtered.append(pa)
        print(f'Total number of template patterns: {len(self.templates_filtered)}')

        self.model, self.indices = self.build_model(infer_config)
        print('Done initializing proposer\n')

    def build_model(self, infer_config: Dict):
         # load model from checkpoint
        checkpoint = torch.load(
            CHECKPOINT_FOLDER / f"{infer_config['expt_name']}.pth.tar",
            map_location=self.device,
        )
        model = TemplateNN_Highway(
            output_size=len(self.templates_filtered),
            size=infer_config['hidden_size'],
            num_layers_body=infer_config['depth'],
            input_size=infer_config['final_fp_size']
        )
        model.load_state_dict(checkpoint["state_dict"])
        model.to(self.device)

        indices = np.loadtxt(DATA_FOLDER / 'variance_indices.txt').astype('int')
        return model, indices

    def propose(self, 
                input_smiles: List[str],
                topk: int = 5,
                **kwargs) -> List[Dict[str, List]]:

        results = {}
        with torch.no_grad():
            for smi in tqdm(input_smiles, desc='Proposing precursors'):
                prod_fp = mol_smi_to_count_fp(smi, infer_config['radius'], infer_config['orig_fp_size'])
                logged = sparse.csr_matrix(np.log(prod_fp.toarray() + 1))
                final_fp = logged[:, self.indices]
                final_fp = torch.as_tensor(final_fp.toarray()).float().to(self.device)

                outputs = self.model(final_fp)
                outputs = nn.Softmax(dim=1)(outputs)
                preds = torch.topk(outputs, k=topk, dim=1)[1].squeeze(dim=0).cpu().numpy()
 
                results[smi] = []
                for idx in preds:
                    score = outputs[0, idx.item()].item()
                    template = self.templates_filtered[idx.item()]
                    rxn = rdchiralReaction(template)
                    prod = rdchiralReactants(smi)
                    try:
                        precs = rdchiralRun(rxn, prod)
                    except:
                        precs = 'N/A'
                    results[smi].append((precs, score)) # Tuple[precs, score] where precs is a List[str]
        return results

if __name__ == '__main__':
    proposer = Proposer(infer_config)
    sample_products = [
        'CCOC(C#Cc1cncc(S(C)(=O)=O)c1)(OCC)OCC',
        'COC(=O)c1cccc2[nH]c(NCC3CCNCC3)nc12',
        'CC(C)(C)OC(=O)N1CC[C@H](NC(=O)C(F)(F)F)C1',
    ]
    results = proposer.propose(sample_products, topk=5)
    print(results)

    # should see this, key = product_smi, value = List of Tuple[precs, score]
    # {'CCOC(C#Cc1cncc(S(C)(=O)=O)c1)(OCC)OCC': [(['C#CC(OCC)(OCC)OCC.CS(=O)(=O)c1cncc(Br)c1'], 0.2964268922805786), 
    # ([], 0.05189067870378494), ([], 0.022873425856232643), ([], 0.0173777025192976), ([], 0.01714753918349743)], 
    
    # 'COC(=O)c1cccc2[nH]c(NCC3CCNCC3)nc12': [(['COC(=O)c1cccc2[nH]c(NCC3CCN(C(=O)OC(C)(C)C)CC3)nc12'], 0.9992377758026123), 
    # (['COC(=O)c1cccc2[nH]c(NCC3CCN(C(=O)OCc4ccccc4)CC3)nc12'], 0.0002514408261049539), (['COC(=O)c1cccc2[nH]c(NCC3CCN(C(=O)C(F)(F)F)CC3)nc12'], 0.00024452427169308066), 
    # (['COC(=O)c1cccc2[nH]c(NCC3CCN(Cc4ccccc4)CC3)nc12'], 0.00012763732229359448), (['COC(=O)c1cccc2[nH]c(NCc3ccncc3)nc12'], 4.4018081098329276e-05)], 
    
    # 'CC(C)(C)OC(=O)N1CC[C@H](NC(=O)C(F)(F)F)C1': [(['CC(C)(C)OC(=O)N1CC[C@H](N)C1.O=C(OC(=O)C(F)(F)F)C(F)(F)F'], 0.7076814770698547), 
    # (['CC(C)(C)OC(=O)N1CC[C@H](N)C1.O=C(Br)C(F)(F)F'], 0.039315130561590195), ([], 0.031778812408447266), 
    # ([], 0.030593203380703926), ([], 0.01709393411874771)]}