app.py

from e_smiles import *
import gradio as gr
from rdkit import Chem
from rdkit.Chem import Draw
import os


def remove_atom_mapping_and_isotopes(smiles):
    mol = Chem.MolFromSmiles(smiles)
    
    if mol is None:
        raise ValueError("Invalid SMILES string")
    
    for atom in mol.GetAtoms():
        atom.SetIsotope(0)
        atom.SetAtomMapNum(0)
    
    clean_smiles = Chem.MolToSmiles(Chem.RemoveHs(mol))
    return clean_smiles

def get_bondtype_and_convert2string(mol, atom1_idx, atom2_idx):
    atom1_idx = int(atom1_idx)
    atom2_idx = int(atom2_idx)
    bond = mol.GetBondBetweenAtoms(atom1_idx - 1, atom2_idx - 1)
    if bond is None:
        raise ValueError("Atoms are not bonded")
    
    if bond.GetBondType() == Chem.BondType.SINGLE:
        return '1.0'
    elif bond.GetBondType() == Chem.BondType.DOUBLE:
        return '2.0'
    elif bond.GetBondType() == Chem.BondType.TRIPLE:
        return '3.0'
    else:
        raise ValueError("Unsupported bond type")
    
def main_interface(smiles):
    img, mapped_smiles, secondary_inputs_visible1, secondary_inputs_visible2, error = process_smiles(smiles)
    return img, mapped_smiles, gr.update(visible=True), secondary_inputs_visible2, error

def process_smiles(smiles):
    try:
        smiles = remove_atom_mapping_and_isotopes(smiles)
        mol = Chem.MolFromSmiles(smiles)
        if mol:
            Chem.Kekulize(mol, clearAromaticFlags=True)
            for atom in mol.GetAtoms():
                atom.SetAtomMapNum(atom.GetIdx() + 1)

            mapped_smiles = Chem.MolToSmiles(mol, canonical=False, kekuleSmiles=True)
        else:
            raise ValueError("Invalid SMILES string")

        mol = Chem.MolFromSmiles(mapped_smiles)
        img = Draw.MolToImage(mol, size=(450, 450))
        return img, mapped_smiles, gr.update(visible=True), gr.update(visible=True), gr.update(visible=False)  
    except Exception as e:
        return None, None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=True, value=str(e))

def process_bond_operation(smiles, atom1_idx, atom2_idx, operation):
    try:
        mol = Chem.MolFromSmiles(smiles)
        if mol is None:
            raise ValueError("Invalid SMILES string")
        
        if atom1_idx > atom2_idx: 
            atom1_idx, atom2_idx = atom2_idx, atom1_idx

        first = str(atom1_idx)
        second = str(atom2_idx)
        third = get_bondtype_and_convert2string(mol, atom1_idx, atom2_idx)
        
        if operation == "break this bond":
            prompt = [first + ':' + second + ':' + third + ':' + '0.0']
        elif operation == "change this bond to single bond":
            prompt = [first + ':' + second + ':' + third + ':' + '1.0']
        elif operation == "change this bond to double bond":
            prompt = [first + ':' + second + ':' + third + ':' + '2.0']
        elif operation == "change this bond to triple bond":
            prompt = [first + ':' + second + ':' + third + ':' + '3.0']
        else:
            raise ValueError("Unsupported operation")
        
        prompt = get_b_smiles_check([smiles, prompt, [], [], [], [], []])

        return prompt, gr.update(visible=False)  
    except Exception as e:
        return None, gr.update(visible=True, value=str(e))

def retrosynthesis(reactant, prompt):
    # 1. generate src file
    with open('tmp_data/src.txt', 'w') as f:
        f.write(" ".join(prompt))

    # 2. inference
    os.system('bash infer.sh')


    # 3. readout reactseq
    with open('./tmp_data/tgt.txt', 'r') as f:
        lines = f.readlines()

    reactseq_list = []
    reactant_smiles = []
    for line in lines:
        line = line.replace(" ", "").rstrip("\n")
        reactseq_list.append(line)
        rxn = reactant + ">>>" + line

    # 4. merge reactseq
        reactant_smiles.append(merge_smiles_with_mapping(rxn))
        
    # 5. draw reactants
    reactants = []
    reactant_images = []
    for smiles in reactant_smiles:
        mol = Chem.MolFromSmiles(smiles)
        reactants.append(smiles)
        reactant_images.append(Draw.MolToImage(mol, size=(600, 600)))

    return reactseq_list, reactants, reactant_images

prompt_output = gr.Textbox(label="Prompt", interactive=False, visible=True) 

