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 = []
    for line in lines:
        line = line.replace(" ", "").rstrip("\n")
        reactseq_list.append(reactant + ">>>" + line)

    # 4. merge reactseq
    reactant_smiles = []
    for seq in reactseq_list:
        reactant_smiles.append(merge_smiles_with_mapping(seq))

    # 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

def secondary_interface(smiles, atom1_idx, atom2_idx, operation):
    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, visible=True)]  # prompt_output

            # Add the reactant SMILES and images
            for i in range(len(reactant_smiles)):
                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))

            # Fill remaining hidden outputs
            for j in range(len(reactant_smiles), 10):
                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]  # Assuming 10 pairs of outputs + prompt_output + error_output
    
    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),  # 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=False),  # 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.")
    
    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=2)
        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=[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()