def secondary_interface(smiles, atom1_idx, atom2_idx, operation, progress=gr.Progress()):
    try:
        prompt, error = process_bond_operation(smiles, atom1_idx, atom2_idx, operation)
        
        if prompt:
            reactseq_list, reactant_smiles, reactant_images = retrosynthesis(smiles, prompt)
            output_components = [gr.update(value=prompt)] 

            for i in range(len(reactant_smiles)):
                if reactant_smiles[i] != "":
                    output_components.append(gr.update(value=reactseq_list[i], visible=True))
                    output_components.append(gr.update(value=reactant_smiles[i], visible=True))
                    output_components.append(gr.update(value=reactant_images[i], visible=True))
                else:
                    output_components.append(gr.update(visible=False))
                    output_components.append(gr.update(visible=False))
                    output_components.append(gr.update(visible=False))

            output_components.append(gr.update(visible=False))  # error_output
            return output_components
        
        else:
            return [gr.update(visible=False)] * 31 + [error]
    
    except Exception as e:
        return [gr.update(visible=False)] * 31 + [gr.update(visible=True, value=str(e))]



def clear_interface():
    return (
        gr.update(value=None),  # smiles_input
        gr.update(value=None),  # img_output
        gr.update(value=None),  # smiles_output
        gr.update(visible=False),  # secondary_inputs
        gr.update(visible=False),  # secondary_inputs
        gr.update(visible=False, value=None),  # error_output
        gr.update(value=None, visible=True),  # prompt_output
    ) + (  # Reactant SMILES and Images
        tuple(gr.update(value=None, visible=False) for _ in range(30))
    )

# 示例数据
examples = [
    ["N#CC1=CC=C(C(N2N=CN=C2)C(O)CC3=CC=C(F)C=C3)C=C1", 7, 18, "break this bond"],
    ["N#CC1=CC=C(C(N2N=CN=C2)C(O)CC3=CC=C(F)C=C3)C=C1", 8, 9, "change this bond to double bond"],
    ["N#CC1=CC=C(C(N2N=CN=C2)C(O)CC3=CC=C(F)C=C3)C=C1", 1, 2, "change this bond to single bond"],
]

with gr.Blocks() as interface:
    gr.Markdown("# ReactSeq")
    gr.Markdown("Please input a SMILES string and two atom indices between the bond which you want to perform an operation.")
    gr.Markdown("The operation can be one of the following: break this bond, change this bond to single bond, change this bond to double bond, change this bond to triple bond.")
    gr.Markdown("The molecule image and SMILES with atom mapping will be displayed.")
    gr.Markdown("After you input the operation, the prompt will be displayed.")
    gr.Markdown("The reactants SMILES and images will be displayed after the retrosynthesis with **default** ranking, with the invalid ones deleted.")
    
    smiles_input = gr.Textbox(placeholder="Please input your SMILES string", label="SMILES input")
    submit_smiles_button = gr.Button("Submit SMILES")

    gr.Markdown("The molecule image and SMILES with atom mapping:")
    img_output = gr.Image(label="Molecule image")
    smiles_output = gr.Textbox(label="SMILES with atom mapping", interactive=False)
    
    with gr.Row(visible=False) as secondary_inputs:
        atom1_idx = gr.Number(label="The first atom index", minimum=1)
        atom2_idx = gr.Number(label="The second atom index", minimum=1)
        operation = gr.Dropdown(choices=["break this bond", "change this bond to single bond", "change this bond to double bond", "change this bond to triple bond"], label="Operation")
        submit_operation_button = gr.Button("Submit Operation")
    
    prompt_output = gr.Textbox(label="Prompt", interactive=False, visible=True)
    
    result_row = []
    for i in range(10):
        result_row.append(gr.Textbox(label=f"Reactseq {i+1} ", interactive=False, visible=False))
        result_row.append(gr.Textbox(label=f"Reactant {i+1} SMILES", interactive=False, visible=False))
        result_row.append(gr.Image(label=f"Reactant {i+1} Image", visible=False))

    error_output = gr.Textbox(label="Error Message", interactive=False, visible=False)
    
    clear_button = gr.Button("Clear")

    submit_smiles_button.click(main_interface, inputs=[smiles_input], outputs=[img_output, smiles_output, secondary_inputs, secondary_inputs, error_output])
    submit_operation_button.click(
        secondary_interface, 
        inputs=[smiles_output, atom1_idx, atom2_idx, operation], 
        outputs=[prompt_output] + result_row + [error_output]
    )
    
    clear_button.click(
        clear_interface, 
        inputs=[], 
        outputs=[smiles_input, img_output, smiles_output, secondary_inputs, secondary_inputs, error_output, prompt_output] + result_row
    )

    gr.Examples(examples=examples, inputs=[smiles_input, atom1_idx, atom2_idx, operation])

interface.launch()